perf: Vectorize get_chunk_slice for faster sharded writes

[mkitti]

Summary

This PR adds vectorized methods to _ShardIndex and _ShardReader for batch chunk slice lookups, significantly reducing per-chunk function call overhead when writing to shards.

Changes

New Methods

_ShardIndex.get_chunk_slices_vectorized: Batch lookup of chunk slices using NumPy vectorized operations instead of per-chunk Python calls.

_ShardReader.to_dict_vectorized: Build a chunk dictionary using vectorized lookup instead of iterating with individual get() calls.

Modified Code Path

In _encode_partial_single, replaced:

shard_dict = {k: shard_reader.get(k) for k in morton_order_iter(chunks_per_shard)}

With vectorized approach:

morton_coords = _morton_order(chunks_per_shard)
chunk_coords_array = np.array(morton_coords, dtype=np.uint64)
shard_dict = shard_reader.to_dict_vectorized(chunk_coords_array, morton_coords)

Benchmark Results

Single Chunk Write to Large Shard

Writing a single 1x1x1 chunk to a shard with 32³ chunks (using test_sharded_morton_write_single_chunk from PR #3712):

Optimization	Time	Speedup vs Main
Main branch (original)	422ms	-
+ Morton optimization (PR #3708)	261ms	1.6x
+ Vectorized get_chunk_slice	95ms	4.4x

Profile Breakdown

Function	Before	After
get_chunk_slice + _localize_chunk	215ms	3ms
to_dict_vectorized loop	81ms	9ms
Total function calls	299k	37k

Checklist

• 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)

[mkitti]

@d-v-b add the benchmark tag here as well please.

[codspeed-hq]

Merging this PR will improve performance by ×7.2

⚡ 4 improved benchmarks
✅ 52 untouched benchmarks
⏩ 6 skipped benchmarks¹

Performance Changes

	Mode	Benchmark	BASE	HEAD	Efficiency
⚡	WallTime	test_morton_order_iter[(8, 8, 8)]	6.2 ms	1.1 ms	×5.5
⚡	WallTime	test_morton_order_iter[(16, 16, 16)]	56.2 ms	8.1 ms	×6.9
⚡	WallTime	test_morton_order_iter[(32, 32, 32)]	502.6 ms	70.2 ms	×7.2
⚡	WallTime	test_sharded_morton_write_single_chunk[(32, 32, 32)-memory]	953.3 ms	164.4 ms	×5.8

---

_{Comparing mkitti:mkitti-get-chunk-slice-vectorization (157d283) with main (36caf1f)}

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. ↩

[mkitti]

Benchmark results after latest commits

Compared against main with cache cleared before each iteration (_morton_order and _morton_order_keys):

Benchmark	main	This PR	Speedup
test_sharded_morton_write_single_chunk[(32,32,32)]	76,235 µs	6,802 µs	11.2×
test_morton_order_iter[(16,16,16)]	2,092 µs	< threshold	>8×
test_morton_order_iter[(32,32,32)]	17,032 µs	< threshold	>8×
test_sharded_morton_indexing[(16,16,16)]	21,925 µs	21,788 µs	~1× (unaffected)
test_sharded_morton_indexing[(32,32,32)]	175,284 µs	177,995 µs	~1× (unaffected)

Profiling breakdown (test_sharded_morton_write_single_chunk, 32³ shard, 32,768 chunks):

Cost	main	This PR
_morton_order computation	18ms (scalar loop × 32,768)	1ms (vectorized)
get_chunk_slice / _localize_chunk	67ms (32,768 Python calls)	1ms (get_chunk_slices_vectorized)
Total function calls	297,531	166,468

The 11× speedup comes from two changes: vectorizing _morton_order (eliminating 32,768 scalar decode_morton calls) and vectorizing the chunk slice lookup (eliminating 32,768 per-chunk get_chunk_slice + _localize_chunk calls). Read paths are unaffected.

Claude generated this report.

[d-v-b]

this looks great, just 1 nit about narrowing the return type of a function and I think this is good to go

[mkitti]

We seem to be stuck due to codspeed running out of time.

The usage of Macro Runners on zarr-developers exceeded the limit of 600 minutes.

The current period ends on February 28, 2026.

This may need to be merged manually.

Also, I have mainly optimized the power-of-2 case. I am preparing additonal benchmarks beyond a power of 2.