fix(iterator): keep exhausted object iterators from fetching again

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)

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,17 +108,21 @@ 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
 
     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()