Skip to content

Commit

Permalink
Merge branch 'main' into feature/group-setitem
Browse files Browse the repository at this point in the history
  • Loading branch information
jhamman authored Oct 24, 2024
2 parents 2dafa7e + 87ca150 commit 255f384
Show file tree
Hide file tree
Showing 4 changed files with 27 additions and 8 deletions.
4 changes: 3 additions & 1 deletion src/zarr/core/indexing.py
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,8 @@ def __iter__(self) -> Iterator[ChunkProjection]: ...


def ceildiv(a: float, b: float) -> int:
if a == 0:
return 0
return math.ceil(a / b)


Expand Down Expand Up @@ -374,7 +376,7 @@ def __init__(self, dim_sel: slice, dim_len: int, dim_chunk_len: int) -> None:

def __iter__(self) -> Iterator[ChunkDimProjection]:
# figure out the range of chunks we need to visit
dim_chunk_ix_from = self.start // self.dim_chunk_len
dim_chunk_ix_from = 0 if self.start == 0 else self.start // self.dim_chunk_len
dim_chunk_ix_to = ceildiv(self.stop, self.dim_chunk_len)

# iterate over chunks in range
Expand Down
15 changes: 11 additions & 4 deletions src/zarr/testing/strategies.py
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ def v2_dtypes() -> st.SearchStrategy[np.dtype]:
stores = st.builds(MemoryStore, st.just({}), mode=st.just("w"))
compressors = st.sampled_from([None, "default"])
zarr_formats: st.SearchStrategy[Literal[2, 3]] = st.sampled_from([2, 3])
array_shapes = npst.array_shapes(max_dims=4)
array_shapes = npst.array_shapes(max_dims=4, min_side=0)


@st.composite # type: ignore[misc]
Expand All @@ -85,17 +85,24 @@ def numpy_arrays(
@st.composite # type: ignore[misc]
def np_array_and_chunks(
draw: st.DrawFn, *, arrays: st.SearchStrategy[np.ndarray] = numpy_arrays
) -> tuple[np.ndarray, tuple[int]]: # type: ignore[type-arg]
) -> tuple[np.ndarray, tuple[int, ...]]: # type: ignore[type-arg]
"""A hypothesis strategy to generate small sized random arrays.
Returns: a tuple of the array and a suitable random chunking for it.
"""
array = draw(arrays)
# We want this strategy to shrink towards arrays with smaller number of chunks
# 1. st.integers() shrinks towards smaller values. So we use that to generate number of chunks
numchunks = draw(st.tuples(*[st.integers(min_value=1, max_value=size) for size in array.shape]))
numchunks = draw(
st.tuples(
*[st.integers(min_value=0 if size == 0 else 1, max_value=size) for size in array.shape]
)
)
# 2. and now generate the chunks tuple
chunks = tuple(size // nchunks for size, nchunks in zip(array.shape, numchunks, strict=True))
chunks = tuple(
size // nchunks if nchunks > 0 else 0
for size, nchunks in zip(array.shape, numchunks, strict=True)
)
return (array, chunks)


Expand Down
10 changes: 9 additions & 1 deletion tests/test_indexing.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from numpy.testing import assert_array_equal

import zarr
from zarr import Array
from zarr.core.buffer import BufferPrototype, default_buffer_prototype
from zarr.core.indexing import (
BasicSelection,
Expand All @@ -31,7 +32,6 @@
if TYPE_CHECKING:
from collections.abc import AsyncGenerator

from zarr.core.array import Array
from zarr.core.buffer.core import Buffer
from zarr.core.common import ChunkCoords

Expand Down Expand Up @@ -1927,3 +1927,11 @@ def test_indexing_with_zarr_array(store: StorePath) -> None:

assert_array_equal(a[ii], za[zii])
assert_array_equal(a[ii], za.oindex[zii])


@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"])
@pytest.mark.parametrize("shape", [(0, 2, 3), (0), (3, 0)])
def test_zero_sized_chunks(store: StorePath, shape: list[int]) -> None:
z = Array.create(store=store, shape=shape, chunk_shape=shape, zarr_format=3, dtype="f8")
z[...] = 42
assert_array_equal(z[...], np.zeros(shape, dtype="f8"))
6 changes: 4 additions & 2 deletions tests/test_properties.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@

import hypothesis.extra.numpy as npst # noqa: E402
import hypothesis.strategies as st # noqa: E402
from hypothesis import given # noqa: E402
from hypothesis import assume, given # noqa: E402

from zarr.testing.strategies import arrays, basic_indices, numpy_arrays, zarr_formats # noqa: E402

Expand Down Expand Up @@ -35,11 +35,13 @@ def test_basic_indexing(data: st.DataObject) -> None:
@given(data=st.data())
def test_vindex(data: st.DataObject) -> None:
zarray = data.draw(arrays())
# integer_array_indices can't handle 0-size dimensions.
assume(all(s > 0 for s in zarray.shape))
nparray = zarray[:]

indexer = data.draw(
npst.integer_array_indices(
shape=nparray.shape, result_shape=npst.array_shapes(max_dims=None)
shape=nparray.shape, result_shape=npst.array_shapes(min_side=1, max_dims=None)
)
)
actual = zarray.vindex[indexer]
Expand Down

0 comments on commit 255f384

Please sign in to comment.