Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 55 additions & 0 deletions redisvl/index/index.py
Original file line number Diff line number Diff line change
Expand Up @@ -2003,6 +2003,33 @@ def paginate(self, query: BaseQuery, page_size: int = 30) -> Generator:
# Increment the offset for the next batch of pagination
offset += page_size

def iter(
self,
filter_expression: str | FilterExpression | None = None,
batch_size: int = DEFAULT_BULK_BATCH_SIZE,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

batch_size wants the same guards paginate has fifteen lines up: TypeError for a non-int, ValueError for anything below 1.

Measured on Redis 8.2.7. batch_size=0 silently returns every key, because redis-py drops a falsy COUNT and the server picks its own page size, which makes "Defaults to 500" in the docstring untrue. batch_size="5" sails straight through. batch_size=-1 surfaces the raw server text Bad arguments for COUNT: Value is outside acceptable bounds inside a RedisSearchError.

Matching paginate is enough. Worth knowing that neither will raise at call time, since validation inside a generator function doesn't run until the first next(). If you'd rather it fail eagerly, the checks have to live in a non-generator wrapper that returns the generator. Your call.

) -> 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. See that method's
docstring for why keys are de-duplicated and why memory is
``O(match count)`` rather than truly streaming.
Comment on lines +2013 to +2017

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These caveats point at something the published docs don't contain. docs/api/searchindex.rst uses autoclass ... :members: with no :private-members:, so iter appears on docs.redisvl.com but _iter_keys_by_filter doesn't. The :meth: reference is dangling and the memory and completeness caveats are invisible to exactly the reader who needs them. The async method is a longer chain still, because the async helper's body is "See the sync counterpart".

Please inline the load-bearing sentences: memory is proportional to the match count, so very large scans should partition the filter; a batch can come back smaller than batch_size; and a drained cursor isn't proof every match was seen. That last one matters most, because a method named for iteration reads as exhaustive. Keep the seen set as it is — dropping it would hand callers the same document twice, and since RediSearch reindexes an updated document under a new, higher id, writing during iteration then provokes further repeats.

Two more worth adding while you're in here. drop_by_filter and update_by_filter carry the only "never string-concatenate untrusted input into a filter" warning in the package, and these accept the same parameter type without it. And the cursor's idle timeout only resets when a page is read, so a caller doing per-key embedding or network work can have its cursor reaped after keys have already been yielded, which fails mid-stream rather than up front.

One softening, too. The MAXSEARCHRESULTS framing is accurate but reads scarier than it is. I measured the default at 1,000,000 on Redis 8.2.7, so it only bites above a million matches. Naming the figure would help a reader judge whether it applies to them.


Args:
filter_expression (Union[str, FilterExpression, None]): Selects the
documents to iterate over. Defaults to None (all documents).
batch_size (int): Number of keys fetched per cursor page. Defaults to 500.

Yields:
str: Document key matching the filter.
"""
filter_expr = (
FilterExpression("*") if filter_expression is None else filter_expression
)
Comment on lines +2027 to +2029

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This wrapper is a no-op, and bypassing the existing helper introduces a divergence. str(FilterExpression("*")) is just "*", and _iter_keys_by_filter already accepts a plain str.

_is_match_all_filter at index.py:236 is what drop_by_filter and update_by_filter use for this normalisation, and it treats "" and whitespace as match-all as well as None. Measured: iter("") and iter(" ") return zero keys here, while drop_by_filter("") matches the whole index. So a filter that arrives empty from config means nothing to one method and everything to its neighbour. iter(FilterExpression()) also raises the bare ValueError("Improperly initialized FilterExpression") that the helper exists to absorb.

Routing through it fixes all three cases and deletes these three lines.

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.

Expand Down Expand Up @@ -3246,6 +3273,34 @@ async def paginate(self, query: BaseQuery, page_size: int = 30) -> AsyncGenerato
yield results
first += page_size

async def aiter(
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. See that method's
docstring for why keys are de-duplicated and why memory is
``O(match count)`` rather than truly streaming.

Args:
filter_expression (Union[str, FilterExpression, None]): Selects the
documents to iterate over. Defaults to None (all documents).
batch_size (int): Number of keys fetched per cursor page. Defaults to 500.

Yields:
str: Document key matching the filter.
"""
filter_expr = (
FilterExpression("*") if filter_expression is None else filter_expression
)
async for batch in self._iter_keys_by_filter(filter_expr, batch_size):
for key in batch:
yield key
Comment on lines +3300 to +3302

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The server-side cursor can outlive this generator. Measured on Redis 8.2.7: break out of async for key in index.aiter() early, then await index.disconnect(), and FT.INFO ... cursor_stats shows index_total climbing 1, 2, 3 across three asyncio.run lifecycles. The cursor is never released. The event loop's async-generator finalisation tries to send FT.CURSOR DEL after the client has gone, and the except RedisError: pass in the helper swallows the failure. Each leaked cursor is held for MAXIDLE 300 seconds, against a per-shard capacity of 128.

