fix: Add embedding_metadata to MatchNeighbor from_embedding.

PiperOrigin-RevId: 879842103

Author: lingyinw

google/cloud/aiplatform/matching_engine/matching_engine_index_endpoint.py

@@ -151,7 +151,7 @@ def __post_init__(self):
 @dataclass
 class HybridQuery:
     """
-    Hyrbid query. Could be used for dense-only or sparse-only or hybrid queries.
+    Hybrid query. Could be used for dense-only or sparse-only or hybrid queries.
 
     dense_embedding (List[float]):
         Optional. The dense part of the hybrid queries.
@@ -328,6 +328,10 @@ def from_embedding(self, embedding: match_service_pb2.Embedding) -> "MatchNeighb
         if embedding.sparse_embedding:
             self.sparse_embedding_values = embedding.sparse_embedding.float_val
             self.sparse_embedding_dimensions = embedding.sparse_embedding.dimension
+
+        # retrieve embedding metadata
+        if embedding.embedding_metadata:
+            self.embedding_metadata = embedding.embedding_metadata
         return self
 
 
@@ -1883,7 +1887,7 @@ def find_neighbors(
             [
                 MatchNeighbor(
                     id=neighbor.datapoint.datapoint_id,
-                    distance=neighbor.distance,
+                    distance=neighbor.distance if neighbor.distance else None,
                     sparse_distance=(
                         neighbor.sparse_distance if neighbor.sparse_distance else None
                     ),
@@ -2219,19 +2223,18 @@ def match(
         # Wrap the results in MatchNeighbor objects and return
         match_neighbors_response = []
         for resp in response.responses[0].responses:
-            match_neighbors_id_map = {}
+            embedding_map = {embedding.id: embedding for embedding in resp.embeddings}
+            neighbors_list = []
             for neighbor in resp.neighbor:
-                match_neighbors_id_map[neighbor.id] = MatchNeighbor(
+                match_neighbor = MatchNeighbor(
                     id=neighbor.id,
-                    distance=neighbor.distance,
+                    distance=neighbor.distance if neighbor.distance else None,
                     sparse_distance=(
                         neighbor.sparse_distance if neighbor.sparse_distance else None
                     ),
                 )
-            for embedding in resp.embeddings:
-                if embedding.id in match_neighbors_id_map:
-                    match_neighbors_id_map[embedding.id] = match_neighbors_id_map[
-                        embedding.id
-                    ].from_embedding(embedding=embedding)
-            match_neighbors_response.append(list(match_neighbors_id_map.values()))
+                if neighbor.id in embedding_map:
+                    match_neighbor.from_embedding(embedding=embedding_map[neighbor.id])
+                neighbors_list.append(match_neighbor)
+            match_neighbors_response.append(neighbors_list)
         return match_neighbors_response

tests/unit/aiplatform/test_matching_engine_index_endpoint.py

@@ -51,6 +51,7 @@
 import constants as test_constants
 
 from google.protobuf import field_mask_pb2
+from google.protobuf import struct_pb2
 
 import grpc
 
@@ -2445,6 +2446,8 @@ def test_from_index_datapoint(self):
         assert result.numeric_restricts[0].value_double is None
 
     def test_from_embedding(self):
+        embedding_metadata_struct = struct_pb2.Struct()
+        embedding_metadata_struct["key"] = "value"
         embedding = match_service_pb2.Embedding(
             id="test_embedding_id",
             float_val=[1.0, 2.0, 3.0],
@@ -2459,6 +2462,7 @@ def test_from_embedding(self):
                     name="namespace2", value_int=10, value_float=None, value_double=None
                 )
             ],
+            embedding_metadata=embedding_metadata_struct,
         )
 
         result = MatchNeighbor(id="embedding_id", distance=0.3).from_embedding(
@@ -2476,3 +2480,4 @@ def test_from_embedding(self):
         assert result.numeric_restricts[0].value_int == 10
         assert not result.numeric_restricts[0].value_float
         assert not result.numeric_restricts[0].value_double
+        assert result.embedding_metadata == {"key": "value"}