feat: Add query by id for MatchingEngineIndexEndpoint `find_neighbo…

…rs()` public endpoint query.

PiperOrigin-RevId: 599497930

google/cloud/aiplatform/matching_engine/matching_engine_index_endpoint.py

@@ -51,10 +51,13 @@ class MatchNeighbor:
  Required. The id of the neighbor.
  distance (float):
  Required. The distance to the query embedding.
+ feature_vector (List(float)):
+ Optional. The feature vector of the matching datapoint.
  """
 
  id: str
  distance: float
+ feature_vector: Optional[List[float]] = None
 
 
 @dataclass
@@ -1185,14 +1188,15 @@ def find_neighbors(
  self,
  *,
  deployed_index_id: str,
- queries: List[List[float]],
+ queries: Optional[List[List[float]]] = None,
  num_neighbors: int = 10,
  filter: Optional[List[Namespace]] = None,
  per_crowding_attribute_neighbor_count: Optional[int] = None,
  approx_num_neighbors: Optional[int] = None,
  fraction_leaf_nodes_to_search_override: Optional[float] = None,
  return_full_datapoint: bool = False,
  numeric_filter: Optional[List[NumericNamespace]] = None,
+ embedding_ids: Optional[List[str]] = None,
  ) -> List[List[MatchNeighbor]]:
  """Retrieves nearest neighbors for the given embedding queries on the
  specified deployed index which is deployed to either public or private
