feat: add iter()/aiter() for lazy filter-based key iteration

redisvl/index/index.py

@@ -4,6 +4,7 @@
 import time
 import warnings
 import weakref
+from contextlib import aclosing
 from dataclasses import dataclass
 from math import ceil
 from typing import (
@@ -2003,6 +2004,65 @@ def paginate(self, query: BaseQuery, page_size: int = 30) -> Generator:
             # Increment the offset for the next batch of pagination
             offset += page_size
 
+    def iter_keys(
+            self,
+            filter_expression: str | FilterExpression | None = None,
+            batch_size: int = DEFAULT_BULK_BATCH_SIZE,
+        ) -> Generator[str, None, None]:
+            """Iterate lazily over document keys matching a filter expression.
+
+            Delegates to :meth:`_iter_keys_by_filter`, which pages with
+            ``FT.AGGREGATE ... WITHCURSOR`` rather than ``FT.SEARCH`` + ``LIMIT``, so
+            this is not subject to the ``MAXSEARCHRESULTS`` limit (1,000,000 on
+            Redis 8.2.7 by default). See that method's docstring for why keys are
+            de-duplicated and why memory is ``O(match count)`` rather than truly
+            streaming.
+
+            A batch can come back smaller than ``batch_size`` (an all-repeat cursor
+            page yields nothing), and a full drain is no proof every match was seen,
+            since the cursor is a position in an ascending walk of internal
+            document ids, not a snapshot. Callers that need completeness should
+            reconcile against a ``CountQuery`` rather than trusting a drained cursor.
+
+            The server-side cursor's idle timeout only resets when a page is read,
+            so a caller doing per-key work (embedding, network) can have its cursor
+            reaped after keys have already been yielded, which fails mid-stream
+            rather than up front. The cursor is released on normal exit, early
+            break, and error; the abandon-and-never-close case cannot be fixed from
+            in here, so callers who break out early should wrap the iterator in
+            ``contextlib.aclosing()``.
+
+            Args:
+                filter_expression (Union[str, FilterExpression, None]): Selects the
+                    documents to iterate over. Defaults to None (all documents).
+                    ``None``, ``""``, whitespace, and a default ``FilterExpression``
+                    are all treated as match-all, consistent with ``drop_by_filter``
+                    and ``update_by_filter``.
+                batch_size (int): Number of keys fetched per cursor page.
+                    Defaults to 500.
+
+            Raises:
+                TypeError: If ``batch_size`` is not an int.
+                ValueError: If ``batch_size`` is less than 1.
+
+            Yields:
+                str: Document key matching the filter.
+
+            Note:
+                Do not concatenate untrusted input into ``filter_expression`` —
+                it is rendered into the ``FT.AGGREGATE`` request verbatim, the same
+                constraint ``drop_by_filter`` and ``update_by_filter`` carry.
+            """
+            if not isinstance(batch_size, int):
+                raise TypeError("batch_size must be an integer")
+            if batch_size < 1:
+                raise ValueError("batch_size must be greater than 0")
+            filter_expr = (
+                "*" if _is_match_all_filter(filter_expression) else filter_expression
+            )
+            for batch in self._iter_keys_by_filter(filter_expr, batch_size):
+                yield from batch
+
     def listall(self) -> list[str]:
         """List all search indices in Redis database.
 
@@ -3246,6 +3306,67 @@ async def paginate(self, query: BaseQuery, page_size: int = 30) -> AsyncGenerato
             yield results
             first += page_size
 
+    async def aiter_keys(
+            self,
+            filter_expression: str | FilterExpression | None = None,
+            batch_size: int = DEFAULT_BULK_BATCH_SIZE,
+        ) -> AsyncGenerator[str, None]:
+            """Iterate lazily over document keys matching a filter expression asynchronously.
+
+            Delegates to :meth:`_iter_keys_by_filter`, which pages with
+            ``FT.AGGREGATE ... WITHCURSOR`` rather than ``FT.SEARCH`` + ``LIMIT``, so
+            this is not subject to the ``MAXSEARCHRESULTS`` limit (1,000,000 on
+            Redis 8.2.7 by default). See that method's docstring for why keys are
+            de-duplicated and why memory is ``O(match count)`` rather than truly
+            streaming.
+
+            A batch can come back smaller than ``batch_size`` (an all-repeat cursor
+            page yields nothing), and a full drain is no proof every match was seen,
+            since the cursor is a position in an ascending walk of internal
+            document ids, not a snapshot. Callers that need completeness should
+            reconcile against a ``CountQuery`` rather than trusting a drained cursor.
+
+            The server-side cursor's idle timeout only resets when a page is read,
+            so a caller doing per-key work (embedding, network) can have its cursor
+            reaped after keys have already been yielded, which fails mid-stream
+            rather than up front. The cursor is released on normal exit, early
+            break, and error; the abandon-and-never-close case cannot be fixed from
+            in here, so callers who break out early should wrap the iterator in
+            ``contextlib.aclosing()``.
+
+            Args:
+                filter_expression (Union[str, FilterExpression, None]): Selects the
+                    documents to iterate over. Defaults to None (all documents).
+                    ``None``, ``""``, whitespace, and a default ``FilterExpression``
+                    are all treated as match-all, consistent with ``drop_by_filter``
+                    and ``update_by_filter``.
+                batch_size (int): Number of keys fetched per cursor page.
+                    Defaults to 500.
+
+            Raises:
+                TypeError: If ``batch_size`` is not an int.
+                ValueError: If ``batch_size`` is less than 1.
+
+            Yields:
+                str: Document key matching the filter.
+
+            Note:
+                Do not concatenate untrusted input into ``filter_expression`` —
+                it is rendered into the ``FT.AGGREGATE`` request verbatim, the same
+                constraint ``drop_by_filter`` and ``update_by_filter`` carry.
+            """
+            if not isinstance(batch_size, int):
+                raise TypeError("batch_size must be an integer")
+            if batch_size < 1:
+                raise ValueError("batch_size must be greater than 0")
+            filter_expr = (
+                "*" if _is_match_all_filter(filter_expression) else filter_expression
+            )
+            async for batch in self._iter_keys_by_filter(filter_expr, batch_size):
+                async with aclosing(batch):
+                    for key in batch:
+                        yield key
+
     async def listall(self) -> list[str]:
         """List all search indices in Redis database.
 

tests/integration/test_index_iteration.py

@@ -0,0 +1,109 @@
+import pytest
+
+from redisvl.index import AsyncSearchIndex, SearchIndex
+from redisvl.query.filter import Tag
+from contextlib import aclosing
+
+DOCS = [
+    {"id": "1", "category": "A"},
+    {"id": "2", "category": "B"},
+    {"id": "3", "category": "A"},
+    {"id": "4", "category": "C"},
+]
+
+
+@pytest.fixture
+def sample_index(redis_url, redis_test_name):
+    index_name = redis_test_name("iter_index")
+    prefix = redis_test_name("iter_doc")
+    index = SearchIndex.from_dict(
+        {
+            "index": {"name": index_name, "prefix": prefix, "storage_type": "hash"},
+            "fields": [{"name": "category", "type": "tag"}],
+        },
+        redis_url=redis_url,
+    )
+    index.create(overwrite=True)
+    index.load(DOCS, id_field="id")
+    yield index
+    index.delete(drop=True)
+
+
+@pytest.fixture
+async def async_sample_index(redis_url, redis_test_name):
+    index_name = redis_test_name("async_iter_index")
+    prefix = redis_test_name("async_iter_doc")
+    index = AsyncSearchIndex.from_dict(
+        {
+            "index": {"name": index_name, "prefix": prefix, "storage_type": "hash"},
+            "fields": [{"name": "category", "type": "tag"}],
+        },
+        redis_url=redis_url,
+    )
+    await index.create(overwrite=True)
+    await index.load(DOCS, id_field="id")
+    yield index
+    await index.delete(drop=True)
+
+
+def test_iter_keys_yields_every_key(sample_index):
+    keys = list(sample_index.iter_keys())
+
+    assert len(keys) == 4
+    assert set(keys) == {f"{sample_index.prefix}:{i}" for i in range(1, 5)}
+
+
+def test_iter_keys_respects_filter_expression(sample_index):
+    keys = list(sample_index.iter_keys(filter_expression=Tag("category") == "A"))
+
+    assert set(keys) == {f"{sample_index.prefix}:1", f"{sample_index.prefix}:3"}
+
+
+def test_iter_keys_is_lazy(sample_index):
+    iterator = sample_index.iter_keys()
+
+    assert next(iterator) is not None
+
+
+def test_iter_keys_pages_when_batch_size_is_smaller_than_the_index(sample_index):
+    keys = list(sample_index.iter_keys(batch_size=2))
+
+    assert sorted(keys) == sorted(f"{sample_index.prefix}:{i}" for i in range(1, 5))
+
+
+@pytest.mark.asyncio
+async def test_aiter_keys_yields_every_key(async_sample_index):
+    keys = [key async for key in async_sample_index.aiter_keys()]
+
+    assert len(keys) == 4
+    assert set(keys) == {f"{async_sample_index.prefix}:{i}" for i in range(1, 5)}
+
+
+@pytest.mark.asyncio
+async def test_aiter_keys_respects_filter_expression(async_sample_index):
+    keys = [
+        key
+        async for key in async_sample_index.aiter_keys(
+            filter_expression=Tag("category") == "A"
+        )
+    ]
+
+    assert set(keys) == {
+        f"{async_sample_index.prefix}:1",
+        f"{async_sample_index.prefix}:3",
+    }
+
+
+def test_iter_keys_raises_with_non_int_batch_size(sample_index):
+    with pytest.raises(TypeError):
+        list(sample_index.iter_keys(batch_size="5"))
+
+
+def test_iter_keys_raises_with_zero_batch_size(sample_index):
+    with pytest.raises(ValueError):
+        list(sample_index.iter_keys(batch_size=0))
+
+
+def test_iter_keys_raises_with_negative_batch_size(sample_index):
+    with pytest.raises(ValueError):
+        list(sample_index.iter_keys(batch_size=-1))