developmentseed / developmentseed/zarrista

Extra memory copy in uncompressed writes

Open
#186 3 comments 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.