fix bigtable schema caching (#23)

pyalex · web-flow · commit 8e5bad73dbb5 · 2021-04-07T16:53:21.000+08:00
Signed-off-by: Oleksii Moskalenko &lt;moskalenko.alexey@gmail.com&gt;
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
@@ -35,7 +35,6 @@
 import java.util.stream.Collectors;
 import java.util.stream.StreamSupport;
 import org.apache.avro.AvroRuntimeException;
-import org.apache.avro.Schema;
 import org.apache.avro.generic.GenericDatumReader;
 import org.apache.avro.generic.GenericRecord;
 import org.apache.avro.io.*;
@@ -138,6 +137,7 @@ private List<Feature> decodeFeatures(
       String tableName,
       ByteString value,
       List<FeatureReferenceV2> featureReferences,
+      BinaryDecoder reusedDecoder,
       long timestamp)
       throws IOException {
     ByteString schemaReferenceBytes = value.substring(0, 4);
@@ -146,11 +146,10 @@ private List<Feature> decodeFeatures(
     BigTableSchemaRegistry.SchemaReference schemaReference =
         new BigTableSchemaRegistry.SchemaReference(tableName, schemaReferenceBytes);
 
-    Schema schema = schemaRegistry.getSchema(schemaReference);
+    GenericDatumReader<GenericRecord> reader = schemaRegistry.getReader(schemaReference);
 
-    GenericDatumReader<GenericRecord> reader = new GenericDatumReader<>(schema);
-    Decoder decoder = DecoderFactory.get().binaryDecoder(featureValueBytes, null);
-    GenericRecord record = reader.read(null, decoder);
+    reusedDecoder = DecoderFactory.get().binaryDecoder(featureValueBytes, reusedDecoder);
+    GenericRecord record = reader.read(null, reusedDecoder);
 
     return featureReferences.stream()
         .map(
@@ -237,6 +236,8 @@ private List<List<Feature>> convertRowToFeature(
       Map<ByteString, Row> rows,
       List<FeatureReferenceV2> featureReferences) {
 
+    BinaryDecoder reusedDecoder = DecoderFactory.get().binaryDecoder(new byte[0], null);
+
     return rowKeys.stream()
         .map(
             rowKey -> {
@@ -269,6 +270,7 @@ private List<List<Feature>> convertRowToFeature(
                                     tableName,
                                     value,
                                     localFeatureReferences,
+                                    reusedDecoder,
                                     rowCell.getTimestamp());
                           } catch (IOException e) {
                             throw new RuntimeException("Failed to decode features from BigTable");
diff --git a/storage/connectors/bigtable/src/main/java/feast/storage/connectors/bigtable/retriever/BigTableSchemaRegistry.java b/storage/connectors/bigtable/src/main/java/feast/storage/connectors/bigtable/retriever/BigTableSchemaRegistry.java
@@ -27,10 +27,12 @@
 import com.google.protobuf.ByteString;
 import java.util.concurrent.ExecutionException;
 import org.apache.avro.Schema;
+import org.apache.avro.generic.GenericDatumReader;
+import org.apache.avro.generic.GenericRecord;
 
 public class BigTableSchemaRegistry {
   private final BigtableDataClient client;
-  private final LoadingCache<SchemaReference, Schema> cache;
+  private final LoadingCache<SchemaReference, GenericDatumReader<GenericRecord>> cache;
 
   private static String COLUMN_FAMILY = "metadata";
   private static String QUALIFIER = "avro";
@@ -52,34 +54,54 @@ public String getTableName() {
     public ByteString getSchemaHash() {
       return schemaHash;
     }
+
+    @Override
+    public int hashCode() {
+      int result = tableName.hashCode();
+      result = 31 * result + schemaHash.hashCode();
+      return result;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) return true;
+      if (o == null || getClass() != o.getClass()) return false;
+
+      SchemaReference that = (SchemaReference) o;
+
+      if (!tableName.equals(that.tableName)) return false;
+      return schemaHash.equals(that.schemaHash);
+    }
   }
 
   public BigTableSchemaRegistry(BigtableDataClient client) {
     this.client = client;
 
-    CacheLoader<SchemaReference, Schema> schemaCacheLoader = CacheLoader.from(this::loadSchema);
+    CacheLoader<SchemaReference, GenericDatumReader<GenericRecord>> schemaCacheLoader =
+        CacheLoader.from(this::loadReader);
 
     cache = CacheBuilder.newBuilder().build(schemaCacheLoader);
   }
 
-  public Schema getSchema(SchemaReference reference) {
-    Schema schema;
+  public GenericDatumReader<GenericRecord> getReader(SchemaReference reference) {
+    GenericDatumReader<GenericRecord> reader;
     try {
-      schema = this.cache.get(reference);
+      reader = this.cache.get(reference);
     } catch (ExecutionException | CacheLoader.InvalidCacheLoadException e) {
       throw new RuntimeException(String.format("Unable to find Schema"), e);
     }
-    return schema;
+    return reader;
   }
 
-  private Schema loadSchema(SchemaReference reference) {
+  private GenericDatumReader<GenericRecord> loadReader(SchemaReference reference) {
     Row row =
         client.readRow(
             reference.getTableName(),
             ByteString.copyFrom(KEY_PREFIX.getBytes()).concat(reference.getSchemaHash()),
             Filters.FILTERS.family().exactMatch(COLUMN_FAMILY));
     RowCell last = Iterables.getLast(row.getCells(COLUMN_FAMILY, QUALIFIER));
 
-    return new Schema.Parser().parse(last.getValue().toStringUtf8());
+    Schema schema = new Schema.Parser().parse(last.getValue().toStringUtf8());
+    return new GenericDatumReader<>(schema);
   }
 }
