zarr-chunk-grid-tests

Integration tests and benchmarks for rectilinear chunk grid support in zarr-python.

Two implementation approaches are being explored:

• separate-chunk-grids (zarr-developers/zarr-python#3534) -- keeps RegularChunkGrid intact and adds a new RectilinearChunkGrid class alongside it
• unified-chunk-grid (zarr-developers/zarr-python#3802) -- builds on #3534, replacing both grid classes with a single polymorphic ChunkGrid that uses FixedDimension (regular) and VaryingDimension (rectilinear) per axis

Both approaches use the same user-facing API (chunks=[[10, 20, 30], [25, 25]]) and the same core algorithm (binary search on cumulative sums for rectilinear lookup). They differ in internal architecture, which affects performance characteristics at scale.

Tests

The test suite covers three layers:

• Direct zarr indexing (test_zarr_basic_indexing.py) -- integer, slice, fancy, boolean, and block indexing for both regular and rectilinear chunk grids
• xarray roundtrip and indexing (test_xarray_roundtrip.py, test_xarray_indexing.py) -- .isel(), .sel(), .where(), dtype roundtrips, attribute preservation
• Dask integration (test_dask_integration.py) -- lazy loading, chunking strategies, compute correctness, reductions
• Edge cases (test_edge_cases.py) -- empty selections, chunk boundaries, zarr v2 compat, .chunks deprecation warnings

Running tests

cd zarr-chunk-grid-tests
uv run --group test pytest -v

To run only rectilinear-specific tests:

uv run --group test pytest -k rectilinear -v

Benchmarks

Performance comparison across three zarr-python branches:

Branch	Description
main	Current zarr-python main (baseline, regular chunks only)
unified-chunk-grid	Single ChunkGrid class with per-dimension polymorphism
separate-chunk-grids	Separate RegularChunkGrid + RectilinearChunkGrid classes

Each benchmark script is a standalone uv script that installs zarr from a specific branch via inline metadata. The shared benchmarking logic lives in benchmark.py.

Benchmark suites

Suite	Array	What it tests
Array creation	1D (10M) + 2D (5000x5000)	Chunk grid construction overhead
Scalar indexing	2D (5000x5000, 50x50 chunks)	Per-lookup dispatch cost (500 accesses)
Slice indexing	2D (5000x5000, 50x50 chunks)	Full array, partial, strided, thin row/col
Fancy / oindex	1D (10M, 1K chunks)	Sorted + unsorted integer arrays (1K--1M indices)
Boolean masking	1D (10M) + 2D (5000x5000)	Various mask densities (0.1%--50%)
Block indexing	2D (5000x5000, 50x50 chunks)	Single block through full array
Coordinate / vindex	1D (10M, 1K chunks)	Sorted + unsorted (1K--1M indices)
Chunk count scaling	1D (10M, 100--10K chunks)	Same operations as chunk count grows

Array sizes are tuned for an M1 Mac with 16GB RAM.

Running benchmarks

Run all three branches and generate a comparison chart:

bash benchmarks/run.sh

Or run individual branches:

uv run benchmarks/test_main.py
uv run benchmarks/test_rectilinear.py
uv run benchmarks/test_separate.py

Regenerate the chart from existing results:

uv run benchmarks/plot_results.py

Results are saved as JSON files (results_*.json) and the chart as benchmarks/comparison.png.

Results

All ratios below are relative to main on regular chunks.

Regular chunks (no regressions expected)

Both implementations preserve regular chunk performance well. separate-chunk-grids is at exact parity with main (it doesn't touch RegularChunkGrid). unified-chunk-grid shows a small overhead in some slice operations (~5--10%) from routing through its polymorphic DimensionGrid protocol, but is faster for array creation.

Category	unified	separate
Slice	1.21x	1.00x
Fancy / oindex	1.09x	1.05x
Boolean	1.05x	1.01x
Scalar	1.03x	0.98x
Coordinate / vindex	1.02x	0.97x
Block	1.02x	0.96x
Chunk scaling	0.94x	0.96x
Creation	0.78x	0.90x

Rectilinear chunks

For typical workloads (moderate chunk counts, standard indexing), both implementations perform similarly -- generally within 5--10% of main's regular chunk performance.

Category	unified	separate
Block	0.99x	1.02x
Coordinate / vindex	0.99x	0.98x
Boolean	1.03x	1.09x
Fancy / oindex	1.06x	1.09x
Scalar	1.29x	1.22x
Creation	2.48x	2.45x

Chunk count scaling

The most architecturally revealing benchmark. As the number of rectilinear chunks increases from 100 to 10,000:

Chunks	unified (slice)	separate (slice)	unified (oindex)	separate (oindex)
100	1.42x	1.47x	0.88x	0.91x
500	0.82x	0.76x	0.93x	0.97x
1,000	1.14x	1.11x	0.92x	0.97x
5,000	0.62x	1.08x	0.80x	1.17x
10,000	0.98x	2.04x	0.82x	1.67x

unified-chunk-grid maintains near-parity with main at all chunk counts. separate-chunk-grids degrades at high chunk counts (5K+), likely due to per-chunk coordinate tuple construction in get_chunk_start/get_chunk_shape.

Array creation

Both approaches are ~6--7x slower than main when creating rectilinear arrays with 10K chunks (cumulative sum precomputation + metadata construction). At typical chunk counts (<1K) the overhead is negligible. This is a one-time cost per array.

Setup

The test suite uses the local zarr-python checkout (configured in pyproject.toml):

[tool.uv.sources]
zarr = { path = "../zarr-python", editable = true }

The benchmark scripts each install zarr from their respective git branches independently.