Skip to content
This repository was archived by the owner on Sep 9, 2026. It is now read-only.

Commit e2d4d19

Browse files
JohannesMessnerNicholasDunhamsamsjaalaeddine-13
authored
feat: milvus document store (#587)
* fix(plot): be robust against non-existing subindices * chore: add file structure * feat: first draft of backend implementation * feat: docarray can now connect to running milvus database * feat: implement basics of getsetdel and seqlike * fix: type hint * fix: saving and loading offset2ids * feat: first implementation of vector search * refactor: declare static methods * feat: add consistency level as a configuration parameter * feat: change default consistency to session * refactor: in clear_storage, drop and re-create collection * feat: implement filter and hybrid search * feat: implement columns feature * feat: allow consistency level to be passed to extend, append, insert * fix: fix columns feature in milvus * test: add test for milvus columns feature * fix: import types for type hints * fix: add pymilvus requirement * test: fix docker compose for milvus * test: add milvus to test construct * fix: backend naming * fix: sort returned docs when accessing by id * test: add milvus to sequence tests * fix: raise keyerror on non existing id * refactor: use context managar to laod and release collections * fix: map embeddings to np before inserting * test: add milvus to advanced indexing tests * test: add milvus to test pull out * fix: deleting to size zero * fix: when accessing with no ids, return empty da instead of rasing * test: add milvus to content tests * test: add milvus to del tests * test: add milvus to embed tests * test: add milvus to test empty * test: add milvus to test eval class * test: add milvus to test find * test: add milvus to getset tests * test: add milvus to all remaining tests * feat: load and release collection in context manager * test: use context manager to speed up milvus tests * test: add milvus to test plot * test: fix plot tests for milvus * test: fix multimodal find test * test: fix test embed * test: use context manager to speed up milvus * fix: implement state methods to enable pickling * test: fix more tests * feat: add overloaded milvus init * test: fix even more tests * refactor: remove some comments * test: use context manager even more * test: moar context manager usage * test: fix test fixture input * test: remove milvus from test that it can't handle * test: remove milvus from one more test * refactor: better mechanism for automatic collection loading * docs: add docs section for milvus * chore: remove comment * fix: enforce limit for filter only queries * docs: fix typo * docs: add milvus to comparison and add docstring * feat: bulk extend * docs: document advanced milvus options * feat: allo passing of kwargs to insert * chore: remove comment * chore: update accepted array types * test: add milvus specific tests * docs: apply changes from code review Co-authored-by: Nicholas Dunham <[email protected]> Signed-off-by: Johannes Messner <[email protected]> * docs: update docs/advanced/document-store/milvus.md Co-authored-by: Nicholas Dunham <[email protected]> Signed-off-by: Johannes Messner <[email protected]> * refactor: update type hint Co-authored-by: samsja <[email protected]> Signed-off-by: Johannes Messner <[email protected]> * refactor: apply suggestions from code review * docs: clarify docstring Signed-off-by: Johannes Messner <[email protected]> * docs: fix fautly docstring Signed-off-by: Johannes Messner <[email protected]> * refactor: another round of review changes * fix: set consistency level for offset id loading * fix: set stricter defaut consistency level * perf: optimize sorting of retrieved documents * docs: document loading context manager * refactor: find unboxing done by base class * feat: add batching * test: add test for batch size * test: restart milvus if it breaks * test: dont remove-orphans where it is not needed * feat: add ability to disble list like behaviour Signed-off-by: Johannes Messner <[email protected]> * ci: increase timeout to see how long it really takes * docs: apply suggestions from code review Co-authored-by: AlaeddineAbdessalem <[email protected]> Signed-off-by: Johannes Messner <[email protected]> * ci: change timeouts Signed-off-by: Johannes Messner <[email protected]> Signed-off-by: Johannes Messner <[email protected]> Signed-off-by: Johannes Messner <[email protected]> Co-authored-by: Nicholas Dunham <[email protected]> Co-authored-by: samsja <[email protected]> Co-authored-by: AlaeddineAbdessalem <[email protected]>
1 parent 227b2e5 commit e2d4d19

39 files changed

Lines changed: 1934 additions & 223 deletions

‎.github/workflows/ci.yml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -187,7 +187,7 @@ jobs:
187187
pytest --suppress-no-test-exit-code --cov=docarray --cov-report=xml \
188188
-v -s -m "not gpu" ${{ matrix.test-path }}
189189
echo "codecov_flag=docarray" >> $GITHUB_OUTPUT
190-
timeout-minutes: 45
190+
timeout-minutes: 60
191191
env:
192192
JINA_AUTH_TOKEN: "${{ secrets.JINA_AUTH_TOKEN }}"
193193
- name: Check codecov file
@@ -238,7 +238,7 @@ jobs:
238238
pytest --suppress-no-test-exit-code --cov=docarray --cov-report=xml \
239239
-v -s -m "not gpu" ${{ matrix.test-path }}
240240
echo "::set-output name=codecov_flag::docarray"
241-
timeout-minutes: 40
241+
timeout-minutes: 60
242242
env:
243243
JINA_AUTH_TOKEN: "${{ secrets.JINA_AUTH_TOKEN }}"
244244
- name: Check codecov file

‎docarray/array/document.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,11 +12,13 @@
1212
from docarray.array.weaviate import DocumentArrayWeaviate
1313
from docarray.array.elastic import DocumentArrayElastic
1414
from docarray.array.redis import DocumentArrayRedis
15+
from docarray.array.milvus import DocumentArrayMilvus
1516
from docarray.array.storage.sqlite import SqliteConfig
1617
from docarray.array.storage.annlite import AnnliteConfig
1718
from docarray.array.storage.weaviate import WeaviateConfig
1819
from docarray.array.storage.elastic import ElasticConfig
1920
from docarray.array.storage.redis import RedisConfig
21+
from docarray.array.storage.milvus import MilvusConfig
2022

2123

2224
class DocumentArray(AllMixins, BaseDocumentArray):
@@ -140,6 +142,16 @@ def __new__(
140142
"""Create a Redis-powered DocumentArray object."""
141143
...
142144

145+
@overload
146+
def __new__(
147+
cls,
148+
_docs: Optional['DocumentArraySourceType'] = None,
149+
storage: str = 'milvus',
150+
config: Optional[Union['MilvusConfig', Dict]] = None,
151+
) -> 'DocumentArrayMilvus':
152+
"""Create a Milvus-powered DocumentArray object."""
153+
...
154+
143155
def __enter__(self):
144156
self._exit_stack = ExitStack()
145157
# Ensure that we sync the data to the storage backend when exiting the context manager
@@ -184,6 +196,10 @@ def __new__(cls, *args, storage: str = 'memory', **kwargs):
184196
from .redis import DocumentArrayRedis
185197

186198
instance = super().__new__(DocumentArrayRedis)
199+
elif storage == 'milvus':
200+
from .milvus import DocumentArrayMilvus
201+
202+
instance = super().__new__(DocumentArrayMilvus)
187203

188204
else:
189205
raise ValueError(f'storage=`{storage}` is not supported.')

‎docarray/array/milvus.py‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
from .document import DocumentArray
2+
3+
from .storage.milvus import StorageMixins, MilvusConfig
4+
5+
__all__ = ['MilvusConfig', 'DocumentArrayMilvus']
6+
7+
8+
class DocumentArrayMilvus(StorageMixins, DocumentArray):
9+
"""
10+
DocumentArray that stores Documents in a `Milvus <https://milvus.io//>`_ vector search engine.
11+
12+
.. note::
13+
This DocumentArray requires `pymilvus`. You can install it via `pip install "docarray[milvus]"`.
14+
15+
To use Milvus as storage backend, a Milvus service needs to be running on your machine.
16+
17+
With this implementation, :meth:`match` and :meth:`find` perform fast (approximate) vector search.
18+
Additionally, search with filters is supported.
19+
20+
Example usage:
21+
22+
.. code-block:: python
23+
24+
from docarray import DocumentArray
25+
26+
# connect to running Milvus service with default configuration (address: http://localhost:19530)
27+
da = DocumentArray(storage='milvus', config={'n_dim': 10})
28+
29+
# connect to a previously persisted DocumentArrayMilvus by specifying collection_name, host, and port
30+
da = DocumentArray(
31+
storage='milvus',
32+
config={
33+
'collection_name': 'persisted',
34+
'host': 'localhost',
35+
'port': '19530',
36+
'n_dim': 10,
37+
},
38+
)
39+
40+
41+
.. seealso::
42+
For further details, see our :ref:`user guide <milvus>`.
43+
"""
44+
45+
def __new__(cls, *args, **kwargs):
46+
return super().__new__(cls)

‎docarray/array/mixins/find.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ def find(
9696
limit: Optional[Union[int, float]] = 20,
9797
metric_name: Optional[str] = None,
9898
exclude_self: bool = False,
99-
filter: Optional[Dict] = None,
99+
filter: Union[Dict, str, None] = None,
100100
only_id: bool = False,
101101
index: str = 'text',
102102
on: Optional[str] = None,

‎docarray/array/storage/base/seqlike.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,13 +14,14 @@ def _update_subindices_append_extend(self, value):
1414
if len(docs_selector) > 0:
1515
da.extend(docs_selector)
1616

17-
def insert(self, index: int, value: 'Document'):
17+
def insert(self, index: int, value: 'Document', **kwargs):
1818
"""Insert `doc` at `index`.
1919
2020
:param index: Position of the insertion.
2121
:param value: The doc needs to be inserted.
22+
:param kwargs: Additional Arguments that are passed to the Document Store. This has no effect for in-memory DocumentArray.
2223
"""
23-
self._set_doc_by_id(value.id, value)
24+
self._set_doc_by_id(value.id, value, **kwargs)
2425
self._offset2ids.insert(index, value.id)
2526

2627
def append(self, value: 'Document', **kwargs):

‎docarray/array/storage/milvus/__init__.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
from abc import ABC
2+
3+
from .backend import BackendMixin, MilvusConfig
4+
from .find import FindMixin
5+
from .getsetdel import GetSetDelMixin
6+
from .seqlike import SequenceLikeMixin
7+
8+
__all__ = ['StorageMixins', 'MilvusConfig']
9+
10+
11+
class StorageMixins(FindMixin, BackendMixin, GetSetDelMixin, SequenceLikeMixin, ABC):
12+
...

0 commit comments

Comments
 (0)