From 52809ec35fb3c66a213d2ebc6edef77609fb84be Mon Sep 17 00:00:00 2001 From: dong xiang Date: Thu, 24 Nov 2022 12:42:08 +0800 Subject: [PATCH] fix: add support independent len calculation Signed-off-by: dong xiang --- docarray/array/storage/milvus/getsetdel.py | 2 ++ docarray/array/storage/redis/seqlike.py | 6 +++++- tests/unit/document/test_disable_offset2id.py | 9 +++++++-- 3 files changed, 14 insertions(+), 3 deletions(-) diff --git a/docarray/array/storage/milvus/getsetdel.py b/docarray/array/storage/milvus/getsetdel.py index 08757f0f715..3c0354d2783 100644 --- a/docarray/array/storage/milvus/getsetdel.py +++ b/docarray/array/storage/milvus/getsetdel.py @@ -40,6 +40,8 @@ def _load_offset2ids(self): ) sorted_res = sorted(res, key=lambda k: int(k['offset'])) self._offset2ids = Offset2ID([r['document_id'] for r in sorted_res]) + else: + self._offset2ids = Offset2ID([], list_like=self._list_like) def _save_offset2ids(self): if self._list_like: diff --git a/docarray/array/storage/redis/seqlike.py b/docarray/array/storage/redis/seqlike.py index ad60024434e..607a1fffc01 100644 --- a/docarray/array/storage/redis/seqlike.py +++ b/docarray/array/storage/redis/seqlike.py @@ -26,8 +26,12 @@ def __len__(self): :return: the length of this :class:`DocumentArrayRedis` object """ - try: + if self._list_like: return len(self._offset2ids) + try: + lua_script = f'return #redis.pcall("keys", "{self._config.index_name}:*")' + cmd = self._client.register_script(lua_script) + return cmd() except: return 0 diff --git a/tests/unit/document/test_disable_offset2id.py b/tests/unit/document/test_disable_offset2id.py index db7fab0e7c5..b5d4104b3b3 100644 --- a/tests/unit/document/test_disable_offset2id.py +++ b/tests/unit/document/test_disable_offset2id.py @@ -25,9 +25,14 @@ def docs(): ) def test_disable_offset2id(docs, storage, config, start_storage): if config: - da = DocumentArray(docs, storage=storage, config=config) + da = DocumentArray(storage=storage, config=config) else: - da = DocumentArray(docs, storage=storage) + da = DocumentArray(storage=storage) + + assert len(da) == 0 + + da.extend(docs) + assert len(da) == 3 with pytest.raises(ValueError): da[0]