Skip to content

perf: use sync methods for chunk encoding / decoding#3885

Merged
d-v-b merged 113 commits into
zarr-developers:mainfrom
d-v-b:perf/prepared-write-v2
Jul 13, 2026
Merged

perf: use sync methods for chunk encoding / decoding#3885
d-v-b merged 113 commits into
zarr-developers:mainfrom
d-v-b:perf/prepared-write-v2

Conversation

@d-v-b

@d-v-b d-v-b commented Apr 8, 2026

Copy link
Copy Markdown
Contributor

perf: synchronous codec pipeline (FusedCodecPipeline), opt-in

Summary

This PR adds FusedCodecPipeline, an opt-in codec pipeline. It removes the per-chunk async scheduling overhead of the default BatchedCodecPipeline — roughly one coroutine per chunk, which on real workloads dominates the actual codec work — by running codec compute synchronously and in bulk.

The default pipeline is unchanged for standard configurations: existing users see no behavior change unless they opt in. One caveat — sharded arrays whose inner codec chain changes the dtype now evolve that chain correctly (the #2179 fix under Notable internal changes); that is a bugfix and it also runs on the default pipeline. Standard inner chains ([BytesCodec], [BytesCodec, ZstdCodec/BloscCodec], transpose + bytes) are byte-identical to before. We want to move slowly here — let early adopters exercise the new pipeline before considering it as a default.

Note on the original framing. This PR started as PhasedCodecPipeline, pitched around "separate IO from compute." That pitch did not pan out: indexing / Python-object-creation overhead swamped the benefit. The win that remained — and what this PR now delivers — is simpler: avoid async for storage backends that don't need it. See this comment for the history.

What it does

  1. Synchronous codec chain. When every codec implements SupportsSyncCodec, a ChunkTransform runs the codec chain synchronously — no event loop, no per-chunk coroutine — optionally across a thread pool for CPU-heavy decode/encode.
  2. Sharded reads use synchronous IO. Byte-range reads are coalesced via a new Store.get_ranges_sync, and dense, fixed-size, uncompressed shards take a vectorized whole-shard bulk decode. Sharded writes go through the codec's synchronous full-shard-rewrite path.
  3. Graceful fallback. When the store lacks synchronous IO (e.g. ZipStore), or a codec isn't sync-capable, the pipeline falls back to the async path — behavior-equivalent to BatchedCodecPipeline.

IO ownership is unchanged: the sharding codec still holds the byte getter/setter and reads/writes storage directly (the same model as before; this PR does not move toward storage-free codecs).

Performance

Large speedups for sharded arrays — up to ~24× writes / ~14× reads on many-chunks-per-shard layouts, more with compression — and no regressions on compute-bound workloads.

Opting in

  • The default codec_pipeline.path is unchanged (BatchedCodecPipeline). Enable the new pipeline with:
    zarr.config.set({"codec_pipeline.path": "zarr.core.codec_pipeline.FusedCodecPipeline"})
  • codec_pipeline.max_workers defaults to None, which means FusedCodecPipeline runs fully threaded (os.cpu_count() workers). It is the only pipeline that reads this setting — BatchedCodecPipeline ignores it entirely. Set it to 1 for sequential, single-threaded execution (often faster for memory-backed stores, where store reads are GIL-locked and thread-pool overhead dominates), or to a fixed integer to cap the pool.

These config additions are non-breaking: the new pipeline and its threading only take effect once you opt in via codec_pipeline.path. The default pipeline never reads max_workers and never spins up a thread pool.

Public API additions

  • SyncByteGetter / SyncByteSetter — runtime-checkable protocols letting custom byte getters/setters opt into the sync fast path for in-memory IO (used by the sharding codec for its inner chunks).
  • Store.get_ranges_sync — synchronous, coalescing counterpart of get_ranges, with the same coalescing policy.

Correctness

Both pipelines are held to the same behavior, with new equivalence/parity suites asserting that each fast path produces results identical to the general path it bypasses:

  • tests/test_pipeline_parity.py — Fused vs Batched parity across sharded/unsharded, compressed/uncompressed, read/write/overwrite scenarios.
  • tests/test_codec_pipeline_suite.py — a shared suite both pipelines must pass (incl. zarr v2 + nested sharding).
  • tests/test_fastpath_equivalence.py — property tests pinning each fast path (complete-chunk merge view, bulk shard decode, scalar-broadcast write, byte-range coalescing) equal to the general path.
  • tests/test_fused_pipeline.py — Fused-specific behavior incl. thread-pool and async-fallback paths.

Three fast-path correctness issues were found, reproduced, and fixed with regression tests: bulk shard decode could serve a reordering vindex/oindex selection in natural order on uncompressed crc-free shards (now gated on sel_shape); the ChunkTransform spec cache was not torn-read-safe under max_workers > 1 (now a single atomically-assigned entry); and sharded reads with variable-length / numcodecs inner codecs crashed because Codec.is_fixed_size had no default (now defaults to False).

Notable internal changes

  • Shared single-chunk helpers in src/zarr/core/chunk_utils.py (merge/encode, decode/scatter, empty-chunk normalization, spec evolution) used by both pipelines and the sharding codec, so the sync and async paths can't drift.
  • V2Codec and the numcodecs wrappers gained _decode_sync/_encode_sync cores; the async methods are now thin asyncio.to_thread wrappers.
  • The sharding codec's inner sub-chunk pipeline now follows the configured codec_pipeline.path (restoring Narrow JSON type, ensure that to_dict always returns a dict, and v2 filter / compressor parsing  #2179 behavior) and is evolved/memoized per spec.

Lineage

Builds on prior work in #3719. Depends on changes from #3907 and #3908. Mostly authored with Claude.

d-v-b added 4 commits April 7, 2026 10:38
`PreparedWrite` models a set of per-chunk changes that would be applied to a stored chunk. `SupportsChunkPacking`
is a protocol for array -> bytes codecs that can use `PreparedWrite` objects to update an existing chunk.
@github-actions github-actions Bot added the needs release notes Automatically applied to PRs which haven't added release notes label Apr 8, 2026
@codecov

codecov Bot commented Apr 8, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 94.19753% with 47 lines in your changes missing coverage. Please review.
✅ Project coverage is 93.87%. Comparing base (b4179e4) to head (9fe6a9a).

Files with missing lines Patch % Lines
src/zarr/codecs/sharding.py 94.57% 16 Missing ⚠️
src/zarr/core/codec_pipeline.py 96.05% 11 Missing ⚠️
src/zarr/testing/buffer.py 18.18% 9 Missing ⚠️
src/zarr/codecs/numcodecs/_codecs.py 83.33% 4 Missing ⚠️
src/zarr/abc/store.py 91.89% 3 Missing ⚠️
src/zarr/core/chunk_utils.py 98.27% 2 Missing ⚠️
src/zarr/testing/store.py 87.50% 2 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #3885      +/-   ##
==========================================
+ Coverage   93.62%   93.87%   +0.25%     
==========================================
  Files          90       91       +1     
  Lines       11936    12560     +624     
==========================================
+ Hits        11175    11791     +616     
- Misses        761      769       +8     
Files with missing lines Coverage Δ
src/zarr/abc/codec.py 98.76% <100.00%> (ø)
src/zarr/codecs/_v2.py 94.11% <100.00%> (+0.50%) ⬆️
src/zarr/core/array.py 97.85% <100.00%> (+0.01%) ⬆️
src/zarr/core/common.py 88.75% <100.00%> (+0.06%) ⬆️
src/zarr/core/config.py 100.00% <ø> (ø)
src/zarr/storage/_fsspec.py 91.32% <ø> (ø)
src/zarr/testing/strategies.py 95.42% <100.00%> (+0.02%) ⬆️
src/zarr/core/chunk_utils.py 98.27% <98.27%> (ø)
src/zarr/testing/store.py 98.91% <87.50%> (-0.53%) ⬇️
src/zarr/abc/store.py 95.36% <91.89%> (-0.82%) ⬇️
... and 4 more

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@d-v-b

d-v-b commented Apr 9, 2026

Copy link
Copy Markdown
Contributor Author

@TomAugspurger how would this design work with CUDA codecs?

@d-v-b d-v-b force-pushed the perf/prepared-write-v2 branch from 5d3064e to b67a5a0 Compare April 15, 2026 09:51
@github-actions github-actions Bot removed the needs release notes Automatically applied to PRs which haven't added release notes label Apr 15, 2026
@d-v-b d-v-b force-pushed the perf/prepared-write-v2 branch 2 times, most recently from a84a15a to 68a7cdc Compare April 17, 2026 10:41
Comment thread src/zarr/core/codec_pipeline.py Outdated
Comment on lines +943 to +962
# Phase 1: fetch all chunks (IO, sequential)
raw_buffers: list[Buffer | None] = [
bg.get_sync(prototype=cs.prototype) # type: ignore[attr-defined]
for bg, cs, *_ in batch
]

# Phase 2: decode (compute, optionally threaded)
def _decode_one(raw: Buffer | None, chunk_spec: ArraySpec) -> NDBuffer | None:
if raw is None:
return None
return transform.decode_chunk(raw, chunk_spec)

specs = [cs for _, cs, *_ in batch]
if n_workers > 0 and len(batch) > 1:
with ThreadPoolExecutor(max_workers=n_workers) as pool:
decoded_list = list(pool.map(_decode_one, raw_buffers, specs))
else:
decoded_list = [
_decode_one(raw, spec) for raw, spec in zip(raw_buffers, specs, strict=True)
]

@ilan-gold ilan-gold Apr 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Why isn't this all multi-threaded i.e., the I/O as well?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

I should benchmark this, but my expectation was that IO against memory storage and local storage is not compute-limited, and so threads wouldn't remove a real bottleneck. for memory storage i'm sure this is true, not sure about local storage though

d-v-b and others added 6 commits April 17, 2026 22:51
Adds a SupportsSetRange protocol to zarr.abc.store for stores that
allow overwriting a byte range within an existing value. Implementations
are added for LocalStore (using file-handle seek+write) and MemoryStore
(in-memory bytearray slice assignment).

This is the prerequisite for the partial-shard write fast path in
ShardingCodec, which can patch individual inner-chunk slots without
rewriting the entire shard blob when the inner codec chain is fixed-size.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
V2Codec, BytesCodec, BloscCodec, etc. previously only implemented the
async _decode_single / _encode_single methods. Add their sync
counterparts (_decode_sync / _encode_sync) so that the upcoming
SyncCodecPipeline can dispatch through them without spinning up an
event loop.

For codecs that wrap external compressors (numcodecs.Zstd, numcodecs.Blosc,
the V2 fallback chain), the sync versions just call the underlying
compressor's blocking API directly instead of routing through
asyncio.to_thread.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…arallelism

Adds SyncCodecPipeline alongside BatchedCodecPipeline. The new pipeline
runs codecs through their sync entry points (_decode_sync / _encode_sync)
and dispatches per-chunk work to a module-level thread pool sized by
the codec_pipeline.max_workers config (default = os.cpu_count()).

Each chunk's full lifecycle (fetch + decode + scatter for reads;
get-existing + merge + encode + set/delete for writes) runs as one
pool task — overlapping IO of one chunk with compute of another.
Scatter into the shared output buffer is thread-safe because chunks
have non-overlapping output selections.

