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

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 93.56%. Comparing base (8932bbe) to head (2b8b095).

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #4012   +/-   ##
=======================================
  Coverage   93.55%   93.56%           
=======================================
  Files          88       88           
  Lines       11896    11912   +16     
=======================================
+ Hits        11129    11145   +16     
  Misses        767      767           

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

🚀 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 76.71%

⚡ 3 improved benchmarks
✅ 63 untouched benchmarks
⏩ 6 skipped benchmarks¹

Performance Changes

	Mode	Benchmark	BASE	HEAD	Efficiency
⚡	WallTime	test_sharded_morton_write_single_chunk[(32, 32, 32)-memory]	348.6 ms	198.7 ms	+75.45%
⚡	WallTime	test_sharded_morton_write_single_chunk[(33, 33, 33)-memory]	386.8 ms	217.2 ms	+78.06%
⚡	WallTime	test_sharded_morton_write_single_chunk[(30, 30, 30)-memory]	292.6 ms	165.7 ms	+76.62%

Tip

Curious why this is faster? Comment @codspeedbot explain why this is faster on this PR, or directly use the CodSpeed MCP with your agent.

---

_{Comparing d-v-b:perf-sharding-coord-cache (3372b05) with main (c0e2afa)²}

Footnotes

1. 6 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 (1cda981) during the generation of this report, so c0e2afa 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!