Skip to content

Commit ee112b9

Browse files
authored
fix reading partial shards (#2397)
1 parent e35c4fc commit ee112b9

File tree

2 files changed

+38
-2
lines changed

2 files changed

+38
-2
lines changed

src/zarr/codecs/sharding.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -129,7 +129,7 @@ def get_chunk_slice(self, chunk_coords: ChunkCoords) -> tuple[int, int] | None:
129129
if (chunk_start, chunk_len) == (MAX_UINT_64, MAX_UINT_64):
130130
return None
131131
else:
132-
return (int(chunk_start), int(chunk_start) + int(chunk_len))
132+
return (int(chunk_start), int(chunk_len))
133133

134134
def set_chunk_slice(self, chunk_coords: ChunkCoords, chunk_slice: slice | None) -> None:
135135
localized_chunk = self._localize_chunk(chunk_coords)
@@ -203,7 +203,7 @@ def create_empty(
203203
def __getitem__(self, chunk_coords: ChunkCoords) -> Buffer:
204204
chunk_byte_slice = self.index.get_chunk_slice(chunk_coords)
205205
if chunk_byte_slice:
206-
return self.buf[chunk_byte_slice[0] : chunk_byte_slice[1]]
206+
return self.buf[chunk_byte_slice[0] : (chunk_byte_slice[0] + chunk_byte_slice[1])]
207207
raise KeyError
208208

209209
def __len__(self) -> int:

tests/test_codecs/test_sharding.py

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,42 @@ def test_sharding_partial(
118118
assert np.array_equal(data, read_data)
119119

120120

121+
@pytest.mark.parametrize("index_location", ["start", "end"])
122+
@pytest.mark.parametrize("store", ["local", "memory", "zip"], indirect=["store"])
123+
@pytest.mark.parametrize(
124+
"array_fixture",
125+
[
126+
ArrayRequest(shape=(128,) * 3, dtype="uint16", order="F"),
127+
],
128+
indirect=["array_fixture"],
129+
)
130+
def test_sharding_partial_readwrite(
131+
store: Store, array_fixture: npt.NDArray[Any], index_location: ShardingCodecIndexLocation
132+
) -> None:
133+
data = array_fixture
134+
spath = StorePath(store)
135+
a = Array.create(
136+
spath,
137+
shape=data.shape,
138+
chunk_shape=data.shape,
139+
dtype=data.dtype,
140+
fill_value=0,
141+
codecs=[
142+
ShardingCodec(
143+
chunk_shape=(1, data.shape[1], data.shape[2]),
144+
codecs=[BytesCodec()],
145+
index_location=index_location,
146+
)
147+
],
148+
)
149+
150+
a[:] = data
151+
152+
for x in range(data.shape[0]):
153+
read_data = a[x, :, :]
154+
assert np.array_equal(data[x], read_data)
155+
156+
121157
@pytest.mark.parametrize(
122158
"array_fixture",
123159
[

0 commit comments

Comments
 (0)