From 17760da3d376d8b988c68d99af0f7f6ed4c3367d Mon Sep 17 00:00:00 2001 From: JackyHsu-X Date: Thu, 22 Oct 2026 18:29:03 +0300 Subject: [PATCH] Keep object iterators exhausted until restarted --- test/collection/test_iterator.py | 100 +++++++++++++++++++++++++++++++ weaviate/collections/iterator.py | 10 ++++ 2 files changed, 110 insertions(+) create mode 100644 test/collection/test_iterator.py diff --git a/test/collection/test_iterator.py b/test/collection/test_iterator.py new file mode 100644 index 000000000..43b3a3f9b --- /dev/null +++ b/test/collection/test_iterator.py @@ -0,0 +1,100 @@ +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock, call +from uuid import UUID + +import pytest + +from weaviate.collections.iterator import _IteratorInputs, _ObjectAIterator, _ObjectIterator + + +def _inputs(): + return _IteratorInputs( + include_vector=True, + return_metadata=None, + return_properties=["title"], + return_references=None, + after=str(UUID(int=1)), + ) + + +@pytest.mark.parametrize("initially_empty", [False, True]) +def test_iterator_stays_exhausted_until_restarted(initially_empty): + obj = SimpleNamespace(uuid=UUID(int=2)) + earlier_obj = SimpleNamespace(uuid=UUID(int=3)) + first_page = [] if initially_empty else [earlier_obj, obj] + pages = [SimpleNamespace(objects=first_page)] + if not initially_empty: + pages.append(SimpleNamespace(objects=[])) + pages.append(SimpleNamespace(objects=[obj])) + query = Mock() + query.fetch_objects.side_effect = pages + iterator = _ObjectIterator(query, _inputs(), cache_size=5) + iter(iterator) + + if not initially_empty: + assert next(iterator) is earlier_obj + assert next(iterator) is obj + assert query.fetch_objects.call_count == 1 + with pytest.raises(StopIteration): + next(iterator) + terminal_call_count = query.fetch_objects.call_count + for _ in range(2): + with pytest.raises(StopIteration): + next(iterator) + assert query.fetch_objects.call_count == terminal_call_count + assert query.fetch_objects.call_args_list[0] == call( + limit=5, + after=UUID(int=1), + include_vector=True, + return_metadata=None, + return_properties=["title"], + return_references=None, + ) + if not initially_empty: + assert query.fetch_objects.call_args.kwargs["after"] == obj.uuid + + iter(iterator) + assert next(iterator) is obj + assert query.fetch_objects.call_args.kwargs["after"] == UUID(int=1) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("initially_empty", [False, True]) +async def test_async_iterator_stays_exhausted_until_restarted(initially_empty): + obj = SimpleNamespace(uuid=UUID(int=2)) + earlier_obj = SimpleNamespace(uuid=UUID(int=3)) + first_page = [] if initially_empty else [earlier_obj, obj] + pages = [SimpleNamespace(objects=first_page)] + if not initially_empty: + pages.append(SimpleNamespace(objects=[])) + pages.append(SimpleNamespace(objects=[obj])) + query = Mock() + query.fetch_objects = AsyncMock(side_effect=pages) + iterator = _ObjectAIterator(query, _inputs(), cache_size=5) + iterator.__aiter__() + + if not initially_empty: + assert await iterator.__anext__() is earlier_obj + assert await iterator.__anext__() is obj + assert query.fetch_objects.await_count == 1 + with pytest.raises(StopAsyncIteration): + await iterator.__anext__() + terminal_call_count = query.fetch_objects.await_count + for _ in range(2): + with pytest.raises(StopAsyncIteration): + await iterator.__anext__() + assert query.fetch_objects.await_count == terminal_call_count + assert query.fetch_objects.call_args_list[0] == call( + limit=5, + after=UUID(int=1), + include_vector=True, + return_metadata=None, + return_properties=["title"], + return_references=None, + ) + if not initially_empty: + assert query.fetch_objects.call_args.kwargs["after"] == obj.uuid + + iterator.__aiter__() + assert await iterator.__anext__() is obj + assert query.fetch_objects.call_args.kwargs["after"] == UUID(int=1) diff --git a/weaviate/collections/iterator.py b/weaviate/collections/iterator.py index 326000f9a..c2e9af74e 100644 --- a/weaviate/collections/iterator.py +++ b/weaviate/collections/iterator.py @@ -57,15 +57,19 @@ def __init__( self.__iter_object_cache: deque[Object[TProperties, TReferences]] = deque() self.__iter_object_last_uuid: Optional[UUID] = _parse_after(self.__inputs.after) self.__iter_cache_size = cache_size or ITERATOR_CACHE_SIZE + self.__iter_object_exhausted = False def __iter__( self, ) -> Iterator[Object[TProperties, TReferences]]: + self.__iter_object_exhausted = False self.__iter_object_cache = deque() self.__iter_object_last_uuid = _parse_after(self.__inputs.after) return self def __next__(self) -> Object[TProperties, TReferences]: + if self.__iter_object_exhausted: + raise StopIteration if len(self.__iter_object_cache) == 0: res = self.__query.fetch_objects( limit=self.__iter_cache_size, @@ -77,6 +81,7 @@ def __next__(self) -> Object[TProperties, TReferences]: ) self.__iter_object_cache = deque(res.objects) # type: ignore if len(self.__iter_object_cache) == 0: + self.__iter_object_exhausted = True raise StopIteration ret_object = self.__iter_object_cache.popleft() @@ -103,10 +108,12 @@ def __init__( self.__iter_object_cache: deque[Object[TProperties, TReferences]] = deque() self.__iter_object_last_uuid: Optional[UUID] = _parse_after(self.__inputs.after) self.__iter_cache_size = cache_size or ITERATOR_CACHE_SIZE + self.__iter_object_exhausted = False def __aiter__( self, ) -> AsyncIterator[Object[TProperties, TReferences]]: + self.__iter_object_exhausted = False self.__iter_object_cache = deque() self.__iter_object_last_uuid = _parse_after(self.__inputs.after) return self @@ -114,6 +121,8 @@ def __aiter__( async def __anext__( self, ) -> Object[TProperties, TReferences]: + if self.__iter_object_exhausted: + raise StopAsyncIteration if len(self.__iter_object_cache) == 0: res = await self.__query.fetch_objects( limit=self.__iter_cache_size, @@ -125,6 +134,7 @@ async def __anext__( ) self.__iter_object_cache = deque(res.objects) # type: ignore if len(self.__iter_object_cache) == 0: + self.__iter_object_exhausted = True raise StopAsyncIteration ret_object = self.__iter_object_cache.popleft()