perf: cache lexicographic chunk coords in sharding codec

[d-v-b]

Fixes a performance regression introduced in #3826 that made morton-order-indexing slower.

for details, see claude's summary below.

The subchunk_write_order feature (#3826) regressed sharded write performance: _encode_partial_single rebuilt the full per-shard chunk coordinate grid on every write via
np.array(list(_subchunk_order_iter(..., "lexicographic"))), and to_dict_vectorized rebuilt a tuple key per row with tuple(coords.ravel()). For a single-chunk write into a shard with tens of thousands of chunks this roughly doubled write time (~22ms -> ~40ms on test_sharded_morton_write_single_chunk, matching the -44% CodSpeed regression).

Add cached _lexicographic_order (array) and _lexicographic_order_keys (tuples) helpers in indexing.py, mirroring _morton_order/_morton_order_keys, and pass the cached keys into to_dict_vectorized instead of deriving them row-by-row. This restores write throughput to the pre-#3826 baseline while preserving identical chunk ordering (verified equal to np.ndindex across shapes including 0-d and empty).

[Description of PR]

TODO:

• Add unit tests and/or doctests in docstrings
• Add docstrings and API docs for any new/modified user-facing classes and functions
• New/modified features documented in docs/user-guide/*.md
• Changes documented as a new file in changes/
• GitHub Actions have all passed
• Test coverage is 100% (Codecov passes)

[codecov]

Codecov Report

❌ Patch coverage is 97.43590% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 93.55%. Comparing base (fe22910) to head (aa2bb30).

Files with missing lines	Patch %	Lines
src/zarr/codecs/sharding.py	93.33%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #4012   +/-   ##
=======================================
  Coverage   93.55%   93.55%           
=======================================
  Files          88       88           
  Lines       11896    11909   +13     
=======================================
+ Hits        11129    11142   +13     
  Misses        767      767           

Files with missing lines	Coverage Δ
src/zarr/core/indexing.py	96.29% <100.00%> (+0.04%)	⬆️
src/zarr/codecs/sharding.py	92.20% <93.33%> (+0.06%)	⬆️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[codspeed-hq]

Merging this PR will improve performance by 40.32%

⚡ 3 improved benchmarks
❌ 3 regressed benchmarks
✅ 53 untouched benchmarks
🆕 10 new benchmarks
⏩ 13 skipped benchmarks¹

Warning

Please fix the performance issues or acknowledge them on CodSpeed.

Performance Changes

	Mode	Benchmark	BASE	HEAD	Efficiency
❌	WallTime	test_sharded_morton_single_chunk[(33, 33, 33)-memory]	2.2 ms	7.7 ms	-71.7%
❌	WallTime	test_sharded_morton_single_chunk[(32, 32, 32)-memory]	2.1 ms	7 ms	-69.24%
❌	WallTime	test_sharded_morton_single_chunk[(30, 30, 30)-memory]	2.1 ms	6.2 ms	-65.96%
⚡	WallTime	test_sharded_morton_write_single_chunk[(32, 32, 32)-memory]	350.3 ms	54 ms	×6.5
⚡	WallTime	test_sharded_morton_write_single_chunk[(33, 33, 33)-memory]	391.3 ms	61.7 ms	×6.3
⚡	WallTime	test_sharded_morton_write_single_chunk[(30, 30, 30)-memory]	292.7 ms	46.7 ms	×6.3
🆕	WallTime	test_morton_order[(8, 8, 8)]	N/A	469.3 µs	N/A
🆕	WallTime	test_morton_order[(10, 10, 10)]	N/A	861.5 µs	N/A
🆕	WallTime	test_morton_order[(20, 20, 20)]	N/A	5.6 ms	N/A
🆕	WallTime	test_morton_order[(30, 30, 30)]	N/A	18.3 ms	N/A
🆕	WallTime	test_morton_order[(16, 16, 16)]	N/A	2.5 ms	N/A
🆕	WallTime	test_morton_order[(32, 32, 32)]	N/A	20.8 ms	N/A
🆕	WallTime	test_sharded_morton_write_single_chunk_warm_cache[(33, 33, 33)-memory]	N/A	37.4 ms	N/A
🆕	WallTime	test_morton_order[(33, 33, 33)]	N/A	24.6 ms	N/A
🆕	WallTime	test_sharded_morton_write_single_chunk_warm_cache[(30, 30, 30)-memory]	N/A	28.6 ms	N/A
🆕	WallTime	test_sharded_morton_write_single_chunk_warm_cache[(32, 32, 32)-memory]	N/A	33.1 ms	N/A

Tip

Investigate this regression by commenting @codspeedbot fix this regression on this PR, or directly use the CodSpeed MCP with your agent.

---

_{Comparing d-v-b:perf-sharding-coord-cache (2742897) with main (6ce787d)²}

Footnotes

1. 13 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩

2. No successful run was found on main (fe22910) during the generation of this report, so 6ce787d was used instead as the comparison base. There might be some changes unrelated to this pull request in this report. ↩

[maxrjones]

I think the following reads cleaner than adding another parameter. The reader already knows its own chunks_per_shard, so it can fetch both cached structures itself instead of having them threaded in. With this, the call site simplifies to just shard_reader.to_dict_vectorized().

I tested it still fixes the regression — same ~1.6–1.8× speedup over main on test_sharded_morton_write_single_chunk, and the difference vs the committed version is within run-to-run noise.

    def to_dict_vectorized(self) -dict[tuple[int, ...], Buffer | None]:
        """Build a dict of chunk coordinates to buffers using vectorized lookup.

        The full per-shard chunk coordinate grid (both the array used for the
        vectorized index lookup and the plain tuples used as dict keys) is
        cached on ``chunks_per_shard``, so neither is rebuilt on every call.
        For a shard with tens of thousands of chunks this avoids reconstructing
        that many tuples on every partial write.

        Returns
        -------
        dict mapping chunk coordinate tuples to Buffer or None
        """
        chunks_per_shard = self.index.chunks_per_shard
        chunk_coords_array = _lexicographic_order(chunks_per_shard)
        chunk_coords_keys = _lexicographic_order_keys(chunks_per_shard)
        starts, ends, valid = self.index.get_chunk_slices_vectorized(chunk_coords_array)

        result: dict[tuple[int, ...], Buffer | None] = {}
        for i, coords in enumerate(chunk_coords_keys):
            if valid[i]:
                result[coords] = self.buf[int(starts[i]) : int(ends[i])]
            else:
                result[coords] = None

        return result

[d-v-b]

that's great feedback! have a look at 6e02c8f and see if things look better to you

[ilan-gold]

Apologies for the regression!

[chuckwondo]

Not sure this makes a lot of sense. Why would we provide a "convenience" function for producing an "lazy" iterator over something that has already been eagerly constructed? Shouldn't we expose the lazy option to the user and then the user can eagerly collect the results, if that's what they want. Otherwise, we simply provide no means for the caller to get lazy results.

[chuckwondo]

Although, in this case, does being lazy even matter? I'm not familiar enough with what's going on here to tell if chunk_shape is always something relatively small, in which case laziness is likely unnecessary.

[ilan-gold]

I didn't want to open this can of worms, but I agree things are funny here. I don't think the iter function here is super important but

user can eagerly collect the results, if that's what they want.

The idea here is definitely to cache the eager results and then let the user create an iterator. But I don't even think an iterator is necessary _lexicographic_order_keys should satisfy the APIs that lexicographic_order_iter is meant for e.e.,g

dict.fromkeys((1, 2, 3))

works just fine (which is one place where the iterator is used).

[mkitti]

It really should be the other way around. We should define the lazy version and the eager version should just be "collect" the lazy version.

The problem is that pay an enormous penalty for generating the entire list of all chunks when we might only just want the first one.

[d-v-b]

agreed! this pattern is what you get when you optimize for benchmarks. we might need to relax that concern, or add a separate benchmark that checks "time to fetch 1 chunk"

[d-v-b]

I think @chuckwondo's point indicates something more basic -- there's no reason to expose an iteration interface, lazy or not, for something that has a defined size, like a dense, finite, coordinate grid. We should be exposing an indexable data structure from the start.

[mkitti]

Suggested change

	def _lexicographic_order_keys(chunk_shape: tuple[int, ...]) -> tuple[tuple[int, ...], ...]:
	return tuple(tuple(int(x) for x in row) for row in _lexicographic_order(chunk_shape))
	def lexicographic_order_iter(chunk_shape: tuple[int, ...]) -> Iterator[tuple[int, ...]]:
	# Iterate the chunk grid in lexicographic (row-major / C) order. The full
	# coordinate tuple is built once and cached on `chunk_shape` (see
	# _lexicographic_order_keys), so repeated calls for the same shape — e.g. one
	# per shard write — reuse it rather than rebuilding tens of thousands of
	# tuples each time.
	return iter(_lexicographic_order_keys(chunk_shape))
	def _lexicographic_order_keys(chunk_shape: tuple[int, ...]) -> tuple[tuple[int, ...], ...]:
	return tuple(lexicographic_order_iter(chunk_shape))
	def lexicographic_order_iter(chunk_shape: tuple[int, ...]) -> Iterator[tuple[int, ...]]:
	# Iterate the chunk grid in lexicographic (row-major / C) order. The full
	# coordinate tuple is built once and cached on `chunk_shape` (see
	# _lexicographic_order_keys), so repeated calls for the same shape — e.g. one
	# per shard write — reuse it rather than rebuilding tens of thousands of
	# tuples each time.
	return (tuple(int(x) for x in row) for row in _lexicographic_order(chunk_shape))

[mkitti]

The lazy "lexicographic_order_iter" uses a generator comprehension. The eager "_lexicographic_order_keys " just pulls that into a tuple.

[Copilot]

Pull request overview

This PR addresses a sharded-write performance regression introduced with subchunk_write_order by caching lexicographic (C-order) chunk coordinates and precomputed tuple keys, avoiding repeated rebuilding of large per-shard coordinate grids during partial writes.

Changes:

• Added cached lexicographic coordinate helpers in zarr.core.indexing and a new lexicographic_order_iter.
• Updated the sharding codec to reuse cached lexicographic coordinate arrays/keys and simplified _ShardReader.to_dict_vectorized() to be shape-derived.
• Added a unit test to validate to_dict_vectorized() behavior and included a changelog entry.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
tests/test_codecs/test_sharding.py	Adds coverage for _ShardReader.to_dict_vectorized() and switches tests to lexicographic_order_iter.
src/zarr/core/indexing.py	Introduces cached lexicographic coordinate/key generators and aliases c_order_iter to the new iterator.
src/zarr/codecs/sharding.py	Uses cached lexicographic coords/keys to eliminate per-write reconstruction overhead.
changes/4001.misc.md	Documents the performance restoration in the changelog.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[d-v-b]

something I added here: I realized that we already had a c_order_iter function, but lexicographic_order_iter is a much better name for this functionality. so c_order_iter is now just a wrapper for lexicographic_order_iter.

[d-v-b]

thanks for the careful review all! I think I've addressed everything

[mkitti]

Suggested change

	return iter(_morton_order_keys(chunk_shape))
	@lru_cache(maxsize=16)
	def _morton_order_keys(chunk_shape: tuple[int, ...]) -> tuple[tuple[int, ...], ...]:
	return tuple(morton_order_iter)
	def morton_order_iter(chunk_shape: tuple[int, ...]) -> Iterator[tuple[int, ...]]:
	return (tuple(int(x) for x in row) for row in _morton_order(chunk_shape))

This is also backwards.

[mkitti]

The suggestion is not applying to the correct lines. Manual intervention is likely required.

[d-v-b]

done in 6840249

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

[mkitti]

@codspeedbot explain why this is faster

[mkitti]

@codspeedbot explain why this is faster

[codspeed-hq]

@codspeedbot explain why this is faster

Why this PR is faster

I compared the flamegraphs between main and this PR for test_sharded_morton_write_single_chunk[(33, 33, 33)-memory] (the benchmark with the largest improvement: 391 ms → 226 ms, +72.9%). The speedup comes from eliminating two expensive, redundant computations in the sharded write hot path:

1. Cached coordinate keys eliminate per-write grid rebuilding (~180 ms saved)

On main, every call to _encode_partial_single triggered _encode_shard_dict → _subchunk_order_iter → _morton_order_keys, which rebuilt the full per-shard chunk coordinate grid from scratch. The nested generators in _morton_order_keys consumed 181 ms (23% of total time), dominated by expensive per-element array_item numpy scalar extraction (~115 ms in the inner <genexpr> alone).

With this PR, the morton/lexicographic order arrays and their tuple keys are cached, so _encode_partial_single's coordinate work drops from ~181 ms to <1 ms on subsequent writes to the same shard shape.

2. Pre-computed keys eliminate row-by-row tuple construction in to_dict_vectorized (~139 ms saved)

On main, _ShardReader.to_dict_vectorized spent 138.6 ms (17.6% of total time) building a dictionary key per row by calling tuple(coords.ravel()) in a loop. Each iteration hit array_item (63 ms), array_ravel (10.5 ms), and PyArray_Scalar (10 ms) — all numpy C-level operations that are fast individually but extremely costly when repeated tens of thousands of times (once per chunk in the shard).

With this PR, to_dict_vectorized receives pre-computed cached tuple keys directly, completely bypassing the per-row numpy extraction. This function no longer appears in the top hotspots.

Net effect

Path	main	This PR
_subchunk_order_iter → _morton_order_keys (coordinate grid)	181 ms	< 1 ms
to_dict_vectorized (tuple key construction)	138.6 ms	eliminated from hotspots
Total benchmark	786 ms	458 ms

The combined ~320 ms savings from these two paths accounts for the full improvement, restoring write throughput to the pre-#3826 baseline.

Summary

• Warming up the performance wizard...
• Analyzing performance improvement — Flamegraph analysis complete — identified two root causes of the speedup
• Handling comment — Comment handling finished successfully

_{Was this response helpful? React with 👍 or 👎 to help us improve CodSpeedBot.}

[mkitti]

Claude review incoming...

[mkitti]

This review was authored by Claude (Claude Code), at the request of @mkitti.

Reviewed the diff critically. The core fix is sound — caching _lexicographic_order/_lexicographic_order_keys and feeding the cached keys straight into to_dict_vectorized/dict.fromkeys removes exactly the two per-write costs #3826 introduced (np.array(list(...)) and per-row tuple(coords.ravel())). The to_dict_vectorized() signature simplification is a genuine improvement (the reader already knows chunks_per_shard), order equivalence with np.ndindex holds across 0-d/empty/N-d, and chunk_coords_array[i]/chunk_coords_keys[i] are guaranteed aligned.

My concern is that the PR has accreted a lazy/eager + naming-symmetry refactor that overshoots the goal, and one piece works against it. Details inline, summarized by severity:

1. (high) The morton "symmetry" commit de-optimizes a hot write path — see comment on morton_order_iter.
2. (medium) _is_total_shard is actually made worse than the itertools.product it replaced, and the commit rationale claims the reverse — see comment on _is_total_shard.
3. (low-med) c_order_iter quietly changes operational characteristics (lazy/no-alloc → eager + cached numpy grid) — see comment on c_order_iter.

Minor: dropping tuple(chunk_shape) coercion in morton_order_iter means a non-tuple arg now fails inside lru_cache rather than being coerced (acceptable per type hints); _lexicographic_order caches a non-contiguous transposed view that pins its larger np.indices base alive; and the rename/soft-deprecation is scope creep relative to a perf-regression fix — worth a line in the PR description so the c_order_iter behavior change is intentional and visible.

Net: I'd approve the lexicographic caching + to_dict_vectorized() change as-is, but push back on item 1 and trim/correct the item 2 narrative before merge.

[mkitti]

Issue 1 (high): this de-optimizes a hot write path — the opposite of this PR's goal.

Previously morton_order_iter returned iter(_morton_order_keys(...)), i.e. it handed back an lru_cached tuple-of-tuples, so repeated calls were essentially free. Now it rebuilds every coordinate tuple from the (cached) array on every call.

morton_order_iter is consumed in _encode_shard_dict, the per-shard serialization path that runs on every shard write — so this reintroduces exactly the per-write tuple-reconstruction cost this PR removes for the lexicographic case. Two knock-on effects:

• _morton_order_keys is now dead in src/ (only the benchmark's cache_clear references it); its cache is never populated in production.
• The asymmetry is backwards: the lexicographic hot path uses cached _lexicographic_order_keys, while the morton hot path uses the uncached morton_order_iter.

Suggest reverting the morton inversion (keep morton_order_iter delegating to cached _morton_order_keys), or route the _encode_shard_dict call site through _morton_order_keys. Source-shape symmetry isn't worth a hot-path de-optimization in a perf PR.

— Claude (Claude Code)

[mkitti]

Issue 2 (medium): this is strictly worse than the code it replaced for the early-exit case, and the commit message claims the reverse.

This previously iterated c_order_iter = itertools.product — genuinely lazy, O(1) memory, true early-exit. lexicographic_order_iter calls _lexicographic_order, which eagerly materializes the whole np.indices grid before yielding the first tuple. So the "laziness" only skips tuple conversion of unconsumed rows, not the expensive grid build.

Concretely, commit f6d91e6's claim ("on a cold cache for a 32³ shard … ~15.8ms → ~24µs, the lazy generator builds one coordinate and bails") can't hold here — np.indices((32,32,32)) runs in full on the first next(). For the pure early-exit case on a cold cache, the original itertools.product was strictly better and used less memory.

In practice this is mostly masked because the dict.fromkeys hot path warms _lexicographic_order for the same shape. But if the early-exit win is actually wanted, make lexicographic_order_iter a true generator (e.g. itertools.product) and have _lexicographic_order_keys collect that — otherwise drop the early-exit framing from the commit history.

— Claude (Claude Code)

[mkitti]

Issue 3 (low-med): the alias quietly changes operational characteristics.

c_order_iter went from pure-lazy itertools.product (no allocation, no cache) to an alias that materializes and lru_caches a full numpy grid. For any large-chunks_per_shard caller this is a cold-start + memory regression, and it now retains up to 16 large arrays. In practice chunks_per_shard is modest so it's likely fine — but it's a real semantic change presented as a no-op "soft-deprecated alias." Worth calling out explicitly in the PR description.

— Claude (Claude Code)

[d-v-b]

thanks for everyone's patience here. Some recent changes:

• I added explanatory comments for why we need two representations (tuples of ints and numpy arrays of ints) of the same data in to_dict_vectorized. Just using Numpy arrays of ints would be more performant, but fully exploiting that requires changes to the calling code, which models chunk indices as tuples of ints. IMO we need to consider the chunk index representation a target for performance optimization.
• I removed the premise of laziness from the *_order routines. Now they are not iterators, they return full collections of coordinates, and they have names that make this clear: morton_order_coords and lexicographic_order_coords. I also changed the name of the parameters from chunk_shape to just shape, because these functions are not just for chunks, and chunk_shape is actually rather ambiguous.

@mkitti feel free to throw your agents at the latest changes.

[d-v-b]

@claude please review