[OPIK-8253] [SDK] [DOCS] feat: parallel chunked dataset item reads - #8156
Conversation
…em reads get_items() downloads every item and rebuilds each one as a pydantic model before returning, which dominates the wait on large datasets. stream_items() reads the paginated items endpoint through the SDK's own httpx client, fetches pages across a thread pool with a bounded look-ahead, and yields raw dicts in dataset order. Measured against a local backend on 20k items: 3.3x faster at 8 threads. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
get_items() now flattens stream_items() instead of walking the cursor-chained typed stream, so it inherits the parallel paginated read. Its public signature and payload are unchanged: both read paths share one projection, which drops data keys shadowing DatasetItem fields and warns once per read, exactly as building a DatasetItem did. Verified byte-identical to the previous implementation (ordered comparison) across nested containers, every JSON scalar, empties, unicode, all seven shadowing keys, filter_string, nb_samples, and 2-50 pages at 1 and 8 threads. The one behavior change: nb_samples=0 or negative now raises ValueError rather than returning a single item. The backend-fetch factories seed Dataset.id from the response they already hold, so switching to the id-addressed endpoint costs no extra round trip. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
stream_items() defaulted to 1000-item chunks, which doubled the request count against the paginated endpoint relative to the typed stream and made a sequential read materially slower than before. Reuse constants.DATASET_STREAM_BATCH_SIZE so both read paths issue the same number of requests, and point the tests and the e2e helper at that constant so they cannot drift from the default. Measured on 100k items, num_threads=1: 4.01s at 2000 against 4.43s for the previous implementation. The default 4 threads gives 1.35s (2.8x). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
github-actions
Bot
added
documentation
⏱️ pre-commit per-hook timing
⏭️ 39 skipped (no matching files changed)
|
|
🌿 Preview your docs: https://opik-preview-01a06c74-11db-70dc-a2c4-e1da57063f66.docs.buildwithfern.com/docs/opik No broken links found Unverified links (timeout / rate-limited / server error — not failing the check)• https://aistudio.google.com/apikey (401) 📌 Results for commit 0e0bb10 |
get_items() read at a fixed 4 threads, so reaching the faster end of the range meant flattening stream_items() by hand. Accept num_threads and forward it. Appended to the signature, so existing positional callers are unaffected; validation is inherited from stream_items(). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
![baz-reviewer[bot]](https://avatars.githubusercontent.com/in/933528?s=60&v=4){kind=small}
- Validate the page response instead of defaulting `total` to zero. A missing, non-numeric or negative `total` capped the read at page 1 and returned part of the dataset as if it were all of it; it now raises. - Cap `chunk_size` at DATASET_STREAM_BATCH_SIZE. It was forwarded straight to the endpoint's `size`, so oversized values could materialize an unbounded page and values above int32 failed only after the request went out. Rejects rather than clamps: quietly returning smaller pages than asked for would look like the argument had no effect. - Defer the dataset-id lookup in `stream_items()` to first iteration by making the Dataset hook a generator. The laziness test only checked httpx and passed despite the eager by-name lookup; it now asserts the lookup too. - Document the `nb_samples` valid range and its ValueError. - Cover the `Dataset.from_public` id seeding, including the no-id fallback. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
petrotiurin
left a comment
petrotiurin
left a comment
There was a problem hiding this comment.
Two findings on the new reader: a read-consistency regression versus the cursor stream, and a blocking early-exit. Details inline.
Pages are addressed by offset and the backend sorts newest id first, so an item inserted mid-read landed at offset 0 and shifted every unfetched page: the read returned one item twice and skipped another. The cursor stream got snapshot semantics for free from its `id < last_retrieved_id` seek, so this was a regression against it, and it landed on callers that match items by exact equality. Every page of a read now carries one version hash, resolved once. Backends without dataset versioning fall back to a live read, which is documented on stream_items(). The thread pool also moves out of a `with` block: the yields happen inside its scope, so closing the generator ran shutdown(wait=True) and blocked the caller until every outstanding page returned -- at GC time, on an unrelated line. It now shuts down with cancel_futures=True and wait=False. Adds get_items(chunk_size=...) so the page size is tunable from the simple API as well. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
Already covered by a test in this PR. Python-SDK-only change (plus docs): get_items() now reads through the paginated items endpoint in parallel instead of the cursor stream, pins the read to a dataset version, and a public stream_items() is added. You shipped the tests for it. tests/unit/api_objects/dataset/test_stream_item_chunks.py covers page fan-out and single-thread order, nb_samples trimming, filter serialization, chunk_size/num_threads validation before any request, the malformed-'total' guard, and the version-pinned vs live-state read including the concurrent-insert shift; tests/e2e/test_dataset.py reads 10k items over 8 threads and asserts every item exactly once in the same order as num_threads=1, that get_items() equals stream_items() flattened, and that a version view reads its own snapshot. Those fail if the read drops, duplicates or reorders items, so we're not proposing anything on top. Our Playwright estate reads dataset items over REST/the TS client and never exercises the Python read path, so there was no e2e gap here to begin with. also touches Python SDK Advisory, from the QA test radar. Nothing here blocks this PR, and anything it proposes is a draft for review. |
| page_delay = 0.5 | ||
| endpoint = FakeItemsEndpoint(_rest_items(200), delay_seconds=page_delay) | ||
| dataset = _build_dataset(endpoint) | ||
|
|
||
| stream = dataset.stream_items(chunk_size=10, num_threads=4) | ||
| for _ in stream: | ||
| break # abandon after the first chunk, with pages still in flight | ||
|
|
||
| started = time.perf_counter() | ||
| stream.close() | ||
| elapsed = time.perf_counter() - started | ||
|
|
There was a problem hiding this comment.
Wall-clock assertion flakes in CI
The test uses a real 0.5-second delay and requires stream.close() to finish within that wall-clock window, so scheduler or CI contention can make correct shutdown(wait=False) behavior flaky — should we use event-based synchronization to hold in-flight workers, then release and join them during teardown?
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
`sdks/python/tests/unit/api_objects/dataset/test_stream_item_chunks.py` around lines
679-692, refactor
`test_stream_items__abandoned_early__does_not_block_on_in_flight_pages` to remove the
real 0.5-second delay and wall-clock assertion. Use synchronization events to hold
in-flight requests, assert that `stream.close()` returns before those events are
released, then release the workers and join or otherwise clean them up during teardown.
| version_info = self.get_version_info() | ||
| version_hash = version_info.version_hash if version_info else None |
There was a problem hiding this comment.
Nit: get_version_info() isn't cached, so every read pays an extra round trip to resolve the pin. evaluate() already calls it separately (evaluator.py:335), so an evaluation run resolves it twice, and the PR seeds Dataset.id elsewhere specifically to avoid a round trip like this one. Memoizing it on the instance would keep the pin free for the many small get_items() calls (cli/imports/utils.py, evaluation/rest_operations.py).
Not blocking, and not free either: a cache needs an invalidation story for a version committed after the dataset object was created, which is why resolving per read is a defensible default.
3 participants
Details
get_items()walked the cursor-chained/items/streamendpoint and rebuilt every row into a typedDatasetItembefore returning, so a large read was one serial request chain plus per-item pydantic validation. This addsDataset.stream_items()— a chunked reader that fetches pages of the id-addressedGET /v1/private/datasets/{id}/itemsendpoint concurrently through the SDK's own httpx client, bypassing the Fern-generated typed layer — and rewiresget_items()on top of it.stream_items(chunk_size, num_threads, filter_string, nb_samples)yields lists of item dicts in dataset order. The read is pinned to a single dataset version, so it is a snapshot: items inserted or deleted while it runs do not affect it. Page 1 is fetched first for itstotal; the rest fan out over a thread pool behind an ordered queue with a bounded look-ahead of2 * num_threadschunks, so a slow consumer cannot make the reader buffer the whole dataset.get_items()keeps its payload exactly — it is now a flatten ofstream_items(). Both share a single projection, so data keys that shadowDatasetItemfields are still dropped and still warned about once per read. It also gains optionalnum_threadsandchunk_size(appended to the signature, so existing positional callers are unaffected), since it would otherwise be stuck on the defaults.DatasetVersiongets the same method, pinned to its version hash.Dataset.from_public,rest_operations.get_datasets) now seedDataset.idfrom the response they already hold, so moving to the id-addressed endpoint costs no extra round trip. This also removes that lookup fromdelete(),get_version_info()anddataset_items_count.ORDER BY (workspace_id, dataset_id, id) DESC) is equivalent to the stream's (id DESC) within a workspace and dataset, verified empirically below.Behavior change.
get_items(nb_samples=0)(or a negative value) previously returned a single item: the stream parser'slen(result) == nb_samplescheck never fires, and the caller loop then stops after the first yield. It now raisesValueError. Everything else about the payload is unchanged.Follow-up worth a separate ticket (no action needed here).
DatasetItemDAO.getItemsrunsMono.zip(countMono, columnsMono)on every page, so each page request costs three ClickHouse queries — the items query plus a fullCOUNTover the filtered dataset and a column-types scan — where the stream endpoint runs one. A 100k read at the default chunk size therefore issues roughly 150 queries instead of 50, and all but the firstCOUNTis unused (the reader only needstotalfrom page 1). Measured per-request cost is 22% higher than the stream endpoint at identical page size, which is why the sequential case does not improve. Atotal=false/ skip-count query parameter on that endpoint would make this path strictly faster than the cursor stream at every thread count.Review round 2 (2 findings from @petrotiurin, 3 from
baz-reviewer[bot]):id < last_retrieved_idseek, so this was a genuine regression against it, and it landed on callers that match items by exact equality. Every page now carries one version hash, resolved once per read; backends without dataset versioning fall back to a live read, documented onstream_items(). Regression tests drive an insert at offset 0 between pages and assert the unversioned read duplicates (so a regression in the pinning surfaces as that test passing) while the pinned read does not.withscope, so closing the generator ranshutdown(wait=True)and waited for every outstanding page, potentially at GC time on an unrelated line. Nowtry/finallywithcancel_futures=True, wait=False; test assertsclose()returns in under one page delay with pages in flight.HttpClientviahttpx.MockTransport(worth doing SDK-wide rather than in this PR — the plumbing is covered e2e today), and tightening the cap test's assertions (correct, but the same count and content are asserted by neighbouring tests).Because version pinning changes which backend query runs, the payload parity check was re-run against the previous implementation: byte-identical, and still identical across 3-50 pages at 1 and 8 threads.
Review round 1 (5 findings from
baz-reviewer[bot], all addressed):totalin a page response silently capped the read at page 1 and returned part of the dataset as if it were all of it.totalandcontentare now validated, raisingOpikExceptionthat says why continuing would be wrong.chunk_sizewas unbounded and forwarded straight to the endpoint'ssize, so oversized values could materialize an unbounded page and values above int32 failed only after the request went out. Now capped atDATASET_STREAM_BATCH_SIZE. This one rejects rather than clamps, unlike the thread ceiling — handing back smaller pages than asked for would look like the argument had no effect.stream_items()resolved the dataset id eagerly, doing a by-name REST lookup before iteration began. TheDatasethook is now a generator, so it defers to first iteration; the laziness test previously only checked httpx and would have passed regardless, and now asserts the lookup itself.nb_samplesValueErrorwas undocumented. Both docstrings and the guide now state the valid range and what to pass for an unbounded read.Dataset.from_publicid seeding had no test. Added, plus the fallback case where the response carries no id.Change checklist
Issues
AI-WATERMARK
AI-WATERMARK: yes
chunk_sizedefault and API shape chosen by the author against the measurements belowTesting
Automated
New unit coverage (61 tests) drives a fake items endpoint that records every call and tracks peak in-flight requests: chunk ordering, real concurrency at
num_threads > 1versus strictly serial at1,nb_sampleslimiting both items and pages fetched, filter serialization, laziness, error mapping, and argument validation. Separate tests pinget_items()to the new path, assert it equalsstream_items()output, and cover itsnum_threadsargument — that it is forwarded (real concurrency observed), that1is strictly serial, that the thread count never changes the result, that invalid values raise, and that positional calls likeget_items(10)still meannb_samples.New e2e coverage: 10k items over 10 pages at 8 threads, asserting every item exactly once and that the 8-thread id order matches the 1-thread order; plus
nb_samples,filter_string, and aDatasetVersionsnapshot read.test_cli_import_export.pyis included deliberately — the CLI import path matches dataset items by exact dict equality minusid, so any drift in the returned payload fails it.Payload parity with the previous implementation
__internal_api__stream_items_as_dataclasses__is unchanged on this branch, so the oldget_items()expression could be run in the same process as the new one and compared directly. Ordered (not set) comparison, all byte-identical:None, empty dict/list/string, zero, unicode + emoji, list-of-dicts.id,source,description,trace_id,span_id,evaluators,execution_policy) inside itsdata— both paths reduce it identically and both warn once.filter_stringmatching some rows and none;nb_samplesat 1/2/3 and at the exact item count.Performance
100k items, local Docker backend, small rows (~480 bytes). Baseline is
mainchecked out in a separate worktree, 3+ repeats each, fresh client andget_dataset()per repeat.get_items()get_items()(default: 2000 x 4 threads)stream_items()by thread count,chunk_size=2000:Read honestly: the gain here is parallelism, not the Fern bypass. Decomposing a sequential read shows the bypass does work as intended — client-side parsing drops from 0.64 s to 0.21 s, a 3x saving — but it is outweighed by the paginated endpoint's 22% higher per-request server cost described under Details, which is why
num_threads=1lands slightly behind the baseline. Larger pages amortize it (measured at 5000: 2.43 s sequential, 1.8x faster than baseline, and 5.5x at 4 threads), butchunk_sizeis deliberately capped atDATASET_STREAM_BATCH_SIZEso a read can never request a bigger page than the typed stream's own batch and peak memory stays bounded as before.chunk_sizeis therefore a downward knob, for datasets whose individual items are large; the sequential case is left to the backend follow-up above rather than bought with unbounded pages. Two caveats on the absolute numbers: a single local backend container understates network latency (which parallelism helps more with) and overstates backend contention (which it helps less with), and these rows are small, so this measures request overhead rather than parse throughput — heavier per-item payloads should favor the Fern bypass more.Not run
Backend, frontend and TypeScript SDK suites — no changes in those areas.
Documentation
apps/opik-documentation/documentation/fern/docs-v2/evaluation/advanced/manage_datasets.mdx— new "Downloading large datasets faster" section leading withget_items(num_threads=8), thenstream_items()for datasets too large to hold in memory, afilter_string/nb_samplesexample, and a note on thechunk_sizetradeoff.get_items()andstream_items()covernum_threads, the returned shape, the2 * num_threadslook-ahead, and which direction to movechunk_size.