Add shard cursor support to filtered iterators

Feed shard_cursors from each SearchReply back into the next SearchRequest,
enabling per-shard pagination state for filtered iterators. Also updates
protobuf definitions to include the shard_cursors field.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: Michelangelo Partipilo

weaviate/collections/classes/internal.py

@@ -269,9 +269,15 @@ class GroupByReturn(Generic[P, R]):
 
 @dataclass
 class QueryReturn(Generic[P, R]):
-    """The return type of a query within the `.query` namespace of a collection."""
+    """The return type of a query within the `.query` namespace of a collection.
+
+    Attributes:
+        objects: List of objects returned from the query.
+        shard_cursors: Optional dictionary mapping shard IDs to cursor positions for paginated queries.
+    """
 
     objects: List[Object[P, R]]
+    shard_cursors: Optional[Dict[str, str]] = None
 
 
 _GQLEntryReturnType: TypeAlias = Dict[str, List[Dict[str, Any]]]

weaviate/collections/grpc/query.py

@@ -121,6 +121,7 @@ def get(
         return_references: Optional[REFERENCES] = None,
         generative: Optional[_Generative] = None,
         rerank: Optional[Rerank] = None,
+        shard_cursors: Optional[Dict[str, str]] = None,
     ) -> search_get_pb2.SearchRequest:
         if self._validate_arguments:
             _validate_input(_ValidateArgument([_Sorting, None], "sort", sort))
