Skip to content

fix(sharding): flatten any multi-dim value for coordinate selections in partial writes - #320

Open
d-v-b wants to merge 6 commits into
mainfrom
fix/sharding-oindex-mixed-int-arrays
Open

fix(sharding): flatten any multi-dim value for coordinate selections in partial writes#320
d-v-b wants to merge 6 commits into
mainfrom
fix/sharding-oindex-mixed-int-arrays

Conversation

@d-v-b

@d-v-b d-v-b commented Sep 4, 2026 •
edited
Loading

Copy link
Copy Markdown
Owner

🤖 AI text below 🤖

Summary

Follow-up to zarr-developers#4284, fixing zarr-developers#4315.

An orthogonal write on a sharded array that mixes an integer index with two or more array indices still failed after zarr-developers#4284:

a = zarr.create_array(MemoryStore(), shape=(4, 4, 4), chunks=(2, 4, 4), dtype="int32",
                      serializer=ShardingCodec(chunk_shape=(2, 2, 2), codecs=(BytesCodec(),)))
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,)

The guard from zarr-developers#4284 reshaped the value only when value.shape == sel_shape. OrthogonalIndexer drops the integer axis from the value it hands down, but np.ix_ keeps that axis as length 1 in the chunk selection, so the re-derived CoordinateIndexer.sel_shape is (2, 1, 2) while the value is (2, 2). Same element count, different rank, reshape skipped.

Changes

  • The invariant is "a coordinate indexer addresses the value flat", so any multi-dimensional value is ravelled to the indexer's shape. Scalars still pass through. Dropping a length-1 axis does not change C order, so this stays consistent with how CoordinateIndexer ravels the broadcast coordinates.
  • Both partial-encode paths (_encode_partial_sync and _encode_partial_single) now share one helper, _get_shard_indexer_and_value, instead of two copies of the block. The check is isinstance(indexer, CoordinateIndexer) rather than getattr(..., "sel_shape"), so mypy types the attribute.
  • test_sharding_orthogonal_set_multiple_array_dims is parametrized over selections: the original 2-D pair, an unsorted pair spanning two shards, an integer axis in each of the three positions, and three array axes. The twelve integer-axis cases fail against the fix: shape the value buffer for coordinate selections in sharded writes zarr-developers/zarr-python#4284 guard and pass here.

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/

🤖 Generated with Claude Code

d-v-b added a commit that referenced this pull request Sep 4, 2026
@d-v-b @claude
docs: add changelog entry for #320
8989750 
Assisted-by: ClaudeCode:claude-fable-5-1
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
d-v-b and others added 2 commits September 4, 2026 11:08
@d-v-b @claude
fix(sharding): flatten any multi-dim value for coordinate selections …
4ccec6a 
…in partial writes

The guard added in zarr-developers#4284 only reshaped the value when its shape equalled
the re-derived CoordinateIndexer's sel_shape. An orthogonal selection that
mixes an integer index with two or more array indices defeats that:
OrthogonalIndexer drops the integer axis from the value but np.ix_ keeps
it as a length-1 axis in the chunk selection, so the shapes differ in rank
while agreeing in element count, the reshape was skipped, and the write
still raised the shape-mismatch ValueError.

The invariant is that a coordinate indexer addresses the value flat, so
ravel any multi-dimensional value instead. Both partial-encode paths now
share one helper for deriving the shard indexer and shaping the value, and
the check is an isinstance on CoordinateIndexer so mypy types sel_shape.

The regression test is parametrized over selections with an integer axis
in each position, three array axes, and an unsorted selection spanning two
shards.

Closes zarr-developers#4315

Assisted-by: ClaudeCode:claude-fable-5-1
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
@d-v-b @claude
docs: add changelog entry for #320
ee25884 
Assisted-by: ClaudeCode:claude-fable-5-1
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
@d-v-b
d-v-b force-pushed the fix/sharding-oindex-mixed-int-arrays branch from 8989750 to ee25884 Compare September 4, 2026 09:08
@d-v-b d-v-b mentioned this pull request Sep 4, 2026
Merged
7 tasks
d-v-b and others added 4 commits September 4, 2026 11:28
@d-v-b
Rename 320.bugfix.md to 4316.bugfix.md
8a61445
@d-v-b @claude
fix(sharding): ravel only a value shaped like the selection minus its…
49f8872 
… unit axes

Ravelling every multi-dimensional value for a coordinate selection was too
lenient. A mask write with a (2, 2) value for four selected elements, or an
orthogonal write with a spurious trailing axis, raises on an unsharded array
but was silently accepted on a sharded one, because the element count
matched and the shard-level selection cannot tell orthogonal from mask
indexing.

The value shape can. An np.ix_ selection has an N-D sel_shape and the
caller's value is that shape minus the integer-indexed axes, which np.ix_
keeps as length-1 axes. Ravel exactly that shape and leave any other rank
alone, so an invalid write fails the same way it does without sharding.

Adds an error test for both leniencies and a positive case with a length-1
array axis next to an integer axis.

Assisted-by: ClaudeCode:claude-fable-5-1
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
@d-v-b @claude
test: exercise sharded writes and bare integer axes in the indexing p…
aaf30c8 
…roperty tests

The property tests could not have found the sharded orthogonal-write bugs:

- test_oindex, test_mask_indexing and test_block_indexing skipped their set
  half on sharded arrays with assume(zarray.shards is None), added in zarr-developers#2825
  when the bug was first seen and never lifted. test_vindex had its set half
  commented out.
