Skip to content

Commit 50b7e01

Browse files
authored
Fix/memorystore buffer aliasing (#4157)
* fix: byte-order handling for structured dtypes in the bytes codec (#220) * fix: byte-order handling for structured dtypes in the bytes codec The bytes codec neither byte-swapped structured-dtype fields to its configured endian on encode (numpy reports byteorder '|' for void dtypes, so the top-level byteorder comparison never detected a mismatch) nor honored its endian when decoding, silently corrupting any structured data whose field byte order differed from the stored one (e.g. virtual references to external big-endian data). Encode now detects byte-order mismatches by comparing full dtypes via newbyteorder, and decode reinterprets raw bytes in the stored byte order before converting to the data type's declared byte order, so the stored layout (codec state) and the in-memory layout (array data type) are independent. Closes #4141 Assisted-by: ClaudeCode:claude-fable-5 * test: fold structured byte-order cases into existing bytes codec tests Extend test_endian's parametrization with structured dtypes and test_bytes_codec_sync_roundtrip with endian/dtype parametrization plus stored-layout and decoded-dtype assertions, instead of adding parallel test functions for the same properties. Assisted-by: ClaudeCode:claude-fable-5 * refactor: rename stored_dtype to view_dtype in BytesCodec decode The variable is the dtype used to view the raw chunk bytes (byte order from the codec's endian configuration), not a property of the stored data or of the returned buffer, which always carries the array's declared dtype. Assisted-by: ClaudeCode:claude-fable-5 * docs: note that the decode-side byte-order conversion copies the chunk Assisted-by: ClaudeCode:claude-fable-5 * fix(store): copy buffers on write in MemoryStore 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 #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 * docs: add changelog entry for MemoryStore buffer copy Assisted-by: ClaudeCode:claude-opus-4.8 * docs: correct changelog
1 parent 3265452 commit 50b7e01

4 files changed

Lines changed: 84 additions & 4 deletions

File tree

‎changes/4157.bugfix.md‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
`MemoryStore` now copies buffers as they are written, so it never retains the
2+
caller's memory. Previously an uncompressed write handed the store a zero-copy
3+
view of the user's array, and mutating that array afterwards would silently
4+
rewrite chunks already committed to the store.
5+
6+
Only `MemoryStore` is affected: stores that serialize on write, such as
7+
`LocalStore` and `ZipStore`, never aliased the caller's memory. Uncompressed
8+
writes to a `MemoryStore` are correspondingly slower, since the copy that makes
9+
the stored data independent is now actually performed; compressed writes are
10+
unchanged. Buffers supplied through the `store_dict` argument remain the
11+
caller's responsibility and are stored as-is.

‎src/zarr/core/buffer/core.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,8 @@ def __getitem__(self, key: slice) -> Self: ...
4545

4646
def __setitem__(self, key: slice, value: Any) -> None: ...
4747

48+
def copy(self) -> Self: ...
49+
4850

4951
@runtime_checkable
5052
class NDArrayLike(Protocol):

‎src/zarr/storage/_memory.py‎

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,18 @@
2626
logger = getLogger(__name__)
2727

2828

29+
def _copy_buffer(value: Buffer) -> Buffer:
30+
"""Copy `value` so the store does not retain the caller's memory.
31+
32+
Encoding a chunk can hand the store a zero-copy view of the user's array
33+
(an uncompressed write is the common case), and unlike stores that
34+
serialize on write, this one keeps whatever it is given alive in a dict.
35+
Without this copy a later mutation of the user's array would rewrite
36+
chunks already committed to the store.
37+
"""
38+
return type(value).from_array_like(value.as_array_like().copy())
39+
40+
2941
class MemoryStore(Store):
3042
"""
3143
Store for local memory.
@@ -42,6 +54,12 @@ class MemoryStore(Store):
4254
supports_writes
4355
supports_deletes
4456
supports_listing
57+
58+
Notes
59+
-----
60+
Writes copy the buffer they are given, so the store never aliases the
61+
caller's memory. Buffers passed via `store_dict` are the caller's
62+
responsibility and are stored as-is.
4563
"""
4664

4765
supports_writes: bool = True
@@ -117,7 +135,7 @@ def set_sync(self, key: str, value: Buffer) -> None:
117135
raise TypeError(
118136
f"MemoryStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead."
119137
)
120-
self._store_dict[key] = value
138+
self._store_dict[key] = _copy_buffer(value)
121139

122140
def delete_sync(self, key: str) -> None:
123141
self._check_writable()
@@ -178,13 +196,13 @@ async def set(self, key: str, value: Buffer, byte_range: tuple[int, int] | None
178196
buf[byte_range[0] : byte_range[1]] = value
179197
self._store_dict[key] = buf
180198
else:
181-
self._store_dict[key] = value
199+
self._store_dict[key] = _copy_buffer(value)
182200

183201
async def set_if_not_exists(self, key: str, value: Buffer) -> None:
184202
# docstring inherited
185203
self._check_writable()
186204
await self._ensure_open()
187-
self._store_dict.setdefault(key, value)
205+
self._store_dict.setdefault(key, _copy_buffer(value))
188206

189207
async def delete(self, key: str) -> None:
190208
# docstring inherited

‎tests/test_store/test_memory.py‎

Lines changed: 50 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
import pytest
99

1010
import zarr
11-
from zarr.core.buffer import Buffer, cpu, gpu
11+
from zarr.core.buffer import Buffer, cpu, default_buffer_prototype, gpu
1212
from zarr.errors import ZarrUserWarning
1313
from zarr.storage import GpuMemoryStore, ManagedMemoryStore, MemoryStore
1414
from zarr.testing.store import StoreTests
@@ -76,6 +76,55 @@ async def test_deterministic_size(
7676
np.testing.assert_array_equal(a[:3], 1)
7777
np.testing.assert_array_equal(a[3:], 0)
7878

79+
@pytest.mark.parametrize("method", ["set", "set_sync", "set_if_not_exists"])
80+
async def test_set_does_not_retain_caller_buffer(self, store: MemoryStore, method: str) -> None:
81+
"""Writing a buffer must not alias the caller's memory.
82+
83+
MemoryStore keeps whatever it is handed alive in a dict, so retaining
84+
the caller's buffer lets a later mutation of that buffer rewrite data
85+
already committed to the store.
86+
"""
87+
source = np.frombuffer(bytearray(b"\x01\x02\x03\x04"), dtype="B")
88+
value = cpu.Buffer.from_array_like(source)
89+
90+
if method == "set_sync":
91+
store.set_sync("k", value)
92+
else:
93+
await getattr(store, method)("k", value)
94+
95+
source[:] = 0xF # mutate the caller's memory after the write
96+
stored = await store.get("k", prototype=default_buffer_prototype())
97+
assert stored is not None
98+
assert stored.to_bytes() == b"\x01\x02\x03\x04"
99+
100+
@pytest.mark.parametrize(
101+
"pipeline",
102+
[
103+
"zarr.core.codec_pipeline.BatchedCodecPipeline",
104+
"zarr.core.codec_pipeline.FusedCodecPipeline",
105+
],
106+
)
107+
@pytest.mark.parametrize(("shape", "chunks"), [((30,), (10,)), ((8,), (4,)), ((4,), (4,))])
108+
def test_write_does_not_alias_source_array(
109+
self, pipeline: str, shape: tuple[int], chunks: tuple[int]
110+
) -> None:
111+
"""Mutating the source array after a write must not corrupt stored chunks.
112+
113+
Without compression the encoded buffer is a zero-copy view of the
114+
caller's array all the way down to the store, so this covers both the
115+
single-chunk and multi-chunk write paths.
116+
"""
117+
with zarr.config.set({"codec_pipeline.path": pipeline}):
118+
array = zarr.create_array(
119+
store=MemoryStore(), shape=shape, chunks=chunks, dtype="i4", compressors=None
120+
)
121+
source = np.arange(shape[0], dtype="i4")
122+
expected = source.copy()
123+
array[:] = source
124+
source[:] = -1
125+
126+
np.testing.assert_array_equal(array[:], expected)
127+
79128
# --- byte-range-write tests: disabled ---
80129
# Byte-range-write support (set_range / set_range_sync / SupportsSetRange)
81130
# was removed from this PR pending a decision on the store interface. These

0 commit comments

Comments
 (0)