@@ -144,6 +145,7 @@ def get(
             generative=generative,
             rerank=rerank,
             sort_by=sort_by,
+            shard_cursors=shard_cursors,
         )
 
     def hybrid(
@@ -416,6 +418,7 @@ def __create_request(
         near_imu: Optional[base_search_pb2.NearIMUSearch] = None,
         near_thermal: Optional[base_search_pb2.NearThermalSearch] = None,
         near_video: Optional[base_search_pb2.NearVideoSearch] = None,
+        shard_cursors: Optional[Dict[str, str]] = None,
     ) -> search_get_pb2.SearchRequest:
         if self._validate_arguments:
             _validate_input(
@@ -507,6 +510,7 @@ def __create_request(
             near_imu=near_imu,
             near_thermal=near_thermal,
             near_video=near_video,
+            shard_cursors=shard_cursors if shard_cursors else {},
         )
 
     def _metadata_to_grpc(self, metadata: _MetadataQuery) -> search_get_pb2.MetadataRequest:

weaviate/collections/iterator.py

@@ -3,6 +3,7 @@
     Any,
     AsyncIterable,
     AsyncIterator,
+    Dict,
     Generic,
     Iterable,
     Iterator,
@@ -62,16 +63,24 @@ def __init__(
         self.__iter_object_cache: List[Object[TProperties, TReferences]] = []
         self.__iter_object_last_uuid: Optional[UUIDorStr] = _parse_after(self.__inputs.after)
         self.__iter_cache_size = cache_size or ITERATOR_CACHE_SIZE
+        self.__iter_shard_cursors: Optional[Dict[str, str]] = None
 
     def __iter__(
         self,
     ) -> Iterator[Object[TProperties, TReferences]]:
         self.__iter_object_cache = []
         self.__iter_object_last_uuid = _parse_after(self.__inputs.after)
+        self.__iter_shard_cursors = None
         return self
 
     def __next__(self) -> Object[TProperties, TReferences]:
         if len(self.__iter_object_cache) == 0:
+            # Shard cursor pagination:
+            # - First call uses shard_cursors=None; subsequent calls include the cursors
+            #   returned by the server in the previous SearchReply.
+            # - The server returns updated shard_cursors in each response, which are fed
+            #   back into the next request so each shard resumes from the right position.
+            # - Iteration ends when the server returns an empty result set.
             res = self.__query.fetch_objects(
                 limit=self.__iter_cache_size,
                 after=self.__iter_object_last_uuid,
@@ -80,8 +89,10 @@ def __next__(self) -> Object[TProperties, TReferences]:
                 return_properties=self.__inputs.return_properties,
                 return_references=self.__inputs.return_references,
                 filters=self.__inputs.filters,
+                shard_cursors=self.__iter_shard_cursors,
             )
             self.__iter_object_cache = res.objects  # type: ignore
+            self.__iter_shard_cursors = res.shard_cursors
             if len(self.__iter_object_cache) == 0:
                 raise StopIteration
 
@@ -109,18 +120,26 @@ def __init__(
         self.__iter_object_cache: List[Object[TProperties, TReferences]] = []
         self.__iter_object_last_uuid: UUIDorStr = _parse_after(self.__inputs.after)
         self.__iter_cache_size = cache_size or ITERATOR_CACHE_SIZE
+        self.__iter_shard_cursors: Optional[Dict[str, str]] = None
 
     def __aiter__(
         self,
     ) -> AsyncIterator[Object[TProperties, TReferences]]:
         self.__iter_object_cache = []
         self.__iter_object_last_uuid = _parse_after(self.__inputs.after)
+        self.__iter_shard_cursors = None
         return self
 
     async def __anext__(
         self,
     ) -> Object[TProperties, TReferences]:
         if len(self.__iter_object_cache) == 0:
+            # Shard cursor pagination:
+            # - First call uses shard_cursors=None; subsequent calls include the cursors
+            #   returned by the server in the previous SearchReply.
+            # - The server returns updated shard_cursors in each response, which are fed
+            #   back into the next request so each shard resumes from the right position.
+            # - Iteration ends when the server returns an empty result set.
             res = await self.__query.fetch_objects(
                 limit=self.__iter_cache_size,
                 after=self.__iter_object_last_uuid,
@@ -129,8 +148,10 @@ async def __anext__(
                 return_properties=self.__inputs.return_properties,
                 return_references=self.__inputs.return_references,
                 filters=self.__inputs.filters,
+                shard_cursors=self.__iter_shard_cursors,
             )
             self.__iter_object_cache = res.objects  # type: ignore
+            self.__iter_shard_cursors = res.shard_cursors
             if len(self.__iter_object_cache) == 0:
                 raise StopAsyncIteration
 

weaviate/collections/queries/base_executor.py

@@ -457,11 +457,14 @@ def _result_to_query_return(
         res: search_get_pb2.SearchReply,
         options: _QueryOptions,
     ) -> QueryReturn[WeaviateProperties, CrossReferences]:
+        # Extract shard_cursors from protobuf map and convert to Python dict
+        shard_cursors = dict(res.shard_cursors) if res.shard_cursors else None
         return QueryReturn(
             objects=[
                 self.__result_to_query_object(obj.properties, obj.metadata, options)
                 for obj in res.results
-            ]
+            ],
+            shard_cursors=shard_cursors,
         )
 
     def _result_to_generative_query_return(

weaviate/collections/queries/fetch_objects/query/executor.py

@@ -1,4 +1,4 @@
-from typing import Any, Generic, Literal, Optional, Type, Union, cast, overload
+from typing import Any, Dict, Generic, Literal, Optional, Type, Union, cast, overload
 
 from weaviate.collections.classes.filters import _Filters
 from weaviate.collections.classes.grpc import METADATA, PROPERTIES, REFERENCES, Sorting
@@ -39,6 +39,7 @@ def fetch_objects(
         return_metadata: Optional[METADATA] = None,
         return_properties: Union[PROPERTIES, bool, None] = None,
         return_references: Literal[None] = None,
+        shard_cursors: Optional[Dict[str, str]] = None,
     ) -> executor.Result[QueryReturn[Properties, References]]: ...
 
     @overload
@@ -54,6 +55,7 @@ def fetch_objects(
         return_metadata: Optional[METADATA] = None,
         return_properties: Union[PROPERTIES, bool, None] = None,
         return_references: REFERENCES,
+        shard_cursors: Optional[Dict[str, str]] = None,
     ) -> executor.Result[QueryReturn[Properties, CrossReferences]]: ...
 
     @overload
@@ -69,6 +71,7 @@ def fetch_objects(
         return_metadata: Optional[METADATA] = None,
         return_properties: Union[PROPERTIES, bool, None] = None,
         return_references: Type[TReferences],
+        shard_cursors: Optional[Dict[str, str]] = None,
     ) -> executor.Result[QueryReturn[Properties, TReferences]]: ...
 
     @overload
@@ -84,6 +87,7 @@ def fetch_objects(
         return_metadata: Optional[METADATA] = None,
         return_properties: Type[TProperties],
         return_references: Literal[None] = None,
+        shard_cursors: Optional[Dict[str, str]] = None,
     ) -> executor.Result[QueryReturn[TProperties, References]]: ...
 
     @overload
@@ -99,6 +103,7 @@ def fetch_objects(
         return_metadata: Optional[METADATA] = None,
         return_properties: Type[TProperties],
         return_references: REFERENCES,
+        shard_cursors: Optional[Dict[str, str]] = None,
     ) -> executor.Result[QueryReturn[TProperties, CrossReferences]]: ...
 
     @overload
@@ -114,6 +119,7 @@ def fetch_objects(
         return_metadata: Optional[METADATA] = None,
         return_properties: Type[TProperties],
         return_references: Type[TReferences],
+        shard_cursors: Optional[Dict[str, str]] = None,
     ) -> executor.Result[QueryReturn[TProperties, TReferences]]: ...
 
     @overload
@@ -129,6 +135,7 @@ def fetch_objects(
         return_metadata: Optional[METADATA] = None,
         return_properties: Optional[ReturnProperties[TProperties]] = None,
         return_references: Optional[ReturnReferences[TReferences]] = None,
+        shard_cursors: Optional[Dict[str, str]] = None,
     ) -> executor.Result[QueryReturnType[Properties, References, TProperties, TReferences]]: ...
 
     def fetch_objects(
@@ -143,6 +150,7 @@ def fetch_objects(
         return_metadata: Optional[METADATA] = None,
         return_properties: Optional[ReturnProperties[TProperties]] = None,
         return_references: Optional[ReturnReferences[TReferences]] = None,
+        shard_cursors: Optional[Dict[str, str]] = None,
     ) -> executor.Result[QueryReturnType[Properties, References, TProperties, TReferences]]:
         """Retrieve the objects in this collection without any search.
 
@@ -156,6 +164,7 @@ def fetch_objects(
             return_metadata: The metadata to return for each object, defaults to `None`.
             return_properties: The properties to return for each object.
             return_references: The references to return for each object.
+            shard_cursors: The shard cursors from the previous response for paginated queries.
 
         NOTE:
             - If `return_properties` is not provided then all properties are returned except for blob properties.
@@ -195,6 +204,7 @@ def resp(
             return_metadata=self._parse_return_metadata(return_metadata, include_vector),
             return_properties=self._parse_return_properties(return_properties),
             return_references=self._parse_return_references(cast(Any, return_references)),
+            shard_cursors=shard_cursors,
         )
         return executor.execute(
             response_callback=resp,