Skip to content

Commit a2bc98c

Browse files
perf: cut the fill_value check on the chunk write path
`NDBuffer.all_equal` runs on every chunk write, because that is how `write_empty_chunks=False` (the default) decides a chunk is uniformly `fill_value` and can be skipped. Instrumented on an uncompressed float32 write it was 62.6% of total write time. Two independent causes. The bit-pattern comparison for a zero `fill_value`, which exists so that `-0.0` stays distinct from `0.0` (zarr-developers#3144), viewed both sides as `V{itemsize}`. NumPy has vectorised integer comparison loops and no void one, so that falls back to a generic elementwise path: 27x slower than the equivalent unsigned integer view on float32 and 33x on float64, measured on numpy 2.5.3. Viewing as an unsigned integer of the same width is equally bit-exact. Widths with no same-size integer, longdouble and complex128, keep the void form. The scan also had no early exit, so a buffer that is not uniformly `fill_value` still had to be read in full before `np.array_equal` could report the first mismatch, which on the write path is the common case. It now walks slices of the leading axis, doubling in length, and returns on the first slice that differs. The leading axis is used rather than a flattened view because a chunk is normally a strided view into a larger array and cannot be flattened without a copy; gating the early exit on contiguity makes it silently never engage. A buffer that really is uniform is still covered in the same number of element comparisons. Both `fill_value` regimes benefit, for different reasons: zarr's own default is `0.0` and takes the first change, while xarray writes `fill_value=NaN` and takes only the second. Behaviour is unchanged, including for `-0.0`, NaN, and non-contiguous buffers. Assisted-by: ClaudeCode:claude-opus-5
1 parent 7291469 commit a2bc98c

3 files changed

Lines changed: 167 additions & 18 deletions

File tree

‎changes/4379.bugfix.md‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
Made the `fill_value` check on the write path substantially cheaper.
2+
`NDBuffer.all_equal` runs on every chunk write when `write_empty_chunks` is
3+
false, which is the default, and it dominated the cost of writing chunks.
4+
5+
Two changes. The bit-pattern comparison used for a zero `fill_value`, which
6+
keeps `-0.0` distinct from `0.0`, now views the buffer as an unsigned integer
7+
of the same width instead of as a void dtype; NumPy has vectorised integer
8+
comparison loops but no void one, so the void form fell back to a generic
9+
elementwise path. Separately, the scan now walks slices of the leading axis,
10+
doubling in length, and returns on the first slice that differs, so a buffer
11+
that is not uniformly `fill_value` no longer has to be read in full. The
12+
leading axis is used rather than a flattened view because a chunk is normally
13+
a strided view into a larger array and cannot be flattened without a copy. A
14+
buffer that really is uniform is still covered in the same number of element
15+
comparisons.
16+
17+
Results are unchanged, including for `-0.0`, NaN, and non-contiguous buffers.

‎src/zarr/core/buffer/core.py‎

Lines changed: 61 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,15 @@
2828
__all__: list[str] = []
2929

3030

31+
# Integer dtypes used to compare raw bytes when `fill_value` is zero, keyed by
32+
# itemsize. A view to a dtype of the same width is always permitted by numpy.
33+
_BITWISE_DTYPE = {1: "u1", 2: "u2", 4: "u4", 8: "u8"}
34+
35+
# Elements in the first slice examined by `NDBuffer.all_equal`. Large enough to
36+
# amortise one numpy call, small enough that a mismatch in it is nearly free.
37+
_ALL_EQUAL_BLOCK = 1 << 14
38+
39+
3140
@runtime_checkable
3241
class ArrayLike(Protocol):
3342
"""Protocol for the array-like type that underlie Buffer"""
@@ -531,32 +540,66 @@ def __len__(self) -> int:
531540
def __repr__(self) -> str:
532541
return f"<NDBuffer shape={self.shape} dtype={self.dtype} {self._data!r}>"
533542

534-
def all_equal(self, other: Any, equal_nan: bool = True) -> bool:
535-
"""Compare to `other` using np.array_equal."""
536-
if other is None:
537-
# Handle None fill_value for Zarr V2
538-
return False
543+
@staticmethod
544+
def _compare_all(data: NDArrayLike, other: Any, equal_nan: bool) -> bool:
545+
"""Whether every element of `data` equals `other`, inspecting all of it."""
539546
# Handle positive and negative zero by comparing bit patterns:
540547
if (
541548
np.asarray(other).dtype.kind == "f"
542549
and other == 0.0
543-
and self._data.dtype.kind not in ("U", "S", "T", "O", "V")
550+
and data.dtype.kind not in ("U", "S", "T", "O", "V")
544551
):
545-
_data, other = np.broadcast_arrays(self._data, np.asarray(other, self._data.dtype))
546-
void_dtype = f"V{_data.dtype.itemsize}"
547-
return np.array_equal(_data.view(void_dtype), other.view(void_dtype))
552+
_data, other = np.broadcast_arrays(data, np.asarray(other, data.dtype))
553+
# Read the bytes as unsigned integers rather than as a void dtype.
554+
# Both compare bit patterns, so both keep -0.0 distinct from 0.0,
555+
# but numpy has vectorised integer comparison loops and no void
556+
# one: the void form falls back to a generic elementwise path that
557+
# is more than an order of magnitude slower. Widths with no integer
558+
# of the same size (longdouble, complex128) keep the void form.
559+
bitwise_dtype = _BITWISE_DTYPE.get(_data.dtype.itemsize, f"V{_data.dtype.itemsize}")
560+
return bool(np.array_equal(_data.view(bitwise_dtype), other.view(bitwise_dtype)))
548561
# use array_equal to obtain equal_nan=True functionality
549-
# Since fill-value is a scalar, isn't there a faster path than allocating a new array for fill value
550-
# every single time we have to write data?
551-
_data, other = np.broadcast_arrays(self._data, other)
552-
return np.array_equal(
553-
_data,
554-
other,
555-
equal_nan=equal_nan
556-
if self._data.dtype.kind not in ("U", "S", "T", "O", "V")
557-
else False,
562+
_data, other = np.broadcast_arrays(data, other)
563+
return bool(
564+
np.array_equal(
565+
_data,
566+
other,
567+
equal_nan=equal_nan if data.dtype.kind not in ("U", "S", "T", "O", "V") else False,
568+
)
558569
)
559570

571+
def all_equal(self, other: Any, equal_nan: bool = True) -> bool:
572+
"""Compare to `other` using np.array_equal."""
573+
if other is None:
574+
# Handle None fill_value for Zarr V2
575+
return False
576+
data = self._data
577+
# A buffer that is not uniformly `other` still has to be read in full
578+
# before a whole-buffer comparison can report the first mismatch, and
579+
# on the write path that is the common case. Walking it in slices of
580+
# the leading axis returns as soon as one slice differs. Slices double
581+
# in length, so a buffer that really is uniform is covered in O(log n)
582+
# comparisons and costs the same as the single scan this replaces.
583+
#
584+
# The leading axis is used rather than a flattened view because a chunk
585+
# is usually a strided view into a larger array, which cannot be
586+
# flattened without copying it. Slicing axis 0 is a view whatever the
587+
# layout.
588+
if data.ndim > 0 and np.ndim(other) == 0 and data.size > _ALL_EQUAL_BLOCK:
589+
leading = data.shape[0]
590+
per_index = data.size // leading if leading else 0
591+
if leading > 1 and per_index:
592+
step = max(1, _ALL_EQUAL_BLOCK // per_index)
593+
start = 0
594+
while start < leading:
595+
stop = min(start + step, leading)
596+
if not self._compare_all(data[start:stop], other, equal_nan):
597+
return False
598+
start = stop
599+
step *= 2
600+
return True
601+
return self._compare_all(data, other, equal_nan)
602+
560603
def fill(self, value: Any) -> None:
561604
self._data.fill(value)
562605

‎tests/test_buffer.py‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -238,3 +238,92 @@ def test_empty(
238238
assert result.flags.c_contiguous # type: ignore[attr-defined]
239239
else:
240240
assert result.flags.f_contiguous # type: ignore[attr-defined]
241+
242+
243+
@pytest.mark.parametrize("dtype", ["float16", "float32", "float64", "complex64", "int32"])
244+
@pytest.mark.parametrize("order", ["C", "F"])
245+
def test_all_equal_matches_full_scan(dtype: str, order: str) -> None:
246+
"""The blocked scan must agree with a whole-buffer comparison.
247+
248+
Buffers are deliberately larger than one block and not a whole number of
249+
blocks, and the mismatching element is placed at the very end so an
250+
early-exit bug cannot pass by luck.
251+
"""
252+
n = (1 << 14) * 3 + 7
253+
side = int(np.sqrt(n))
254+
for base in (np.zeros(n, dtype=dtype), np.zeros((side, side), dtype=dtype)):
255+
arr = np.asfortranarray(base) if order == "F" and base.ndim > 1 else base
256+
uniform = cpu.NDBuffer.from_numpy_array(arr)
257+
differs = arr.copy()
258+
differs.reshape(-1)[-1] = 1
259+
mixed = cpu.NDBuffer.from_numpy_array(differs)
260+
for fill in (0, 0.0, 1):
261+
assert uniform.all_equal(fill) == cpu.NDBuffer._compare_all(arr, fill, True)
262+
assert mixed.all_equal(fill) == cpu.NDBuffer._compare_all(differs, fill, True)
263+
264+
265+
@pytest.mark.parametrize("dtype", ["float16", "float32", "float64"])
266+
def test_all_equal_distinguishes_negative_zero(dtype: str) -> None:
267+
"""Regression test for #3144: -0.0 is not the same chunk as 0.0."""
268+
n = (1 << 14) * 2 + 3
269+
negative = cpu.NDBuffer.from_numpy_array(np.full(n, -0.0, dtype=dtype))
270+
positive = cpu.NDBuffer.from_numpy_array(np.zeros(n, dtype=dtype))
271+
assert positive.all_equal(0.0)
272+
assert not negative.all_equal(0.0)
273+
assert negative.all_equal(-0.0)
274+
assert not positive.all_equal(-0.0)
275+
# a single -0.0 in an otherwise +0.0 buffer, past the first block
276+
mixed_data = np.zeros(n, dtype=dtype)
277+
mixed_data[-1] = np.array(-0.0, dtype=dtype)
278+
assert not cpu.NDBuffer.from_numpy_array(mixed_data).all_equal(0.0)
279+
280+
281+
def test_all_equal_non_contiguous() -> None:
282+
"""A strided buffer cannot be flattened without copying; it must still be correct."""
283+
n = (1 << 14) * 4
284+
base = np.zeros(n * 2, dtype="float32")
285+
strided = base[::2]
286+
assert not strided.flags.c_contiguous
287+
assert cpu.NDBuffer.from_numpy_array(strided).all_equal(0.0)
288+
strided2 = base[::2].copy()
289+
strided2[-1] = 1
290+
assert not cpu.NDBuffer.from_numpy_array(strided2).all_equal(0.0)
291+
292+
293+
@pytest.mark.parametrize("fill", [0.0, np.nan, 1.0])
294+
def test_all_equal_strided_chunk_view(fill: float) -> None:
295+
"""A chunk carved out of a larger array is strided, which is the shape the
296+
write path actually hands this method. Slicing the leading axis has to work
297+
for those, not just for contiguous buffers."""
298+
whole = np.full((512, 512), fill, dtype="float32")
299+
chunk = whole[64:192, 64:192]
300+
assert not chunk.flags.c_contiguous
301+
assert not chunk.flags.f_contiguous
302+
assert cpu.NDBuffer.from_numpy_array(chunk).all_equal(fill)
303+
# a single differing element, placed last so an early exit cannot pass by luck
304+
whole2 = whole.copy()
305+
whole2[191, 191] = 12345.0
306+
chunk2 = whole2[64:192, 64:192]
307+
assert not cpu.NDBuffer.from_numpy_array(chunk2).all_equal(fill)
308+
# and placed first
309+
whole3 = whole.copy()
310+
whole3[64, 64] = 12345.0
311+
assert not cpu.NDBuffer.from_numpy_array(whole3[64:192, 64:192]).all_equal(fill)
312+
313+
314+
def test_all_equal_single_leading_row() -> None:
315+
"""shape[0] == 1 has no leading axis to slice; it must fall back, not break."""
316+
row = np.zeros((1, (1 << 14) * 2), dtype="float32")
317+
assert cpu.NDBuffer.from_numpy_array(row).all_equal(0.0)
318+
row2 = row.copy(); row2[0, -1] = 1
319+
assert not cpu.NDBuffer.from_numpy_array(row2).all_equal(0.0)
320+
321+
322+
def test_all_equal_nan() -> None:
323+
n = (1 << 14) * 2 + 1
324+
nans = cpu.NDBuffer.from_numpy_array(np.full(n, np.nan, dtype="float64"))
325+
assert nans.all_equal(np.nan)
326+
assert not nans.all_equal(0.0)
327+
one_nan = np.zeros(n, dtype="float64")
328+
one_nan[-1] = np.nan
329+
assert not cpu.NDBuffer.from_numpy_array(one_nan).all_equal(0.0)

0 commit comments

Comments
 (0)