fix: shape the value buffer for coordinate selections in sharded writes

[dylanpulver]

Summary

Nightly Slow Hypothesis CI on main filed #4280 on 2026-08-22 and hit it again on 2026-08-25 (run 32792180919): ValueError: shape mismatch: value array of shape (3,1) could not be broadcast to indexing result of shape (3,). An orthogonal set on a sharded array with two array-indexed dimensions reproduces it:

a = zarr.create_array(MemoryStore(), shape=(4, 4), chunks=(2, 4), dtype="int32",
                      serializer=ShardingCodec(chunk_shape=(2, 2), codecs=(BytesCodec(),)))
a.oindex[np.array([3, 1, 2]), np.array([0, 2])] = np.arange(6).reshape(3, 2)

OrthogonalIndexer hands such a chunk selection down as an np.ix_ pair (indexing.py:989). get_indexer reads it back as a coordinate selection, whose projections address shard_array flat while the caller shaped it like sel_shape. _decode_partial_single reshapes out to sel_shape on the way out (sharding.py:1085); this does the same on the way in. Reads were unaffected.

For reviewers

Guarded on shard_array.shape == sel_shape, so a flat value passes through. Checked against a numpy np.ix_ oracle over 1568 combinations of chunk grid, sharding nesting, per-dimension selector: 96 failures on main, 0 after, every one a write with two array-indexed dimensions.

Author attestation

• I am a human, these are my changes, and I have reviewed and understood every change and can explain why each is correct.

TODO

• Add unit tests and/or doctests in docstrings
• Changes documented as a new file in changes/

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 94.21%. Comparing base (45ead4f) to head (e6478da).

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #4284   +/-   ##
=======================================
  Coverage   94.21%   94.21%           
=======================================
  Files          92       92           
  Lines       12863    12871    +8     
=======================================
+ Hits        12119    12127    +8     
  Misses        744      744           

Files with missing lines	Coverage Δ
src/zarr/codecs/sharding.py	96.22% <100.00%> (+0.04%)	⬆️

🚀 New features to boost your workflow:

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

[d-v-b]

🤖 AI text below 🤖

Code review

Found 1 issue:

1. The fix is applied to _encode_partial_single (async, BatchedCodecPipeline) but not to its sync twin _encode_partial_sync (used by FusedCodecPipeline), which derives its indexer the same way and still fails with the same ValueError for this PR's own regression scenario when codec_pipeline.path is set to FusedCodecPipeline. The decode side applies the sel_shape reshape in both twins (_decode_partial_single and _decode_partial_sync), so the sync encode path needs the matching reshape too. Parametrizing the new test over both pipelines (as test_sharding_vlen_inner_codec_roundtrip does) would cover it.

Missing-fix location:

zarr-python/src/zarr/codecs/sharding.py

Lines 798 to 806 in 01aac1f

	indexer = list(
	get_indexer(
	selection,
	shape=shard_shape,
	chunk_grid=ChunkGrid.from_sizes(shard_shape, self.chunk_shape),
	)
	)

The reshape added on the async side:

zarr-python/src/zarr/codecs/sharding.py

Lines 1361 to 1374 in 01aac1f

	shard_indexer = get_indexer(
	selection,
	shape=shard_shape,
	chunk_grid=ChunkGrid.from_sizes(shard_shape, chunk_shape),
	)
	# A coordinate indexer flattens the selection, so its projections address
	# `shard_array` as 1-D while the caller shaped it like `sel_shape`. This
	# mirrors the reshape `_decode_partial_single` applies on the way out.
	sel_shape = getattr(shard_indexer, "sel_shape", None)
	if sel_shape is not None and shard_array.shape == sel_shape:
	shard_array = shard_array.reshape(shard_indexer.shape)
	indexer = list(shard_indexer)

Sync decode twin already carrying the mirrored reshape:

zarr-python/src/zarr/codecs/sharding.py

Lines 1306 to 1309 in 01aac1f

	if hasattr(indexer, "sel_shape"):
	return out.reshape(indexer.sel_shape)
	return out

🤖 Generated with Claude Code

_{- If this code review was useful, please react with 👍. Otherwise, react with 👎.}

[dylanpulver]

Good catch, that was a real gap. _encode_partial_sync derives its indexer the same way and hit the same ValueError, which I reproduced before fixing:

ValueError: shape mismatch: value array of shape (1,2) could not be broadcast to indexing result of shape (1,)

Same reshape applied there, and the regression test is now parametrized over both pipelines. All four partial paths handle sel_shape: _decode_partial_single and _decode_partial_sync on the way out, _encode_partial_single and _encode_partial_sync on the way in.

I confirmed the two Fused variants fail on the previous commit and pass on this one, so the parametrization is doing work rather than duplicating a passing case.

tests/test_codecs/test_sharding.py is 207 passed, 1 skipped. ruff and mypy are clean.

[li-em]

We hit this issue in icechunk CI too, and are currently sidestepping it

There is also an extra fix in src/zarr/core/indexing.py by using a stable sort for np.argsort, but I can open a PR later after this is merged.

[dylanpulver]

Thanks for the corroboration — useful to know it reproduces in icechunk CI and not just in a constructed case.

The stable-sort change in indexing.py looks like a separate defect on the same path rather than part of this one, so keeping it to its own PR seems right to me. Happy either way if a maintainer would rather see them together.

[d-v-b]

this looks good, i'm going to approve + merge. there's a remaining bugged case that I will fix in a follow-up PR:

import numpy as np, zarr
from zarr.storage import MemoryStore
from zarr.codecs import ShardingCodec, BytesCodec

a = zarr.create_array(
    MemoryStore(), shape=(4, 4, 4), chunks=(2, 4, 4), dtype="int32",
    serializer=ShardingCodec(chunk_shape=(2, 2, 2), codecs=(BytesCodec(),)),
    compressors=None, fill_value=0,
)
a[:] = np.arange(64, dtype="int32").reshape(4, 4, 4)
a.oindex[np.array([3, 1, 2]), 1, np.array([0, 2])] = np.arange(6, dtype="int32").reshape(3, 2)
# ValueError: shape mismatch: value array of shape (1,2) could not be broadcast to indexing result of shape (1,)

[dylanpulver]

Thanks for confirming the bug hits icechunk too, that's useful beyond our regression case.

Agreed the stable-sort change in indexing.py is separate from this fix. A follow-up PR after this one merges sounds right.