- orthogonal_indices wrapped every bare integer as a one-element array, so
  zarr never received an integer index and OrthogonalIndexer's dropped-axis
  path was unreachable. basic_indices(min_dims=1) never yields an integer
  either, so that branch was dead.
- arrays() only drew a shard shape when every axis had a chunk strictly
  between 1 and the axis length, on top of the v3 and regular-grid draws:
  2 of 500 test_oindex examples were sharded.

Lift the skips, draw integers explicitly and give the numpy indexer the
same dropped-axis result, enable the vindex write with a duplicate-point
filter, and let any chunk that fits the array be sharded (33 of 500 now).
With these changes test_oindex fails against the code before this PR with
the mixed-integer shape mismatch, and passes with it.

Assisted-by: ClaudeCode:claude-fable-5-1
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
@d-v-b @claude
test: assert invalid value ranks are rejected on chunked and sharded …
23a0003 
…arrays alike

Replace the sharded-only wrong-rank test in test_sharding.py and the GH2469
one-off in test_indexing.py with one parametrized error test: a coordinate
write with twice the elements, a mask write with a 2-D value, and an
orthogonal write with an extra axis each raise ValueError on chunked and
sharded arrays under both codec pipelines. The property under test is that
storage layout does not change which writes are rejected, which a
sharded-only test could not state. zarr_array_from_numpy_array grows a
shards argument for it.

Only the rejection is asserted; a write that fails inside the chunk merge
may already have touched other chunks on a chunked array.

Assisted-by: ClaudeCode:claude-fable-5-1
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
d-v-b added a commit that referenced this pull request Sep 5, 2026
@d-v-b @claude
fix(sharding): flatten any multi-dim value for coordinate selections …
a50d791 
…in partial writes (zarr-developers#4316)

* fix(sharding): flatten any multi-dim value for coordinate selections in partial writes

The guard added in zarr-developers#4284 only reshaped the value when its shape equalled
the re-derived CoordinateIndexer's sel_shape. An orthogonal selection that
mixes an integer index with two or more array indices defeats that:
OrthogonalIndexer drops the integer axis from the value but np.ix_ keeps
it as a length-1 axis in the chunk selection, so the shapes differ in rank
while agreeing in element count, the reshape was skipped, and the write
still raised the shape-mismatch ValueError.

The invariant is that a coordinate indexer addresses the value flat, so
ravel any multi-dimensional value instead. Both partial-encode paths now
share one helper for deriving the shard indexer and shaping the value, and
the check is an isinstance on CoordinateIndexer so mypy types sel_shape.

The regression test is parametrized over selections with an integer axis
in each position, three array axes, and an unsorted selection spanning two
shards.

Closes zarr-developers#4315

Assisted-by: ClaudeCode:claude-fable-5-1
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* docs: add changelog entry for #320

Assisted-by: ClaudeCode:claude-fable-5-1
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* Rename 320.bugfix.md to 4316.bugfix.md

* fix(sharding): ravel only a value shaped like the selection minus its unit axes

Ravelling every multi-dimensional value for a coordinate selection was too
lenient. A mask write with a (2, 2) value for four selected elements, or an
orthogonal write with a spurious trailing axis, raises on an unsharded array
but was silently accepted on a sharded one, because the element count
matched and the shard-level selection cannot tell orthogonal from mask
indexing.

The value shape can. An np.ix_ selection has an N-D sel_shape and the
caller's value is that shape minus the integer-indexed axes, which np.ix_
keeps as length-1 axes. Ravel exactly that shape and leave any other rank
alone, so an invalid write fails the same way it does without sharding.

Adds an error test for both leniencies and a positive case with a length-1
array axis next to an integer axis.

Assisted-by: ClaudeCode:claude-fable-5-1
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* test: exercise sharded writes and bare integer axes in the indexing property tests

The property tests could not have found the sharded orthogonal-write bugs:

- test_oindex, test_mask_indexing and test_block_indexing skipped their set
  half on sharded arrays with assume(zarray.shards is None), added in zarr-developers#2825
  when the bug was first seen and never lifted. test_vindex had its set half
  commented out.
- orthogonal_indices wrapped every bare integer as a one-element array, so
  zarr never received an integer index and OrthogonalIndexer's dropped-axis
  path was unreachable. basic_indices(min_dims=1) never yields an integer
  either, so that branch was dead.
- arrays() only drew a shard shape when every axis had a chunk strictly
  between 1 and the axis length, on top of the v3 and regular-grid draws:
  2 of 500 test_oindex examples were sharded.

Lift the skips, draw integers explicitly and give the numpy indexer the
same dropped-axis result, enable the vindex write with a duplicate-point
filter, and let any chunk that fits the array be sharded (33 of 500 now).
With these changes test_oindex fails against the code before this PR with
the mixed-integer shape mismatch, and passes with it.

Assisted-by: ClaudeCode:claude-fable-5-1
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* test: assert invalid value ranks are rejected on chunked and sharded arrays alike

Replace the sharded-only wrong-rank test in test_sharding.py and the GH2469
one-off in test_indexing.py with one parametrized error test: a coordinate
write with twice the elements, a mask write with a 2-D value, and an
orthogonal write with an extra axis each raise ValueError on chunked and
sharded arrays under both codec pipelines. The property under test is that
storage layout does not change which writes are rejected, which a
sharded-only test could not state. zarr_array_from_numpy_array grows a
shards argument for it.

Only the rejection is asserted; a write that fails inside the chunk merge
may already have touched other chunks on a chunked array.

Assisted-by: ClaudeCode:claude-fable-5-1
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

---------

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@d-v-b