diff --git a/virtualizarr/manifests/store.py b/virtualizarr/manifests/store.py index 6933bd90d..1076c7716 100644 --- a/virtualizarr/manifests/store.py +++ b/virtualizarr/manifests/store.py @@ -1,9 +1,10 @@ from __future__ import annotations +import asyncio import warnings -from collections.abc import AsyncGenerator, Iterable +from collections.abc import AsyncGenerator, Iterable, Sequence from dataclasses import dataclass -from typing import TYPE_CHECKING, Literal, TypeAlias +from typing import TYPE_CHECKING, Literal, NamedTuple, TypeAlias from urllib.parse import urlparse from obspec_utils.registry import ObjectStoreRegistry @@ -44,6 +45,76 @@ class StoreRequest: """The key within the store to request.""" +@dataclass +class _ChunkRef: + """A single chunk request resolved to an absolute byte range in a source file.""" + + store: ObjectStore + path: str + start: int + end: int + + +class _Member(NamedTuple): + """One chunk request within a file group.""" + + request_index: int + """Position of this request in the original ``get_many`` ``requests``.""" + start: int + """Start of the chunk's byte range, absolute within the source file.""" + end: int + """End (exclusive) of the chunk's byte range, absolute within the source file.""" + + +@dataclass +class _FileGroup: + """A set of chunk requests that all read from the same source file.""" + + store: ObjectStore + path: str + members: list[_Member] + + +def _coalesce_members( + members: list[_Member], *, max_gap: int, max_bytes: int +) -> list[list[_Member]]: + """Group members (all in the same file) into runs, each served by one read. + + Two members join the same run when the gap between them is at most + ``max_gap`` *and* the resulting run span stays at most ``max_bytes``. + + The two knobs play distinct roles: + + - ``max_gap`` decides *whether* to bridge a gap between references, i.e. how + much unwanted data coalescing is allowed to pull in. This is what governs + strided access (a large gap bridges the stride of a 2D grid and over-reads; + ``0`` merges only adjacent references and never over-reads). + - ``max_bytes`` caps the *size* of any single read, so a long run of + (near-)adjacent members is split into several bounded reads that the caller + fetches concurrently. It bounds per-read memory and lets a large contiguous + span be fetched over several parallel connections rather than one. It does + *not* affect over-read. + + Runs are returned sorted by start offset and cover every member exactly + once. A member larger than ``max_bytes`` still forms its own (single) run. + """ + ordered = sorted(members, key=lambda member: member.start) + runs: list[list[_Member]] = [] + run_start = run_end = 0 + for member in ordered: + if ( + runs + and member.start - run_end <= max_gap + and max(run_end, member.end) - run_start <= max_bytes + ): + runs[-1].append(member) + run_end = max(run_end, member.end) + else: + runs.append([member]) + run_start, run_end = member.start, member.end + return runs + + def get_store_prefix(url: str) -> str: """ Get a logical prefix to use for a url in an ObjectStoreRegistry @@ -112,7 +183,12 @@ def __eq__(self, value: object): NotImplementedError def __init__( - self, group: ManifestGroup, *, registry: ObjectStoreRegistry | None = None + self, + group: ManifestGroup, + *, + registry: ObjectStoreRegistry | None = None, + coalesce_max_gap_bytes: int = 0, + coalesce_max_bytes: int = 8 * 1024 * 1024, ) -> None: """Instantiate a new ManifestStore. @@ -123,6 +199,27 @@ def __init__( registry A registry mapping the URL scheme and netloc to [ObjectStore][obstore.store.ObjectStore] instances, allowing [ManifestStores][virtualizarr.manifests.ManifestStore] to read from different [ObjectStore][obstore.store.ObjectStore] instances. + coalesce_max_gap_bytes + When multiple chunks requested together (via ``get_many``) refer to the + same source file, virtual references separated by at most this many bytes + are coalesced into a single, larger read. Defaults to ``0``, which merges + only references that are exactly adjacent in the file - a pure win, since + no unwanted bytes are read. A larger value additionally bridges gaps up to + that size, trading some wasted bytes for fewer requests; this can backfire + for strided access patterns (e.g. a 2D spatial box, whose chunks are + contiguous along one axis but far apart along another), where bridging the + stride pulls in the intervening chunks. See + [ManifestStore.get_many][virtualizarr.manifests.ManifestStore.get_many]. + coalesce_max_bytes + Upper bound on the size of a single coalesced read. A run of adjacent (or, + with a non-zero gap, near-adjacent) references longer than this is split + into several reads that are fetched concurrently, bounding per-read memory + and letting a large contiguous span be pulled over several parallel + connections instead of one. This is the equivalent of icechunk's + ``ideal_concurrent_request_size``, provided here because obstore does not + split a single ranged read on its own. It does not affect over-read (that + is ``coalesce_max_gap_bytes``), and is inert unless a query coalesces into a + run larger than this. Defaults to 8 MiB. """ if not isinstance(group, ManifestGroup): @@ -131,6 +228,8 @@ def __init__( super().__init__(read_only=True) self._registry = ObjectStoreRegistry() if registry is None else registry self._group = group + self._coalesce_max_gap_bytes = coalesce_max_gap_bytes + self._coalesce_max_bytes = coalesce_max_bytes def __str__(self) -> str: return f"ManifestStore(group={self._group}, registry={self._registry})" @@ -182,16 +281,30 @@ async def get( entry = manifest.get_entry(chunk_indexes) if entry is None: return None - path = entry["path"] - offset = entry["offset"] - length = entry["length"] - # Get the configured object store instance that matches the path - store, path_after_prefix = self._registry.resolve(path) + store, path_in_store = self._resolve_store_and_path(entry["path"]) + # Transform the input byte range to account for the chunk location in the file + byte_range = _transform_byte_range( + byte_range, + chunk_start=entry["offset"], + chunk_end_exclusive=entry["offset"] + entry["length"], + ) + + # Actually get the bytes + bytes = await store.get_range_async( + path_in_store, + start=byte_range.start, + end=byte_range.end, + ) + return prototype.buffer.from_bytes(bytes) # type: ignore[arg-type] + + def _resolve_store_and_path(self, path: str) -> tuple[ObjectStore, str]: + """Resolve a manifest entry ``path`` to its ObjectStore and the path + within that store (with any store prefix stripped).""" + store, _ = self._registry.resolve(path) if not store: raise ValueError( f"Could not find a store to use for {path} in the store registry" ) - path_in_store = urlparse(path).path if hasattr(store, "prefix") and store.prefix: prefix = str(store.prefix).lstrip("/") @@ -199,20 +312,39 @@ async def get( prefix = urlparse(store.url).path.lstrip("/") else: prefix = "" - path_in_store = path_in_store.lstrip("/").removeprefix(prefix).lstrip("/") - # Transform the input byte range to account for the chunk location in the file - chunk_end_exclusive = offset + length - byte_range = _transform_byte_range( - byte_range, chunk_start=offset, chunk_end_exclusive=chunk_end_exclusive - ) + return store, path_in_store.lstrip("/").removeprefix(prefix).lstrip("/") - # Actually get the bytes - bytes = await store.get_range_async( - path_in_store, - start=byte_range.start, - end=byte_range.end, + def _resolve_chunk_ref( + self, key: str, byte_range: ByteRequest | None + ) -> _ChunkRef | None: + """Resolve a chunk key to an absolute byte range within its source file. + + Returns ``None`` when the key is not a plain, manifest-backed chunk that + can participate in coalescing - i.e. metadata documents, inlined chunks, + missing chunks, or group keys. Those are handled individually by ``get``. + """ + node, suffix = _get_deepest_group_or_array(self._group, key) + if suffix.endswith( + ("zarr.json", ".zattrs", ".zgroup", ".zarray", ".zmetadata") + ) or isinstance(node, ManifestGroup): + return None + manifest = node.manifest + separator: Literal[".", "/"] = getattr( + node.metadata.chunk_key_encoding, "separator", "." ) - return prototype.buffer.from_bytes(bytes) # type: ignore[arg-type] + chunk_indexes = parse_manifest_index(key, separator, expand_pattern=True) + if chunk_indexes in manifest._inlined: + return None + entry = manifest.get_entry(chunk_indexes) + if entry is None: + return None + store, path_in_store = self._resolve_store_and_path(entry["path"]) + rng = _transform_byte_range( + byte_range, + chunk_start=entry["offset"], + chunk_end_exclusive=entry["offset"] + entry["length"], + ) + return _ChunkRef(store=store, path=path_in_store, start=rng.start, end=rng.end) async def get_partial_values( self, @@ -223,6 +355,95 @@ async def get_partial_values( # TODO: Implement using private functions from the upstream Zarr obstore integration raise NotImplementedError + async def get_many( + self, + requests: Sequence[tuple[str, ByteRequest | None] | str], + *, + prototype: BufferPrototype, + ) -> AsyncGenerator[Sequence[tuple[int, Buffer | None]], None]: + """Retrieve many chunks at once, coalescing reads from the same file. + + Overrides [Store.get_many][zarr.abc.store.Store.get_many]. Requested + chunks are resolved through the manifests to ``(source file, byte + range)`` and grouped by source file. Within each file, references closer + together than ``coalesce_max_gap_bytes`` are coalesced into runs; each run + is then split into reads of at most ``coalesce_max_bytes`` and served by a + ranged read that is sliced back into per-chunk buffers. The gap controls + how much unwanted data coalescing may pull in - ``0`` (the default) merges + only adjacent references and never over-reads, which matters for strided + access such as a 2D spatial box - while the size cap only bounds how large + any single read may get (splitting big contiguous runs into parallel + reads); it does not affect over-read. + This is the same technique object_store and async-tiff use to read many + tiles efficiently, applied here to virtual chunk references. + + Keys that are not plain manifest-backed chunks (metadata documents, + inlined chunks, or missing chunks) are served individually via ``get``. + Results are yielded as ``(request_index, Buffer | None)`` batches in + completion order, per the ``Store.get_many`` contract. + """ + # Local import so importing this module doesn't require the zarr config. + from zarr.core.config import config + + # Partition requests into coalescable chunk reads (grouped by source + # file) and everything else (served one-by-one via ``get``). + groups: dict[tuple[int, str], _FileGroup] = {} + singletons: list[tuple[int, str, ByteRequest | None]] = [] + for index, request in enumerate(requests): + key, byte_range = (request, None) if isinstance(request, str) else request + ref = self._resolve_chunk_ref(key, byte_range) + if ref is None: + singletons.append((index, key, byte_range)) + continue + group = groups.setdefault( + (id(ref.store), ref.path), _FileGroup(ref.store, ref.path, []) + ) + group.members.append(_Member(index, ref.start, ref.end)) + + semaphore = asyncio.Semaphore(config.get("async.concurrency")) + + async def fetch_run( + store: ObjectStore, path: str, run: list[_Member] + ) -> Sequence[tuple[int, Buffer | None]]: + # One ranged read spanning the whole run, sliced back per chunk. + run_start = run[0].start + run_end = max(member.end for member in run) + async with semaphore: + data = await store.get_range_async(path, start=run_start, end=run_end) + view = memoryview(data) # type: ignore[arg-type] + return [ + ( + member.request_index, + prototype.buffer.from_bytes( + view[member.start - run_start : member.end - run_start] + ), + ) + for member in run + ] + + async def fetch_single( + index: int, key: str, byte_range: ByteRequest | None + ) -> Sequence[tuple[int, Buffer | None]]: + async with semaphore: + buffer = await self.get(key, prototype, byte_range) + return ((index, buffer),) + + tasks = [ + asyncio.ensure_future(fetch_run(group.store, group.path, run)) + for group in groups.values() + for run in _coalesce_members( + group.members, + max_gap=self._coalesce_max_gap_bytes, + max_bytes=self._coalesce_max_bytes, + ) + ] + tasks += [ + asyncio.ensure_future(fetch_single(index, key, byte_range)) + for index, key, byte_range in singletons + ] + for coro in asyncio.as_completed(tasks): + yield await coro + async def exists(self, key: str) -> bool: # docstring inherited raise NotImplementedError diff --git a/virtualizarr/tests/test_manifests/test_store_get_many.py b/virtualizarr/tests/test_manifests/test_store_get_many.py new file mode 100644 index 000000000..38ae9552d --- /dev/null +++ b/virtualizarr/tests/test_manifests/test_store_get_many.py @@ -0,0 +1,292 @@ +from __future__ import annotations + +import numpy as np +import pytest +from obspec_utils.registry import ObjectStoreRegistry +from zarr.core.buffer import default_buffer_prototype + +from virtualizarr.manifests import ( + ChunkManifest, + ManifestArray, + ManifestGroup, + ManifestStore, +) +from virtualizarr.manifests.store import _coalesce_members, _Member +from virtualizarr.manifests.utils import create_v3_array_metadata +from virtualizarr.tests import requires_obstore + +# 16 distinct bytes, enough for four contiguous 4-byte chunks. +FILE_BYTES = bytes(range(1, 17)) +_CODECS = [{"configuration": {"endian": "little"}, "name": "bytes"}] + + +def _build_store( + tmpdir, + *, + files: dict[str, bytes] | None = None, + entries: dict[str, dict] | None = None, + coalesce_max_gap_bytes: int = 1024 * 1024, + coalesce_max_bytes: int = 8 * 1024 * 1024, +) -> ManifestStore: + """Build a ManifestStore over a LocalStore with one array ``foo`` (4x4, 2x2 chunks).""" + import obstore as obs + + store = obs.store.LocalStore() + prefix = "file://" + files = files or {f"{tmpdir}/data.tmp": FILE_BYTES} + for filepath, data in files.items(): + obs.put(store, filepath, data) + + if entries is None: + fp = next(iter(files)) + entries = { + "0.0": {"path": f"{prefix}/{fp}", "offset": 0, "length": 4}, + "0.1": {"path": f"{prefix}/{fp}", "offset": 4, "length": 4}, + "1.0": {"path": f"{prefix}/{fp}", "offset": 8, "length": 4}, + "1.1": {"path": f"{prefix}/{fp}", "offset": 12, "length": 4}, + } + manifest = ChunkManifest(entries=entries) + metadata = create_v3_array_metadata( + shape=(4, 4), + chunk_shape=(2, 2), + data_type=np.dtype("int32"), + codecs=_CODECS, + chunk_key_encoding={"name": "default", "separator": "."}, + fill_value=0, + ) + marr = ManifestArray(metadata=metadata, chunkmanifest=manifest) + group = ManifestGroup(arrays={"foo": marr}) + registry = ObjectStoreRegistry({prefix: store}) + return ManifestStore( + group=group, + registry=registry, + coalesce_max_gap_bytes=coalesce_max_gap_bytes, + coalesce_max_bytes=coalesce_max_bytes, + ) + + +async def _collect(store: ManifestStore, requests) -> dict[int, object]: + """Drain get_many into an index -> buffer mapping, asserting each index appears once.""" + collected: dict[int, object] = {} + async for batch in store.get_many(requests, prototype=default_buffer_prototype()): + for index, buffer in batch: + assert index not in collected, "each request must be reported exactly once" + collected[index] = buffer + return collected + + +class _RecordingStore: + """Proxy that records each ranged read (``get_range_async``) and forwards + everything else. + + Installed by monkeypatching ``ManifestStore._resolve_store_and_path`` so we + can see exactly which byte ranges get_many issues after coalescing. + """ + + def __init__(self, inner, log: list[dict]) -> None: + self._inner = inner + self._log = log + + def __getattr__(self, name): + return getattr(self._inner, name) + + async def get_range_async(self, path, *, start, end): + self._log.append({"path": path, "start": start, "end": end}) + return await self._inner.get_range_async(path, start=start, end=end) + + +@pytest.fixture() +def record_range_calls(monkeypatch): + """Record every get_ranges_async call ManifestStore.get_many issues.""" + log: list[dict] = [] + original = ManifestStore._resolve_store_and_path + proxies: dict[int, _RecordingStore] = {} + + def patched(self, path): + inner, path_in_store = original(self, path) + proxy = proxies.get(id(inner)) + if proxy is None: + proxy = _RecordingStore(inner, log) + proxies[id(inner)] = proxy + return proxy, path_in_store + + monkeypatch.setattr(ManifestStore, "_resolve_store_and_path", patched) + return log + + +@requires_obstore +class TestGetMany: + pytestmark = pytest.mark.asyncio + + async def test_matches_individual_get(self, tmpdir): + store = _build_store(tmpdir) + keys = ["foo/c.0.0", "foo/c.0.1", "foo/c.1.0", "foo/c.1.1"] + collected = await _collect(store, keys) + + assert set(collected) == set(range(len(keys))) + for index, key in enumerate(keys): + expected = await store.get(key, prototype=default_buffer_prototype()) + assert collected[index].to_bytes() == expected.to_bytes() + # sanity: the four chunks are the four consecutive 4-byte slices + assert collected[0].to_bytes() == FILE_BYTES[0:4] + assert collected[3].to_bytes() == FILE_BYTES[12:16] + + async def test_accepts_key_range_tuples(self, tmpdir): + from zarr.abc.store import RangeByteRequest + + store = _build_store(tmpdir) + requests = [ + "foo/c.0.0", + ("foo/c.0.1", None), + ("foo/c.1.0", RangeByteRequest(1, 3)), # middle two bytes of that chunk + ] + collected = await _collect(store, requests) + assert collected[0].to_bytes() == FILE_BYTES[0:4] + assert collected[1].to_bytes() == FILE_BYTES[4:8] + assert collected[2].to_bytes() == FILE_BYTES[8:12][1:3] + + async def test_missing_chunk_is_none(self, tmpdir): + prefix = "file://" + fp = f"{tmpdir}/data.tmp" + # omit chunk "1.1" so it has no manifest entry + entries = { + "0.0": {"path": f"{prefix}/{fp}", "offset": 0, "length": 4}, + "0.1": {"path": f"{prefix}/{fp}", "offset": 4, "length": 4}, + "1.0": {"path": f"{prefix}/{fp}", "offset": 8, "length": 4}, + } + store = _build_store(tmpdir, entries=entries) + collected = await _collect(store, ["foo/c.0.0", "foo/c.1.1"]) + assert collected[0].to_bytes() == FILE_BYTES[0:4] + assert collected[1] is None # missing -> None, not omitted + + async def test_metadata_key_served_individually(self, tmpdir): + store = _build_store(tmpdir) + collected = await _collect(store, ["foo/zarr.json", "foo/c.0.0"]) + assert collected[0] is not None # zarr.json metadata document + assert b'"zarr_format"' in collected[0].to_bytes() + assert collected[1].to_bytes() == FILE_BYTES[0:4] + + async def test_empty_requests(self, tmpdir): + store = _build_store(tmpdir) + collected = await _collect(store, []) + assert collected == {} + + async def test_coalesces_contiguous_into_one_read(self, tmpdir, record_range_calls): + store = _build_store(tmpdir) + keys = ["foo/c.0.0", "foo/c.0.1", "foo/c.1.0", "foo/c.1.1"] + collected = await _collect(store, keys) + + # the four contiguous chunks are served by a single ranged read [0, 16) + assert [(c["start"], c["end"]) for c in record_range_calls] == [(0, 16)] + # and the bytes were sliced back to the right chunks + assert [collected[i].to_bytes() for i in range(4)] == [ + FILE_BYTES[0:4], + FILE_BYTES[4:8], + FILE_BYTES[8:12], + FILE_BYTES[12:16], + ] + + async def test_groups_requests_by_file(self, tmpdir, record_range_calls): + prefix = "file://" + fp_a = f"{tmpdir}/a.tmp" + fp_b = f"{tmpdir}/b.tmp" + entries = { + "0.0": {"path": f"{prefix}/{fp_a}", "offset": 0, "length": 4}, + "0.1": {"path": f"{prefix}/{fp_a}", "offset": 4, "length": 4}, + "1.0": {"path": f"{prefix}/{fp_b}", "offset": 0, "length": 4}, + "1.1": {"path": f"{prefix}/{fp_b}", "offset": 4, "length": 4}, + } + store = _build_store( + tmpdir, files={fp_a: FILE_BYTES, fp_b: FILE_BYTES}, entries=entries + ) + await _collect(store, ["foo/c.0.0", "foo/c.0.1", "foo/c.1.0", "foo/c.1.1"]) + + # one read per distinct source file, each spanning that file's two chunks + assert len(record_range_calls) == 2 + for call in record_range_calls: + assert (call["start"], call["end"]) == (0, 8) + + async def test_gap_larger_than_max_gap_is_not_merged( + self, tmpdir, record_range_calls + ): + prefix = "file://" + fp = f"{tmpdir}/data.tmp" + # two chunks 996 bytes apart in the file + entries = { + "0.0": {"path": f"{prefix}/{fp}", "offset": 0, "length": 4}, + "0.1": {"path": f"{prefix}/{fp}", "offset": 1000, "length": 4}, + } + store = _build_store( + tmpdir, + files={fp: bytes(2048)}, + entries=entries, + coalesce_max_gap_bytes=8, # < 996 gap -> keep separate + ) + await _collect(store, ["foo/c.0.0", "foo/c.0.1"]) + assert sorted((c["start"], c["end"]) for c in record_range_calls) == [ + (0, 4), + (1000, 1004), + ] + + async def test_max_bytes_caps_a_merge(self, tmpdir, record_range_calls): + prefix = "file://" + fp = f"{tmpdir}/data.tmp" + entries = { + "0.0": {"path": f"{prefix}/{fp}", "offset": 0, "length": 4}, + "0.1": {"path": f"{prefix}/{fp}", "offset": 500, "length": 4}, + } + # the 500-byte gap is within max_gap, but merging would make a 504-byte + # read which exceeds max_bytes, so the two chunks stay separate. + store = _build_store( + tmpdir, + files={fp: bytes(2048)}, + entries=entries, + coalesce_max_gap_bytes=1000, + coalesce_max_bytes=100, + ) + await _collect(store, ["foo/c.0.0", "foo/c.0.1"]) + assert sorted((c["start"], c["end"]) for c in record_range_calls) == [ + (0, 4), + (500, 504), + ] + + +@requires_obstore +class TestCoalesceMembers: + def _members(self, *ranges): + return [_Member(i, start, end) for i, (start, end) in enumerate(ranges)] + + def test_merges_contiguous_and_near(self): + # 0-4 and 4-8 touch; 8-12 is 2 bytes after -> all merge with gap>=2 + runs = _coalesce_members( + self._members((0, 4), (6, 10)), max_gap=2, max_bytes=1 << 20 + ) + assert len(runs) == 1 + assert [(m.start, m.end) for m in runs[0]] == [(0, 4), (6, 10)] + + def test_splits_on_large_gap(self): + runs = _coalesce_members( + self._members((0, 4), (100, 104)), max_gap=8, max_bytes=1 << 20 + ) + assert [[(m.start, m.end) for m in run] for run in runs] == [ + [(0, 4)], + [(100, 104)], + ] + + def test_splits_when_span_exceeds_max_bytes(self): + runs = _coalesce_members( + self._members((0, 4), (50, 54)), max_gap=1000, max_bytes=40 + ) + assert len(runs) == 2 + + def test_oversized_single_member_is_its_own_run(self): + runs = _coalesce_members(self._members((0, 1000)), max_gap=8, max_bytes=100) + assert len(runs) == 1 + assert runs[0][0].end == 1000 + + def test_sorts_by_start(self): + runs = _coalesce_members( + self._members((8, 12), (0, 4), (4, 8)), max_gap=0, max_bytes=1 << 20 + ) + # all adjacent -> one run, ordered by start + assert [(m.start, m.end) for m in runs[0]] == [(0, 4), (4, 8), (8, 12)]