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

[d-v-b]

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.

If you need fine-grained control over writing your array, I would not do it through xarray. IIRC when you write an xarray data structure to zarr, you can declare whether you want to write the chunks or just prepare the array / group metadata documents. In your case, I would only use xarray for writing the metadata, and take control over writing the chunks with low-level control.

More broadly, the numpy API is not stateful w.r.t. iterating over elements of an array, and therefore the array API is not stateful wr.r.t array iteration, but your data is stateful, and in a surprising way (x[0], x[0] is more than 2X slower than x[0] for your data, iirc). I worry that this detail will cause a lot of friction / surprised users, and maybe a chunked array API (where each chunk conforms to the array API without any surprises, but fetching a chunk is slow) would be a better fit for your data.

[hmaarrfk]

If you need fine-grained control over writing your array, I would not do it through xarray.

Thats fair, i'm just explaining my personal constraints.

More broadly, the numpy API is not stateful w.r.t.

The truth about memory access, and hardware, is that any kind of memory access is stateful, and no truely "random access".

Cache, pre-fetching, memory alignment all play a role in performance.

[d-v-b]

No matter what objective truth of memory is, the array API is broadly designed around the assumption that random access is fine. Your data, and IMO data stored in Zarr arrays, violates that assumption. So for these situations I think relying on the array API is questionable.

[hmaarrfk]

i'm not sure i really follow why the conversation has gone down this path.

I would really like to bring it back to the issue at hand:

I'm trying to bring up a usability issue i've noticed in dask, and now in zarr.

1. Lazy arrays are very useful for out of memory operations.
2. Reading lazy arrays can violate the assumption of fast random access.
3. I noticed what I think is a "convenience" in zarr where "checking with a zero" can be quite costly.
4. If you agree, lets continue our conversation on how to address it (the idea of moving the check to the encoder is a good idea).
5. If you disagree, i'm ok with "we do not wish to spend time on this edge case". and please close the issue.