fix: Add embedding_metadata to MatchNeighbor from_embedding.

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,16 @@ 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,
-                    sparse_distance=(
-                        neighbor.sparse_distance if neighbor.sparse_distance else None
-                    ),
+                    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

@@ -2459,6 +2459,7 @@ def test_from_embedding(self):
                     name="namespace2", value_int=10, value_float=None, value_double=None
                 )
             ],
+            embedding_metadata={"key": "value"},
         )
 
         result = MatchNeighbor(id="embedding_id", distance=0.3).from_embedding(
@@ -2476,3 +2477,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"}