Skip to content

Commit f078bd3

Browse files
d-v-bclaude
andauthored
fix(sharding): reject a stored shard too short to hold its index on every path (#4475)
* refactor(sharding): test shard readers and byte slices with `is None`, not truthiness `_ShardReader` defines `__len__`, so `shard_reader or create_empty(...)` replaced a loaded zero-chunk shard with an empty reader; the outcome was equivalent but the test meant "was a shard loaded". Likewise `get_chunk_slice` returns `tuple[int, int] | None`, and `_load_full_shard_maybe` relied on `Buffer.__len__` to treat a missing and a zero-length value alike; that conflation is now spelled out. No behaviour changes. Assisted-by: ClaudeCode:claude-opus-5-5 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(sharding): reject a stored shard too short to hold its index on every path A zero-length shard was read as missing by the async full-shard read and partial write, and failed the index checksum on every other path. The index decoder now checks the index length first, so every read and partial-write path, under either pipeline, raises the same ValueError. zarrs and tensorstore also reject such a shard: https://github.com/zarrs/zarrs/blob/b8ad3ab45082b7d55ac0dd653594953a9d9e6e7d/zarrs/src/array/codec/array_to_bytes/sharding/sharding_codec.rs#L1273 https://github.com/google/tensorstore/blob/f47412a006d7b3da888f3dbb7fc001d1ac6553e1/tensorstore/kvstore/zarr3_sharding_indexed/shard_format.cc#L113 Assisted-by: ClaudeCode:claude-opus-5-5 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * docs(changes): name the fragment after PR 388 Assisted-by: ClaudeCode:claude-opus-5-5 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * docs(changes): name the fragment after upstream PR 4475 Assisted-by: ClaudeCode:claude-opus-5-5 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
1 parent 24ea324 commit f078bd3

3 files changed

Lines changed: 66 additions & 4 deletions

File tree

‎changes/4475.bugfix.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
A stored shard too short to hold its shard index, including a zero-length value, now raises a `ValueError` saying the shard is truncated or corrupt, on every read and partial-write path under either codec pipeline. Previously the asynchronous full-shard read and partial write treated a zero-length shard as missing (reading the fill value, or overwriting it), while every other path failed with a checksum mismatch or a decode error. zarrs and tensorstore also reject such a shard.

‎src/zarr/codecs/sharding.py‎

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -355,7 +355,7 @@ def create_empty(
355355

356356
def __getitem__(self, chunk_coords: tuple[int, ...]) -> Buffer:
357357
chunk_byte_slice = self.index.get_chunk_slice(chunk_coords)
358-
if chunk_byte_slice:
358+
if chunk_byte_slice is not None:
359359
return self.buf[chunk_byte_slice[0] : chunk_byte_slice[1]]
360360
raise KeyError
361361

@@ -667,6 +667,7 @@ def _decode_shard_index_sync(
667667
self, index_bytes: Buffer, chunks_per_shard: tuple[int, ...]
668668
) -> _ShardIndex:
669669
"""Decode shard index synchronously using ChunkTransform."""
670+
self._check_shard_index_size(index_bytes, chunks_per_shard)
670671
index_transform = self._get_index_chunk_transform(chunks_per_shard)
671672
index_spec = self._get_index_chunk_spec(chunks_per_shard)
672673
index_array = index_transform.decode_chunk(index_bytes, index_spec)
@@ -1403,7 +1404,8 @@ async def _encode_partial_single(
14031404
prototype=chunk_spec.prototype,
14041405
chunks_per_shard=chunks_per_shard,
14051406
)
1406-
shard_reader = shard_reader or _ShardReader.create_empty(chunks_per_shard)
1407+
if shard_reader is None:
1408+
shard_reader = _ShardReader.create_empty(chunks_per_shard)
14071409
# Use vectorized lookup for better performance. The lexicographic
14081410
# coordinate array and keys are cached, so neither is rebuilt on
14091411
# every write.
@@ -1547,6 +1549,7 @@ async def _decode_shard_index(
15471549
# the synchronous read paths cannot use but this async path still can.
15481550
if self._index_codecs_sync_capable():
15491551
return self._decode_shard_index_sync(index_bytes, chunks_per_shard)
1552+
self._check_shard_index_size(index_bytes, chunks_per_shard)
15501553
index_array = next(
15511554
iter(
15521555
await get_pipeline_class()
@@ -1581,6 +1584,23 @@ async def _encode_shard_index(self, index: _ShardIndex) -> Buffer:
15811584
raise RuntimeError("Encoding the shard index produced no bytes.")
15821585
return index_bytes
15831586

1587+
def _check_shard_index_size(
1588+
self, index_bytes: Buffer, chunks_per_shard: tuple[int, ...]
1589+
) -> None:
1590+
"""Raise if `index_bytes` is shorter than the encoded shard index.
1591+
1592+
A stored shard always holds its index, so a shorter read means the
1593+
stored value is truncated (a zero-length value included). Every read
1594+
and write path decodes the index through here, so they all reject it
1595+
alike. zarrs and tensorstore also reject such a shard.
1596+
"""
1597+
expected = self._shard_index_size(chunks_per_shard)
1598+
if len(index_bytes) < expected:
1599+
raise ValueError(
1600+
f"The stored shard is too short to hold its index: read {len(index_bytes)} "
1601+
f"bytes of index, expected {expected}. The shard is truncated or corrupt."
1602+
)
1603+
15841604
def _shard_index_size(self, chunks_per_shard: tuple[int, ...]) -> int:
15851605
return (
15861606
get_pipeline_class()
@@ -1669,7 +1689,7 @@ async def _load_full_shard_maybe(
16691689

16701690
return (
16711691
await _ShardReader.from_bytes(shard_bytes, self, chunks_per_shard)
1672-
if shard_bytes
1692+
if shard_bytes is not None
16731693
else None
16741694
)
16751695

‎tests/test_codecs/test_sharding.py‎

Lines changed: 42 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
import math
33
import pickle
44
import warnings
5-
from typing import Any, cast, get_args
5+
from typing import TYPE_CHECKING, Any, cast, get_args
66
from unittest.mock import AsyncMock
77

88
import numpy as np
@@ -40,6 +40,9 @@
4040
from ..conftest import ArrayRequest
4141
from .test_codecs import _AsyncArrayProxy, order_from_dim
4242

43+
if TYPE_CHECKING:
44+
from collections.abc import Callable
45+
4346

4447
def _reads_are_sync(store_mock: AsyncMock) -> bool:
4548
"""True when the partial-shard read for this store+pipeline goes through the
@@ -1395,3 +1398,41 @@ def test_sharding_orthogonal_set_multiple_array_dims(
13951398
expected[ix] = value.reshape(expected[ix].shape)
13961399
assert np.array_equal(a[:], expected)
13971400
assert np.array_equal(a.oindex[selection], value)
1401+
1402+
1403+
@pytest.mark.parametrize(
1404+
"pipeline_path",
1405+
[
1406+
"zarr.core.codec_pipeline.FusedCodecPipeline",
1407+
"zarr.core.codec_pipeline.BatchedCodecPipeline",
1408+
],
1409+
)
1410+
@pytest.mark.parametrize("store", ["memory", "local"], indirect=True)
1411+
@pytest.mark.parametrize("index_location", ["start", "end"])
1412+
@pytest.mark.parametrize("stored", [b"", b"abc"])
1413+
@pytest.mark.parametrize("op", ["full read", "partial read", "partial write"])
1414+
def test_sharding_truncated_shard_raises(
1415+
pipeline_path: str, store: Store, index_location: IndexLocation, stored: bytes, op: str
1416+
) -> None:
1417+
"""A stored shard shorter than its index (zero-length included) is
1418+
rejected by every read and partial-write path, under either pipeline,
1419+
rather than read as missing by some paths and failing the checksum in
1420+
others."""
1421+
with zarr.config.set({"codec_pipeline.path": pipeline_path}):
1422+
arr = zarr.create_array(
1423+
store,
1424+
shape=(8,),
1425+
chunks=(2,),
1426+
shards={"shape": (4,), "index_location": index_location},
1427+
dtype="i4",
1428+
compressors=None,
1429+
fill_value=-1,
1430+
)
1431+
zarr.core.sync.sync(store.set("c/0", default_buffer_prototype().buffer.from_bytes(stored)))
1432+
ops: dict[str, Callable[[], object]] = {
1433+
"full read": lambda: arr[0:4],
1434+
"partial read": lambda: arr[0:2],
1435+
"partial write": lambda: arr.__setitem__(slice(0, 2), 7),
1436+
}
1437+
with pytest.raises(ValueError, match="too short to hold its index"):
1438+
ops[op]()

0 commit comments

Comments
 (0)