@@ -1243,11 +1247,18 @@ def find_neighbors(
  Note that returning full datapoint will significantly increase the
  latency and cost of the query.
 
- numeric_filter (Optional[list[NumericNamespace]]):
+ numeric_filter (list[NumericNamespace]):
  Optional. A list of NumericNamespaces for filtering the matching
  results. For example:
  [NumericNamespace(name="cost", value_int=5, op="GREATER")]
  will match datapoints that its cost is greater than 5.
+
+ embedding_ids (str):
+ Optional. If `queries` is set, will use `queries` to do nearest
+ neighbor search. If `queries` isn't set, will first use
+ `embedding_ids` to lookup embedding values from dataset, if embedding
+ with `embedding_ids` exists in the dataset, do nearest neighbor search.
+
  Returns:
  List[List[MatchNeighbor]] - A list of nearest neighbors for each query.
  """
@@ -1262,7 +1273,6 @@ def find_neighbors(
  per_crowding_attribute_num_neighbors=per_crowding_attribute_neighbor_count,
  approx_num_neighbors=approx_num_neighbors,
  fraction_leaf_nodes_to_search_override=fraction_leaf_nodes_to_search_override,
- return_full_datapoint=return_full_datapoint,
  )
 
  # Create the FindNeighbors request
@@ -1271,50 +1281,60 @@ def find_neighbors(
  find_neighbors_request.deployed_index_id = deployed_index_id
  find_neighbors_request.return_full_datapoint = return_full_datapoint
 
- for query in queries:
- find_neighbors_query = (
- gca_match_service_v1beta1.FindNeighborsRequest.Query()
- )
- find_neighbors_query.neighbor_count = num_neighbors
- find_neighbors_query.per_crowding_attribute_neighbor_count = (
- per_crowding_attribute_neighbor_count
- )
- find_neighbors_query.approximate_neighbor_count = approx_num_neighbors
- find_neighbors_query.fraction_leaf_nodes_to_search_override = (
- fraction_leaf_nodes_to_search_override
+ # Token restricts
+ restricts = []
+ if filter:
+ for namespace in filter:
+ restrict = gca_index_v1beta1.IndexDatapoint.Restriction()
+ restrict.namespace = namespace.name
+ restrict.allow_list.extend(namespace.allow_tokens)
+ restrict.deny_list.extend(namespace.deny_tokens)
+ restricts.append(restrict)
+ # Numeric restricts
+ numeric_restricts = []
+ if numeric_filter:
+ for numeric_namespace in numeric_filter:
+ numeric_restrict = gca_index_v1beta1.IndexDatapoint.NumericRestriction()
+ numeric_restrict.namespace = numeric_namespace.name
+ numeric_restrict.op = numeric_namespace.op
+ numeric_restrict.value_int = numeric_namespace.value_int
+ numeric_restrict.value_float = numeric_namespace.value_float
+ numeric_restrict.value_double = numeric_namespace.value_double
+ numeric_restricts.append(numeric_restrict)
+ # Queries
+ query_by_id = False if queries else True
+ queries = queries if queries else embedding_ids
+ if queries:
+ for query in queries:
+ find_neighbors_query = gca_match_service_v1beta1.FindNeighborsRequest.Query(
+ neighbor_count=num_neighbors,
+ per_crowding_attribute_neighbor_count=per_crowding_attribute_neighbor_count,
+ approximate_neighbor_count=approx_num_neighbors,
+ fraction_leaf_nodes_to_search_override=fraction_leaf_nodes_to_search_override,
+ )
+ datapoint = gca_index_v1beta1.IndexDatapoint(
+ datapoint_id=query if query_by_id else None,
+ feature_vector=None if query_by_id else query,
+ )
+ datapoint.restricts.extend(restricts)
+ datapoint.numeric_restricts.extend(numeric_restricts)
+ find_neighbors_query.datapoint = datapoint
+ find_neighbors_request.queries.append(find_neighbors_query)
+ else:
+ raise ValueError(
+ "To find neighbors using matching engine,"
+ "please specify `queries` or `embedding_ids`"
  )
- datapoint = gca_index_v1beta1.IndexDatapoint(feature_vector=query)
- # Token restricts
- if filter:
- for namespace in filter:
- restrict = gca_index_v1beta1.IndexDatapoint.Restriction()
- restrict.namespace = namespace.name
- restrict.allow_list.extend(namespace.allow_tokens)
- restrict.deny_list.extend(namespace.deny_tokens)
- datapoint.restricts.append(restrict)
- # Numeric restricts
- if numeric_filter:
- for numeric_namespace in numeric_filter:
- numeric_restrict = (
- gca_index_v1beta1.IndexDatapoint.NumericRestriction()
- )
- numeric_restrict.namespace = numeric_namespace.name
- numeric_restrict.op = numeric_namespace.op
- numeric_restrict.value_int = numeric_namespace.value_int
- numeric_restrict.value_float = numeric_namespace.value_float
- numeric_restrict.value_double = numeric_namespace.value_double
- datapoint.numeric_restricts.append(numeric_restrict)
-
- find_neighbors_query.datapoint = datapoint
- find_neighbors_request.queries.append(find_neighbors_query)
 
  response = self._public_match_client.find_neighbors(find_neighbors_request)
 
  # Wrap the results in MatchNeighbor objects and return
  return [
  [
  MatchNeighbor(
- id=neighbor.datapoint.datapoint_id, distance=neighbor.distance
+ id=neighbor.datapoint.datapoint_id,
+ distance=neighbor.distance,
+ feature_vector=neighbor.datapoint.feature_vector,
  )
  for neighbor in embedding_neighbors.neighbors
  ]
@@ -1429,13 +1449,12 @@ def _batch_get_embeddings(
  def match(
  self,
  deployed_index_id: str,
- queries: Optional[List[List[float]]] = None,
+ queries: List[List[float]] = None,
  num_neighbors: int = 1,
  filter: Optional[List[Namespace]] = None,
  per_crowding_attribute_num_neighbors: Optional[int] = None,
  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
@@ -1468,11 +1487,6 @@ def match(
  query time allows user to tune search performance. This value
  increase result in both search accuracy and latency increase.
  The value should be between 0.0 and 1.0.
- return_full_datapoint (bool):
- Optional. If set to true, the full datapoints (including all
- 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
@@ -1518,9 +1532,13 @@ def match(
  per_crowding_attribute_num_neighbors=per_crowding_attribute_num_neighbors,
  approx_num_neighbors=approx_num_neighbors,
  fraction_leaf_nodes_to_search_override=fraction_leaf_nodes_to_search_override,
- embedding_enabled=return_full_datapoint,
  )
  requests.append(request)
+ else:
+ raise ValueError(
+ "To find neighbors using matching engine,"
+ "please specify `queries` or `embedding_ids`"
+ )
 
  batch_request_for_index.requests.extend(requests)
  batch_request.requests.append(batch_request_for_index)
@@ -1531,8 +1549,11 @@ def match(
  # Wrap the results in MatchNeighbor objects and return
  return [
  [
- MatchNeighbor(id=neighbor.id, distance=neighbor.distance)
- for neighbor in embedding_neighbors.neighbor
+ MatchNeighbor(
+ id=embedding_neighbors.neighbor[i].id,
+ distance=embedding_neighbors.neighbor[i].distance,
+ )
+ for i in range(len(embedding_neighbors.neighbor))
  ]
  for embedding_neighbors in response.responses[0].responses
  ]

tests/unit/aiplatform/test_matching_engine_index_endpoint.py

@@ -230,6 +230,7 @@
  -0.021106,
  ]
 ]
+_TEST_QUERY_IDS = ["1", "2"]
 _TEST_NUM_NEIGHBOURS = 1
 _TEST_FILTER = [
  Namespace(name="class", allow_tokens=["token_1"], deny_tokens=["token_2"])
@@ -1044,7 +1045,7 @@ def test_index_endpoint_match_queries_backward_compatibility(
  index_endpoint_match_queries_mock.assert_called_with(batch_request)
 
  @pytest.mark.usefixtures("get_index_endpoint_mock")
- def test_private_index_endpoint_match_queries(
+ def test_private_service_access_index_endpoint_match_queries(
  self, index_endpoint_match_queries_mock
  ):
  aiplatform.init(project=_TEST_PROJECT)
@@ -1061,7 +1062,6 @@ def test_private_index_endpoint_match_queries(
  per_crowding_attribute_num_neighbors=_TEST_PER_CROWDING_ATTRIBUTE_NUM_NEIGHBOURS,
  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,
  )
 
@@ -1085,7 +1085,6 @@ def test_private_index_endpoint_match_queries(
  per_crowding_attribute_num_neighbors=_TEST_PER_CROWDING_ATTRIBUTE_NUM_NEIGHBOURS,
  approx_num_neighbors=_TEST_APPROX_NUM_NEIGHBORS,
  fraction_leaf_nodes_to_search_override=_TEST_FRACTION_LEAF_NODES_TO_SEARCH_OVERRIDE,
- embedding_enabled=_TEST_RETURN_FULL_DATAPOINT,
  )
  for i in range(len(_TEST_QUERIES))
  ],
@@ -1135,7 +1134,6 @@ def test_private_index_endpoint_find_neighbor_queries(
  per_crowding_attribute_num_neighbors=_TEST_PER_CROWDING_ATTRIBUTE_NUM_NEIGHBOURS,
  approx_num_neighbors=_TEST_APPROX_NUM_NEIGHBORS,
  fraction_leaf_nodes_to_search_override=_TEST_FRACTION_LEAF_NODES_TO_SEARCH_OVERRIDE,
- embedding_enabled=_TEST_RETURN_FULL_DATAPOINT,
  )
  for test_query in _TEST_QUERIES
  ],
@@ -1241,6 +1239,56 @@ def test_index_public_endpoint_find_neighbors_queries(
  find_neighbors_request
  )
 
+ @pytest.mark.usefixtures("get_index_public_endpoint_mock")
+ def test_index_public_endpoint_find_neiggbor_query_by_id(
+ self, index_public_endpoint_match_queries_mock
+ ):
+ aiplatform.init(project=_TEST_PROJECT)
+
+ my_pubic_index_endpoint = aiplatform.MatchingEngineIndexEndpoint(
+ index_endpoint_name=_TEST_INDEX_ENDPOINT_ID
+ )
+
+ my_pubic_index_endpoint.find_neighbors(
+ deployed_index_id=_TEST_DEPLOYED_INDEX_ID,
+ num_neighbors=_TEST_NUM_NEIGHBOURS,
+ filter=_TEST_FILTER,
+ per_crowding_attribute_neighbor_count=_TEST_PER_CROWDING_ATTRIBUTE_NUM_NEIGHBOURS,
+ 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,
+ embedding_ids=_TEST_QUERY_IDS,
+ )
+
+ find_neighbors_request = gca_match_service_v1beta1.FindNeighborsRequest(
+ index_endpoint=my_pubic_index_endpoint.resource_name,
+ deployed_index_id=_TEST_DEPLOYED_INDEX_ID,
+ queries=[
+ gca_match_service_v1beta1.FindNeighborsRequest.Query(
+ neighbor_count=_TEST_NUM_NEIGHBOURS,
+ datapoint=gca_index_v1beta1.IndexDatapoint(
+ datapoint_id=_TEST_QUERY_IDS[i],
+ restricts=[
+ gca_index_v1beta1.IndexDatapoint.Restriction(
+ namespace="class",
+ allow_list=["token_1"],
+ deny_list=["token_2"],
+ )
+ ],
+ ),
+ per_crowding_attribute_neighbor_count=_TEST_PER_CROWDING_ATTRIBUTE_NUM_NEIGHBOURS,
+ approximate_neighbor_count=_TEST_APPROX_NUM_NEIGHBORS,
+ fraction_leaf_nodes_to_search_override=_TEST_FRACTION_LEAF_NODES_TO_SEARCH_OVERRIDE,
+ )
+ for i in range(len(_TEST_QUERY_IDS))
+ ],
+ return_full_datapoint=_TEST_RETURN_FULL_DATAPOINT,
+ )
+
+ index_public_endpoint_match_queries_mock.assert_called_with(
+ find_neighbors_request
+ )
+
  @pytest.mark.usefixtures("get_index_public_endpoint_mock")
  def test_index_public_endpoint_match_queries_with_numeric_filtering(
  self, index_public_endpoint_match_queries_mock