fix chunk/shard iteration #3299

d-v-b · 2025-07-25T16:40:06Z

In main there are some routines for iterating over the chunks of an array, but these routines do not distinguish between chunks and shards (i.e., stored objects) for arrays with sharding.

This PR adds a separate set of shard-specific iteration routines to complement our chunk-specific iteration routines. Various bugs related to iterating over chunks, when shards were the intended iteration target, have been fixed by these changes, notably bugs causing memory races when creating arrays via create_array (xref ##3169)

I think this supersedes #3217, @bojidar-bg I credited you as a co-author on one of these commits because your idea to change the iteration from chunks to shards was correct.

…nd add a failing test

Co-authored-by: Bojidar Marinov <bojidar.marinov.bg@gmail.com>

d-v-b · 2025-07-25T16:42:29Z

src/zarr/core/array.py

+        return self.chunk_grid_shape
+
+    @property
+    def chunk_grid_shape(self) -> ChunkCoords:


I added a new property called chunk_grid_shape because cdata_shape is ambiguous. cdata_shape is still around, but it now uses chunk_grid_shape

d-v-b · 2025-07-25T16:42:45Z

src/zarr/core/array.py

+        return tuple(starmap(ceildiv, zip(self.shape, self.chunks, strict=True)))
+
+    @property
+    def shard_grid_shape(self) -> ChunkCoords:


this complements chunk_grid_shape.

src/zarr/core/array.py

…nto fix/_iter_chunk_keys

bojidar-bg · 2025-07-25T16:45:02Z

Ooh, that looks like a much more complete implementation of what I did in that other PR! Kudos; I wouldn't have been able to push things this far ✨✨

d-v-b · 2025-07-25T16:46:05Z

Ooh, that looks like a much more complete implementation of what I did in that other PR! Kudos; I wouldn't have been able to push things this far ✨✨

It turned out to be more than I expected 😅 .

d-v-b · 2025-07-25T16:47:22Z

src/zarr/core/indexing.py


    else:
        msg = f"Indexing order {order} is not supported at this time."  # type: ignore[unreachable]
        raise NotImplementedError(msg)


+def iter_regions(


this is a new function for iterating over contiguous regions. When we support irregular chunking, we can overload the type of region_shape accordingly.

codecov · 2025-07-25T16:58:02Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 94.69%. Comparing base (9498336) to head (dfb4ee5).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3299      +/-   ##
==========================================
+ Coverage   94.62%   94.69%   +0.07%     
==========================================
  Files          79       79              
  Lines        9468     9520      +52     
==========================================
+ Hits         8959     9015      +56     
+ Misses        509      505       -4

Files with missing lines	Coverage Δ
src/zarr/core/array.py	`97.44% <100.00%> (+0.33%)`	⬆️
src/zarr/core/indexing.py	`96.40% <100.00%> (+0.31%)`	⬆️

🚀 New features to boost your workflow:

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

src/zarr/core/array.py

dstansby · 2025-07-28T20:30:34Z

Just to make sure we're on the same page here, this is where we're heading in my mind, as something that can go in the user guide:

In zarr-python, a chunk of data represents a single file in storage. There is one exception: when a single sharding codec is part of the codecs. In this case a shard represents a single file in storage, and a chunk represents a unit of data the same size or smaller that can be read independently of other chunks in the shard. This distinction is made to allow zarr-python to optimize the creation and processing of sharded arrays.

dstansby · 2025-08-11T12:17:06Z

I think my earlier comments still needs addressing here:

It is a bit confusing that .shards now returns None with just chunks, where all the other shard releated properties return the equivalent chunk values. I would suggest returning self.chunks from .shards if there's no sharding to match the other new properties, and adding a new .is_sharded: bool property to determine whether shards are being used (previously one would have done self.shards is None).

and

I like the concept that even when shards are not set, the shard properties still return the chunk values, but since this is a new concept it should be accomapanied with appropriate explanation in the Sharding user guide section, and in the release notes (which could just be a link to the sharding user guide section).

Happy to try and review in full once we work these points out (especially the second one).

d-v-b · 2025-08-11T15:12:06Z

It is a bit confusing that .shards now returns None with just chunks, where all the other shard releated properties return the equivalent chunk values. I would suggest returning self.chunks from .shards if there's no sharding to match the other new properties, and adding a new .is_sharded: bool property to determine whether shards are being used (previously one would have done self.shards is None).

I agree that it's confusing, and I also agree that the Array.shards attribute should make sense in light of the other various shard-related properties. But an is_sharded property might also be confusing, especially for an array where is_sharded == False but shards = <some tuple>. What about naming this property something like uses_sharding_codec? IMO we should be normalizing "all arrays are sharded" as much as possible across the entire stack, and to me this argues for treating the way an array is sharded (via the sharding codec, or not) as an implementation detail. But I don't feel strongly here except that whatever we do should not be confusing.

And sorry for all the pings recently but @zarr-developers/python-core-devs it would be great to get some other POVs here, since we are talking about some potential changes to the user-facing Array object.

For this PR, i'm totally fine making all the new shard methods private until we can decide on a good story for the array API. We can still use these methods to fix the various bugs related to mixing up chunks and shards. This would also address your second point, because we wouldn't need to add any new user-facing docs at this time.

dcherian · 2025-08-13T19:04:10Z

What about naming this property something like uses_sharding_codec?

This is better.

IMO we should be normalizing "all arrays are sharded" as much as possible across the entire stack, and to me this argues for treating the way an array is sharded (via the sharding codec, or not) as an implementation detail.

I'm 50/50 on this, given that historically Zarr has only had chunks, not shards; and that these are also kwargs to the constructor. It's going to be confusing regardless. Should we just call it shards_or_chunks :D

i'm totally fine making all the new shard methods private until we can decide on a good story for the array API.

Let's do this ASAP, and punt on the changing the meaning of the .shards property.

dstansby · 2025-08-13T19:22:11Z

Perhaps the recommendation to replace array.shards == None could be array.shards == array.chunks instead?

Certainly .shards shouldn't be changed in code that goes in 3.1.x now, so again I'd advocate for any new API being private, and only the minimum API added that's needed to fix the original issue.

d-v-b · 2025-08-13T19:54:46Z

the new API is all private, let me know if we need to do anything else

…x/_iter_chunk_keys

d-v-b · 2025-08-21T08:05:41Z

ping @dstansby @dcherian, let me know if we need anything else here.

dcherian · 2025-08-22T14:58:17Z

src/zarr/core/indexing.py

-    >>> tuple(iter_grid((2,3)), origin=(1,1))
-    ((1, 1), (1, 2), (1, 3), (2, 1), (2, 2), (2, 3))
+    >>> tuple(iter_grid((2,3), origin=(1,1)))
+    ((1, 1), (1, 2))

-    >>> tuple(iter_grid((2,3)), origin=(1,1), selection_shape=(2,2))
-    ((1, 1), (1, 2), (1, 3), (2, 1))
+    >>> tuple(iter_grid((2,3), origin=(0,0), selection_shape=(2,2)))
+    ((0, 0), (0, 1), (1, 0), (1, 1))


hah that was so wrong earlier!

src/zarr/core/array.py

Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>

src/zarr/core/indexing.py

Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>

dcherian · 2025-08-22T15:03:13Z

src/zarr/core/indexing.py

+        for g_pos, r_shape, d_shape in zip(grid_position, region_shape, domain_shape, strict=True):
+            start = g_pos * r_shape
+            stop = start + r_shape
+            if trim_excess:


just curious, is there a reason to not trim excess?

the chunks / shards written to storage are always fully-sized, no matter the shape of the array. trim excess lets you control whether iteration is done over array-indexing space (trim_excess=True) or stored objects indexing space (trim_excess=False), for the same shape

dcherian · 2025-08-22T15:08:51Z

tests/test_array.py


    # write chunks one at a time
-    for idx, region in enumerate(arr._iter_chunk_regions()):
+    for idx, region in enumerate(arr._iter_shard_regions()):


these tests are a bit circular in that the n*_initialized property uses this function under the hood.

Not a blocker, but would be good to think about making this test more independent.

that's a good point, maybe later we can simplify this function and make each shard 1x1x1, then iterate over array elements directly

dcherian · 2025-08-22T15:17:42Z

tests/test_indexing.py

+        )
+    else:
+        selection_shape_parsed = selection_shape
+    for d_s, r_s, o, ss in zip(


Again here this is reimplementing the logic being tested.

Would be good to add some property tests in the future, for example,

I think we know what the last index of the last slice should be, depending on trim_excess being True or False.

we never receive a slice with start < corresponding index in origin.

output shape == selection_shape (with some trim_excess dependency perhaps)

dcherian

LGTM. I have some concerns about better testing these functions, but they can be addressed later.

lumberbot-app · 2025-08-22T15:21:00Z

Owee, I'm MrMeeseeks, Look at me.

There seem to be a conflict, please backport manually. Here are approximate instructions:

Checkout backport branch and update it.

git checkout 3.1.x
git pull

Cherry pick the first parent branch of the this PR on top of the older branch:

git cherry-pick -x -m1 c9eefe663b34b019e9673bcaa37a88aae6158280

You will likely have some merge/cherry-pick conflict here, fix them and commit:

git commit -am 'Backport PR #3299: fix chunk/shard iteration'

Push to a named branch:

git push YOURFORK 3.1.x:auto-backport-of-pr-3299-on-3.1.x

Create a PR against branch 3.1.x, I would have named this PR:

"Backport PR #3299 on branch 3.1.x (fix chunk/shard iteration)"

And apply the correct labels and milestones.

Congratulations — you did some good work! Hopefully your backport PR will be tested by the continuous integration and merged soon!

Remember to remove the Still Needs Manual Backport label once the PR gets merged.

If these instructions are inaccurate, feel free to suggest an improvement.

d-v-b and others added 6 commits July 25, 2025 11:04

factor array element iteration routines into stand-alone functions, a…

cc3698b

…nd add a failing test

add shard_grid_shape, chunk_grid_shape

39b5c07

Co-authored-by: Bojidar Marinov <bojidar.marinov.bg@gmail.com>

docstrings

70a4876

handle null shards

35b4891

use shard_grid_shape instead of cdata_shape in nchunks

8846c96

add improved set of low-level iteration routines for arrays

31698d4

github-actions bot added the needs release notes Automatically applied to PRs which haven't added release notes label Jul 25, 2025

Merge branch 'main' into fix/_iter_chunk_keys

a258780

d-v-b commented Jul 25, 2025

View reviewed changes

src/zarr/core/array.py Outdated Show resolved Hide resolved

d-v-b added 2 commits July 25, 2025 18:44

correct deprecation message

5c34e4a

Merge branch 'fix/_iter_chunk_keys' of github.com:d-v-b/zarr-python i…

943af28

…nto fix/_iter_chunk_keys

d-v-b commented Jul 25, 2025

View reviewed changes

bojidar-bg mentioned this pull request Jul 25, 2025

Fix race condition in from_array for arrays with shards #3217

Closed

6 tasks

d-v-b mentioned this pull request Jul 25, 2025

_iter_chunk_keys iterates over the wrong shape #3293

Open

d-v-b added 3 commits July 25, 2025 19:56

check for deprecation warnings around iter_chunk_keys

54d7fd7

plug tiny coverage holes

82d9bdc

lint

ec5e0f1

d-v-b mentioned this pull request Jul 26, 2025

array.nchunks is ambiguous #3296

Open

d-v-b requested a review from a team July 26, 2025 21:02

dcherian reviewed Jul 28, 2025

View reviewed changes

src/zarr/core/array.py Show resolved Hide resolved

dcherian reviewed Jul 28, 2025

View reviewed changes

src/zarr/core/array.py Outdated Show resolved Hide resolved

dcherian reviewed Jul 28, 2025

View reviewed changes

src/zarr/core/array.py Outdated Show resolved Hide resolved

d-v-b added 2 commits July 28, 2025 18:24

rename chunks_initialized to shards_initialized

83a3ea5

add nshards, nshards_initialized

29c1cdd

Merge branch 'main' into fix/_iter_chunk_keys

1c2c6cd

Merge branch 'main' into fix/_iter_chunk_keys

70f9f4f

make new API private

aa9d0b7

d-v-b added 3 commits August 21, 2025 09:52

Merge branch 'main' of github.com:zarr-developers/zarr-python into fi…

6a79fb4

…x/_iter_chunk_keys

Merge branch 'main' of github.com:zarr-developers/zarr-python into fi…

c8bad03

…x/_iter_chunk_keys

docstrings

4b12879

d-v-b added 2 commits August 21, 2025 10:08

remove references to private API

952c64b

Merge branch 'main' into fix/_iter_chunk_keys

7b62410

d-v-b requested a review from a team August 22, 2025 12:30

maxrjones mentioned this pull request Aug 22, 2025

Release Zarr-Python v3.1.2 #3397

Open

25 tasks

dcherian reviewed Aug 22, 2025

View reviewed changes

src/zarr/core/array.py Outdated Show resolved Hide resolved

Update src/zarr/core/array.py

cdecba8

Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>

dcherian reviewed Aug 22, 2025

View reviewed changes

src/zarr/core/indexing.py Outdated Show resolved Hide resolved

Update src/zarr/core/indexing.py

09bd11c

Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>

dcherian reviewed Aug 22, 2025

View reviewed changes

fix docstring

dfb4ee5

dcherian reviewed Aug 22, 2025

View reviewed changes

dcherian approved these changes Aug 22, 2025

View reviewed changes

d-v-b merged commit c9eefe6 into zarr-developers:main Aug 22, 2025
31 checks passed

lumberbot-app bot added the Still Needs Manual Backport label Aug 22, 2025

d-v-b deleted the fix/_iter_chunk_keys branch August 22, 2025 15:21

Uh oh!

fix chunk/shard iteration #3299

fix chunk/shard iteration #3299

Conversation

d-v-b commented Jul 25, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

bojidar-bg commented Jul 25, 2025

Uh oh!

d-v-b commented Jul 25, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov bot commented Jul 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dstansby commented Jul 28, 2025

Uh oh!

dstansby commented Aug 11, 2025

Uh oh!

d-v-b commented Aug 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

dcherian commented Aug 13, 2025

Uh oh!

dstansby commented Aug 13, 2025

Uh oh!

d-v-b commented Aug 13, 2025

Uh oh!

d-v-b commented Aug 21, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dcherian Aug 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dcherian left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

lumberbot-app bot commented Aug 22, 2025

Uh oh!

Uh oh!

codecov bot commented Jul 25, 2025 •

edited

Loading

d-v-b commented Aug 11, 2025 •

edited

Loading

dcherian Aug 22, 2025 •

edited

Loading