refactor: unify embedding map methods naming

hanxiao · Mar 2, 2022 · Mar 1, 2022 · Mar 1, 2022 · Mar 1, 2022 · Mar 2, 2022
commit be4c734c107997cb22810eba4ad3a12ed44c17c4
diff --git a/docarray/array/storage/base/backend.py b/docarray/array/storage/base/backend.py
@@ -18,10 +18,10 @@ def _init_storage(
     def _get_storage_infos(self) -> Optional[Dict]:
         ...
 
-    @staticmethod
-    def _map_id(_id: str) -> str:
+    def _map_id(self, _id: str) -> str:
         return _id
 
-    @staticmethod
-    def _map_embedding(embedding: 'ArrayType') -> 'ArrayType':
-        return embedding
+    def _map_embedding(self, embedding: 'ArrayType') -> 'ArrayType':
+        from ....math.ndarray import to_numpy_array
+
+        return to_numpy_array(embedding)
diff --git a/docarray/array/storage/pqlite/__init__.py b/docarray/array/storage/pqlite/__init__.py
@@ -16,9 +16,3 @@
 
 class StorageMixins(FindMixin, BackendMixin, GetSetDelMixin, SequenceLikeMixin, ABC):
     ...
-
-    def _to_numpy_embedding(self, doc: 'Document'):
-        if doc.embedding is None:
-            doc.embedding = np.zeros(self._pqlite.dim, dtype=np.float32)
-        elif isinstance(doc.embedding, list):
-            doc.embedding = np.array(doc.embedding, dtype=np.float32)
diff --git a/docarray/array/storage/pqlite/backend.py b/docarray/array/storage/pqlite/backend.py
@@ -9,6 +9,8 @@
     Generator,
     Iterator,
 )
+
+import numpy as np
 from pqlite import PQLite
 
 from ..base.backend import BaseBackendMixin
@@ -96,3 +98,10 @@ def _get_storage_infos(self) -> Dict:
             'Data Path': self._config.data_path,
             'Serialization Protocol': self._config.serialize_config.get('protocol'),
         }
+
+    def _map_embedding(self, embedding: 'ArrayType') -> 'ArrayType':
+        if embedding is None:
+            embedding = np.zeros(self._pqlite.dim, dtype=np.float32)
+        elif isinstance(embedding, list):
+            embedding = np.array(embedding, dtype=np.float32)
+        return embedding
diff --git a/docarray/array/storage/pqlite/getsetdel.py b/docarray/array/storage/pqlite/getsetdel.py
@@ -19,7 +19,7 @@ def _get_doc_by_id(self, _id: str) -> 'Document':
         return doc
 
     def _set_doc_by_id(self, _id: str, value: 'Document'):
-        self._to_numpy_embedding(value)
+        value.embedding = self._map_embedding(value.embedding)
         docs = DocumentArrayInMemory([value])
         self._pqlite.update(docs)
 
@@ -33,7 +33,7 @@ def _set_docs_by_ids(self, ids, docs: Iterable['Document'], mismatch_ids: Dict):
         self._pqlite.delete(list(mismatch_ids.keys()))
         docs = DocumentArrayInMemory(docs)
         for doc in docs:
-            self._to_numpy_embedding(doc)
+            doc.embedding = self._map_embedding(doc.embedding)
         self._pqlite.update(docs)
 
     def _del_docs_by_ids(self, ids):

diff --git a/docarray/array/storage/pqlite/seqlike.py b/docarray/array/storage/pqlite/seqlike.py
@@ -15,7 +15,7 @@ def extend(self, values: Iterable['Document']) -> None:
             return
 
         for doc in docs:
-            self._to_numpy_embedding(doc)
+            doc.embedding = self._map_embedding(doc.embedding)
 
         self._pqlite.index(docs)
         self._offset2ids.extend([doc.id for doc in docs])

diff --git a/docarray/array/storage/qdrant/backend.py b/docarray/array/storage/qdrant/backend.py
@@ -12,6 +12,7 @@
     List,
 )
 
