Merge pull request #991 from weaviate/add-after-to-iterator

Add after parameter to collection.iterator()

integration/test_iterator.py

@@ -203,3 +203,20 @@ def test_iterator(collection_factory: CollectionFactory, count: int) -> None:
             assert first_order == ret
 
         assert sorted(ret) == expected
+
+
+def test_iterator_with_after(collection_factory: CollectionFactory) -> None:
+    collection = collection_factory(
+        properties=[Property(name="data", data_type=DataType.INT)],
+        vectorizer_config=Configure.Vectorizer.none(),
+        data_model_properties=Dict[str, int],
+    )
+
+    collection.data.insert_many([DataObject(properties={"data": i}) for i in range(10)])
+
+    uuids = [obj.uuid for obj in collection.iterator()]
+    iterator = collection.iterator(after=uuids[5])
+    assert (
+        next(iterator).properties["data"]
+        == collection.query.fetch_object_by_id(uuids[6]).properties["data"]
+    )

weaviate/collections/collection.py

@@ -1,4 +1,5 @@
 import json
+import uuid as uuid_package
 from dataclasses import asdict
 from typing import Generic, Literal, Optional, Type, Union, overload
 
@@ -25,6 +26,7 @@
 from weaviate.collections.query import _GenerateCollection, _QueryCollection
 from weaviate.collections.tenants import _Tenants
 from weaviate.connect import ConnectionV4
+from weaviate.types import UUID
 from weaviate.validator import _validate_input, _ValidateArgument
 
 
@@ -192,6 +194,7 @@ def iterator(
         *,
         return_properties: Optional[PROPERTIES] = None,
         return_references: Literal[None] = None,
+        after: Optional[UUID] = None,
     ) -> _ObjectIterator[Properties, References]:
         ...
 
@@ -203,6 +206,7 @@ def iterator(
         *,
         return_properties: Optional[PROPERTIES] = None,
         return_references: REFERENCES,
+        after: Optional[UUID] = None,
     ) -> _ObjectIterator[Properties, CrossReferences]:
         ...
 
@@ -214,6 +218,7 @@ def iterator(
         *,
         return_properties: Optional[PROPERTIES] = None,
         return_references: Type[TReferences],
+        after: Optional[UUID] = None,
     ) -> _ObjectIterator[Properties, TReferences]:
         ...
 
@@ -225,6 +230,7 @@ def iterator(
         *,
         return_properties: Type[TProperties],
         return_references: Literal[None] = None,
+        after: Optional[UUID] = None,
     ) -> _ObjectIterator[TProperties, References]:
         ...
 
@@ -236,6 +242,7 @@ def iterator(
         *,
         return_properties: Type[TProperties],
         return_references: REFERENCES,
+        after: Optional[UUID] = None,
     ) -> _ObjectIterator[TProperties, CrossReferences]:
         ...
 
@@ -247,6 +254,7 @@ def iterator(
         *,
         return_properties: Type[TProperties],
         return_references: Type[TReferences],
+        after: Optional[UUID] = None,
     ) -> _ObjectIterator[TProperties, TReferences]:
         ...
 
@@ -256,6 +264,7 @@ def iterator(
         return_metadata: Optional[METADATA] = None,
         return_properties: Optional[ReturnProperties[TProperties]] = None,
         return_references: Optional[ReturnReferences[TReferences]] = None,
+        after: Optional[UUID] = None,
     ) -> Union[
         _ObjectIterator[Properties, References],
         _ObjectIterator[Properties, CrossReferences],
@@ -283,6 +292,8 @@ def iterator(
                 The properties to return with each object.
             `return_references`
                 The references to return with each object.
+            `after`
+                The cursor to use to mark the initial starting point of the iterator in the collection.
 
         Raises:
             `weaviate.exceptions.WeaviateGRPCQueryError`:
@@ -296,5 +307,8 @@ def iterator(
                 return_metadata=return_metadata,
                 return_properties=return_properties,
                 return_references=return_references,
-            ).objects
+            ).objects,
+            after
+            if after is None or isinstance(after, uuid_package.UUID)
+            else uuid_package.UUID(after),
         )

weaviate/collections/iterator.py

@@ -10,16 +10,19 @@
 
 class _ObjectIterator(Generic[P, R], Iterable[Object[P, R]]):
     def __init__(
-        self, fetch_objects_query: Callable[[int, Optional[UUID]], List[Object[P, R]]]
+        self,
+        fetch_objects_query: Callable[[int, Optional[UUID]], List[Object[P, R]]],
+        init_after: Optional[UUID],
     ) -> None:
         self.__query = fetch_objects_query
+        self.__init_after = init_after
 
         self.__iter_object_cache: List[Object[P, R]] = []
-        self.__iter_object_last_uuid: Optional[UUID] = None
+        self.__iter_object_last_uuid: Optional[UUID] = init_after
 
     def __iter__(self) -> Iterator[Object[P, R]]:
         self.__iter_object_cache = []
-        self.__iter_object_last_uuid = None
+        self.__iter_object_last_uuid = self.__init_after
         return self
 
     def __next__(self) -> Object[P, R]: