Repository navigation
Avoid multipart buffer copies for in-memory buffers #767
Description
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.
Activity
kylebarron commentedon Aug 20, 2026
Thanks! This was indeed a place for us to improve here. #771 should fix this
tomwhite commentedon Aug 21, 2026
Thanks for fixing this Kyle! I have just confirmed the fix in memray-array by doing a local build.