Add API to `Store` interface to read directly into an output buffer

[TomAugspurger]

Currently, we reading uncompressed data we have the store read the data into a temporary buffer and then copy those bytes into the output buffer.

To get the best performance, we should consider adding an optional get_into API to Store. Instead of taking a prototype, this would take the actual output Buffer to read into. Stores must opt into this, for backwards compatibility, by overriding supports_get_into:

    async def get_into(
        self,
        key: str,
        out: Buffer,
        byte_range: ByteRequest | None = None,
    ) -> bool:
        raise NotImplementedError

    @property
    def supports_get_into(self) -> bool:
        """Does the store support get_into?"""
        return False

For the special case of

• uncompress data, and
• The chunk being read is a contiguous subset of the output ndarray

then the bytes on disk can be interpreted directly as an ndarray (when combined with a shape and itemsize (and maybe endianness?), and we can avoid a memcpy. Some early testing indicates that this might be worth doing. Over in https://github.com/TomAugspurger/zarr-python/blob/tom/zero-copy-alt/simple.py, I see about 7.5x higher throughput for reading uncompressed data with read_into (compared to about 2.5x higher throughput for compressed data, where this get_into optimization isn't an option).

Real world gains will probably be lower, and remote file system APIs typically don't offer APIs to read directly into a user-allocated output buffer like .readinto does.

[d-v-b]

this looks like a good idea! how would we handle the case where key denotes an object bigger than out?

[dcherian]

Is the underlying issue that we can't get numpy to take ownership of a buffer passed back from a different library? That's why the memcpy is needed?

[negin513]

This is a really neat idea. That extra memcpy costs time and bandwidth, especially for unoptimized chunks...

IIUC, most stores just return bytes, which NumPy can’t safely wrap without memcpying, since bytes are immutable....this solution is really nice because numpy just allocates the memory and store directly reads into that buffer...

[TomAugspurger]

NumPy can make make a view of outside data. And as Negin notes, for bytes the result will be immutable, which might cause issues later:

In [3]: a = np.frombuffer(b'\x01\x02', dtype=np.uint8)

In [4]: a
Out[4]: array([1, 2], dtype=uint8)

In [5]: a.flags.writeable
Out[5]: False

There's a few other subtle things going on here:

1. I think it's mentioned somewhere in Codec pipeline performance #2904, but zarr-python always allocates an out buffer (I think some APIs accept a user-provided buffer) but doesn't ever read / decode into it. So we could either remove the out buffer allocation and just use the bytes from the Store, or we can read directly into the out buffer.
2. When reading multiple chunks and when each chunk is a contiguous slice of the output buffer, there's an advantage to reading or decoding directly into the out buffer. Then you can do a single malloc for the output, and give each read / decode task a (zero-copy) slice of the out buffer to read / decode into. Then you have just a single large malloc instead of many small ones, and you avoid the many small copies into the out buffer.

[dcherian]

IMO the immutable piece is an OK tradeoff. Numpy users are somewhat used to it being required after broadcasting.

When reading multiple chunks and when each chunk is a contiguous slice of the output buffer

Good point, but this requires a very specific chunking IIRC. Specifically chunksize=-1 along the last axis for C-order and first axis for F-order.

>>> import numpy as np
>>> np.array([[1, 2, 3], [4, 5, 6]])
array([[1, 2, 3],
       [4, 5, 6]])
>>> array = np.array([[1, 2, 3], [4, 5, 6]])
>>> array.flags
  C_CONTIGUOUS : True
  F_CONTIGUOUS : False

>>> array[slice(2), :].flags
  C_CONTIGUOUS : True
  F_CONTIGUOUS : False

>>> array[:, slice(2)].flags
  C_CONTIGUOUS : False
  F_CONTIGUOUS : False

>>> array[slice(2), slice(2)].flags
  C_CONTIGUOUS : False
  F_CONTIGUOUS : False

I'm only -0.25 for the added complexity of an out kwarg. Are there workloads that would benefit greatly? Is there precedent in the Arrow/parquet ecosystem?