feat: Add support for low_level_batch_size in match() for MatchingEngineIndexEndpoint private query.

lingyinw · copybara-github · commit 67e593b34928 · 2024-01-18T06:03:05.000-08:00
PiperOrigin-RevId: 599483250
diff --git a/google/cloud/aiplatform/matching_engine/matching_engine_index_endpoint.py b/google/cloud/aiplatform/matching_engine/matching_engine_index_endpoint.py
@@ -1436,6 +1436,7 @@ def match(
         approx_num_neighbors: Optional[int] = None,
         fraction_leaf_nodes_to_search_override: Optional[float] = None,
         return_full_datapoint: bool = False,
+        low_level_batch_size: int = 0,
     ) -> List[List[MatchNeighbor]]:
         """Retrieves nearest neighbors for the given embedding queries on the
         specified deployed index for private endpoint only.
@@ -1472,6 +1473,13 @@ def match(
                 vector values and of the nearest neighbors are returned.
                 Note that returning full datapoint will significantly increase the
                 latency and cost of the query.
+            low_level_batch_size (int):
+                Optional. Selects the optimal batch size to use for low-level
+                batching. Queries within each low level batch are executed
+                sequentially while low level batches are executed in parallel.
+                This field is optional, defaults to 0 if not set. A non-positive
+                number disables low level batching, i.e. all queries are
+                executed sequentially.
 
         Returns:
             List[List[MatchNeighbor]] - A list of nearest neighbors for each query.
@@ -1487,6 +1495,7 @@ def match(
             match_service_pb2.BatchMatchRequest.BatchMatchRequestPerIndex()
         )
         batch_request_for_index.deployed_index_id = deployed_index_id
+        batch_request_for_index.low_level_batch_size = low_level_batch_size
 
         # Preprocess restricts to be used for each request
         restricts = []
diff --git a/tests/unit/aiplatform/test_matching_engine_index_endpoint.py b/tests/unit/aiplatform/test_matching_engine_index_endpoint.py
@@ -244,6 +244,7 @@
 _TEST_APPROX_NUM_NEIGHBORS = 2
 _TEST_FRACTION_LEAF_NODES_TO_SEARCH_OVERRIDE = 0.8
 _TEST_RETURN_FULL_DATAPOINT = True
+_TEST_LOW_LEVEL_BATCH_SIZE = 3
 _TEST_ENCRYPTION_SPEC_KEY_NAME = "kms_key_name"
 _TEST_PROJECT_ALLOWLIST = ["project-1", "project-2"]
 _TEST_PRIVATE_SERVICE_CONNECT_IP_ADDRESS = "10.128.0.5"
@@ -1061,12 +1062,14 @@ def test_private_index_endpoint_match_queries(
             approx_num_neighbors=_TEST_APPROX_NUM_NEIGHBORS,
             fraction_leaf_nodes_to_search_override=_TEST_FRACTION_LEAF_NODES_TO_SEARCH_OVERRIDE,
             return_full_datapoint=_TEST_RETURN_FULL_DATAPOINT,
+            low_level_batch_size=_TEST_LOW_LEVEL_BATCH_SIZE,
         )
 
         batch_request = match_service_pb2.BatchMatchRequest(
             requests=[
                 match_service_pb2.BatchMatchRequest.BatchMatchRequestPerIndex(
                     deployed_index_id=_TEST_DEPLOYED_INDEX_ID,
+                    low_level_batch_size=_TEST_LOW_LEVEL_BATCH_SIZE,
                     requests=[
                         match_service_pb2.MatchRequest(
                             num_neighbors=_TEST_NUM_NEIGHBOURS,
