Skip to content

fix(store): copy buffers on write in MemoryStore#224

Open
d-v-b wants to merge 2 commits into
mainfrom
fix/memorystore-buffer-aliasing
Open

fix(store): copy buffers on write in MemoryStore#224
d-v-b wants to merge 2 commits into
mainfrom
fix/memorystore-buffer-aliasing

Conversation

@d-v-b

@d-v-b d-v-b commented Jul 16, 2026

Copy link
Copy Markdown
Owner

🤖 AI text below 🤖

Fixes F1 from the v3.3.0 evaluation — silent wrong data on the default pipeline.

The bug

z = zarr.create_array(store=MemoryStore(), shape=(8,), chunks=(4,), dtype="i4", compressors=None)
data = np.arange(8, dtype="i4")
z[:] = data
data[:] = -1        # mutate the source *after* the write
z[:]                # -> [-1 -1 -1 -1 -1 -1 -1 -1]   the stored chunks changed

No error, no warning. Reproduces on both pipelines.

Root cause

Three zero-copy links in a row, ending in a store that retains what it is handed:

  1. _merge_chunk_array returns a view of the caller's value for complete chunks
  2. BytesCodec._encode_sync wraps that view without copying when there is no compressor
  3. MemoryStore.set stores the buffer by reference

Every other store serializes on write, so the alias dies at the syscall. MemoryStore is the one that keeps the caller's pages alive in a dict.

Why fix it here

The view itself is sound, and the existing contract around it is deliberate — _merge_chunk_array's docstring and test_merge_complete_chunk_returns_view_and_write_does_not_mutate_source both pin it. But that contract only reasoned about zarr's callers treating the view as read-only. The invariant a store needs is stronger: don't retain. So the defect is at the store boundary, and that is where this fixes it.

Copying in MemoryStore.set rather than narrowing the fast path means:

  • _merge_chunk_array is untouched — the pinned test still passes, and perf: use sync methods for chunk encoding / decoding zarr-developers/zarr-python#3885's win is intact for LocalStore/ZipStore/remote stores, which were never affected.
  • The single-chunk case is fixed too. Narrowing the fast path back to v3.2.1's value.shape == chunk_spec.shape would only shrink the blast radius: v3.2.1 aliased on single-chunk writes and silently corrupted just the same. This fixes both.

Cost

Measured on a 34 MB write to MemoryStore:

before after
uncompressed 22.6 ms 39.4 ms
compressed 52.0 ms 52.6 ms

The copy is only material on the uncompressed path — which is exactly the path that was aliasing. Compressed writes hand over a freshly-encoded, smaller buffer and are unchanged. Worth being explicit that part of zarr-developers#3885's uncompressed speedup came from eliminating a copy MemoryStore actually needed; v3.2.1 paid a comparable one via _merge_chunk_array's create/fill/copy.

Changes

  • MemoryStore.set / set_sync / set_if_not_exists copy via a new _copy_buffer helper. The byte_range branch already copies bytes into an existing buffer, so it is left alone.
  • ArrayLike protocol gains copy() (NDArrayLike already declares it; numpy and cupy both implement it).
  • MemoryStore docstring now states the ownership semantics.

Tests

  • test_set_does_not_retain_caller_buffer — store-level, across all three write methods.
  • test_write_does_not_alias_source_array — end-to-end, across both pipelines and the multi-chunk / multi-chunk-exact / single-chunk shapes.

Both fail on main and pass here. Full suite: 6739 passed, 0 failed. mypy clean.

d-v-b added 2 commits July 16, 2026 22:48
Encoding an uncompressed chunk hands the store a zero-copy view of the
caller's array, and MemoryStore keeps whatever it is given alive in a
dict rather than serializing it. Mutating the source array after a write
therefore rewrote chunks already committed to the store, silently.

Stores that serialize on write (LocalStore, ZipStore, remote stores) are
unaffected, so they keep the full benefit of zarr-developers#3885. Only MemoryStore pays
the copy, and only where it was aliasing to begin with: an uncompressed
34 MB write goes from ~23 ms to ~39 ms, while compressed writes are
unchanged.

Copying at the store boundary rather than narrowing the fast path in
_merge_chunk_array also fixes the single-chunk case, which aliased in
v3.2.1 too.

Assisted-by: ClaudeCode:claude-opus-4.8
Assisted-by: ClaudeCode:claude-opus-4.8
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant