diff --git a/src/zarr/codecs/sharding.py b/src/zarr/codecs/sharding.py index 609e32f87d..d0b2cec285 100644 --- a/src/zarr/codecs/sharding.py +++ b/src/zarr/codecs/sharding.py @@ -171,6 +171,16 @@ def get_chunk_slices_vectorized( valid : ndarray of shape (n_chunks,) Boolean mask indicating which chunks are non-empty. """ + # Handle 0-dimensional arrays (n_dims == 0) + if chunk_coords_array.shape[1] == 0: + # offsets_and_lengths has shape (2,) for 0D, reshape to (1, 2) + offsets_and_lengths = self.offsets_and_lengths.reshape(1, 2) + starts = offsets_and_lengths[:, 0] + lengths = offsets_and_lengths[:, 1] + valid = starts != MAX_UINT_64 + ends = starts + lengths + return starts, ends, valid + # Localize coordinates via modulo (vectorized) shard_shape = np.array(self.offsets_and_lengths.shape[:-1], dtype=np.uint64) localized = chunk_coords_array.astype(np.uint64) % shard_shape diff --git a/tests/test_codecs/test_sharding.py b/tests/test_codecs/test_sharding.py index 233cc4cb77..2325069dd0 100644 --- a/tests/test_codecs/test_sharding.py +++ b/tests/test_codecs/test_sharding.py @@ -16,6 +16,7 @@ ShardingCodecIndexLocation, TransposeCodec, ) +from zarr.codecs.sharding import MAX_UINT_64, _ShardIndex from zarr.core.buffer import NDArrayLike, default_buffer_prototype from zarr.storage import StorePath, ZipStore @@ -554,3 +555,30 @@ def test_sharding_mixed_integer_list_indexing(store: Store) -> None: s3 = sharded[0:5, 1, 0:3] assert c3.shape == s3.shape == (5, 3) # type: ignore[union-attr] np.testing.assert_array_equal(c3, s3) + + +def test_sharding_zero_dimensional() -> None: + """Regression test for https://github.com/zarr-developers/zarr-python/issues/3751""" + arr = zarr.create_array({}, shape=(), dtype="f4", chunks=(), shards=()) + arr[()] = 42.0 + assert arr[()] == pytest.approx(42.0) + # Overwriting should also work + arr[()] = 43.0 + assert arr[()] == pytest.approx(43.0) + + +def test_shard_index_get_chunk_slices_vectorized_zero_dimensional() -> None: + """Directly cover the 0-D path in _ShardIndex.get_chunk_slices_vectorized.""" + # For a 0D array offsets_and_lengths has shape (2,) — reshape to (1, 2) inside. + index = _ShardIndex(np.array([10, 4], dtype=np.uint64)) + chunk_coords = np.empty((1, 0), dtype=np.uint64) + starts, ends, valid = index.get_chunk_slices_vectorized(chunk_coords) + np.testing.assert_array_equal(starts, np.array([10], dtype=np.uint64)) + np.testing.assert_array_equal(ends, np.array([14], dtype=np.uint64)) + np.testing.assert_array_equal(valid, np.array([True])) + + # Empty/unwritten chunk case + index_empty = _ShardIndex(np.array([MAX_UINT_64, MAX_UINT_64], dtype=np.uint64)) + starts_e, _ends_e, valid_e = index_empty.get_chunk_slices_vectorized(chunk_coords) + np.testing.assert_array_equal(starts_e, np.array([MAX_UINT_64], dtype=np.uint64)) + np.testing.assert_array_equal(valid_e, np.array([False]))