code: /*
   * SPDX-License-Identifier: Apache-2.0
   * Copyright 2018-2021 The Feast Authors
   *
   * Licensed under the Apache License, Version 2.0 (the "License");
   * you may not use this file except in compliance with the License.
   * You may obtain a copy of the License at
   *
   *     https://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package feast.serving.service;
  
  import com.google.common.collect.Lists;
  import com.google.protobuf.ByteString;
  import com.google.protobuf.Timestamp;
  import feast.proto.core.*;
  import feast.proto.serving.ServingAPIProto;
  import feast.proto.serving.TransformationServiceAPIProto.TransformFeaturesRequest;
  import feast.proto.serving.TransformationServiceAPIProto.TransformFeaturesResponse;
  import feast.proto.serving.TransformationServiceAPIProto.ValueType;
  import feast.proto.serving.TransformationServiceGrpc;
  import feast.proto.types.ValueProto;
  import feast.serving.registry.RegistryRepository;
  import io.grpc.ManagedChannel;
  import io.grpc.ManagedChannelBuilder;
  import io.grpc.Status;
  import java.io.ByteArrayOutputStream;
  import java.io.IOException;
  import java.nio.channels.Channels;
  import java.util.*;
  import org.apache.arrow.memory.BufferAllocator;
  import org.apache.arrow.memory.RootAllocator;
  import org.apache.arrow.vector.*;
  import org.apache.arrow.vector.ipc.ArrowFileReader;
  import org.apache.arrow.vector.ipc.ArrowFileWriter;
  import org.apache.arrow.vector.types.FloatingPointPrecision;
  import org.apache.arrow.vector.types.pojo.ArrowType;
  import org.apache.arrow.vector.types.pojo.Field;
  import org.apache.arrow.vector.types.pojo.Schema;
  import org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel;
  import org.apache.commons.lang3.tuple.Pair;
  import org.slf4j.Logger;
  
  public class OnlineTransformationService implements TransformationService {
  
    private static final Logger log =
        org.slf4j.LoggerFactory.getLogger(OnlineTransformationService.class);
    private final TransformationServiceGrpc.TransformationServiceBlockingStub stub;
    private final RegistryRepository registryRepository;
    static final int INT64_BITWIDTH = 64;
    static final int INT32_BITWIDTH = 32;
  
    public OnlineTransformationService(
        String transformationServiceEndpoint, RegistryRepository registryRepository) {
      if (transformationServiceEndpoint != null) {
        final ManagedChannel channel =
            ManagedChannelBuilder.forTarget(transformationServiceEndpoint).usePlaintext().build();
        this.stub = TransformationServiceGrpc.newBlockingStub(channel);
      } else {
        this.stub = null;
      }
      this.registryRepository = registryRepository;
    }
  
    /** {@inheritDoc} */
    @Override
    public TransformFeaturesResponse transformFeatures(
        TransformFeaturesRequest transformFeaturesRequest) {
      if (this.stub == null) {
        throw new RuntimeException(
            "Transformation service endpoint must be configured to enable this functionality.");
      }
      return this.stub.transformFeatures(transformFeaturesRequest);
    }
  
    /** {@inheritDoc} */
    @Override
    public List<ServingAPIProto.FeatureReferenceV2> extractOnDemandFeaturesDependencies(
        List<ServingAPIProto.FeatureReferenceV2> onDemandFeatureReferences) {
      List<ServingAPIProto.FeatureReferenceV2> onDemandFeatureSources = new ArrayList<>();
      for (ServingAPIProto.FeatureReferenceV2 featureReference : onDemandFeatureReferences) {
        OnDemandFeatureViewProto.OnDemandFeatureViewSpec onDemandFeatureViewSpec =
            this.registryRepository.getOnDemandFeatureViewSpec(featureReference);
        Map<String, OnDemandFeatureViewProto.OnDemandSource> sources =
            onDemandFeatureViewSpec.getSourcesMap();
  
        for (OnDemandFeatureViewProto.OnDemandSource source : sources.values()) {
          OnDemandFeatureViewProto.OnDemandSource.SourceCase sourceCase = source.getSourceCase();
          switch (sourceCase) {
            case REQUEST_DATA_SOURCE:
              // Do nothing. The value should be provided as dedicated request parameter
              break;
            case FEATURE_VIEW_PROJECTION:
              FeatureReferenceProto.FeatureViewProjection projection =
                  source.getFeatureViewProjection();
              for (FeatureProto.FeatureSpecV2 featureSpec : projection.getFeatureColumnsList()) {
                String featureName = featureSpec.getName();
                ServingAPIProto.FeatureReferenceV2 onDemandFeatureSource =
                    ServingAPIProto.FeatureReferenceV2.newBuilder()
                        .setFeatureViewName(projection.getFeatureViewName())
                        .setFeatureName(featureName)
                        .build();
                onDemandFeatureSources.add(onDemandFeatureSource);
              }
              break;
            case FEATURE_VIEW:
              FeatureViewProto.FeatureView featureView = source.getFeatureView();
              FeatureViewProto.FeatureViewSpec featureViewSpec = featureView.getSpec();
              String featureViewName = featureViewSpec.getName();
              for (FeatureProto.FeatureSpecV2 featureSpec : featureViewSpec.getFeaturesList()) {
                String featureName = featureSpec.getName();
                ServingAPIProto.FeatureReferenceV2 onDemandFeatureSource =
                    ServingAPIProto.FeatureReferenceV2.newBuilder()
                        .setFeatureViewName(featureViewName)
                        .setFeatureName(featureName)
                        .build();
                onDemandFeatureSources.add(onDemandFeatureSource);
              }
              break;
            default:
              throw Status.INTERNAL
                  .withDescription(
                      "OnDemandSource proto source field has an unexpected type: " + sourceCase)
                  .asRuntimeException();
          }
        }
      }
      return onDemandFeatureSources;
    }
  
    /** {@inheritDoc} */
    public void processTransformFeaturesResponse(
        feast.proto.serving.TransformationServiceAPIProto.TransformFeaturesResponse
            transformFeaturesResponse,
        String onDemandFeatureViewName,
        Set<String> onDemandFeatureStringReferences,
        ServingAPIProto.GetOnlineFeaturesResponse.Builder responseBuilder) {
      try {
        BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
        ArrowFileReader reader =
            new ArrowFileReader(
                new ByteArrayReadableSeekableByteChannel(
                    transformFeaturesResponse
                        .getTransformationOutput()
                        .getArrowValue()
                        .toByteArray()),
                allocator);
        reader.loadNextBatch();
        VectorSchemaRoot readBatch = reader.getVectorSchemaRoot();
        Schema responseSchema = readBatch.getSchema();
        List<Field> responseFields = responseSchema.getFields();
        Timestamp now = Timestamp.newBuilder().setSeconds(System.currentTimeMillis() / 1000).build();
  
        for (Field field : responseFields) {
          String columnName = field.getName();
          String fullFeatureName = columnName.replace("__", ":");
          ArrowType columnType = field.getType();
  
          // The response will contain all features for the specified ODFV, so we
          // skip the features that were not requested.
          if (!onDemandFeatureStringReferences.contains(fullFeatureName)) {
            continue;
          }
  
          FieldVector fieldVector = readBatch.getVector(field);
          int valueCount = fieldVector.getValueCount();
          ServingAPIProto.GetOnlineFeaturesResponse.FeatureVector.Builder vectorBuilder =
              responseBuilder.addResultsBuilder();
          List<ValueProto.Value> valueList = Lists.newArrayListWithExpectedSize(valueCount);
  
          // TODO: support all Feast types
          // TODO: clean up the switch statement
          if (columnType instanceof ArrowType.Int) {
            int bitWidth = ((ArrowType.Int) columnType).getBitWidth();
            switch (bitWidth) {
              case INT64_BITWIDTH:
                for (int i = 0; i < valueCount; i++) {
                  long int64Value = ((BigIntVector) fieldVector).get(i);
                  valueList.add(ValueProto.Value.newBuilder().setInt64Val(int64Value).build());
                }
                break;
              case INT32_BITWIDTH:
                for (int i = 0; i < valueCount; i++) {
                  int int32Value = ((IntVector) fieldVector).get(i);
                  valueList.add(ValueProto.Value.newBuilder().setInt32Val(int32Value).build());
                }
                break;
              default:
                throw Status.INTERNAL
                    .withDescription(
                        "Column "
                            + columnName
                            + " is of type ArrowType.Int but has bitWidth "
                            + bitWidth
                            + " which cannot be handled.")
                    .asRuntimeException();
            }
          } else if (columnType instanceof ArrowType.FloatingPoint) {
            FloatingPointPrecision precision = ((ArrowType.FloatingPoint) columnType).getPrecision();
            switch (precision) {
              case DOUBLE:
                for (int i = 0; i < valueCount; i++) {
                  double doubleValue = ((Float8Vector) fieldVector).get(i);
                  valueList.add(ValueProto.Value.newBuilder().setDoubleVal(doubleValue).build());
                }
                break;
              case SINGLE:
                for (int i = 0; i < valueCount; i++) {
                  float floatValue = ((Float4Vector) fieldVector).get(i);
                  valueList.add(ValueProto.Value.newBuilder().setFloatVal(floatValue).build());
                }
                break;
              default:
                throw Status.INTERNAL
                    .withDescription(
                        "Column "
                            + columnName
                            + " is of type ArrowType.FloatingPoint but has precision "
                            + precision
                            + " which cannot be handled.")
                    .asRuntimeException();
            }
          }
  
          for (ValueProto.Value v : valueList) {
            vectorBuilder.addValues(v);
            vectorBuilder.addStatuses(ServingAPIProto.FieldStatus.PRESENT);
            vectorBuilder.addEventTimestamps(now);
          }
  
          responseBuilder.getMetadataBuilder().getFeatureNamesBuilder().addVal(fullFeatureName);
        }
      } catch (IOException e) {
        log.info(e.toString());
        throw Status.INTERNAL
            .withDescription("Unable to correctly process transform features response: " + e)
            .asRuntimeException();
      }
    }
  
    /** {@inheritDoc} */
    public ValueType serializeValuesIntoArrowIPC(List<Pair<String, List<ValueProto.Value>>> values) {
      // In order to be serialized correctly, the data must be packaged in a VectorSchemaRoot.
      // We first construct all the columns.
      BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
  
      List<Field> columnFields = new ArrayList<>();
      List<FieldVector> columns = new ArrayList<>();
  
      for (Pair<String, List<ValueProto.Value>> columnEntry : values) {
        // The Python FTS does not expect full feature names, so we extract the feature name.
        String columnName = columnEntry.getKey();
  
        List<ValueProto.Value> columnValues = columnEntry.getValue();
        FieldVector column;
        ValueProto.Value.ValCase valCase = columnValues.get(0).getValCase();
        // TODO: support all Feast types
        switch (valCase) {
          case INT32_VAL:
            column = new IntVector(columnName, allocator);
            column.setValueCount(columnValues.size());
            for (int idx = 0; idx < columnValues.size(); idx++) {
              ((IntVector) column).set(idx, columnValues.get(idx).getInt32Val());
            }
            break;
          case INT64_VAL:
            column = new BigIntVector(columnName, allocator);
            column.setValueCount(columnValues.size());
            for (int idx = 0; idx < columnValues.size(); idx++) {
              ((BigIntVector) column).set(idx, columnValues.get(idx).getInt64Val());
            }
  
            break;
          case DOUBLE_VAL:
            column = new Float8Vector(columnName, allocator);
            column.setValueCount(columnValues.size());
            for (int idx = 0; idx < columnValues.size(); idx++) {
              ((Float8Vector) column).set(idx, columnValues.get(idx).getDoubleVal());
            }
            break;
          case FLOAT_VAL:
            column = new Float4Vector(columnName, allocator);
            column.setValueCount(columnValues.size());
            for (int idx = 0; idx < columnValues.size(); idx++) {
              ((Float4Vector) column).set(idx, columnValues.get(idx).getFloatVal());
            }
            break;
          default:
            throw Status.INTERNAL
                .withDescription(
                    "Column " + columnName + " has a type that is currently not handled: " + valCase)
                .asRuntimeException();
        }
  
        columns.add(column);
        columnFields.add(column.getField());
      }
  
      VectorSchemaRoot schemaRoot = new VectorSchemaRoot(columnFields, columns);
  
      // Serialize the VectorSchemaRoot into Arrow IPC format.
      ByteArrayOutputStream out = new ByteArrayOutputStream();
      ArrowFileWriter writer = new ArrowFileWriter(schemaRoot, null, Channels.newChannel(out));
      try {
        writer.start();
        writer.writeBatch();
        writer.end();
      } catch (IOException e) {
        log.info(e.toString());
        throw Status.INTERNAL
            .withDescription("ArrowFileWriter could not write properly; failed with error: " + e)
            .asRuntimeException();
      }
      byte[] byteData = out.toByteArray();
      ByteString sourceData = ByteString.copyFrom(byteData);
      ValueType transformationInput = ValueType.newBuilder().setArrowValue(sourceData).build();
      return transformationInput;
    }
  }
  