The async wrappers (read/write) detect SupportsGetSync/SupportsSetSync
stores and dispatch to the sync fast path, passing the configured
max_workers. Other stores fall through to the async path, which still
uses asyncio.concurrent_map at async.concurrency.

Notes on perf:
- Default (None → cpu_count) is tuned for chunks ≥ ~512 KB.
- Small chunks (≤ 64 KB) regress 1.5-3x because pool dispatch overhead
  (~30-50 µs/task) dominates per-chunk work. Workaround:
  zarr.config.set({"codec_pipeline.max_workers": 1}).
- For large chunks on local/memory stores, IO+compute parallelism
  yields 1.7-2.5x over BatchedCodecPipeline on direct-API reads and
  ~2.5x on roundtrip.

ChunkTransform encapsulates the sync codec chain. It caches resolved
ArraySpecs across calls with the same chunk_spec — combined with the
constant-ArraySpec optimization in indexing, hot-path overhead is
minimized.

Includes test scaffolding for the new pipeline (test_sync_codec_pipeline)
and config plumbing for the max_workers key.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Adds _encode_partial_sync and _decode_partial_sync to ShardingCodec.
For fixed-size inner codec chains and stores that implement
SupportsSetRange, partial writes patch individual inner-chunk slots
in-place instead of rewriting the whole shard:

  - Reads existing shard index (one byte-range get).
  - For each affected inner chunk: decodes the slot, merges the new
    region, re-encodes.
  - Writes each modified slot at its deterministic byte offset, then
    rewrites just the index.

