Skip to content

feat: add direct async iteration on Dataset and KeyValueStore - #2267

Open
Pijukatel wants to merge 12 commits into
masterfrom
claude/sharp-hopper-yidrmf
Open

Pijukatel wants to merge 12 commits into
masterfrom
claude/sharp-hopper-yidrmf

Conversation

@Pijukatel

@Pijukatel Pijukatel commented Oct 2, 2026 •

Copy link
Copy Markdown
Collaborator

Description

  • Add Dataset.__aiter__, and KeyValueStore.iterate_values, iterate_entries and __aiter__, mirroring the JS async iterators.
  • Add overridable KeyValueStoreClient.iterate_entries with a key-by-key default; the SQL and Redis clients read in batches that are bounded by lenth and memory
  • Explicitly state that the Redis client created with decode_responses=True is not supported for KVS and fix related typing

Issues

Testing

  • Unit tests for the new frontend methods on all storage clients, plus client-level tests for the SQL and Redis overrides.

Checklist

  • CI passed

Pijukatel and others added 3 commits October 1, 2026 08:19
Add `__aiter__` to `Dataset` (delegating to `iterate_items`) and to
`KeyValueStore` (delegating to the new `iterate_entries`), plus
`KeyValueStore.iterate_values` and `KeyValueStore.iterate_entries`.

Values are fetched one at a time as the iteration advances, so only a
single value is held in memory regardless of the record sizes. On the
Apify platform this means one request per record on top of the paginated
key listing, which is the only way the API offers to read values.

The memory key-value store client now skips keys deleted while a key
iteration is suspended instead of raising `KeyError`, matching the other
backends.

Closes #1745

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Wg6jZgyQp7XkvVvxuueJyV
Add a non-abstract `KeyValueStoreClient.iterate_entries` to the storage
client base class, yielding `KeyValueStoreRecord`s. The default
implementation lists the keys with `iterate_keys` and reads each value
with `get_value` as the iteration advances, so existing and third-party
clients need no changes. Backends that can read keys together with their
values more efficiently can override it.

`KeyValueStore.iterate_entries` and `iterate_values` now delegate to the
client method instead of combining `iterate_keys` and `get_value`
themselves.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Wg6jZgyQp7XkvVvxuueJyV
Override `iterate_entries` where the backend can do better than one
value read per key:

- The SQL client selects keys, metadata and values in a single streamed
  query, so a store with N records costs one query instead of N+1 and
  still holds only one row in memory at a time.
- The Redis client fetches values with one HMGET per batch instead of
  two round trips per record. Batches are bounded by key count and by
  the record sizes known from the metadata hash, so large values do not
  pile up in memory. A record larger than the byte limit is fetched
  alone.

Both clients share the value decoding with `get_value` through a new
`_build_record` helper.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Wg6jZgyQp7XkvVvxuueJyV
@codecov

codecov Bot commented Oct 2, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 98.21429% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 93.91%. Comparing base (df1178c) to head (c0c334e).
⚠️ Report is 6 commits behind head on master.

Files with missing lines Patch % Lines
...e/storage_clients/_base/_key_value_store_client.py 83.33% 1 Missing ⚠️
...storage_clients/_memory/_key_value_store_client.py 66.66% 1 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##           master    #2267      +/-   ##
==========================================
+ Coverage   93.87%   93.91%   +0.03%     
==========================================
  Files         182      183       +1     
  Lines       13130    13247     +117     
==========================================
+ Hits        12326    12441     +115     
- Misses        804      806       +2     
Flag Coverage Δ
unit 93.91% <98.21%> (+0.03%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

Pijukatel and others added 7 commits October 2, 2026 07:09
Whether a record deleted mid-iteration is yielded depends on how the
storage client reads values, which is not part of the contract.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Wg6jZgyQp7XkvVvxuueJyV
…e client

`set_value` stores an empty byte string for `None`, so the batched read
can fetch those records like any other and let `_build_record` map the
`application/x-none` content type to `None`. This drops the key filtering
and the per-key lookup dict from the batch fetch.

Also reword the base `iterate_entries` docstring so it does not present
the handling of records deleted mid-iteration as a contract.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Wg6jZgyQp7XkvVvxuueJyV
…ype checker

redis-py types every reply as `bytes | str | None` because a client
created with `decode_responses=True` returns strings. The key-value
store client stores binary values and needs the raw bytes back, so it
narrowed the type with `ty: ignore` and would fail with an
`AttributeError` on such a client.

Add `expect_bytes`, which narrows a reply to bytes and raises a
`TypeError` naming the unsupported option, use it at both value-reading
sites, and document the requirement on `RedisStorageClient`.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Wg6jZgyQp7XkvVvxuueJyV
`expect_bytes` now handles a single reply, which `isinstance` narrows
on its own, and the HMGET result is narrowed with a comprehension.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Wg6jZgyQp7XkvVvxuueJyV
…ient

Datasets use RedisJSON and request queues parse JSON strings, so both
work with a decoding Redis client. Only key-value store values are
binary.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Wg6jZgyQp7XkvVvxuueJyV
`stream_results` buffers up to 1000 rows client side, so iterating a
store of large values held hundreds of megabytes before the first yield,
and the single open transaction lasted for the whole iteration.

The SQL client now reads record metadata in keyset-paginated pages,
splits each page into batches bounded by the record sizes, and fetches
every batch's values with one query. Each query runs in its own short
session. The batching rule is shared with the Redis client through a new
`batch_records_by_size` helper.

Measured with 300 records of 1 MiB: peak memory drops from 300 MiB to
9 MiB. With 2000 small records the iteration issues 41 selects instead
of 2001 and is about 40 times faster than the default implementation.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Wg6jZgyQp7XkvVvxuueJyV

This comment was marked as resolved.

Pijukatel and others added 2 commits October 2, 2026 11:26
The default iteration retries every read through `get_value`, but the
batched HMGET in the Redis override was not retried. It now returns a
list under the same `retry_on_error` as the other Redis reads, matching
the SQL client.

Add a direct unit test for `batch_records_by_size`, covering the count
bound, the size bound, an oversized record alone in its batch and
records of unknown size.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Wg6jZgyQp7XkvVvxuueJyV
The test directories have no `__init__.py`, so pytest cannot import two
modules named `test_utils.py` and aborted the whole unit test collection.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Wg6jZgyQp7XkvVvxuueJyV
@Pijukatel
Pijukatel marked this pull request as ready for review October 2, 2026 11:38
@Pijukatel
Pijukatel requested review from Mantisus and vdusek October 2, 2026 11:38
@apify-service-account apify-service-account added the tested Temporary label used only programatically for some analytics. label Oct 2, 2026
@apify-service-account apify-service-account added the t-tooling Issues with this label are in the ownership of the tooling team. label Oct 2, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

t-tooling Issues with this label are in the ownership of the tooling team. tested Temporary label used only programatically for some analytics.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add direct async iteration on Dataset and KeyValueStore

3 participants