+import numpy as np
 from qdrant_client import QdrantClient
 from qdrant_openapi_client.models.models import (
     Distance,
@@ -24,6 +25,7 @@
 from docarray.array.storage.base.backend import BaseBackendMixin
 from docarray.array.storage.qdrant.helper import DISTANCES
 from docarray.helper import dataclass_from_dict, random_identity
+from docarray.math.helper import EPSILON
 
 if TYPE_CHECKING:
     from docarray.types import (
@@ -181,3 +183,18 @@ def _get_storage_infos(self) -> Dict:
             'Distance': self._config.distance,
             'Serialization Protocol': self._config.serialize_config.get('protocol'),
         }
+
+    def _map_embedding(self, embedding: 'ArrayType') -> List[float]:
+        if embedding is None:
+            embedding = np.random.rand(self.n_dim)
+        else:
+            from ....math.ndarray import to_numpy_array
+
+            embedding = to_numpy_array(embedding)
+
+        if embedding.ndim > 1:
+            embedding = np.asarray(embedding).squeeze()
+
+        if np.all(embedding == 0):
+            embedding = embedding + EPSILON
+        return embedding.tolist()
diff --git a/docarray/array/storage/qdrant/find.py b/docarray/array/storage/qdrant/find.py
@@ -50,7 +50,7 @@ def distance(self) -> 'Distance':
         raise NotImplementedError()
 
     def _find_similar_vectors(self, q: 'QdrantArrayType', limit=10):
-        query_vector = QdrantStorageHelper.embedding_to_array(q, default_dim=0)
+        query_vector = self._map_embedding(q)
 
         search_result = self.client.search(
             self.collection_name,

diff --git a/docarray/array/storage/qdrant/getsetdel.py b/docarray/array/storage/qdrant/getsetdel.py
@@ -69,7 +69,7 @@ def _document_to_qdrant(self, doc: 'Document') -> 'PointStruct':
         return PointStruct(
             id=self._map_id(doc.id),
             payload=dict(_serialized=doc.to_base64(**self.serialization_config)),
-            vector=QdrantStorageHelper.embedding_to_array(doc.embedding, self.n_dim),
+            vector=self._map_embedding(doc.embedding),
         )
 
     def _get_doc_by_id(self, _id: str) -> 'Document':

diff --git a/docarray/array/storage/qdrant/helper.py b/docarray/array/storage/qdrant/helper.py
@@ -1,35 +1,5 @@
-from typing import List, TYPE_CHECKING
-
-import numpy as np
-import scipy.sparse
 from qdrant_openapi_client.models.models import Distance
 
-from docarray.math.helper import EPSILON
-
-if TYPE_CHECKING:
-    from docarray.types import ArrayType
-
-
-class QdrantStorageHelper:
-    @classmethod
-    def embedding_to_array(
-        cls, embedding: 'ArrayType', default_dim: int
-    ) -> List[float]:
-        if embedding is None:
-            embedding = np.random.rand(default_dim)
-        else:
-            from ....math.ndarray import to_numpy_array
-
-            embedding = to_numpy_array(embedding)
-
-        if embedding.ndim > 1:
-            embedding = np.asarray(embedding).squeeze()
-
-        if np.all(embedding == 0):
-            embedding = embedding + EPSILON
-        return embedding.tolist()
-
-
 DISTANCES = {
     'cosine': Distance.COSINE,
     'euclidean': Distance.EUCLID,

diff --git a/docarray/array/storage/weaviate/backend.py b/docarray/array/storage/weaviate/backend.py
@@ -277,27 +277,11 @@ def _doc2weaviate_create_payload(self, value: 'Document'):
         :param value: document to create a payload for
         :return: the payload dictionary
         """
-        if value.embedding is not None:
-            from ....math.ndarray import to_numpy_array
-
-            embedding = to_numpy_array(value.embedding)
-
-            if embedding.ndim > 1:
-                embedding = np.asarray(embedding).squeeze()
-
-            # Weaviate expects vector to have dim 2 at least
-            # or get weaviate.exceptions.UnexpectedStatusCodeException:  models.C11yVector
-            # hence we cast it to list of a single element
-            if len(embedding) == 1:
-                embedding = [embedding[0]]
-        else:
-            embedding = None
-
         return dict(
             data_object={'_serialized': value.to_base64(**self._serialize_config)},
             class_name=self._class_name,
             uuid=self._map_id(value.id),
-            vector=embedding,
+            vector=self._map_embedding(value.embedding),
         )
 
     def _map_id(self, doc_id: str):
@@ -312,6 +296,24 @@ def _map_id(self, doc_id: str):
         # daw2[0, 'text'] == 'hi' # this will be False if we don't append class name
         return str(uuid.uuid5(uuid.NAMESPACE_URL, doc_id + self._class_name))
 
+    def _map_embedding(self, embedding: 'ArrayType'):
+        if embedding is not None:
+            from ....math.ndarray import to_numpy_array
+
+            embedding = to_numpy_array(embedding)
+
+            if embedding.ndim > 1:
+                embedding = np.asarray(embedding).squeeze()
+
+            # Weaviate expects vector to have dim 2 at least
+            # or get weaviate.exceptions.UnexpectedStatusCodeException:  models.C11yVector
+            # hence we cast it to list of a single element
+            if len(embedding) == 1:
+                embedding = [embedding[0]]
+        else:
+            embedding = None
+        return embedding
+
     def _get_storage_infos(self) -> Dict:
         return {
             'Backend': 'Weaviate',