Two parts. Wrapping the inner generator makes an explicit close deterministic:

async with contextlib.aclosing(
    self._iter_keys_by_filter(filter_expr, batch_size)
) as batches:
    async for batch in batches:
        for key in batch:
            yield key

I measured the cursor released the instant aclose() returns with that in place, versus still open without it. The abandon-and-never-close case can't be fixed from in here, so the docstring should tell callers who break out early to wrap the iterator in contextlib.aclosing().

The sync path needs nothing. I watched index_total drop back to 0 on close(), on del plus a collection, and on break.


async def listall(self) -> list[str]:
"""List all search indices in Redis database.

Expand Down
112 changes: 112 additions & 0 deletions tests/integration/test_index_iteration.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,112 @@
import pytest

from redisvl.index import AsyncSearchIndex, SearchIndex
from redisvl.query.filter import Tag

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)
# id_field makes the key deterministic: <prefix>:<id>
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_yields_every_key(sample_index):
"""iter() with no filter must yield every key in the index, once each."""
keys = list(sample_index.iter())

assert len(keys) == 4
assert set(keys) == {f"{sample_index.prefix}:{i}" for i in range(1, 5)}


def test_iter_respects_filter_expression(sample_index):
"""A filter expression must narrow the yielded keys."""
keys = list(sample_index.iter(filter_expression=Tag("category") == "A"))

assert set(keys) == {f"{sample_index.prefix}:1", f"{sample_index.prefix}:3"}


def test_iter_is_lazy(sample_index):
"""Iteration must stream: the first key arrives without draining the index."""
iterator = sample_index.iter()

assert next(iterator) is not None

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This assertion can't fail. Keys are strings, so is not None holds for anything the generator yields, and an empty index would surface as an uncaught StopIteration rather than a failed assertion.

I checked by swapping in an iter() that drains the whole index into a list before yielding anything. The test still passed, so it can't tell a streaming implementation from an eager one, which is the one property it's named for.

What would work is counting FT.AGGREGATE round trips before the first key arrives. tests/unit/test_bulk_cursor_dedup.py already has a _ReplayCursor harness that fakes client.ft(name), so this can be a hermetic unit test with no server and no fixture. While you're there, the async method has no laziness test at all.



def test_iter_pages_when_batch_size_is_smaller_than_the_index(sample_index):
"""A batch_size below the document count must still yield every key exactly once."""
keys = list(sample_index.iter(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_yields_every_key(async_sample_index):
"""aiter() must mirror iter() on the async client."""
keys = [key async for key in async_sample_index.aiter()]

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_respects_filter_expression(async_sample_index):
"""The async iterator must apply the filter the same way the sync one does."""
keys = [
key
async for key in async_sample_index.aiter(
filter_expression=Tag("category") == "A"
)
]

assert set(keys) == {
f"{async_sample_index.prefix}:1",
f"{async_sample_index.prefix}:3",
}


@pytest.mark.asyncio
async def test_aiter_pages_when_batch_size_is_smaller_than_the_index(
async_sample_index,
):
"""A batch_size below the document count must still yield every key exactly once."""
keys = [key async for key in async_sample_index.aiter(batch_size=2)]

assert sorted(keys) == sorted(
f"{async_sample_index.prefix}:{i}" for i in range(1, 5)
)
Loading