Add NDBuffer.empty (#3191)

In #2904, I've
found that `np.empty(...)` can be quite a bit faster than the `np.full`
we use in places
(https://github.com/zarr-developers/zarr-python/blob/baabf08d07e8518e3d37bd83c493a1d46ea7ac1d/src/zarr/core/buffer/cpu.py#L149.
Though note that `np.empty` and `np.zeros` are about the same for the
cases where we use that.)

For the common case of chunk-aligned reads, the memset used by `np.full`
or `np.zeros` should be unnecessary, because the codec pipeline will
overwrite the memory anyway (or overwritten with `fill_value` if the
store is missing the key).

This adds a new method `NDBuffer.empty`, mirroring `np.empty`. To
preserve backwards compatibility, it's *not* abstract. It delegates to
the less efficient `NDBuffer.create`. But I've implemented it for our
`gpu` and `cpu` buffers.

Co-authored-by: Davis Bennett <[email protected]>

Author: TomAugspurger

changes/3191.feature.rst

@@ -0,0 +1 @@
+Added `NDBuffer.empty` method for faster ndbuffer initialization.

src/zarr/core/buffer/core.py

@@ -374,6 +374,41 @@ def create(
             cast("NDArrayLike", None)
         )  # This line will never be reached, but it satisfies the type checker
 
+    @classmethod
+    def empty(
+        cls, shape: ChunkCoords, dtype: npt.DTypeLike, order: Literal["C", "F"] = "C"
+    ) -> Self:
+        """
+        Create an empty buffer with the given shape, dtype, and order.
+
+        This method can be faster than ``NDBuffer.create`` because it doesn't
+        have to initialize the memory used by the underlying ndarray-like
+        object.
+
+        Parameters
+        ----------
+        shape
+            The shape of the buffer and its underlying ndarray-like object
+        dtype
+            The datatype of the buffer and its underlying ndarray-like object
+        order
+            Whether to store multi-dimensional data in row-major (C-style) or
+            column-major (Fortran-style) order in memory.
+
+        Returns
+        -------
+        buffer
+            New buffer representing a new ndarray_like object with empty data.
+
+        See Also
+        --------
+        NDBuffer.create
+            Create a new buffer with some initial fill value.
+        """
+        # Implementations should override this method if they have a faster way
+        # to allocate an empty buffer.
+        return cls.create(shape=shape, dtype=dtype, order=order)
+
     @classmethod
     def from_ndarray_like(cls, ndarray_like: NDArrayLike) -> Self:
         """Create a new buffer of a ndarray-like object

src/zarr/core/buffer/cpu.py

@@ -20,7 +20,7 @@
     from typing import Self
 
     from zarr.core.buffer.core import ArrayLike, NDArrayLike
-    from zarr.core.common import BytesLike
+    from zarr.core.common import BytesLike, ChunkCoords
 
 
 class Buffer(core.Buffer):
@@ -160,6 +160,12 @@ def create(
         else:
             return cls(np.full(shape=tuple(shape), fill_value=fill_value, dtype=dtype, order=order))
 
+    @classmethod
+    def empty(
+        cls, shape: ChunkCoords, dtype: npt.DTypeLike, order: Literal["C", "F"] = "C"
+    ) -> Self:
+        return cls(np.empty(shape=shape, dtype=dtype, order=order))
+
     @classmethod
     def from_numpy_array(cls, array_like: npt.ArrayLike) -> Self:
         return cls.from_ndarray_like(np.asanyarray(array_like))

src/zarr/core/buffer/gpu.py

@@ -22,7 +22,7 @@
     from collections.abc import Iterable
     from typing import Self
 
-    from zarr.core.common import BytesLike
+    from zarr.core.common import BytesLike, ChunkCoords
 
 try:
     import cupy as cp
@@ -178,6 +178,12 @@ def create(
             ret.fill(fill_value)
         return ret
 
+    @classmethod
+    def empty(
+        cls, shape: ChunkCoords, dtype: npt.DTypeLike, order: Literal["C", "F"] = "C"
+    ) -> Self:
+        return cls(cp.empty(shape=shape, dtype=dtype, order=order))
+
     @classmethod
     def from_numpy_array(cls, array_like: npt.ArrayLike) -> Self:
         """Create a new buffer of Numpy array-like object

src/zarr/testing/buffer.py

@@ -13,6 +13,8 @@
     from collections.abc import Iterable
     from typing import Self
 
+    from zarr.core.common import ChunkCoords
+
 
 __all__ = [
     "NDBufferUsingTestNDArrayLike",
@@ -51,6 +53,15 @@ def create(
             ret.fill(fill_value)
         return ret
 
+    @classmethod
+    def empty(
+        cls,
+        shape: ChunkCoords,
+        dtype: npt.DTypeLike,
+        order: Literal["C", "F"] = "C",
+    ) -> Self:
+        return super(cpu.NDBuffer, cls).empty(shape=shape, dtype=dtype, order=order)
+
 
 class StoreExpectingTestBuffer(MemoryStore):
     """Example of a custom Store that expect MyBuffer for all its non-metadata

src/zarr/testing/utils.py

@@ -40,13 +40,10 @@ def has_cupy() -> bool:
 T = TypeVar("T")
 
 
+gpu_mark = pytest.mark.gpu
+skip_if_no_gpu = pytest.mark.skipif(not has_cupy(), reason="CuPy not installed or no GPU available")
+
+
 # Decorator for GPU tests
 def gpu_test(func: T) -> T:
-    return cast(
-        "T",
-        pytest.mark.gpu(
-            pytest.mark.skipif(not has_cupy(), reason="CuPy not installed or no GPU available")(
-                func
-            )
-        ),
-    )
+    return cast("T", gpu_mark(skip_if_no_gpu(func)))

tests/test_buffer.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Literal
 
 import numpy as np
 import pytest
@@ -20,7 +20,7 @@
     TestBuffer,
     TestNDArrayLike,
 )
-from zarr.testing.utils import gpu_test
+from zarr.testing.utils import gpu_mark, gpu_test, skip_if_no_gpu
 
 if TYPE_CHECKING:
     import types
@@ -200,3 +200,39 @@ def test_gpu_buffer_prototype() -> None:
 def test_cpu_buffer_as_scalar() -> None:
     buf = cpu.buffer_prototype.nd_buffer.create(shape=(), dtype="int64")
     assert buf.as_scalar() == buf.as_ndarray_like()[()]  # type: ignore[index]
+
+
+@pytest.mark.parametrize(
+    "prototype",
+    [
+        cpu.buffer_prototype,
+        pytest.param(
+            gpu.buffer_prototype,
+            marks=[gpu_mark, skip_if_no_gpu],
+        ),
+        BufferPrototype(
+            buffer=cpu.Buffer,
+            nd_buffer=NDBufferUsingTestNDArrayLike,
+        ),
+    ],
+)
+@pytest.mark.parametrize(
+    "shape",
+    [
+        (1, 2),
+        (1, 2, 3),
+    ],
+)
+@pytest.mark.parametrize("dtype", ["int32", "float64"])
+@pytest.mark.parametrize("order", ["C", "F"])
+def test_empty(
+    prototype: BufferPrototype, shape: tuple[int, ...], dtype: str, order: Literal["C", "F"]
+) -> None:
+    buf = prototype.nd_buffer.empty(shape=shape, dtype=dtype, order=order)
+    result = buf.as_ndarray_like()
+    assert result.shape == shape
+    assert result.dtype == dtype
+    if order == "C":
+        assert result.flags.c_contiguous  # type: ignore[attr-defined]
+    else:
+        assert result.flags.f_contiguous  # type: ignore[attr-defined]