Skip to content

Commit

Permalink
enable milvus pagination query (#422)
Browse files Browse the repository at this point in the history
  • Loading branch information
dorren002 authored Feb 5, 2025
1 parent 3d854b2 commit 366e302
Showing 1 changed file with 20 additions and 2 deletions.
22 changes: 20 additions & 2 deletions lazyllm/tools/rag/milvus_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
from lazyllm.common import override, obj2str, str2obj

MILVUS_UPSERT_BATCH_SIZE = 500
MILVUS_PAGINATION_OFFSET = 1000

class MilvusStore(StoreBase):
# we define these variables as members so that pymilvus is not imported until MilvusStore is instantiated.
Expand Down Expand Up @@ -246,8 +247,25 @@ def _gen_field_key(self, k: str) -> str:
def _load_all_nodes_to(self, store: StoreBase) -> None:
uid2node = {}
for group_name in self._client.list_collections():
results = self._client.query(collection_name=group_name,
filter=f'{self._primary_key} != ""')
collection_desc = self._client.describe_collection(collection_name=group_name)
field_names = [field.get("name") for field in collection_desc.get('fields', [])]

iterator = self._client.query_iterator(
collection_name=group_name,
batch_size=MILVUS_PAGINATION_OFFSET,
filter=f'{self._primary_key} != ""',
output_fields=field_names
)

results = []
while True:
result = iterator.next()

if not result:
iterator.close()
break
results += result

for result in results:
node = self._deserialize_node_partial(result)
node._group = group_name
Expand Down

0 comments on commit 366e302

Please sign in to comment.