Fix multi cells in bigtable retriever & type bugs (#22)

pyalex · web-flow · commit 8f41e36eafd2 · 2021-04-05T17:41:22.000+08:00
Signed-off-by: Oleksii Moskalenko &lt;moskalenko.alexey@gmail.com&gt;
diff --git a/serving/src/test/java/feast/serving/it/ServingServiceBigTableIT.java b/serving/src/test/java/feast/serving/it/ServingServiceBigTableIT.java
@@ -218,7 +218,7 @@ static void globalSetup() throws IOException {
         SchemaBuilder.record("DriverData")
             .namespace(featureTableName)
             .fields()
-            .requiredInt(feature1Reference.getName())
+            .requiredLong(feature1Reference.getName())
             .requiredDouble(feature2Reference.getName())
             .nullableString(feature3Reference.getName(), "null")
             .requiredString(feature4Reference.getName())
@@ -228,7 +228,7 @@ static void globalSetup() throws IOException {
 
     GenericRecord record =
         new GenericRecordBuilder(ftSchema)
-            .set("trip_cost", 5)
+            .set("trip_cost", 5L)
             .set("trip_distance", 3.5)
             .set("trip_empty", null)
             .set("trip_wrong_type", "test")
@@ -245,7 +245,7 @@ static void globalSetup() throws IOException {
         SchemaBuilder.record("DriverMerchantData")
             .namespace(rideMerchantFeatureTableName)
             .fields()
-            .requiredInt(feature1Reference.getName())
+            .requiredLong(feature1Reference.getName())
             .requiredDouble(feature2Reference.getName())
             .nullableString(feature3Reference.getName(), "null")
             .requiredString(feature4Reference.getName())
@@ -256,7 +256,7 @@ static void globalSetup() throws IOException {
     // Entity-Feature Row
     GenericRecord compoundEntityRecord =
         new GenericRecordBuilder(compoundFtSchema)
-            .set("trip_cost", 10)
+            .set("trip_cost", 10L)
             .set("trip_distance", 5.5)
             .set("trip_empty", null)
             .set("trip_wrong_type", "wrong_type")
diff --git a/storage/api/src/main/java/feast/storage/api/retriever/NativeFeature.java b/storage/api/src/main/java/feast/storage/api/retriever/NativeFeature.java
@@ -58,13 +58,13 @@ public ValueProto.Value getFeatureValue(ValueProto.ValueType.Enum valueType) {
           finalValue = ValueProto.Value.newBuilder().setInt32Val((Integer) featureValue).build();
           break;
         case INT64:
-          finalValue = ValueProto.Value.newBuilder().setInt64Val((Integer) featureValue).build();
+          finalValue = ValueProto.Value.newBuilder().setInt64Val((Long) featureValue).build();
           break;
         case DOUBLE:
           finalValue = ValueProto.Value.newBuilder().setDoubleVal((Double) featureValue).build();
           break;
         case FLOAT:
-          finalValue = ValueProto.Value.newBuilder().setFloatVal((Long) featureValue).build();
+          finalValue = ValueProto.Value.newBuilder().setFloatVal((Float) featureValue).build();
           break;
         case BYTES:
           finalValue = ValueProto.Value.newBuilder().setBytesVal((ByteString) featureValue).build();
diff --git a/storage/connectors/bigtable/src/main/java/feast/storage/connectors/bigtable/retriever/BigTableOnlineRetriever.java b/storage/connectors/bigtable/src/main/java/feast/storage/connectors/bigtable/retriever/BigTableOnlineRetriever.java
@@ -20,6 +20,7 @@
 import com.google.cloud.bigtable.data.v2.models.Filters;
 import com.google.cloud.bigtable.data.v2.models.Query;
 import com.google.cloud.bigtable.data.v2.models.Row;
+import com.google.cloud.bigtable.data.v2.models.RowCell;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.Timestamp;
 import feast.proto.serving.ServingAPIProto.FeatureReferenceV2;
@@ -242,9 +243,15 @@ private List<List<Feature>> convertRowToFeature(
               if (!rows.containsKey(rowKey)) {
                 return Collections.<Feature>emptyList();
               } else {
-                return rows.get(rowKey).getCells().stream()
+                Row row = rows.get(rowKey);
+                return featureReferences.stream()
+                    .map(FeatureReferenceV2::getFeatureTable)
+                    .distinct()
+                    .map(cf -> row.getCells(cf, ""))
+                    .filter(ls -> !ls.isEmpty())
                     .flatMap(
-                        rowCell -> {
+                        rowCells -> {
+                          RowCell rowCell = rowCells.get(0); // Latest cell
                           String family = rowCell.getFamily();
                           ByteString value = rowCell.getValue();
 
