diff --git a/docarray/array/storage/milvus/getsetdel.py b/docarray/array/storage/milvus/getsetdel.py index 3c0354d2783..ec9b188fcf8 100644 --- a/docarray/array/storage/milvus/getsetdel.py +++ b/docarray/array/storage/milvus/getsetdel.py @@ -74,8 +74,7 @@ def _get_docs_by_ids(self, ids: 'Iterable[str]', **kwargs) -> 'DocumentArray': raise KeyError(f'No documents found for ids {ids}') docs.extend(self._docs_from_query_response(res)) # sort output docs according to input id sorting - id_to_index = {id_: i for i, id_ in enumerate(ids)} - return DocumentArray(sorted(docs, key=lambda d: id_to_index[d.id])) + return DocumentArray([docs[d] for d in ids]) def _del_docs_by_ids(self, ids: 'Iterable[str]', **kwargs) -> 'DocumentArray': kwargs = self._update_kwargs_from_config('consistency_level', **kwargs) diff --git a/tests/unit/array/test_advance_indexing.py b/tests/unit/array/test_advance_indexing.py index 86fb1a07f9e..088aa5f3bef 100644 --- a/tests/unit/array/test_advance_indexing.py +++ b/tests/unit/array/test_advance_indexing.py @@ -273,12 +273,13 @@ def test_sequence_str(docs, storage, config, start_storage): else: docs = DocumentArray(docs, storage=storage) # getter - idx = [d.id for d in docs[1, 3, 5, 7, -1, -2]] + idx = [d.id for d in docs[1, 3, 5, 7, -1, -2, 1]] assert len(docs[idx]) == len(idx) assert len(docs[tuple(idx)]) == len(idx) # setter + idx = [d.id for d in docs[1, 3, 5, 7, -1, -2]] docs[idx] = [Document(text='repl') for _ in range(len(idx))] idx = [d.id for d in docs[1, 3, 5, 7, -1, -2]] for _id in idx: