Optimize feature retrieval for Cassandra online storage (#24)

Signed-off-by: Khor Shu Heng <[email protected]>

Co-authored-by: Khor Shu Heng <[email protected]>

Author: khorshuheng

storage/connectors/cassandra/src/main/java/feast/storage/connectors/cassandra/retriever/CassandraOnlineRetriever.java

@@ -16,8 +16,10 @@
  */
 package feast.storage.connectors.cassandra.retriever;
 
+import com.datastax.oss.driver.api.core.AsyncPagingIterable;
 import com.datastax.oss.driver.api.core.CqlSession;
-import com.datastax.oss.driver.api.core.cql.BoundStatement;
+import com.datastax.oss.driver.api.core.cql.AsyncResultSet;
+import com.datastax.oss.driver.api.core.cql.PreparedStatement;
 import com.datastax.oss.driver.api.core.cql.Row;
 import com.datastax.oss.driver.api.querybuilder.QueryBuilder;
 import com.datastax.oss.driver.api.querybuilder.select.Select;
@@ -30,9 +32,11 @@
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.*;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
+import java.util.concurrent.ExecutionException;
 import java.util.function.Function;
 import java.util.stream.Collectors;
-import java.util.stream.StreamSupport;
 import org.apache.avro.AvroRuntimeException;
 import org.apache.avro.generic.GenericDatumReader;
 import org.apache.avro.generic.GenericRecord;
@@ -157,12 +161,40 @@ public Map<ByteBuffer, Row> getFeaturesFromSSTable(
     for (String columnFamily : columnFamilies) {
       query = query.writeTime(columnFamily).as(columnFamily + EVENT_TIMESTAMP_SUFFIX);
     }
-    query = query.whereColumn(ENTITY_KEY).in(QueryBuilder.bindMarker());
+    query = query.whereColumn(ENTITY_KEY).isEqualTo(QueryBuilder.bindMarker());
 
-    BoundStatement statement = session.prepare(query.build()).bind(rowKeys);
+    PreparedStatement preparedStatement = session.prepare(query.build());
 
-    return StreamSupport.stream(session.execute(statement).spliterator(), false)
-        .collect(Collectors.toMap((Row row) -> row.getByteBuffer(ENTITY_KEY), Function.identity()));
+    List<CompletableFuture<AsyncResultSet>> completableAsyncResultSets =
+        rowKeys.stream()
+            .map(preparedStatement::bind)
+            .map(session::executeAsync)
+            .map(CompletionStage::toCompletableFuture)
+            .collect(Collectors.toList());
+
+    CompletableFuture<Void> allResultComputed =
+        CompletableFuture.allOf(completableAsyncResultSets.toArray(new CompletableFuture[0]));
+
+    Map<ByteBuffer, Row> resultMap;
+    try {
+      resultMap =
+          allResultComputed
+              .thenApply(
+                  v ->
+                      completableAsyncResultSets.stream()
+                          .map(CompletableFuture::join)
+                          .filter(result -> result.remaining() != 0)
+                          .map(AsyncPagingIterable::one)
+                          .filter(Objects::nonNull)
+                          .collect(
+                              Collectors.toMap(
+                                  (Row row) -> row.getByteBuffer(ENTITY_KEY), Function.identity())))
+              .get();
+    } catch (InterruptedException | ExecutionException e) {
+      throw new RuntimeException(e.getMessage());
+    }
+
+    return resultMap;
   }
 
   /**

storage/connectors/cassandra/src/main/java/feast/storage/connectors/cassandra/retriever/CassandraSchemaRegistry.java

@@ -18,6 +18,7 @@
 
 import com.datastax.oss.driver.api.core.CqlSession;
 import com.datastax.oss.driver.api.core.cql.BoundStatement;
+import com.datastax.oss.driver.api.core.cql.PreparedStatement;
 import com.datastax.oss.driver.api.core.cql.Row;
 import com.datastax.oss.driver.api.querybuilder.QueryBuilder;
 import com.datastax.oss.driver.api.querybuilder.select.Select;
@@ -34,6 +35,7 @@
 
 public class CassandraSchemaRegistry {
   private final CqlSession session;
+  private final PreparedStatement preparedStatement;
   private final LoadingCache<SchemaReference, GenericDatumReader<GenericRecord>> cache;
 
   private static String SCHEMA_REF_TABLE = "feast_schema_reference";
@@ -67,6 +69,13 @@ public int hashCode() {
 
   public CassandraSchemaRegistry(CqlSession session) {
     this.session = session;
+    String tableName = String.format("\"%s\"", SCHEMA_REF_TABLE);
+    Select query =
+        QueryBuilder.selectFrom(tableName)
+            .column(SCHEMA_COLUMN)
+            .whereColumn(SCHEMA_REF_COLUMN)
+            .isEqualTo(QueryBuilder.bindMarker());
+    this.preparedStatement = session.prepare(query.build());
 
     CacheLoader<SchemaReference, GenericDatumReader<GenericRecord>> schemaCacheLoader =
         CacheLoader.from(this::loadReader);
@@ -85,14 +94,7 @@ public GenericDatumReader<GenericRecord> getReader(SchemaReference reference) {
   }
 
   private GenericDatumReader<GenericRecord> loadReader(SchemaReference reference) {
-    String tableName = String.format("\"%s\"", SCHEMA_REF_TABLE);
-    Select query =
-        QueryBuilder.selectFrom(tableName)
-            .column(SCHEMA_COLUMN)
-            .whereColumn(SCHEMA_REF_COLUMN)
-            .isEqualTo(QueryBuilder.bindMarker());
-
-    BoundStatement statement = session.prepare(query.build()).bind(reference.getSchemaHash());
+    BoundStatement statement = preparedStatement.bind(reference.getSchemaHash());
 
     Row row = session.execute(statement).one();