For variable-size inner codecs (e.g. with compression) or stores that
don't support byte-range writes, falls through to a full-shard rewrite
matching BatchedCodecPipeline semantics.

The partial-decode path computes a ReadPlan from the shard index and
issues one byte-range get per overlapping chunk, decoding only what
the read selection touches.

Both paths are dispatched from SyncCodecPipeline via the existing
supports_partial_decode / supports_partial_encode protocol checks.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Two new test files:

  test_codec_invariants — asserts contract-level properties that every
  codec / shard / buffer combination must satisfy: round-trip exactness,
  prototype propagation, fill-value handling, all-empty shard handling.

  test_pipeline_parity — exhaustive matrix asserting that
  SyncCodecPipeline and BatchedCodecPipeline produce semantically
  identical results across codec configs, layouts (including
  nested sharding), write sequences, and write_empty_chunks settings.
  Three checks per cell:
    1. Same array contents on read.
    2. Same set of store keys after writes.
    3. Each pipeline reads the other's output identically (catches
       layout-divergence bugs).

These tests pinned the design throughout the SyncCodecPipeline +
partial-shard development.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Adds .gitignore entries for .claude/, CLAUDE.md, and docs/superpowers/
so local IDE/agent planning artifacts don't get committed by accident.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@d-v-b d-v-b force-pushed the perf/prepared-write-v2 branch from aa111a2 to 1be5563 Compare April 17, 2026 21:04
Comment thread src/zarr/core/codec_pipeline.py Outdated
selected = decoded[chunk_selection]
if drop_axes:
selected = selected.squeeze(axis=drop_axes)
out[out_selection] = selected

