Merge branch 'refs/heads/main' into return-scalar-for-zero-dim-indexing

Author: brokkoli71

changes/2755.bugfix.rst

@@ -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.

changes/2799.bugfix.rst

@@ -0,0 +1 @@
+Enitialise empty chunks to the default fill value during writing and add default fill values for datetime, timedelta, structured, and other (void* fixed size) data types

src/zarr/api/asynchronous.py

@@ -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
     ----------
@@ -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
     ----------
@@ -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)

src/zarr/api/synchronous.py

@@ -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
     ----------
@@ -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
     ----------
@@ -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)))
 

src/zarr/core/buffer/cpu.py

@@ -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:

src/zarr/core/codec_pipeline.py

@@ -56,6 +56,19 @@ def resolve_batched(codec: Codec, chunk_specs: Iterable[ArraySpec]) -> Iterable[
     return [codec.resolve_metadata(chunk_spec) for chunk_spec in chunk_specs]
 
 
+def fill_value_or_default(chunk_spec: ArraySpec) -> Any:
+    fill_value = chunk_spec.fill_value
+    if fill_value is None:
+        # Zarr V2 allowed `fill_value` to be null in the metadata.
+        # Zarr V3 requires it to be set. This has already been
+        # validated when decoding the metadata, but we support reading
+        # Zarr V2 data and need to support the case where fill_value
+        # is None.
+        return _default_fill_value(dtype=chunk_spec.dtype)
+    else:
+        return fill_value
+
+
 @dataclass(frozen=True)
 class BatchedCodecPipeline(CodecPipeline):
     """Default codec pipeline.
@@ -247,17 +260,7 @@ async def read_batch(
                 if chunk_array is not None:
                     out[out_selection] = chunk_array
                 else:
-                    fill_value = chunk_spec.fill_value
-
-                    if fill_value is None:
-                        # Zarr V2 allowed `fill_value` to be null in the metadata.
-                        # Zarr V3 requires it to be set. This has already been
-                        # validated when decoding the metadata, but we support reading
-                        # Zarr V2 data and need to support the case where fill_value
-                        # is None.
-                        fill_value = _default_fill_value(dtype=chunk_spec.dtype)
-
-                    out[out_selection] = fill_value
+                    out[out_selection] = fill_value_or_default(chunk_spec)
         else:
             chunk_bytes_batch = await concurrent_map(
                 [
@@ -284,10 +287,7 @@ async def read_batch(
                         tmp = tmp.squeeze(axis=drop_axes)
                     out[out_selection] = tmp
                 else:
-                    fill_value = chunk_spec.fill_value
-                    if fill_value is None:
-                        fill_value = _default_fill_value(dtype=chunk_spec.dtype)
-                    out[out_selection] = fill_value
+                    out[out_selection] = fill_value_or_default(chunk_spec)
 
     def _merge_chunk_array(
         self,
@@ -305,7 +305,7 @@ def _merge_chunk_array(
                 shape=chunk_spec.shape,
                 dtype=chunk_spec.dtype,
                 order=chunk_spec.order,
-                fill_value=chunk_spec.fill_value,
+                fill_value=fill_value_or_default(chunk_spec),
             )
         else:
             chunk_array = existing_chunk_array.copy()  # make a writable copy
@@ -394,7 +394,7 @@ async def _read_key(
                     chunk_array_batch.append(None)  # type: ignore[unreachable]
                 else:
                     if not chunk_spec.config.write_empty_chunks and chunk_array.all_equal(
-                        chunk_spec.fill_value
+                        fill_value_or_default(chunk_spec)
                     ):
                         chunk_array_batch.append(None)
                     else:

src/zarr/core/group.py

@@ -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
         ----------
@@ -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(
@@ -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
         ----------
@@ -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
         ----------
@@ -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
         ----------
@@ -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)))
 

src/zarr/core/metadata/v2.py

@@ -349,6 +349,14 @@ def _default_fill_value(dtype: np.dtype[Any]) -> Any:
         return b""
     elif dtype.kind in "UO":
         return ""
+    elif dtype.kind in "Mm":
+        return dtype.type("nat")
+    elif dtype.kind == "V":
+        if dtype.fields is not None:
+            default = tuple([_default_fill_value(field[0]) for field in dtype.fields.values()])
+            return np.array([default], dtype=dtype)
+        else:
+            return np.zeros(1, dtype=dtype)
     else:
         return dtype.type(0)
 

src/zarr/core/sync.py

@@ -3,6 +3,7 @@
 import asyncio
 import atexit
 import logging
+import os
 import threading
 from concurrent.futures import ThreadPoolExecutor, wait
 from typing import TYPE_CHECKING, TypeVar
@@ -89,6 +90,26 @@ def cleanup_resources() -> None:
 atexit.register(cleanup_resources)
 
 
+def reset_resources_after_fork() -> None:
+    """
+    Ensure that global resources are reset after a fork. Without this function,
+    forked processes will retain invalid references to the parent process's resources.
+    """
+    global loop, iothread, _executor
+    # These lines are excluded from coverage because this function only runs in a child process,
+    # which is not observed by the test coverage instrumentation. Despite the apparent lack of
+    # test coverage, this function should be adequately tested by any test that uses Zarr IO with
+    # multiprocessing.
+    loop[0] = None  # pragma: no cover
+    iothread[0] = None  # pragma: no cover
+    _executor = None  # pragma: no cover
+
+
+# this is only available on certain operating systems
+if hasattr(os, "register_at_fork"):
+    os.register_at_fork(after_in_child=reset_resources_after_fork)
+
+
 async def _runner(coro: Coroutine[Any, Any, T]) -> T | BaseException:
     """
     Await a coroutine and return the result of running it. If awaiting the coroutine raises an

tests/test_array.py

@@ -1,8 +1,10 @@
 import dataclasses
 import json
 import math
+import multiprocessing as mp
 import pickle
 import re
+import sys
 from itertools import accumulate
 from typing import TYPE_CHECKING, Any, Literal
 from unittest import mock
@@ -1390,3 +1392,39 @@ def test_roundtrip_numcodecs() -> None:
     metadata = root["test"].metadata.to_dict()
     expected = (*filters, BYTES_CODEC, *compressors)
     assert metadata["codecs"] == expected
+
+
+def _index_array(arr: Array, index: Any) -> Any:
+    return arr[index]
+
+
+@pytest.mark.parametrize(
+    "method",
+    [
+        pytest.param(
+            "fork",
+            marks=pytest.mark.skipif(
+                sys.platform in ("win32", "darwin"), reason="fork not supported on Windows or OSX"
+            ),
+        ),
+        "spawn",
+        pytest.param(
+            "forkserver",
+            marks=pytest.mark.skipif(
+                sys.platform == "win32", reason="forkserver not supported on Windows"
+            ),
+        ),
+    ],
+)
+@pytest.mark.parametrize("store", ["local"], indirect=True)
+def test_multiprocessing(store: Store, method: Literal["fork", "spawn", "forkserver"]) -> None:
+    """
+    Test that arrays can be pickled and indexed in child processes
+    """
+    data = np.arange(100)
+    arr = zarr.create_array(store=store, data=data)
+    ctx = mp.get_context(method)
+    pool = ctx.Pool()
+
+    results = pool.starmap(_index_array, [(arr, slice(len(data)))])
+    assert all(np.array_equal(r, data) for r in results)