fix(store): copy buffers on write in MemoryStore#224
Open
d-v-b wants to merge 2 commits into
Open
Conversation
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
🤖 AI text below 🤖
Fixes F1 from the v3.3.0 evaluation — silent wrong data on the default pipeline.
The bug
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:
_merge_chunk_arrayreturns a view of the caller'svaluefor complete chunksBytesCodec._encode_syncwraps that view without copying when there is no compressorMemoryStore.setstores the buffer by referenceEvery other store serializes on write, so the alias dies at the syscall.
MemoryStoreis 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 andtest_merge_complete_chunk_returns_view_and_write_does_not_mutate_sourceboth 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.setrather than narrowing the fast path means:_merge_chunk_arrayis untouched — the pinned test still passes, and perf: use sync methods for chunk encoding / decoding zarr-developers/zarr-python#3885's win is intact forLocalStore/ZipStore/remote stores, which were never affected.value.shape == chunk_spec.shapewould 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: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
MemoryStoreactually needed; v3.2.1 paid a comparable one via_merge_chunk_array's create/fill/copy.Changes
MemoryStore.set/set_sync/set_if_not_existscopy via a new_copy_bufferhelper. Thebyte_rangebranch already copies bytes into an existing buffer, so it is left alone.ArrayLikeprotocol gainscopy()(NDArrayLikealready declares it; numpy and cupy both implement it).MemoryStoredocstring 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
mainand pass here. Full suite: 6739 passed, 0 failed.mypyclean.