@ilan-gold ilan-gold Apr 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

It might be worth experimenting with moving this setting operation out[out_selection] = selected outside the threadpool execution since, IIRC, it holds the GIL and is probably non-trivial time-wise.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

The memory usage will probably go up a bit though....

@d-v-b

d-v-b commented Jun 19, 2026

Copy link
Copy Markdown
Contributor Author

@ilan-gold some changes for you to note:

  • the batched codec pipeline is now the default
  • I added a documentation page entry explaining the motivation for the new codec pipeline, and how to opt in.
  • i removed the commented tests for the partial write stuff
  • i updated the PR description

@ilan-gold ilan-gold left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Just little things at this point!

Comment thread docs/user-guide/experimental.md Outdated
Comment thread docs/user-guide/experimental.md Outdated
Comment thread docs/user-guide/experimental.md Outdated
Comment on lines +97 to +100
Threading is off by default because it is not universally a win: it slows small chunks (where the
per-task overhead outweighs the work), and on many-core nodes a pool sized to `cpu_count` can
oversubscribe workloads that already parallelize at a higher level (e.g. Dask). `max_workers` only
affects `FusedCodecPipeline`; the default batched pipeline ignores it.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

I could do some work on this. I'll have to check the benchmarks to understand more. But any guidance here would be appreciated

Comment thread src/zarr/codecs/sharding.py Outdated
Comment thread src/zarr/codecs/sharding.py
Comment thread src/zarr/codecs/sharding.py Outdated
d-v-b and others added 8 commits June 22, 2026 11:56
Co-authored-by: Ilan Gold <ilanbassgold@gmail.com>
Co-authored-by: Ilan Gold <ilanbassgold@gmail.com>
Co-authored-by: Ilan Gold <ilanbassgold@gmail.com>
Co-authored-by: Ilan Gold <ilanbassgold@gmail.com>
Co-authored-by: Ilan Gold <ilanbassgold@gmail.com>
@ilan-gold ilan-gold self-assigned this Jun 25, 2026
@ilan-gold

ilan-gold commented Jul 6, 2026

Copy link
Copy Markdown
Contributor

I have found Lachi's https://github.com/zarrs/zarr_benchmarks to be the most reflective benchmark of real-world data. Here are the most relevant results IMO on my 16-core, SSD-backed linux VM from https://www.denbi.de/:

Read-only

single-threaded, new pipeline:

| Image                              |   Time (s)<br>zarrs<br>rust |   <br>zarr<br>python |   <br>zarrs<br>python |   Memory (GB)<br>zarrs<br>rust |   <br>zarr<br>python |   <br>zarrs<br>python |
|:-----------------------------------|----------------------------:|---------------------:|----------------------:|-------------------------------:|---------------------:|----------------------:|
| data/benchmark_compress_shard.zarr |                        0.98 |                 3.88 |                  1.32 |                           2.16 |                 2.20 |                  2.23 |

fully-threaded, new pipeline:

| Image                              |   Time (s)<br>zarrs<br>rust |   <br>zarr<br>python |   <br>zarrs<br>python |   Memory (GB)<br>zarrs<br>rust |   <br>zarr<br>python |   <br>zarrs<br>python |
|:-----------------------------------|----------------------------:|---------------------:|----------------------:|-------------------------------:|---------------------:|----------------------:|
| data/benchmark_compress_shard.zarr |                        1.00 |                 1.60 |                  1.37 |                           2.16 |                 2.76 |                  2.22 |

