developmentseed / developmentseed/zarrista
Extra memory copy in uncompressed writes
- Dominant language
- Python
- Stars
- 45
- Forks
- 2
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 7
Description
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
```python
# /// 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.
```python
# /// 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.
Contributor guide
No contributing guide indexed for this repository
Research direction
Run the supplied memray reproducer with zarrista, focusing on the ArrayBuilder-created array, the z[:] = arr write, and FilesystemStore. Trace the uncompressed write path and compare peak memory with the Zarr baseline; done means the zarrista run is near 100 MB with no unnecessary extra buffer copy.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, rust
- Domain
- backend, performance
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 68/100