Skip to content

Commit

Permalink
Deterministic chunk padding (#2755)
Browse files Browse the repository at this point in the history
* test deterministic memory store

* deterministic memory store

* simplify test

* document changes

* Update src/zarr/core/buffer/cpu.py

Co-authored-by: Joe Hamman <jhamman1@gmail.com>

* lint

* handle fill_value==None

* better test

* improve changes documentation

Co-authored-by: David Stansby <dstansby@gmail.com>

* update docstrings

* document changed `zarr.empty`

* add notes to empty() and empty_like()

---------

Co-authored-by: Norman Rzepka <code@normanrz.com>
Co-authored-by: Joe Hamman <jhamman1@gmail.com>
Co-authored-by: David Stansby <dstansby@gmail.com>
  • Loading branch information
4 people authored Feb 11, 2025
1 parent f4278a5 commit 8b77464
Show file tree
Hide file tree
Showing 6 changed files with 67 additions and 13 deletions.
3 changes: 3 additions & 0 deletions changes/2755.bugfix.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
The array returned by ``zarr.empty`` and an empty ``zarr.core.buffer.cpu.NDBuffer`` will now be filled with the
specified fill value, or with zeros if no fill value is provided.
This fixes a bug where Zarr format 2 data with no fill value was written with un-predictable chunk sizes.
12 changes: 10 additions & 2 deletions src/zarr/api/asynchronous.py
Original file line number Diff line number Diff line change
Expand Up @@ -1065,7 +1065,8 @@ async def create(
async def empty(
shape: ChunkCoords, **kwargs: Any
) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
"""Create an empty array.
"""Create an empty array with the specified shape. The contents will be filled with the
array's fill value or zeros if no fill value is provided.
Parameters
----------
Expand All @@ -1087,7 +1088,8 @@ async def empty(
async def empty_like(
a: ArrayLike, **kwargs: Any
) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
"""Create an empty array like `a`.
"""Create an empty array like `a`. The contents will be filled with the
array's fill value or zeros if no fill value is provided.
Parameters
----------
Expand All @@ -1100,6 +1102,12 @@ async def empty_like(
-------
Array
The new array.
Notes
-----
The contents of an empty Zarr array are not defined. On attempting to
retrieve data from an empty Zarr array, any values may be returned,
and these are not guaranteed to be stable from one access to the next.
"""
like_kwargs = _like_args(a, kwargs)
return await empty(**like_kwargs)
Expand Down
12 changes: 10 additions & 2 deletions src/zarr/api/synchronous.py
Original file line number Diff line number Diff line change
Expand Up @@ -902,7 +902,8 @@ def create_array(

# TODO: add type annotations for kwargs
def empty(shape: ChunkCoords, **kwargs: Any) -> Array:
"""Create an empty array.
"""Create an empty array with the specified shape. The contents will be filled with the
array's fill value or zeros if no fill value is provided.
Parameters
----------
Expand All @@ -928,7 +929,8 @@ def empty(shape: ChunkCoords, **kwargs: Any) -> Array:
# TODO: move ArrayLike to common module
# TODO: add type annotations for kwargs
def empty_like(a: ArrayLike, **kwargs: Any) -> Array:
"""Create an empty array like another array.
"""Create an empty array like another array. The contents will be filled with the
array's fill value or zeros if no fill value is provided.
Parameters
----------
Expand All @@ -941,6 +943,12 @@ def empty_like(a: ArrayLike, **kwargs: Any) -> Array:
-------
Array
The new array.
Notes
-----
The contents of an empty Zarr array are not defined. On attempting to
retrieve data from an empty Zarr array, any values may be returned,
and these are not guaranteed to be stable from one access to the next.
"""
return Array(sync(async_api.empty_like(a, **kwargs)))

Expand Down
8 changes: 4 additions & 4 deletions src/zarr/core/buffer/cpu.py
Original file line number Diff line number Diff line change
Expand Up @@ -154,10 +154,10 @@ def create(
order: Literal["C", "F"] = "C",
fill_value: Any | None = None,
) -> Self:
ret = cls(np.empty(shape=tuple(shape), dtype=dtype, order=order))
if fill_value is not None:
ret.fill(fill_value)
return ret
if fill_value is None:
return cls(np.zeros(shape=tuple(shape), dtype=dtype, order=order))
else:
return cls(np.full(shape=tuple(shape), fill_value=fill_value, dtype=dtype, order=order))

@classmethod
def from_numpy_array(cls, array_like: npt.ArrayLike) -> Self:
Expand Down
19 changes: 14 additions & 5 deletions src/zarr/core/group.py
Original file line number Diff line number Diff line change
Expand Up @@ -1498,7 +1498,8 @@ async def tree(self, expand: bool | None = None, level: int | None = None) -> An
async def empty(
self, *, name: str, shape: ChunkCoords, **kwargs: Any
) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
"""Create an empty array in this Group.
"""Create an empty array with the specified shape in this Group. The contents will
be filled with the array's fill value or zeros if no fill value is provided.
Parameters
----------
Expand All @@ -1515,7 +1516,6 @@ async def empty(
retrieve data from an empty Zarr array, any values may be returned,
and these are not guaranteed to be stable from one access to the next.
"""

return await async_api.empty(shape=shape, store=self.store_path, path=name, **kwargs)

async def zeros(
Expand Down Expand Up @@ -1592,7 +1592,8 @@ async def full(
async def empty_like(
self, *, name: str, data: async_api.ArrayLike, **kwargs: Any
) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
"""Create an empty sub-array like `data`.
"""Create an empty sub-array like `data`. The contents will be filled with
the array's fill value or zeros if no fill value is provided.
Parameters
----------
Expand Down Expand Up @@ -2442,7 +2443,8 @@ def require_array(self, name: str, *, shape: ShapeLike, **kwargs: Any) -> Array:

@_deprecate_positional_args
def empty(self, *, name: str, shape: ChunkCoords, **kwargs: Any) -> Array:
"""Create an empty array in this Group.
"""Create an empty array with the specified shape in this Group. The contents will be filled with
the array's fill value or zeros if no fill value is provided.
Parameters
----------
Expand Down Expand Up @@ -2531,7 +2533,8 @@ def full(

@_deprecate_positional_args
def empty_like(self, *, name: str, data: async_api.ArrayLike, **kwargs: Any) -> Array:
"""Create an empty sub-array like `data`.
"""Create an empty sub-array like `data`. The contents will be filled
with the array's fill value or zeros if no fill value is provided.
Parameters
----------
Expand All @@ -2546,6 +2549,12 @@ def empty_like(self, *, name: str, data: async_api.ArrayLike, **kwargs: Any) ->
-------
Array
The new array.
Notes
-----
The contents of an empty Zarr array are not defined. On attempting to
retrieve data from an empty Zarr array, any values may be returned,
and these are not guaranteed to be stable from one access to the next.
"""
return Array(self._sync(self._async_group.empty_like(name=name, data=data, **kwargs)))

Expand Down
26 changes: 26 additions & 0 deletions tests/test_store/test_memory.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,19 @@
from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np
import pytest

import zarr
from zarr.core.buffer import Buffer, cpu, gpu
from zarr.storage import GpuMemoryStore, MemoryStore
from zarr.testing.store import StoreTests
from zarr.testing.utils import gpu_test

if TYPE_CHECKING:
from zarr.core.common import ZarrFormat


class TestMemoryStore(StoreTests[MemoryStore, cpu.Buffer]):
store_cls = MemoryStore
Expand Down Expand Up @@ -46,6 +53,25 @@ def test_store_supports_partial_writes(self, store: MemoryStore) -> None:
def test_list_prefix(self, store: MemoryStore) -> None:
assert True

@pytest.mark.parametrize("dtype", ["uint8", "float32", "int64"])
@pytest.mark.parametrize("zarr_format", [2, 3])
async def test_deterministic_size(
self, store: MemoryStore, dtype, zarr_format: ZarrFormat
) -> None:
a = zarr.empty(
store=store,
shape=(3,),
chunks=(1000,),
dtype=dtype,
zarr_format=zarr_format,
overwrite=True,
)
a[...] = 1
a.resize((1000,))

np.testing.assert_array_equal(a[:3], 1)
np.testing.assert_array_equal(a[3:], 0)


@gpu_test
class TestGpuMemoryStore(StoreTests[GpuMemoryStore, gpu.Buffer]):
Expand Down

0 comments on commit 8b77464

Please sign in to comment.