Extra memory copy in uncompressed writes

[tomwhite]

With the release of Zarrista 0.1.0 (congrats!) I have been revisiting my benchmarks for measuring zero-copy in Zarr (https://github.com/tomwhite/memray-array). I'm very excited about having a true zero-copy Zarr python library.

Compressed reads and writes incur only a single extra buffer copy (as expected), and uncompressed reads don't incur any extra copies. I don't think that any Zarr python library has achieved zero copy for reads before, so this is great to see!

Only uncompressed writes incur an extra copy - which shouldn't be needed - hence this issue.

To reproduce, run the following

Details

# /// script
# requires-python = ">=3.11"
# dependencies = ["zarrista==0.1.0", "memray", "numpy"]
# ///
import tempfile
import time
from pathlib import Path

import memray
import numpy as np
from memray import FileReader
from zarrista import ArrayBuilder, ChunkGrid, DataType, FillValue
from zarrista.store import FilesystemStore

with tempfile.TemporaryDirectory() as tmp:
    profile = str(Path(tmp) / "write.bin")
    store = FilesystemStore(tmp)

    with memray.Tracker(profile, native_traces=True):
        arr = np.random.default_rng().random((5000, 5000), dtype=np.float32)
        grid = ChunkGrid.regular(arr.shape, chunk_shape=arr.shape)
        dtype = DataType.from_string(np.dtype(arr.dtype).name)
        fill_value = FillValue(np.zeros((), dtype=arr.dtype).tobytes())
        z = ArrayBuilder(grid, dtype, fill_value).create(store, "/a.zarr")
        z[:] = arr
        del arr
        del z
        time.sleep(1)

    peak = FileReader(profile).metadata.peak_memory
    print(f"peak {peak / 1e6:.1f} MB")

peak 200.7 MB

This should be around 100MB (a single allocation). Compare to Zarr with the default local store which is around 100MB.

Details

# /// script
# requires-python = ">=3.11"
# dependencies = ["zarr==3.3.0", "numcodecs>=0.16.0", "memray", "numpy"]
# ///
import tempfile
import time
from pathlib import Path

import memray
import numpy as np
import zarr
from memray import FileReader

with tempfile.TemporaryDirectory() as tmp:
    profile = str(Path(tmp) / "write.bin")
    store = str(Path(tmp) / "a.zarr")

    with memray.Tracker(profile, native_traces=True):
        arr = np.random.default_rng().random((5000, 5000), dtype=np.float32)
        z = zarr.create_array(
            store=store,
            shape=arr.shape,
            dtype=arr.dtype,
            chunks=arr.shape,
            compressors=None,
            overwrite=True,
            config={"write_empty_chunks": True},
        )
        z[:] = arr
        del arr
        del z
        time.sleep(1)

    peak = FileReader(profile).metadata.peak_memory
    print(f"peak {peak / 1e6:.1f} MB")

peak 100.8 MB

cc @LDeakin since this might be getting into Zarrs territory.

[LDeakin]

There is an extra copy at the below line if the data coming in is borrowed and the codecs are pass-through.

https://github.com/zarrs/zarrs/blob/f91f1431f324e34c3f5a02c7a5a220e6b775da38/zarrs/src/array/array_ops/array_write_ops_array.rs#L204

Scanning over zarrista, if you had an into_array_bytes in some spots rather than as_array_bytes that doesn't borrow the underlying data it might be possible to avoid this. However, I should probably make this footgun less of a problem in zarrs itself as most stores are happy to store borrowed data without copying. The in-memory store is actually the only one that really wants to own the data!

[kylebarron]

Thank you Tom for digging into this.

I asked Claude to dig into it a bit and it also suggested that this is a place where we can improve Zarrs.

Scanning over zarrista, if you had an into_array_bytes in some spots rather than as_array_bytes that doesn't borrow the underlying data it might be possible to avoid this

ArrayBytes::Fixed is defined with Cow<'a, [u8]>. Thus, for us to provide owned data, we would need to supply a Vec<u8>.

It's not possible for us to provide a Vec<u8> because the data is all owned externally by Python. It'll never be possible for us to have a zero-copy Vec<u8> from a numpy array.

The best we can do in Zarrista is have an owned Bytes object, created with Bytes::from_owner, where we have an owned Rust object that is pointing to a reference-counted Python allocation.

[kylebarron]

I spent a little time brainstorming back and forth with Claude thinking about fixes. Here's a verbose summary.

The tl;dr is I think a suggestion for ArrayBytesRaw to add a Shared(Bytes) variant.

Claude-generated text

Confirmed, and profiling narrows it to one line.

Root cause. memray with native traces puts the second 100 MB allocation inside Array::store_chunk_opt, at Bytes::from(chunk_encoded.into_owned()) — the line @LDeakin pointed at. z[:] = arr reaches it through the single-chunk fast path in store_array_subset_opt. As a control, writing the same buffer with Array.store_encoded_chunk, which skips encoding, peaks at 100.9 MB instead of 200.7 MB. So nothing on the zarrista side copies.

Both ends already speak bytes::Bytes.

• pyo3_bytes wraps the numpy buffer with Bytes::from_owner — no copy.
• FilesystemStore::set writes &[u8]. AsyncObjectStore::set does put(path, value.into()), and PutPayload::from(Bytes) is Self(Arc::new([value])) — no copy.

Only the middle cannot carry it. ArrayBytesRaw<'a> = Cow<'a, [u8]> has two states: borrow for a lifetime, or own a Vec. Neither says "a refcounted handle to a buffer someone else owns," so ownership is laundered through a Vec and the copy is the laundering.

That also means the into_array_bytes suggestion can't help here: reaching Cow::Owned from a Python buffer costs exactly the copy we are trying to remove.

Reads are zero-copy today, but the mechanism may not survive S3. The measured result is right, and the analysis agrees with it: FilesystemStore reads into a fresh Vec, which becomes Cow::Owned for free. But array_sync_readable.rs:499 does Cow::Owned(chunk_encoded.into()), and Vec::from(Bytes) reclaims the allocation only when the Bytes uniquely owns a Vec. A from_owner or from_static buffer calls slice.to_vec() unconditionally, and a Bytes sliced out of a shared buffer fails the uniqueness check. object_store's LocalStore also reads into a fresh Vec, so obstore lands in the lucky case too — but an HTTP-backed store handing back a slice of a network buffer would not.

That makes the README's note that "for obstore the numbers for S3 and local should be the same" the assumption most worth testing here. I have not measured it.

Layers, simplest first. Each is a place someone could reasonably stop:

Layer	Goal	Today
0	uncompressed whole-chunk read, local store or obstore	already zero-copy
1	uncompressed whole-chunk write, local store	copies
2	uncompressed whole-chunk write, object store	copies
3	uncompressed whole-chunk read, remote object store	probably copies (untested)
4	read one chunk out of a shard	copies
5	partial / multi-chunk writes	copies, and must — assembly needs a buffer

Compressed data sits outside this: the compressor allocates its output, which is the one copy @tomwhite already counts as correct.

What reaches which layer. A borrowed-write escape hatch on the storage traits (set_ref(&[u8]) or similar) reaches layer 1, but object_store needs an owned payload, so it stops there. Recovering ownership at the store call with Bytes::slice_ref reaches layer 2, but only by testing at runtime whether the encoded output happens to alias the input — nothing in the types keeps that true.

Layers 3 and 4 are different in kind. They are about carrying ownership in from the store through decode, and there is no bolt-on for that — the pipeline type has to hold it. It would also make layer 0 hold by construction rather than by luck:

pub enum ArrayBytesRaw<'a> {
    Borrowed(&'a [u8]),   // as today
    Owned(Vec<u8>),       // as today
    Shared(Bytes),        // refcounted, transferable to any store
}

Then store_chunk_opt stops guessing — Shared(b) => b — and the fast path becomes visible in the types rather than dependent on a runtime check. Keeping Borrowed matters: a Rust caller passing &[T] into a compressed chain pays nothing today, and forcing Bytes at the entry would add a copy. Bytes::slice would also give layer 4 zero-copy subranges of a shard.

Costs, so this isn't undersold. ArrayBytesRaw is a public type alias with 228 mentions across 49 files in zarrs/src + zarrs_codec/src, so this is breaking. It looks wide but shallow: with Deref<Target = [u8]>, From<Vec<u8>>, From<&[u8]>, From<Bytes>, and an inherent into_owned(), most call sites should compile unchanged. Bytes::from_owner requires Send + 'static, so zarrista has its own work to do on the DLPack path — I haven't verified that part compiles.

Happy to measure layer 3 against a real S3 bucket, or to prototype the type change, if either would help.