3.2.1:

| Image                              |   Time (s)<br>zarrs<br>rust |   <br>zarr<br>python |   <br>zarrs<br>python |   Memory (GB)<br>zarrs<br>rust |   <br>zarr<br>python |   <br>zarrs<br>python |
|:-----------------------------------|----------------------------:|---------------------:|----------------------:|-------------------------------:|---------------------:|----------------------:|
| data/benchmark_compress_shard.zarr |                        0.98 |                 4.00 |                  1.37 |                           2.16 |                 3.97 |                  2.24 |

Roundtrip

single-threaded, new pipeline:

| Image                              |   Time (s)<br>zarrs<br>rust |   <br>zarr<br>python |   <br>zarrs<br>python |   Memory (GB)<br>zarrs<br>rust |   <br>zarr<br>python |   <br>zarrs<br>python |
|:-----------------------------------|----------------------------:|---------------------:|----------------------:|-------------------------------:|---------------------:|----------------------:|
| data/benchmark_compress_shard.zarr |                        1.49 |                16.32 |                  1.94 |                           0.22 |                 2.20 |                  2.24 |

fully-threaded, new pipeline:

| Image                              |   Time (s)<br>zarrs<br>rust |   <br>zarr<br>python |   <br>zarrs<br>python |   Memory (GB)<br>zarrs<br>rust |   <br>zarr<br>python |   <br>zarrs<br>python |
|:-----------------------------------|----------------------------:|---------------------:|----------------------:|-------------------------------:|---------------------:|----------------------:|
| data/benchmark_compress_shard.zarr |                        1.52 |                 6.07 |                  2.06 |                           0.22 |                 2.79 |                  2.24 |

3.2.1:

| Image                              |   Time (s)<br>zarrs<br>rust |   <br>zarr<br>python |   <br>zarrs<br>python |   Memory (GB)<br>zarrs<br>rust |   <br>zarr<br>python |   <br>zarrs<br>python |
|:-----------------------------------|----------------------------:|---------------------:|----------------------:|-------------------------------:|---------------------:|----------------------:|
| data/benchmark_compress_shard.zarr |                        1.44 |                14.55 |                  1.97 |                           0.22 |                 5.35 |                  2.24 |

Smaller subchunks than Presented Here (i.e., <64x64x64)

If I alter the benchmarks to have even smaller subchunks in the shards, say 32x32x32, the difference between rust and python grows, but the difference between the new pipeline and the old one also grows by a similar order of magnitude

Conclusion

I think we should turn threading on by-default. Most people are going to be seeing the benefit, if I had to guess, on the local-file-system level. That being said, it would be great if someone could benchmark this on cloud-backed data!

d-v-b and others added 9 commits July 6, 2026 13:55
* refactor codecs

* threading + benchmark cleanups
Soften the overstated 'default pipeline is unchanged' claim to note it
holds for standard configurations, and add a bugfix fragment for the
sharding inner-codec spec-evolution fix (zarr-developers#2179) that also runs on the
default BatchedCodecPipeline.

Assisted-by: ClaudeCode:claude-opus-4.8
@d-v-b d-v-b requested a review from ilan-gold July 13, 2026 16:34

@ilan-gold ilan-gold left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

This PR should contain no breaking changes and simply add a new pipeline option that is more performant to the current stack. That pipeline includes a bunch of desired features that have been sitting around in our heads for a while and are now materialized, centered around hte notion that "sync should be sync". It has been tested downstream in the CI here and in scverse/anndata#2529 where everything is passing that would be expected to pass (we have a flag that detects pre-releases so that's what's failing there).

In short, I am confident this PR is ready for prime-time!

@d-v-b

d-v-b commented Jul 13, 2026

Copy link
Copy Markdown
Contributor Author

thanks for your help here ilan! this is getting merged. I will follow-up with some docs / QOL changes in separate PRs

@d-v-b d-v-b merged commit fe50301 into zarr-developers:main Jul 13, 2026
30 checks passed
@d-v-b d-v-b deleted the perf/prepared-write-v2 branch July 13, 2026 16:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

benchmark Code will be benchmarked in a CI job. performance Potential issues with Zarr performance (I/O, memory, etc.) run-downstream Run the tests of downstream libraries (e.g., xarray) against zarr

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants