Merge branch 'main' into dependabot/github_actions/actions-c8837d376d

Author: d-v-b

changes/4392.misc.md

@@ -0,0 +1 @@
+Moved the early-exit `fill_value` scan from the `NDBuffer` base class to the CPU buffer, so GPU buffers keep a single reduction.

pyproject.toml

@@ -47,7 +47,7 @@ requires-python = ">=3.12"
 dependencies = [
     'packaging>=22.0',
     'numpy>=2',
-    'numcodecs>=0.14',
+    'numcodecs>=0.16',
     'google-crc32c>=1.5',
     'typing_extensions>=4.14',
     'donfig>=0.8',
@@ -109,19 +109,19 @@ homepage = "https://github.com/zarr-developers/zarr-python"
 # pins deliberately, e.g. via dependabot or `uv lock --upgrade`.
 [dependency-groups]
 test = [
-    "coverage==7.16.0",
+    "coverage==7.16.1",
     "pytest==9.1.1",
     "pytest-asyncio==1.4.0",
     "pytest-cov==7.1.0",
     "pytest-accept==0.3.0",
-    "numpydoc==1.10.0",
+    "numpydoc==1.11.0",
     "hypothesis==6.168.0",
     "pytest-reportlog==1.0.0",
     "pytest-xdist==3.8.0",
     "pytest-benchmark==5.3.0",
     "pytest-codspeed==5.0.3",
     "tomlkit==0.15.1",
-    "uv==0.12.12",
+    "uv==0.12.15",
 ]
 remote-tests = [
     {include-group = "test"},
@@ -145,7 +145,7 @@ docs = [
     "mkdocs-redirects==1.2.3",
     "markdown-exec[ansi]==1.12.3",
     "griffe-inherited-docstrings==1.1.3",
-    "ruff==0.16.6",
+    "ruff==0.16.8",
     # Changelog generation
     {include-group = "release"},
     # Optional dependencies to run examples
@@ -252,7 +252,7 @@ dependency-groups = ["remote-tests"]
 extra-dependencies = [
     'packaging==22.*',
     'numpy==2.0.*',
-    'numcodecs==0.14.*',  # 0.14 needed for zarr3 codecs
+    'numcodecs==0.16.*',
     'fsspec==2023.10.0',
     's3fs==2023.10.0',
     'universal_pathlib==0.2.0',

src/zarr/core/buffer/core.py

@@ -32,10 +32,6 @@
 # itemsize. A view to a dtype of the same width is always permitted by numpy.
 _BITWISE_DTYPE = {1: "u1", 2: "u2", 4: "u4", 8: "u8"}
 
-# Elements in the first slice examined by `NDBuffer.all_equal`. Large enough to
-# amortise one numpy call, small enough that a mismatch in it is nearly free.
-_ALL_EQUAL_BLOCK = 1 << 14
-
 
 @runtime_checkable
 class ArrayLike(Protocol):
@@ -540,65 +536,12 @@ def __len__(self) -> int:
     def __repr__(self) -> str:
         return f"<NDBuffer shape={self.shape} dtype={self.dtype} {self._data!r}>"
 
-    @staticmethod
-    def _compare_all(data: NDArrayLike, other: Any, equal_nan: bool) -> bool:
-        """Whether every element of `data` equals `other`, inspecting all of it."""
-        # Handle positive and negative zero by comparing bit patterns:
-        if (
-            np.asarray(other).dtype.kind == "f"
-            and other == 0.0
-            and data.dtype.kind not in ("U", "S", "T", "O", "V")
-        ):
-            _data, other = np.broadcast_arrays(data, np.asarray(other, data.dtype))
-            # Read the bytes as unsigned integers rather than as a void dtype.
-            # Both compare bit patterns, so both keep -0.0 distinct from 0.0,
-            # but numpy has vectorised integer comparison loops and no void
-            # one: the void form falls back to a generic elementwise path that
-            # is more than an order of magnitude slower. Widths with no integer
-            # of the same size (longdouble, complex128) keep the void form.
-            bitwise_dtype = _BITWISE_DTYPE.get(_data.dtype.itemsize, f"V{_data.dtype.itemsize}")
-            return bool(np.array_equal(_data.view(bitwise_dtype), other.view(bitwise_dtype)))
-        # use array_equal to obtain equal_nan=True functionality
-        _data, other = np.broadcast_arrays(data, other)
-        return bool(
-            np.array_equal(
-                _data,
-                other,
-                equal_nan=equal_nan if data.dtype.kind not in ("U", "S", "T", "O", "V") else False,
-            )
-        )
-
     def all_equal(self, other: Any, equal_nan: bool = True) -> bool:
-        """Compare to `other` using np.array_equal."""
+        """Whether every element of this buffer equals `other`."""
         if other is None:
             # Handle None fill_value for Zarr V2
             return False
-        data = self._data
-        # A buffer that is not uniformly `other` still has to be read in full
-        # before a whole-buffer comparison can report the first mismatch, and
-        # on the write path that is the common case. Walking it in slices of
-        # the leading axis returns as soon as one slice differs. Slices double
-        # in length, so a buffer that really is uniform is covered in O(log n)
-        # comparisons and costs the same as the single scan this replaces.
-        #
-        # The leading axis is used rather than a flattened view because a chunk
-        # is usually a strided view into a larger array, which cannot be
-        # flattened without copying it. Slicing axis 0 is a view whatever the
-        # layout.
-        if data.ndim > 0 and np.ndim(other) == 0 and data.size > _ALL_EQUAL_BLOCK:
-            leading = data.shape[0]
-            per_index = data.size // leading if leading else 0
-            if leading > 1 and per_index:
-                step = max(1, _ALL_EQUAL_BLOCK // per_index)
-                start = 0
-                while start < leading:
-                    stop = min(start + step, leading)
-                    if not self._compare_all(data[start:stop], other, equal_nan):
-                        return False
-                    start = stop
-                    step *= 2
-                return True
-        return self._compare_all(data, other, equal_nan)
+        return _array_all_equal(self._data, other, equal_nan)
 
     def fill(self, value: Any) -> None:
         self._data.fill(value)
@@ -610,6 +553,34 @@ def transpose(self, axes: SupportsIndex | Sequence[SupportsIndex] | None) -> Sel
         return self.__class__(self._data.transpose(axes))
 
 
+def _array_all_equal(data: NDArrayLike, other: Any, equal_nan: bool) -> bool:
+    """Whether every element of `data` equals `other`, inspecting all of it."""
+    # Handle positive and negative zero by comparing bit patterns:
+    if (
+        np.asarray(other).dtype.kind == "f"
+        and other == 0.0
+        and data.dtype.kind not in ("U", "S", "T", "O", "V")
+    ):
+        _data, other = np.broadcast_arrays(data, np.asarray(other, data.dtype))
+        # Read the bytes as unsigned integers rather than as a void dtype.
+        # Both compare bit patterns, so both keep -0.0 distinct from 0.0,
+        # but numpy has vectorised integer comparison loops and no void
+        # one: the void form falls back to a generic elementwise path that
+        # is more than an order of magnitude slower. Widths with no integer
+        # of the same size (longdouble, complex128) keep the void form.
+        bitwise_dtype = _BITWISE_DTYPE.get(_data.dtype.itemsize, f"V{_data.dtype.itemsize}")
+        return bool(np.array_equal(_data.view(bitwise_dtype), other.view(bitwise_dtype)))
+    # use array_equal to obtain equal_nan=True functionality
+    _data, other = np.broadcast_arrays(data, other)
+    return bool(
+        np.array_equal(
+            _data,
+            other,
+            equal_nan=equal_nan if data.dtype.kind not in ("U", "S", "T", "O", "V") else False,
+        )
+    )
+
+
 class BufferPrototype(NamedTuple):
     """Prototype of the Buffer and NDBuffer class
 

src/zarr/core/buffer/cpu.py

@@ -22,6 +22,10 @@
     from zarr.core.buffer.core import ArrayLike, NDArrayLike
     from zarr.core.common import BytesLike
 
+# Elements in the first slice examined by `NDBuffer.all_equal`. Large enough to
+# amortise one numpy call, small enough that a mismatch in it is nearly free.
+_ALL_EQUAL_BLOCK = 1 << 14
+
 
 class Buffer(core.Buffer):
     """A flat contiguous memory block
@@ -189,6 +193,42 @@ def __setitem__(self, key: Any, value: Any) -> None:
             value = value._data
         self._data.__setitem__(key, value)
 
+    def all_equal(self, other: Any, equal_nan: bool = True) -> bool:
+        """Whether every element of this buffer equals `other`.
+
+        Returns as soon as part of the buffer is found to differ, rather than
+        reading all of it.
+        """
+        if other is None:
+            # Handle None fill_value for Zarr V2
+            return False
+        data = self._data
+        if data.size == 0 or data.ndim == 0 or np.ndim(other) > 0:
+            # nothing to slice, or `other` has to broadcast against all of `data`
+            return core._array_all_equal(data, other, equal_nan)
+        # A buffer that is not uniformly `other` still has to be read in full
+        # before a whole-buffer comparison can report the first mismatch, and
+        # on the write path that is the common case. Walking it in slices of
+        # the leading axis returns as soon as one slice differs. Slices double
+        # in length, so a buffer that really is uniform is covered in O(log n)
+        # comparisons and costs the same as a single scan. A buffer no larger
+        # than one block is covered by the first slice.
+        #
+        # The leading axis is used rather than a flattened view because a chunk
+        # is usually a strided view into a larger array, which cannot be
+        # flattened without copying it. Slicing axis 0 is a view whatever the
+        # layout.
+        leading = data.shape[0]
+        step = max(1, _ALL_EQUAL_BLOCK // (data.size // leading))
+        start = 0
+        while start < leading:
+            stop = min(start + step, leading)
+            if not core._array_all_equal(data[start:stop], other, equal_nan):
+                return False
+            start = stop
+            step *= 2
+        return True
+
 
 def as_numpy_array_wrapper(
     func: Callable[[npt.NDArray[Any]], bytes], buf: core.Buffer, prototype: core.BufferPrototype

tests/test_buffer.py

@@ -240,91 +240,59 @@ def test_empty(
         assert result.flags.f_contiguous  # type: ignore[attr-defined]
 
 
-@pytest.mark.parametrize("dtype", ["float16", "float32", "float64", "complex64", "int32"])
-@pytest.mark.parametrize("order", ["C", "F"])
-def test_all_equal_matches_full_scan(dtype: str, order: str) -> None:
-    """The blocked scan must agree with a whole-buffer comparison.
-
-    Buffers are deliberately larger than one block and not a whole number of
-    blocks, and the mismatching element is placed at the very end so an
-    early-exit bug cannot pass by luck.
-    """
-    n = (1 << 14) * 3 + 7
-    side = int(np.sqrt(n))
-    for base in (np.zeros(n, dtype=dtype), np.zeros((side, side), dtype=dtype)):
-        arr = np.asfortranarray(base) if order == "F" and base.ndim > 1 else base
-        uniform = cpu.NDBuffer.from_numpy_array(arr)
-        differs = arr.copy()
-        differs.reshape(-1)[-1] = 1
-        mixed = cpu.NDBuffer.from_numpy_array(differs)
-        for fill in (0, 0.0, 1):
-            assert uniform.all_equal(fill) == cpu.NDBuffer._compare_all(arr, fill, True)
-            assert mixed.all_equal(fill) == cpu.NDBuffer._compare_all(differs, fill, True)
+def _all_equal_reference(data: np.ndarray, fill: float) -> bool:
+    """`NDBuffer.all_equal` as it was before the fast paths: one full scan,
+    comparing a zero `fill` bitwise through a void view."""
+    if np.asarray(fill).dtype.kind == "f" and fill == 0.0:
+        void = f"V{data.dtype.itemsize}"
+        expected = np.broadcast_to(np.asarray(fill, data.dtype), data.shape)
+        return bool(np.array_equal(data.view(void), expected.view(void)))
+    return bool(np.array_equal(data, np.broadcast_to(fill, data.shape), equal_nan=True))
+
+
+_BLOCK = 1 << 14
+
+
+def _layouts(dtype: str) -> dict[str, np.ndarray]:
+    """Buffers larger than one scan block and not a whole number of blocks."""
+    side = 211
+    return {
+        "1d": np.zeros(_BLOCK * 3 + 7, dtype=dtype),
+        "C": np.zeros((side, side), dtype=dtype),
+        "F": np.zeros((side, side), dtype=dtype, order="F"),
+        # a chunk carved out of a larger array, which is what the write path passes
+        "strided": np.zeros((2 * side, 2 * side), dtype=dtype)[5 : side + 5, 7 : side + 7],
+        "leading-1": np.zeros((1, _BLOCK * 2), dtype=dtype),
+    }
+
+
+@pytest.mark.parametrize(
+    "dtype", ["float16", "float32", "float64", "complex64", "longdouble", "int32"]
+)
+@pytest.mark.parametrize("layout", ["1d", "C", "F", "strided", "leading-1"])
+@pytest.mark.parametrize("contents", [0.0, -0.0, np.nan, 1.0])
+@pytest.mark.parametrize("mismatch", [None, "first", "last"])
+@pytest.mark.parametrize("fill", [0.0, -0.0, np.nan, 1.0, 0])
+def test_all_equal(
+    dtype: str, layout: str, contents: float, mismatch: str | None, fill: float
+) -> None:
+    """`all_equal` agrees with a full scan, whatever the layout or where a mismatch sits."""
+    data = _layouts(dtype)[layout]
+    if np.dtype(dtype).kind == "i" and contents != 1.0:
+        contents = 0
+    data[...] = contents
+    if mismatch is not None:
+        index = tuple(0 if mismatch == "first" else n - 1 for n in data.shape)
+        data[index] = 7
+    assert cpu.NDBuffer.from_numpy_array(data).all_equal(fill) == _all_equal_reference(data, fill)
 
 
 @pytest.mark.parametrize("dtype", ["float16", "float32", "float64"])
 def test_all_equal_distinguishes_negative_zero(dtype: str) -> None:
     """Regression test for #3144: -0.0 is not the same chunk as 0.0."""
-    n = (1 << 14) * 2 + 3
-    negative = cpu.NDBuffer.from_numpy_array(np.full(n, -0.0, dtype=dtype))
-    positive = cpu.NDBuffer.from_numpy_array(np.zeros(n, dtype=dtype))
-    assert positive.all_equal(0.0)
-    assert not negative.all_equal(0.0)
-    assert negative.all_equal(-0.0)
-    assert not positive.all_equal(-0.0)
-    # a single -0.0 in an otherwise +0.0 buffer, past the first block
-    mixed_data = np.zeros(n, dtype=dtype)
-    mixed_data[-1] = np.array(-0.0, dtype=dtype)
-    assert not cpu.NDBuffer.from_numpy_array(mixed_data).all_equal(0.0)
-
-
-def test_all_equal_non_contiguous() -> None:
-    """A strided buffer cannot be flattened without copying; it must still be correct."""
-    n = (1 << 14) * 4
-    base = np.zeros(n * 2, dtype="float32")
-    strided = base[::2]
-    assert not strided.flags.c_contiguous
-    assert cpu.NDBuffer.from_numpy_array(strided).all_equal(0.0)
-    strided2 = base[::2].copy()
-    strided2[-1] = 1
-    assert not cpu.NDBuffer.from_numpy_array(strided2).all_equal(0.0)
-
-
-@pytest.mark.parametrize("fill", [0.0, np.nan, 1.0])
-def test_all_equal_strided_chunk_view(fill: float) -> None:
-    """A chunk carved out of a larger array is strided, which is the shape the
-    write path actually hands this method. Slicing the leading axis has to work
-    for those, not just for contiguous buffers."""
-    whole = np.full((512, 512), fill, dtype="float32")
-    chunk = whole[64:192, 64:192]
-    assert not chunk.flags.c_contiguous
-    assert not chunk.flags.f_contiguous
-    assert cpu.NDBuffer.from_numpy_array(chunk).all_equal(fill)
-    # a single differing element, placed last so an early exit cannot pass by luck
-    whole2 = whole.copy()
-    whole2[191, 191] = 12345.0
-    chunk2 = whole2[64:192, 64:192]
-    assert not cpu.NDBuffer.from_numpy_array(chunk2).all_equal(fill)
-    # and placed first
-    whole3 = whole.copy()
-    whole3[64, 64] = 12345.0
-    assert not cpu.NDBuffer.from_numpy_array(whole3[64:192, 64:192]).all_equal(fill)
-
-
-def test_all_equal_single_leading_row() -> None:
-    """shape[0] == 1 has no leading axis to slice; it must fall back, not break."""
-    row = np.zeros((1, (1 << 14) * 2), dtype="float32")
-    assert cpu.NDBuffer.from_numpy_array(row).all_equal(0.0)
-    row2 = row.copy()
-    row2[0, -1] = 1
-    assert not cpu.NDBuffer.from_numpy_array(row2).all_equal(0.0)
-
-
-def test_all_equal_nan() -> None:
-    n = (1 << 14) * 2 + 1
-    nans = cpu.NDBuffer.from_numpy_array(np.full(n, np.nan, dtype="float64"))
-    assert nans.all_equal(np.nan)
-    assert not nans.all_equal(0.0)
-    one_nan = np.zeros(n, dtype="float64")
-    one_nan[-1] = np.nan
-    assert not cpu.NDBuffer.from_numpy_array(one_nan).all_equal(0.0)
+    positive = np.zeros(_BLOCK * 2 + 3, dtype=dtype)
+    negative = np.full_like(positive, -0.0)
+    assert cpu.NDBuffer.from_numpy_array(positive).all_equal(0.0)
+    assert not cpu.NDBuffer.from_numpy_array(negative).all_equal(0.0)
+    assert cpu.NDBuffer.from_numpy_array(negative).all_equal(-0.0)
+    assert not cpu.NDBuffer.from_numpy_array(positive).all_equal(-0.0)