diff --git a/docarray/index/backends/qdrant.py b/docarray/index/backends/qdrant.py index 0daad346b44..aa0d49095e2 100644 --- a/docarray/index/backends/qdrant.py +++ b/docarray/index/backends/qdrant.py @@ -67,6 +67,10 @@ class QdrantDocumentIndex(BaseDocIndex, Generic[TSchema]): def __init__(self, db_config=None, **kwargs): """Initialize QdrantDocumentIndex""" + if db_config is not None and getattr( + db_config, 'index_name' + ): # this is needed for subindices + db_config.collection_name = db_config.index_name super().__init__(db_config=db_config, **kwargs) self._db_config: QdrantDocumentIndex.DBConfig = cast( QdrantDocumentIndex.DBConfig, self._db_config @@ -98,11 +102,7 @@ def collection_name(self): 'To do so, use the syntax: QdrantDocumentIndex[DocumentType]' ) - return ( - self._db_config.collection_name - or self._db_config.index_name - or default_collection_name - ) + return self._db_config.collection_name or default_collection_name @property def index_name(self): @@ -250,6 +250,12 @@ class DBConfig(BaseDocIndex.DBConfig): } ) + def __post_init__(self): + if self.collection_name is None and self.index_name is not None: + self.collection_name = self.index_name + if self.index_name is None and self.collection_name is not None: + self.index_name = self.collection_name + @dataclass class RuntimeConfig(BaseDocIndex.RuntimeConfig): """Dataclass that contains all "dynamic" configurations of QdrantDocumentIndex.""" diff --git a/tests/index/qdrant/test_subindex.py b/tests/index/qdrant/test_subindex.py index 0af344518ba..9343c8d1d3b 100644 --- a/tests/index/qdrant/test_subindex.py +++ b/tests/index/qdrant/test_subindex.py @@ -222,3 +222,36 @@ def test_subindex_contain(index): # Empty index empty_index = QdrantDocumentIndex[MyDoc]() assert (empty_doc in empty_index) is False + + +def test_subindex_collections(): + from typing import Optional + from docarray.typing.tensor import AnyTensor + from pydantic import Field + + class MetaPathDoc(BaseDoc): + path_id: str + level: int + text: str + embedding: Optional[AnyTensor] = Field(space='cosine', dim=128) + + class MetaCategoryDoc(BaseDoc): + node_id: Optional[str] + node_name: Optional[str] + name: Optional[str] + product_type_definitions: Optional[str] + leaf: bool + paths: Optional[DocList[MetaPathDoc]] + embedding: Optional[AnyTensor] = Field(space='cosine', dim=128) + channel: str + lang: str + + db_config = QdrantDocumentIndex.DBConfig( + host='localhost', + collection_name="channel_category", + ) + + doc_index = QdrantDocumentIndex[MetaCategoryDoc](db_config) + + assert doc_index._subindices["paths"].index_name == 'channel_category__paths' + assert doc_index._subindices["paths"].collection_name == 'channel_category__paths'