Read-modify-write behavior instead of simple write

[akshaysubr]

Here's a piece of code we're running using a DirectoryStore backed array:

subgroup[field][time_index, ...] = t.numpy() # converting pytorch tensor to numpy array

The array is chunked in a way that this should only change one single chunk and change the full chunk contents. I would expect this to only write out data. But we're seeing read-modify-write behavior which adds a lot of overhead. Here is a profile screenshot:

It seems like we only see a write without a read when the chunk doesn't already exist. But when it does exist, it is always a read-modify-write.

Would like to understand why this is happening and how to maybe get around this behavior.

cc @aleckohlhoff

[d-v-b]

I think the culprit will be somewhere in this function:

zarr-python/zarr/core.py

Line 1874 in 4132f36

if (out_is_ndarray and

In particular, if you fail to enter this branch, you are doomed to read the chunk before writing. So I would throw a breakpoint in this function and check which condition is preventing you from taking the fast path

[akshaysubr]

@d-v-b Sorry, I may be misunderstanding, but isn't that function to read data from storage rather than writing to storage?

[d-v-b]

you are right, that's the wrong function to look at. I think this is likely the routine to examine:

zarr-python/zarr/core.py

Line 2119 in 4132f36

def _process_for_setitem(self, ckey, chunk_selection, value, fields=None):

; the second branch of the conditional executes a read from storage, which might be what you are observing

[rabernat]

I would be very surprised to learn that we always read chunks when writing. The code makes it pretty clear that, if we get a full chunk selection, we can skip the read:

zarr-python/zarr/core.py

Lines 2120 to 2124 in 4132f36

	if is_total_slice(chunk_selection, self._chunks) and not fields:
	# totally replace chunk
	# optimization: we are completely replacing the chunk, so no need
	# to access the existing chunk data

@akshaysubr - can you share the shape and chunk shape of the array, plus the indexer that you're using to access it? Is time_index a scalar?

[rabernat]

Just to confirm, here is a little example that verifies that we do not always read chunks

import zarr
from zarr.storage import MemoryStore

class LoggingMemoryStore(MemoryStore):
    
    def __getitem__(self, *args, **kwargs):
        print("__getitem__")
        return super().__getitem__(*args, **kwargs)
        
store = LoggingMemoryStore()

# this calls __getitem__ once
arr = zarr.create(shape=10, chunks=10, store=store)

# these do not call __getitem__ (they fill the full chunk)
arr[:] = 1
arr[slice(0, 10)] = 2

# this DOES call __getitem__
arr[:5] = 3 

[akshaysubr]

@rabernat Yes, time_index is a scalar and the chunk size in that dimension is 1. Looks like the reason we're seeing this RMW behavior is because is_total_slice here doesn't take the fast path when using scalar indexing for 1 length chunks.

Here is an update to your little example that shows this:

import zarr
from zarr.storage import MemoryStore

class LoggingMemoryStore(MemoryStore):

    def __getitem__(self, *args, **kwargs):
        print("__getitem__")
        return super().__getitem__(*args, **kwargs)

store = LoggingMemoryStore()

# this calls __getitem__ once
arr = zarr.create(shape=10, chunks=1, store=store)

# this does not call __getitem__ (they fill the full chunk)
arr[:] = 1

# this DOES call __getitem__ (even though it fills the full chunk)
arr[2] = 2

# no call to __getitem__
arr[:5] = 3

[akshaysubr]

I guess a workaround for now is to use arr[slice(2,3)] = 2 instead of arr[2] = 2. If it is worth fixing this corner case in zarr, we can submit a PR for it.

[d-v-b]

I for one would welcome a PR, and I note that proper type annotations for is_total_slice would have revealed this bug sooner.