Skip to content

Commit ea7aee4

Browse files
michaltakacDimensionLab OSS Sync
authored andcommitted
perf(arrayroute): reuse bounded shard indexes for strided reads
GitOrigin-RevId: 3d410f53747c982d4d4f319d9595484cc80b5d69
1 parent 6475f37 commit ea7aee4

4 files changed

Lines changed: 112 additions & 3 deletions

File tree

‎packages/arrayroute/README.md‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,14 @@ lossy conversion is performed. Decoded-byte estimates assume full inner-chunk de
4343
partial codecs can do better. Transport bytes must be measured by the backend, not inferred
4444
from those estimates. A large logical shape is not evidence of large physical I/O throughput.
4545

46+
Zarrista strided reads gather only touched inner chunks. When partial shard decoding is
47+
available, the adapter reuses one shard index while gathering those chunks and discards
48+
it before moving to another shard or returning from the read. Index bytes and native
49+
bookkeeping are additional to the read-buffer allowance; no field payload is cached by
50+
this mechanism. Reads after a completed external write create fresh indexes. Concurrent
51+
source mutation still requires snapshot isolation supplied by the caller. An outer
52+
shard codec uses whole-shard accounting and the ordinary subset reader.
53+
4654
`arrayroute.statistics.ChannelMoments` accumulates or merges channel-last population
4755
statistics from bounded buffers, with optional per-channel validity masks. It keeps integer
4856
counts and centred float64 moments; missing observations do not contaminate the result.

‎packages/arrayroute/pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[project]
22
name = "arrayroute"
3-
version = "0.1.4"
3+
version = "0.1.5"
44
description = "Measured, bounded selections over scientific chunked arrays"
55
requires-python = ">=3.11"
66
license = "Apache-2.0"

‎packages/arrayroute/src/arrayroute/backends.py‎

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,9 @@ def __init__(self, array: Any):
4040
self.dtype = np.dtype(array.metadata["data_type"])
4141
chunks = array.effective_subchunk_shape or array.chunk_shape([0] * array.ndim)
4242
codecs = array.metadata["codecs"]
43-
if any(c["name"] == "sharding_indexed" for c in codecs[:-1]):
43+
outer_codec = any(c["name"] == "sharding_indexed" for c in codecs[:-1])
44+
self._partial_shards = array.is_sharded and not outer_codec
45+
if outer_codec:
4446
# A codec outside the shard requires decoding the complete shard.
4547
chunks = array.chunk_shape([0] * array.ndim)
4648
self.spec = ArraySpec(tuple(array.shape), tuple(chunks), self.dtype.itemsize)
@@ -59,10 +61,32 @@ def _contiguous(self, selection: Selection) -> np.ndarray:
5961
def _strided(self, selection: Selection) -> np.ndarray:
6062
request = plan(self.spec, selection)
6163
result = np.empty(request.output_shape, dtype=self.dtype)
64+
cache = self.array.shard_cache() if self._partial_shards else None
65+
shard_shape = self.array.chunk_shape([0] * self.array.ndim) if cache is not None else ()
66+
active_shard = None
6267

