-
Notifications
You must be signed in to change notification settings - Fork 99
feat: add iter()/aiter() for lazy filter-based key iteration #682
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
6adf547
441a51d
69dd491
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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, | ||
| ) -> 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
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. These caveats point at something the published docs don't contain. 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 Two more worth adding while you're in here. One softening, too. The |
||
|
|
||
| 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
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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.
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. | ||
|
|
||
|
|
@@ -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
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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 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 keyI measured the cursor released the instant The sync path needs nothing. I watched |
||
|
|
||
| async def listall(self) -> list[str]: | ||
| """List all search indices in Redis database. | ||
|
|
||
|
|
||
| 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 | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This assertion can't fail. Keys are strings, so I checked by swapping in an What would work is counting |
||
|
|
||
|
|
||
| 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) | ||
| ) | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
batch_sizewants the same guardspaginatehas fifteen lines up:TypeErrorfor a non-int,ValueErrorfor anything below 1.Measured on Redis 8.2.7.
batch_size=0silently returns every key, because redis-py drops a falsyCOUNTand the server picks its own page size, which makes "Defaults to 500" in the docstring untrue.batch_size="5"sails straight through.batch_size=-1surfaces the raw server textBad arguments for COUNT: Value is outside acceptable boundsinside aRedisSearchError.Matching
paginateis enough. Worth knowing that neither will raise at call time, since validation inside a generator function doesn't run until the firstnext(). If you'd rather it fail eagerly, the checks have to live in a non-generator wrapper that returns the generator. Your call.