Merge pull request #30 from googleapis/batch_fetch

Add batch fetch API

Author: pbiswa01

src/main/java/com/google/cloud/aiplatform/fs/BigtableClient.java

@@ -18,17 +18,21 @@
 import static com.google.cloud.bigtable.data.v2.models.Filters.FILTERS;
 
 import com.google.api.gax.core.FixedCredentialsProvider;
+import com.google.api.gax.rpc.ServerStream;
 import com.google.auth.oauth2.AccessToken;
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.cloud.bigtable.data.v2.BigtableDataClient;
 import com.google.cloud.bigtable.data.v2.BigtableDataSettings;
 import com.google.cloud.bigtable.data.v2.models.Filters.Filter;
+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.TableId;
 import com.google.cloud.bigtable.data.v2.stub.metrics.NoopMetricsProvider;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.VerifyException;
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Optional;
 
 /**
@@ -96,27 +100,61 @@ public void close() {
     this.bigtableDataClient.close();
   }
 
-  public Row fetchData(InternalFetchRequest request) throws Exception {
-    Filter filter = null;
+  private Filter buildFilter(InternalFetchRequest request) {
     if (request.featureViewSpec.continuousSyncEnabled) {
-      filter =
-          FILTERS
-              .chain()
-              .filter(FILTERS.family().regex(request.featureViewId))
-              .filter(FILTERS.limit().cellsPerColumn(1));
+      return FILTERS
+          .chain()
+          .filter(FILTERS.family().regex(request.featureViewId))
+          .filter(FILTERS.limit().cellsPerColumn(1));
     } else {
-      filter =
-          FILTERS
-              .chain()
-              .filter(
-                  FILTERS
-                      .qualifier()
-                      .rangeWithinFamily(request.featureViewId)
-                      .startClosed(defaultColumn)
-                      .endClosed(directWriteColumn))
-              .filter(FILTERS.limit().cellsPerColumn(1));
+      return FILTERS
+          .chain()
+          .filter(
+              FILTERS
+                  .qualifier()
+                  .rangeWithinFamily(request.featureViewId)
+                  .startClosed(defaultColumn)
+                  .endClosed(directWriteColumn))
+          .filter(FILTERS.limit().cellsPerColumn(1));
     }
+  }
+
+  public Row fetchData(InternalFetchRequest request) throws Exception {
+    Filter filter = buildFilter(request);
     return bigtableDataClient.readRow(
         TableId.of(request.cloudBigtableSpec.tableId), request.dataKey, filter);
   }
+
+  /**
+   * Fetches multiple rows from Bigtable based on the provided InternalFetchRequest.
+   * This method assumes that {@code request.dataKeys} is populated.
+   *
+   * @param request An {@link InternalFetchRequest} with {@code dataKeys} set.
+   * @return A {@link List} of {@link Row} objects for the keys found in Bigtable. Keys not found
+   *     in Bigtable will not have corresponding entries in the returned list.
+   * @throws IllegalArgumentException if {@code request.dataKeys} is null or empty.
+   * @throws Exception for other underlying issues during the Bigtable read.
+   */
+  public List<Row> batchFetchData(InternalFetchRequest request) throws Exception {
+    if (request.dataKeys == null || request.dataKeys.isEmpty()) {
+      throw new IllegalArgumentException("batchFetchData requires InternalFetchRequest with populated dataKeys.");
+    }
+
+    Filter filter = buildFilter(request);
+    TableId tableId = TableId.of(request.cloudBigtableSpec.tableId);
+
+    Query query = Query.create(tableId);
+    for (String dataKey : request.dataKeys) {
+      query = query.rowKey(dataKey);
+    }
+    query = query.filter(filter);
+
+    ServerStream<Row> rowStream = bigtableDataClient.readRows(query);
+
+    List<Row> rows = new ArrayList<>();
+    for (Row row : rowStream) {
+      rows.add(row);
+    }
+    return rows;
+  }
 }

src/main/java/com/google/cloud/aiplatform/fs/CloudBigtableCache.java

@@ -119,4 +119,4 @@ class CloudBigtableSpec {
     this.instanceId = fos.getBigtable().getBigtableMetadata().getInstanceId();
     this.tableId = fos.getBigtable().getBigtableMetadata().getTableId();
   }
-}
+}

src/main/java/com/google/cloud/aiplatform/fs/Converter.java

@@ -29,6 +29,10 @@
 import com.google.protobuf.InvalidProtocolBufferException;
 import io.grpc.Status;
 import io.grpc.Status.Code;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
@@ -163,4 +167,73 @@ public static FetchFeatureValuesResponse rowToResponse(Row row, InternalFetchReq
     responseBuilder.setKeyValues(internalStorageToKeyValuesList(cell, request));
     return responseBuilder.build();
   }
