Repository navigation
Add API to Store interface to read directly into an output buffer #3429
Description
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.
Activity
d-v-b commentedon Sep 3, 2025
this looks like a good idea! how would we handle the case where
keydenotes an object bigger thanout?dcherian commentedon Sep 5, 2025
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 commentedon Sep 5, 2025
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 commentedon Sep 5, 2025
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:
There's a few other subtle things going on here:
outbuffer (I think some APIs accept a user-provided buffer) but doesn't ever read / decode into it. So we could either remove theoutbuffer allocation and just use the bytes from the Store, or we can read directly into the out buffer.outbuffer. Then you can do a single malloc for the output, and give each read / decode task a (zero-copy) slice of theoutbuffer 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 theoutbuffer.dcherian commentedon Sep 8, 2025
IMO the immutable piece is an OK tradeoff. Numpy users are somewhat used to it being required after broadcasting.
Good point, but this requires a very specific chunking IIRC. Specifically
chunksize=-1along the last axis for C-order and first axis for F-order.I'm only -0.25 for the added complexity of an
outkwarg. Are there workloads that would benefit greatly? Is there precedent in the Arrow/parquet ecosystem?