6368
def visit(axis, sources, targets, strides):
69+
nonlocal active_shard
6470
if axis == len(selection):
65-
result[tuple(targets)] = self._contiguous(tuple(sources))[tuple(strides)]
71+
if cache is None:
72+
value = self._contiguous(tuple(sources))
73+
else:
74+
starts = [s if isinstance(s, int) else s.start for s in sources]
75+
shard = tuple(p // c for p, c in zip(starts, shard_shape, strict=True))
76+
if shard != active_shard:
77+
# Retain at most one index, and never reuse it across read calls.
78+
cache.clear()
79+
active_shard = shard
80+
chunks = self.spec.chunks
81+
coordinates = [p // c for p, c in zip(starts, chunks, strict=True)]
82+
local = tuple(
83+
s % c if isinstance(s, int)
84+
else slice(s.start % c, s.start % c + s.stop - s.start)
85+
for s, c in zip(sources, chunks, strict=True)
86+
)
87+
value = np.asarray(self.array.retrieve_subchunk(
88+
coordinates, shard_cache=cache))[local]
89+
result[tuple(targets)] = value[tuple(strides)]
6690
return
6791
sel, chunk = request.selection[axis], self.spec.chunks[axis]
6892
if isinstance(sel, int):
Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
# SPDX-License-Identifier: Apache-2.0
2+
import numpy as np
3+
import pytest
4+
import zarr
5+
6+
from arrayroute import ArrayReader
7+
from arrayroute.backends import ZarristaBackend
8+
9+
zarrista = pytest.importorskip('zarrista')
10+
11+
12+
@pytest.mark.parametrize('sharded', [False, True])
13+
def test_strides_match_bits_across_partial_shards_and_missing_chunks(tmp_path, sharded):
14+
rng = np.random.default_rng(289)
15+
values = rng.integers(0, 2**32, size=(9, 13, 7), dtype=np.uint32).view(np.float32)
16+
path = tmp_path / 'field.zarr'
17+
options = {'shards': (6, 8, 4)} if sharded else {}
18+
stored = zarr.create_array(path, shape=values.shape, dtype='f4', chunks=(3, 4, 2),
19+
fill_value=-7, **options)
20+
stored[:6] = values[:6]
21+
values[6:] = -7
22+
backend = ZarristaBackend(zarrista.Array.open(zarrista.store.FilesystemStore(str(path))))
23+
for budget in [400, 1 << 20]:
24+
reader = ArrayReader(backend, buffer_bytes=budget)
25+
for selection in [
26+
(slice(1, 9, 2), slice(2, 13, 3), 3),
27+
(7, slice(None, None, 2), slice(1, 7, 3)),
28+
(slice(None, None, 4), 0, slice(None, None, 2)),
29+
(slice(0, 0, 2), slice(None), slice(None)),
30+
]:
31+
assert reader.read(selection).tobytes() == values[selection].tobytes()
32+
33+
34+
def test_shard_index_does_not_outlive_read_or_accumulate(tmp_path):
35+
path = tmp_path / 'field.zarr'
36+
values = np.arange(16 * 24, dtype='f4').reshape(16, 24)
37+
stored = zarr.create_array(path, data=values, chunks=(2, 3), shards=(4, 6))
38+
native = zarrista.Array.open(zarrista.store.FilesystemStore(str(path)))
39+
caches, sizes = [], []
40+
41+
class Observed:
42+
def __getattr__(self, name):
43+
return getattr(native, name)
44+
45+
def shard_cache(self):
46+
cache = native.shard_cache()
47+
caches.append(cache)
48+
return cache
49+
50+
def retrieve_subchunk(self, coordinates, *, shard_cache):
51+
value = native.retrieve_subchunk(coordinates, shard_cache=shard_cache)
52+
sizes.append(shard_cache.size())
53+
return value
54+
55+
reader = ArrayReader(ZarristaBackend(Observed()), buffer_bytes=1 << 20)
56+
selection = (slice(1, 16, 2), slice(0, 24, 2))
57+
assert reader.read(selection).tobytes() == values[selection].tobytes()
58+
stored[:] = values + 7
59+
assert reader.read(selection).tobytes() == (values + 7)[selection].tobytes()
60+
assert len(caches) == 2 and caches[0] is not caches[1]
61+
assert sizes and max(sizes) == 1
62+
63+
64+
def test_outer_shard_codec_strides_use_whole_shard_budget(tmp_path):
65+
from zarr.codecs import BytesCodec, ShardingCodec, ZstdCodec
66+
67+
path = tmp_path / 'field.zarr'
68+
values = np.arange(32 * 32, dtype='f4').reshape(32, 32)
69+
with pytest.warns(UserWarning, match='disables partial'):
70+
stored = zarr.create_array(path, shape=values.shape, chunks=(16, 16), dtype='f4',
71+
serializer=ShardingCodec(chunk_shape=(8, 8), codecs=[BytesCodec()]),
72+
compressors=[ZstdCodec()])
73+
stored[:] = values
74+
backend = ZarristaBackend(zarrista.Array.open(zarrista.store.FilesystemStore(str(path))))
75+
selection = (slice(1, 32, 3), slice(2, 32, 4))
76+
reader = ArrayReader(backend, buffer_bytes=2048)
77+
assert reader.read(selection).tobytes() == values[selection].tobytes()

0 commit comments

Comments
 (0)