diff --git a/src/zarr/core/indexing.py b/src/zarr/core/indexing.py index 3d47f5f183..ca227be094 100644 --- a/src/zarr/core/indexing.py +++ b/src/zarr/core/indexing.py @@ -1346,8 +1346,15 @@ def decode_morton(z: int, chunk_shape: ChunkCoords) -> ChunkCoords: def morton_order_iter(chunk_shape: ChunkCoords) -> Iterator[ChunkCoords]: - for i in range(product(chunk_shape)): - yield decode_morton(i, chunk_shape) + i = 0 + order: list[ChunkCoords] = [] + while len(order) < product(chunk_shape): + m = decode_morton(i, chunk_shape) + if m not in order and all(x < y for x, y in zip(m, chunk_shape, strict=False)): + order.append(m) + i += 1 + for j in range(product(chunk_shape)): + yield order[j] def c_order_iter(chunks_per_shard: ChunkCoords) -> Iterator[ChunkCoords]: diff --git a/tests/test_codecs/test_codecs.py b/tests/test_codecs/test_codecs.py index dfb8e1c595..2025e72937 100644 --- a/tests/test_codecs/test_codecs.py +++ b/tests/test_codecs/test_codecs.py @@ -204,6 +204,26 @@ def test_morton() -> None: ] +@pytest.mark.parametrize( + "shape", + [ + [2, 2, 2], + [5, 2], + [2, 5], + [2, 9, 2], + [3, 2, 12], + [2, 5, 1], + [4, 3, 6, 2, 7], + [3, 2, 1, 6, 4, 5, 2], + ], +) +def test_morton2(shape) -> None: + order = list(morton_order_iter(shape)) + for i, x in enumerate(order): + assert x not in order[:i] # no duplicates + assert all(x[j] < shape[j] for j in range(len(shape))) # all indices are within bounds + + @pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) def test_write_partial_chunks(store: Store) -> None: data = np.arange(0, 256, dtype="uint16").reshape((16, 16)) diff --git a/tests/test_codecs/test_sharding.py b/tests/test_codecs/test_sharding.py index 78f32fef0e..51c82067f3 100644 --- a/tests/test_codecs/test_sharding.py +++ b/tests/test_codecs/test_sharding.py @@ -393,3 +393,32 @@ async def test_sharding_with_empty_inner_chunk( print("read data") data_read = await a.getitem(...) assert np.array_equal(data_read, data) + + +@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) +@pytest.mark.parametrize( + "index_location", + [ShardingCodecIndexLocation.start, ShardingCodecIndexLocation.end], +) +@pytest.mark.parametrize("chunks_per_shard", [(5, 2), (2, 5), (5, 5)]) +async def test_sharding_with_chunks_per_shard( + store: Store, index_location: ShardingCodecIndexLocation, chunks_per_shard: tuple[int] +) -> None: + chunk_shape = (2, 1) + shape = [x * y for x, y in zip(chunks_per_shard, chunk_shape, strict=False)] + data = np.ones(np.prod(shape), dtype="int32").reshape(shape) + fill_value = 42 + + path = f"test_sharding_with_chunks_per_shard_{index_location}" + spath = StorePath(store, path) + a = Array.create( + spath, + shape=shape, + chunk_shape=shape, + dtype="int32", + fill_value=fill_value, + codecs=[ShardingCodec(chunk_shape=chunk_shape, index_location=index_location)], + ) + a[...] = data + data_read = a[...] + assert np.array_equal(data_read, data)