refactor: minor adjustments

Signed-off-by: jupyterjazz <saba.sturua@jina.ai>
docarray · JoanFM · Jul 19, 2023 · Jul 2, 2023 · Jul 2, 2023 · Jul 13, 2023
commit 6dce3166c780729c273613306af687b39c0c4f47
diff --git a/docarray/index/__init__.py b/docarray/index/__init__.py
@@ -16,7 +16,15 @@
     from docarray.index.backends.redis import RedisDocumentIndex  # noqa: F401
     from docarray.index.backends.milvus import MilvusDocumentIndex  # noqa: F401
 
-__all__ = ['InMemoryExactNNIndex']
+__all__ = [
+    'InMemoryExactNNIndex',
+    'ElasticDocIndex',
+    'ElasticV7DocIndex',
+    'QdrantDocumentIndex',
+    'WeaviateDocumentIndex',
+    'RedisDocumentIndex',
+    'MilvusDocumentIndex',
+]
 
 
 def __getattr__(name: str):

diff --git a/docarray/index/backends/milvus.py b/docarray/index/backends/milvus.py
@@ -24,10 +24,7 @@
     _raise_not_supported,
     _raise_not_composable,
 )
-from docarray.index.backends.helper import (
-    _execute_find_and_filter_query,
-    _collect_query_args,
-)
+from docarray.index.backends.helper import _collect_query_args
 from docarray.typing import AnyTensor, NdArray
 from docarray.typing.id import ID
 from docarray.typing.tensor.abstract_tensor import AbstractTensor
@@ -201,7 +198,7 @@ def python_type_to_db_type(self, python_type: Type) -> Any:
             if safe_issubclass(python_type, py_type):
                 return db_type
 
-        return None
+        raise ValueError(f'Unsupported column type for {type(self)}: {python_type}')
 
     def _create_or_load_collection(self) -> Collection:
         """
@@ -778,10 +775,10 @@ def execute_query(self, query: Any, *args, **kwargs) -> Any:
 
         return FindResult(documents=docs, scores=scores)
 
-    def _docs_from_query_response(self, result: Sequence[Dict]) -> Sequence[TSchema]:
-        return DocList[self._schema](
+    def _docs_from_query_response(self, result: Sequence[Dict]) -> DocList[Any]:
+        return DocList[self._schema](  # type: ignore
             [
-                self._schema.from_base64(
+                self._schema.from_base64(  # type: ignore
                     result[i]["serialized"], **self._db_config.serialize_config
                 )
                 for i in range(len(result))
@@ -794,7 +791,7 @@ def _docs_from_find_response(self, result: Hits) -> _FindResult:
         )
 
         return _FindResult(
-            documents=DocList[self.out_schema](
+            documents=DocList[self.out_schema](  # type: ignore
                 [
                     self.out_schema.from_base64(
                         hit.entity.get('serialized'), **self._db_config.serialize_config
@@ -815,7 +812,7 @@ def _always_true_expr(self, primary_key: str) -> str:
         """
         return f'({primary_key} in ["1"]) or ({primary_key} not in ["1"])'
 
-    def _map_embedding(self, embedding: Optional[AnyTensor]) -> Optional[AnyTensor]:
+    def _map_embedding(self, embedding: AnyTensor) -> np.ndarray:
         """
         Milvus exclusively supports one-dimensional vectors. If multi-dimensional
         vectors are provided, they will be automatically flattened to ensure compatibility.
@@ -830,7 +827,7 @@ def _map_embedding(self, embedding: Optional[AnyTensor]) -> Optional[AnyTensor]:
 
         embedding = self._to_numpy(embedding)
         if embedding.ndim > 1:
-            embedding = np.asarray(embedding).squeeze()
+            embedding = np.asarray(embedding).squeeze()  # type: ignore
 
         return embedding
 

diff --git a/docarray/index/backends/redis.py b/docarray/index/backends/redis.py
@@ -1,4 +1,3 @@
-import uuid
 from collections import defaultdict
 from typing import (
     TypeVar,