+
+  /**
+   * Converts a list of Bigtable Rows to a list of FetchFeatureValuesResponse.
+   * The output list is ordered to match the keys in {@code request.dataKeys}.
+   *
+   * @param rows The list of Rows returned from Bigtable. These are only the rows that were found.
+   * @param request An {@link InternalFetchRequest} containing the ordered list of keys in {@code request.dataKeys}.
+   * @return A {@link List<FetchFeatureValuesResponse>} where each element corresponds to a key in {@code request.dataKeys}.
+   *         If a key was not found in Bigtable, the response for that key will contain an empty {@link FeatureNameValuePairList}.
+   * @throws UnimplementedException if {@code request.format} is PROTO_STRUCT.
+   * @throws InternalException if conversion of any *found* row fails due to unexpected data format,
+   *         propagating the underlying issue.
+   */
+  public static List<FetchFeatureValuesResponse> rowsToResponses(List<Row> rows, InternalFetchRequest request) throws Exception {
+    if (request.format == FeatureViewDataFormat.PROTO_STRUCT) {
+      throw new UnimplementedException(
+          new Throwable("PROTO_STRUCT is not supported for batch fetch"),
+          /* statusCode= */ GrpcStatusCode.of(Code.UNIMPLEMENTED),
+          /* retryable= */ false);
+    }
+
+    if (request.dataKeys == null || request.dataKeys.isEmpty()) {
+      return new ArrayList<>(); // Return empty list if no keys were requested.
+    }
+
+    // 1. Create a map for quick lookup of Rows by their key.
+    Map<String, Row> keyToRowMap = new HashMap<>();
+    for (Row row : rows) {
+      // Row.getKey() returns a ByteString, convert to String.
+      keyToRowMap.put(row.getKey().toStringUtf8(), row);
+    }
+
+    // 2. Build the ordered list of responses.
+    List<FetchFeatureValuesResponse> responses = new ArrayList<>(request.dataKeys.size());
+
+    // Iterate through the *ordered* requested keys from InternalFetchRequest.dataKeys.
+    for (String dataKey : request.dataKeys) {
+      Row foundRow = keyToRowMap.get(dataKey);
+      FetchFeatureValuesResponse response;
+
+      if (foundRow != null) {
+        // Key was found in Bigtable. Convert the Row to a FeatureViewCell and then to a Response.
+        try {
+          FeatureViewCell cell = rowToFeatureViewCell(foundRow, request);
+          response = FetchFeatureValuesResponse.newBuilder()
+              .setKeyValues(internalStorageToKeyValuesList(cell, request))
+              .build();
+        } catch (Exception e) {
+          // If conversion of a *found* row fails, it indicates an internal data issue.
+          // Following the pattern of rowToResponse, we throw an InternalException.
+          throw new InternalException(
+              new Throwable(String.format("Failed to convert Bigtable Row for key '%s': %s", dataKey, e.getMessage()), e),
+              /* statusCode= */ GrpcStatusCode.of(Status.Code.INTERNAL),
+              /* retryable= */ false);
+        }
+      } else {
+        // Key was not found in Bigtable. Create a response indicating this.
+        // For KEY_VALUE format, an empty FeatureNameValuePairList signifies that no features
+        // were found for the requested key.
+        response = FetchFeatureValuesResponse.newBuilder()
+            .setKeyValues(FeatureNameValuePairList.getDefaultInstance())
+            .build();
+        logger.log(Level.FINE, String.format("Entity id '%s' not found in Bigtable during batch fetch.", dataKey));
+      }
+      responses.add(response);
+    }
+
+    return responses;
+  }
 }

src/main/java/com/google/cloud/aiplatform/fs/FeatureOnlineStoreDirectClient.java

@@ -25,6 +25,7 @@
 import com.google.cloud.aiplatform.v1.FetchFeatureValuesResponse;
 import com.google.cloud.bigtable.data.v2.models.Row;
 import io.grpc.Status.Code;
+import java.util.List;
 import java.util.Optional;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -99,6 +100,21 @@ public FetchFeatureValuesResponse fetchFeatureValues(FetchFeatureValuesRequest r
     return Converter.rowToResponse(row, internalRequest);
   }
 
+  public List<FetchFeatureValuesResponse> batchFetchFeatureValues(List<FetchFeatureValuesRequest> requests) throws Exception {
+    for (FetchFeatureValuesRequest request : requests) {
+      if (request.getDataFormat().equals(FeatureViewDataFormat.PROTO_STRUCT)) {
+        throw new UnimplementedException(
+            new Throwable("PROTO_STRUCT is not supported yet for batch fetch"),
+            /* statusCode= */ GrpcStatusCode.of(Code.UNIMPLEMENTED),
+            /* retryable= */ false);
+      }
+    }
+
+    InternalFetchRequest internalRequest = new InternalFetchRequest(requests);
+    List<Row> rows = this.bigtableClientManager.getClient().batchFetchData(internalRequest);
+    return Converter.rowsToResponses(rows, internalRequest);
+  }
+
   public void close() {
     this.bigtableClientManager.shutdown();
   }