add new function get_embeddings_by_id to collection.py for MMR support

python/src/alayalite/collection.py

@@ -295,3 +295,36 @@ def get_index_params(self):
         Retrieve the configuration parameters of the index in the collection.
         """
         return self.__index_params
+
+
+
+    # this function is mainly for MMR purpose，add by liangpeiran
+    def get_embeddings_by_id(self, ids: List[str]) -> List[List[float]]:
+        """
+        Retrieve embeddings (vectors) from the underlying index by external ids.
+
+        Args:
+            ids: List of external ids used in Collection (the "id" column).
+
+        Returns:
+            A list of vectors (list of float), aligned with the input order.
+
+        Raises:
+            RuntimeError: if index is not initialized.
+            KeyError: if any id does not exist in collection.
+        """
+        if not ids:
+            return []
+
+        _assert(self.__index_py is not None, "Index is not initialized yet")
+
+        missing = [x for x in ids if x not in self.__outer_inner_map]
+        if missing:
+            raise KeyError(f"Some ids do not exist in collection: {missing}")
+
+        vectors: List[List[float]] = []
+        for outer_id in ids:
+            inner_id = self.__outer_inner_map[outer_id]
+            vec = self.__index_py.get_data_by_id(inner_id)
+            vectors.append(np.asarray(vec, dtype=np.float32).tolist())
+        return vectors

python/tests/test_collection_get_embedding_by_id.py

@@ -0,0 +1,29 @@
+import numpy as np
+import pytest
+
+from alayalite.collection import Collection
+
+
+def test_get_embeddings_by_id_roundtrip():
+    col = Collection("mmr_test_collection")
+
+    items = [
+        ("a", "doc-a", [1.0, 0.0, 0.0], {"k": 1}),
+        ("b", "doc-b", [0.0, 1.0, 0.0], {"k": 2}),
+        ("c", "doc-c", [0.0, 0.0, 1.0], {"k": 3}),
+    ]
+    col.insert(items)
+
+    got = col.get_embeddings_by_id(["b", "a", "c"])
+    assert len(got) == 3
+    assert np.allclose(got[0], np.array([0.0, 1.0, 0.0], dtype=np.float32))
+    assert np.allclose(got[1], np.array([1.0, 0.0, 0.0], dtype=np.float32))
+    assert np.allclose(got[2], np.array([0.0, 0.0, 1.0], dtype=np.float32))
+
+
+def test_get_embeddings_by_id_missing_raises():
+    col = Collection("mmr_test_collection_missing")
+    col.insert([("a", "doc-a", [1.0, 0.0], {})])
+
+    with pytest.raises(KeyError):
+        col.get_embeddings_by_id(["a", "not-exist"])