zero check in encoding less than ideal for lazy arrays

[hmaarrfk]

The zero check in the encoding path is less than ideal

zarr-python/src/zarr/core/codec_pipeline.py

Line 416 in b873691

if not chunk_spec.config.write_empty_chunks and chunk_array.all_equal(

For lazy arrays, it eventually calls:

1. np.broadcast_arrays
2. np.array_equal

where the latter can instantiate a full array in memory.

The problem is that some lazy array implementation will not "keep this instantiation of the interediate full memory array" in memory for the the operation 2 lines down, that will eventually write the array.

So if the instantiation of the lazy array is costly, we are:

1. Instantiating it once for checking for zeros.
2. Instantiating a second time to write it to the store.

I would love to avoid this 'double cost' in a "pythonic way".

I would have loved to suggest
My proposal would be to refactor the "zero check" to use np.array_equiv https://numpy.org/devdocs/reference/generated/numpy.array_equiv.html#numpy.array_equiv

>>> import numpy as np
>>> np.array_equiv(np.zeros((3, 4)), 0)
True

but looking at the source, it seems to recreate the same problem https://github.com/numpy/numpy/blob/main/numpy/_core/numeric.py#L2552-L2598

I'm not sure what the right solution is.

One could cache the local array.

I notice that this operation is actually considered inefficient by the original writers leaving a note like

        # use array_equal to obtain equal_nan=True functionality
        # Since fill-value is a scalar, isn't there a faster path than allocating a new array for fill value
        # every single time we have to write data?

---

Where this is a problem for me:

The numpy broadcast operation is actually really difficult to implement if you want to ensure high performance based on the chunks of the array. My experience with dask from 7 years ago reminds me that it was also difficult for them to implement efficient rechunking.

So implementing the np.broadcast_arrays just feels like "alot of work"....

zarr-python/src/zarr/core/buffer/core.py

Line 543 in b873691

_data, other = np.broadcast_arrays(self._data, other)

The zero check feels like it makes sense.... but its like "alot of work to do" for some array backends.

---

Suggested patch:

diff --git a/src/zarr/core/buffer/core.py b/src/zarr/core/buffer/core.py
index f0d01566..54a2c09f 100644
--- a/src/zarr/core/buffer/core.py
+++ b/src/zarr/core/buffer/core.py
@@ -540,6 +540,11 @@ class NDBuffer:
         # use array_equal to obtain equal_nan=True functionality
         # Since fill-value is a scalar, isn't there a faster path than allocating a new array for fill value
         # every single time we have to write data?
+
+        # The operation array_equal operation below effectively will force the array
+        # into memory.
+        # So we cache it here.
+        self._data = np.asarray(self._data)
         _data, other = np.broadcast_arrays(self._data, other)
         return np.array_equal(
             self._data,

[d-v-b]

I see a few options for working around this limitation:

• improve our user-facing APIs for writing single chunks / shards. That means discouraging my_zarr_array[:] = my_array_like and instead promoting APIs that give more control over the chunk encoding process. Then you could simply encode the chunks you want to encode without extra indirection.
• make the boundary between the user's array-like object and zarr's copy of that array-like object smarter. I think zarr could require that we will always make a chunk-sized copy of a user array using a zarr-owned memory buffer (cpu or gpu memory). That means we fill this buffer by iterating over elements of the user's array, and this loop would be a natural place to inject the "is it fill value" check. But that requires writing routines that are lower-level than pure python.

[hmaarrfk]

improve our user-facing APIs

One of the aspects that has stopped me from pushing forward with the recommendations you made in #3628 is that I currently use zarr through xarray. So i'll have to ensure I understand how to configure zarr through the xarray lens, and I just haven't found the time to do that. I'm also not concentrating on the fully disastrous usecase that I outline below. So for me, the cost is at most "2x" read, and not "100x" read. For the case I outline below, I'll likely just flag "write 0 chunks", and I don't think we will be storing with zarr (at least to start).

chunk-sized copy

this is nice, and may be good compromise, I am worried about "unecessary memory allocation and deallocation" for performance reasons.

---

To help make the problem feel "real", I wanted to provide a minimum reproducible example of the codepath i'm concerned with.

I created a "minimal" array-like class that mimics an array that doesn't support random access.

The array is pretty "inexpensive" to read once, and in order, but once you try to revisit data that has already been written, its pretty disastrous. Fun game, let me know when you stop waiting, I pressed alt-tab like 4 seconds in. My attention span is pretty short....

# Copy past this directly into your favorite shell
test -f bbb_sunflower_2160p_60fps_normal.mp4.zip || \
   wget https://download.blender.org/demo/movies/BBB/bbb_sunflower_2160p_60fps_normal.mp4.zip
test -f bbb_sunflower_2160p_60fps_normal.mp4 || \
   unzip bbb_sunflower_2160p_60fps_normal.mp4.zip

python - <<'EOF'
import imageio.v3 as iio
from tqdm import tqdm

with iio.imopen("bbb_sunflower_2160p_60fps_normal.mp4", "r", plugin="pyav") as video:
    for i in tqdm(range(200), desc="Reading frames one a time"):
        video.read(index=i)

with iio.imopen("bbb_sunflower_2160p_60fps_normal.mp4", "r", plugin="pyav") as video:
    for i in tqdm(range(200), desc="Reading each frame twice"):
        video.read(index=i)
        video.read(index=i)
EOF

Array API version

import imageio.v3 as iio
from tqdm import tqdm
import numpy as np

class VideoFrame:
    def __init__(self, video, frame_index):
        self._video = video
        self._frame_index = frame_index

    def __array__(self, dtype=None):
        assert dtype is None
        return self._video.read(index=self._frame_index)

class Video:
    def __init__(self, filename):
        self.filename = filename
        self._video = iio.imopen(filename, "r", plugin="pyav")
        self._index = None

    def __getitem__(self, index):
        return VideoFrame(self._video, index)

if __name__ == "__main__":
    video = Video("bbb_sunflower_2160p_60fps_normal.mp4")

    for i in tqdm(range(200), desc="Reading frames one a time"):
        np.asarray(video[i])

    for i in tqdm(range(200), desc="Reading each frame twice"):
        np.asarray(video[i])  # zero check
        np.asarray(video[i])  # data writing

Reading frames one a time: 100%|██████████| 200/200 [00:02<00:00, 90.87it/s]
Reading each frame twice: 100%|███████████| 200/200 [03:34<00:00,  1.07s/it]

xref: imageio/imageio#933