Avoid multipart buffer copies for in-memory buffers

[tomwhite]

When using obstore to to write to Zarr I noticed that it was using more memory than expected. For example, to write a 100MB array used a peak of around 160MB (the precise amount varies).

You can see that by running this script:

Details

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

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

with tempfile.TemporaryDirectory() as tmp:
    profile = str(Path(tmp) / "write.bin")
    store = zarr.storage.ObjectStore(
        obstore.store.LocalStore(prefix=str(Path(tmp) / "store"), mkdir=True)
    )

    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")

Compare this to using the default local store, which uses around 100MB, as expected:

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")

Sample output:

peak 100.8 MB

This is happening because multipart uploads are creating new scratch copies, even when the data being uploaded is already in memory. The fix would be to use views of the in-memory data for the multipart uploads.

[kylebarron]

Thanks! This was indeed a place for us to improve here. #771 should fix this

[tomwhite]

Thanks for fixing this Kyle! I have just confirmed the fix in memray-array by doing a local build.