feat: add support for per_crowding_attribute_num_neighbors approx_num_neighborsto MatchingEngineIndexEndpoint match()

PiperOrigin-RevId: 579015694

Author: lingyinw

google/cloud/aiplatform/matching_engine/matching_engine_index_endpoint.py

@@ -1073,6 +1073,8 @@ def match(
         queries: List[List[float]],
         num_neighbors: int = 1,
         filter: Optional[List[Namespace]] = [],
+        per_crowding_attribute_num_neighbors: Optional[int] = None,
+        approx_num_neighbors: Optional[int] = None,
     ) -> List[List[MatchNeighbor]]:
         """Retrieves nearest neighbors for the given embedding queries on the specified deployed index.
 
@@ -1089,6 +1091,15 @@ def match(
                 For example, [Namespace("color", ["red"], []), Namespace("shape", [], ["squared"])] will match datapoints
                 that satisfy "red color" but not include datapoints with "squared shape".
                 Please refer to https://cloud.google.com/vertex-ai/docs/matching-engine/filtering#json for more detail.
+            per_crowding_attribute_num_neighbors (int):
+                Optional. Crowding is a constraint on a neighbor list produced by nearest neighbor
+                search requiring that no more than some value k' of the k neighbors
+                returned have the same value of crowding_attribute.
+                It's used for improving result diversity.
+                This field is the maximum number of matches with the same crowding tag.
+            approx_num_neighbors (int):
+                The number of neighbors to find via approximate search before exact reordering is performed.
+                If not set, the default value from scam config is used; if set, this value must be > 0.
 
         Returns:
             List[List[MatchNeighbor]] - A list of nearest neighbors for each query.
@@ -1123,6 +1134,8 @@ def match(
                 num_neighbors=num_neighbors,
                 deployed_index_id=deployed_index_id,
                 float_val=query,
+                per_crowding_attribute_num_neighbors=per_crowding_attribute_num_neighbors,
+                approx_num_neighbors=approx_num_neighbors,
             )
             for namespace in filter:
                 restrict = match_service_pb2.Namespace()

tests/unit/aiplatform/test_matching_engine_index_endpoint.py

@@ -232,6 +232,8 @@
     Namespace(name="class", allow_tokens=["token_1"], deny_tokens=["token_2"])
 ]
 _TEST_IDS = ["123", "456", "789"]
+_TEST_PER_CROWDING_ATTRIBUTE_NUM_NEIGHBOURS = 3
+_TEST_APPROX_NUM_NEIGHBORS = 2
 
 
 def uuid_mock():
@@ -853,6 +855,47 @@ def test_delete_index_endpoint_with_force(
             name=_TEST_INDEX_ENDPOINT_NAME
         )
 
+    @pytest.mark.usefixtures("get_index_endpoint_mock")
+    def test_index_endpoint_match_queries_backward_compatibility(
+        self, index_endpoint_match_queries_mock
+    ):
+        aiplatform.init(project=_TEST_PROJECT)
+
+        my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint(
+            index_endpoint_name=_TEST_INDEX_ENDPOINT_ID
+        )
+
+        my_index_endpoint.match(
+            _TEST_DEPLOYED_INDEX_ID,
+            _TEST_QUERIES,
+            _TEST_NUM_NEIGHBOURS,
+            _TEST_FILTER,
+        )
+
+        batch_request = match_service_pb2.BatchMatchRequest(
+            requests=[
+                match_service_pb2.BatchMatchRequest.BatchMatchRequestPerIndex(
+                    deployed_index_id=_TEST_DEPLOYED_INDEX_ID,
+                    requests=[
+                        match_service_pb2.MatchRequest(
+                            num_neighbors=_TEST_NUM_NEIGHBOURS,
+                            deployed_index_id=_TEST_DEPLOYED_INDEX_ID,
+                            float_val=_TEST_QUERIES[0],
+                            restricts=[
+                                match_service_pb2.Namespace(
+                                    name="class",
+                                    allow_tokens=["token_1"],
+                                    deny_tokens=["token_2"],
+                                )
+                            ],
+                        )
+                    ],
+                )
+            ]
+        )
+
+        index_endpoint_match_queries_mock.assert_called_with(batch_request)
+
     @pytest.mark.usefixtures("get_index_endpoint_mock")
     def test_index_endpoint_match_queries(self, index_endpoint_match_queries_mock):
         aiplatform.init(project=_TEST_PROJECT)
@@ -866,6 +909,8 @@ def test_index_endpoint_match_queries(self, index_endpoint_match_queries_mock):
             queries=_TEST_QUERIES,
             num_neighbors=_TEST_NUM_NEIGHBOURS,
             filter=_TEST_FILTER,
+            per_crowding_attribute_num_neighbors=_TEST_PER_CROWDING_ATTRIBUTE_NUM_NEIGHBOURS,
+            approx_num_neighbors=_TEST_APPROX_NUM_NEIGHBORS,
         )
 
         batch_request = match_service_pb2.BatchMatchRequest(
@@ -884,6 +929,8 @@ def test_index_endpoint_match_queries(self, index_endpoint_match_queries_mock):
                                     deny_tokens=["token_2"],
                                 )
                             ],
+                            per_crowding_attribute_num_neighbors=_TEST_PER_CROWDING_ATTRIBUTE_NUM_NEIGHBOURS,
+                            approx_num_neighbors=_TEST_APPROX_NUM_NEIGHBORS,
                         )
                     ],
                 )