diff --git a/obstore/python/obstore/fsspec.py b/obstore/python/obstore/fsspec.py index d6c76a92..a936e238 100644 --- a/obstore/python/obstore/fsspec.py +++ b/obstore/python/obstore/fsspec.py @@ -32,12 +32,12 @@ from __future__ import annotations -import asyncio import warnings from collections import defaultdict +from collections.abc import Iterable from functools import cached_property, lru_cache from pathlib import Path -from typing import TYPE_CHECKING, Literal, overload +from typing import TYPE_CHECKING, Literal, cast, overload from urllib.parse import urlparse import fsspec.asyn @@ -45,15 +45,22 @@ from fsspec.implementations.local import make_path_posix from obstore import open_reader, open_writer +from obstore.exceptions import NotSupportedError from obstore.store import from_url if TYPE_CHECKING: import sys - from collections.abc import Coroutine, Iterable + from collections.abc import Coroutine, Sequence from datetime import datetime from typing import Any - from obstore import Attributes, Bytes, ReadableFile, WritableFile + from obstore import ( + Attributes, + Bytes, + GetOptions, + ReadableFile, + WritableFile, + ) from obstore.store import ( AzureConfig, AzureCredentialProvider, @@ -113,6 +120,81 @@ """A type hint for all supported protocols.""" +def _needs_object_size(start: int, end: int | None) -> bool: + """Whether resolving a range requires knowing the size of the object. + + Negative bounds require the size to be known, except a bare negative + `start`, which maps directly to a suffix request. + """ + return end is not None and (end < 0 or start < 0) + + +def _apply_object_size( + start: int, + end: int | None, + size: int, +) -> tuple[int, int | None]: + """Resolve bounds that count back from the end of an object of `size`.""" + return ( + max(0, size + start) if start < 0 else start, + max(0, size + end) if end is not None and end < 0 else end, + ) + + +def _split_requests( + paths: list[str], + bounds: Sequence[tuple[int, int | None]], +) -> tuple[ + dict[str, list[tuple[int, int, int]]], + list[tuple[int, str, int]], +]: + """Split the requested ranges into bounded per-object and open-ended.""" + # `get_ranges` takes only bounded ranges, and merges only within one object, + # so split as zarr's obstore store does. + # Ref: https://github.com/zarr-developers/zarr-python/blob/de2cce1adc41a4d38721bf62b25eb312a52066dd/src/zarr/storage/_obstore.py#L440-L512 + per_file_bounded_requests: dict[str, list[tuple[int, int, int]]] = defaultdict(list) + open_ended_requests: list[tuple[int, str, int]] = [] + for idx, (path, (start, end)) in enumerate(zip(paths, bounds, strict=True)): + if end is not None: + per_file_bounded_requests[path].append((idx, start, end)) + else: + open_ended_requests.append((idx, path, start)) + return per_file_bounded_requests, open_ended_requests + + +async def _get_open_ended( + store: ObjectStore, + path: str, + start: int, +) -> Bytes: + """Get an open-ended range, working around stores without suffix support.""" + options: GetOptions + if start == 0: + # `{"offset": 0}` fails on an empty object, so send no range instead. + options = {} + elif start < 0: + options = {"range": {"suffix": -start}} + else: + options = {"range": {"offset": start}} + try: + resp = await store.get_async(path, options=options) + return await resp.bytes_async() + except NotSupportedError: + if start >= 0: + # The store refused something other than a suffix. + raise + + # Azure rejects suffix ranges, so fall back to a size lookup and a bounded + # read, as zarr's obstore store does. A suffix covering the whole object + # becomes a plain `get`. + suffix = -start + size = (await store.head_async(path))["size"] + if suffix >= size: + resp = await store.get_async(path) + return await resp.bytes_async() + return await store.get_range_async(path, start=size - suffix, length=suffix) + + class FsspecStore(fsspec.asyn.AsyncFileSystem): """An fsspec implementation based on a obstore Store. @@ -382,19 +464,27 @@ async def _cat_file( end: int | None = None, **_kwargs: Any, ) -> bytes: - bucket, path = self._split_path(path) + """Get a byte range, interpreting `start` and `end` like Python slices. + + Zero-length, inverted, and past-the-end ranges raise an error where a + slice would return `b""`. + """ + bucket, path_in_bucket = self._split_path(path) store = self._construct_store(bucket) - if start is None and end is None: - resp = await store.get_async(path) - return (await resp.bytes_async()).to_bytes() + start = 0 if start is None else start - if start is None or end is None: - raise NotImplementedError( - "cat_file not implemented for start=None xor end=None", - ) + if _needs_object_size(start, end): + # No range-header equivalent, so resolve against the object size. + size = (await self._info(path))["size"] + start, end = _apply_object_size(start, end, size) + + if end is None: + buffer = await _get_open_ended(store, path_in_bucket, start) + return buffer.to_bytes() - range_bytes = await store.get_range_async(path, start=start, end=end) + # `get_range` only takes bounded ranges, mirroring the Rust API. + range_bytes = await store.get_range_async(path_in_bucket, start=start, end=end) return range_bytes.to_bytes() async def _cat( # type: ignore (fsspec has bad typing) @@ -426,48 +516,158 @@ async def _cat( # type: ignore (fsspec has bad typing) async def _cat_ranges( # noqa: PLR0913, PLR0917 # type: ignore (fsspec has bad typing) self, paths: list[str], - starts: list[int] | int, - ends: list[int] | int, - max_gap=None, # noqa: ANN001, ARG002 - batch_size=None, # noqa: ANN001, ARG002 - on_error="return", # noqa: ANN001, ARG002 + starts: Sequence[int | None] | int | None, + ends: Sequence[int | None] | int | None, + max_gap: int | None = None, + batch_size: int | None = None, + on_error: str = "return", **_kwargs: Any, - ) -> list[bytes]: - if isinstance(starts, int): + ) -> list[bytes | BaseException]: + """Get ranges whose bounds are interpreted as in `_cat_file`. + + Failures are returned in place of the bytes, or the first is raised + when `on_error` is "raise". + """ + # The base class implementation `AsyncFileSystem._cat_ranges` forwards each + # element to `_cat_file`, which documents negative bounds and `None` for + # either end. + # Ref: https://github.com/fsspec/filesystem_spec/blob/e6668a146cd07b9f50530c49ea3916d8ab13e169/fsspec/spec.py#L790-L800 + + # A non-iterable start or end applies to every path, since `None` is not + # `Iterable`. + if not isinstance(starts, Iterable): starts = [starts] * len(paths) - if isinstance(ends, int): + if not isinstance(ends, Iterable): ends = [ends] * len(paths) if not len(paths) == len(starts) == len(ends): raise ValueError - per_file_requests: dict[str, list[tuple[int, int, int]]] = defaultdict(list) - # When upgrading to Python 3.10, use strict=True - for idx, (path, start, end) in enumerate(zip(paths, starts, ends)): - per_file_requests[path].append((start, end, idx)) + bounds = await self._resolve_inexpressible_bounds( + paths, + starts, + ends, + batch_size=batch_size, + ) + per_file_bounded_requests, open_ended_requests = _split_requests(paths, bounds) - futs: list[Coroutine[Any, Any, list[Bytes]]] = [] - for path, ranges in per_file_requests.items(): - bucket, path_no_bucket = self._split_path(path) - store = self._construct_store(bucket) + futs: list[Coroutine[Any, Any, list[tuple[int, bytes | BaseException]]]] = [ + self._cat_bounded_ranges(path, ranges, max_gap) + for path, ranges in per_file_bounded_requests.items() + ] + futs += [ + self._cat_open_ended_range(idx, path, start) + for idx, path, start in open_ended_requests + ] - offsets = [r[0] for r in ranges] - ends = [r[1] for r in ranges] - fut = store.get_ranges_async(path_no_bucket, starts=offsets, ends=ends) - futs.append(fut) + # Batched, to limit how many requests are in flight at once. + results = cast( + "list[list[tuple[int, bytes | BaseException]]]", + await fsspec.asyn._run_coros_in_chunks( # noqa: SLF001 + futs, + batch_size=batch_size or self.batch_size, + nofiles=True, + ), + ) - result = await asyncio.gather(*futs) + output_buffers: list[bytes | BaseException] = [b""] * len(paths) + for responses in results: + for idx, buffer in responses: + output_buffers[idx] = buffer - output_buffers: list[bytes] = [b""] * len(paths) - # When upgrading to Python 3.10, use strict=True - for per_file_request, buffers in zip(per_file_requests.items(), result): - path, ranges = per_file_request - # When upgrading to Python 3.10, use strict=True - for buffer, ranges_ in zip(buffers, ranges): - initial_index = ranges_[2] - output_buffers[initial_index] = buffer.to_bytes() + # Anything except "raise" returns failures in place. + if on_error == "raise": + for buffer in output_buffers: + if isinstance(buffer, BaseException): + raise buffer return output_buffers + async def _cat_bounded_ranges( + self, + path: str, + ranges: list[tuple[int, int, int]], # (output index, start, end) + max_gap: int | None, + ) -> list[tuple[int, bytes | BaseException]]: + """Get several bounded ranges of one object, merging nearby ones.""" + indices, starts, ends = zip(*ranges, strict=True) + # fsspec's `max_gap` is obstore's `coalesce`. It is left out when unset, so + # that obstore's own default applies. + coalesce: dict[str, int] = {} if max_gap is None else {"coalesce": max_gap} + try: + bucket, path_in_bucket = self._split_path(path) + store = self._construct_store(bucket) + buffers = await store.get_ranges_async( + path_in_bucket, + starts=starts, + ends=ends, + lengths=None, + **coalesce, + ) + except Exception as exc: # noqa: BLE001 + # The ranges share one request, and thus the exception. + return [(idx, exc) for idx in indices] + return [ + (idx, buffer.to_bytes()) + for idx, buffer in zip(indices, buffers, strict=True) + ] + + async def _cat_open_ended_range( + self, + idx: int, + path: str, + start: int, + ) -> list[tuple[int, bytes | BaseException]]: + """Get one open-ended range, which `get_ranges` cannot express. + + Returns a list for symmetry with `_cat_bounded_ranges`, so that both can + be batched together. + """ + try: + bucket, path_in_bucket = self._split_path(path) + store = self._construct_store(bucket) + buffer = await _get_open_ended(store, path_in_bucket, start) + except Exception as exc: # noqa: BLE001 + return [(idx, exc)] + return [(idx, buffer.to_bytes())] + + async def _resolve_inexpressible_bounds( + self, + paths: list[str], + starts: Sequence[int | None], + ends: Sequence[int | None], + *, + batch_size: int | None, + ) -> list[tuple[int, int | None]]: + """Resolve the bounds that obstore cannot express.""" + normalized_starts = [0 if start is None else start for start in starts] + paths_needing_size = list( + { + path + for path, start, end in zip(paths, normalized_starts, ends, strict=True) + if _needs_object_size(start, end) + }, + ) + sizes: dict[str, int] = {} + if paths_needing_size: + # `_sizes` goes through `_info`, so the dircache may serve these. + # Type checkers infer `list[None]` here, but the values are sizes. + sizes = dict( + zip( + paths_needing_size, + cast( + "list[int]", + await self._sizes(paths_needing_size, batch_size=batch_size), + ), + strict=True, + ), + ) + return [ + _apply_object_size(start, end, sizes[path]) + if _needs_object_size(start, end) + else (start, end) + for path, start, end in zip(paths, normalized_starts, ends, strict=True) + ] + async def _put_file( self, lpath: str, diff --git a/tests/test_fsspec.py b/tests/test_fsspec.py index 3da2f5d3..022eba53 100644 --- a/tests/test_fsspec.py +++ b/tests/test_fsspec.py @@ -8,11 +8,14 @@ from unittest.mock import patch import fsspec +import fsspec.asyn import pyarrow.parquet as pq import pytest from fsspec.registry import _registry +from obstore.exceptions import NotSupportedError from obstore.fsspec import FsspecStore, register +from obstore.store import ObjectStoreMethods from tests.conftest import TEST_BUCKET_NAME if TYPE_CHECKING: @@ -572,6 +575,226 @@ def test_multi_file_ops(minio_bucket: tuple[S3Config, ClientConfig]): assert out == [f"{bucket}/afile"] +def test_cat_file(fs: FsspecStore, monkeypatch: pytest.MonkeyPatch): + """Test that `cat_file` accepts every range form and sends the right request.""" + data = os.urandom(10000) + path = f"{TEST_BUCKET_NAME}/data1" + fs.pipe_file(path, data) + + ranges_via_get: list[dict[str, int] | None] = [] + bounds_via_get_range: list[tuple[int, int]] = [] + original_get = ObjectStoreMethods.get_async + original_get_range = ObjectStoreMethods.get_range_async + + async def spy_get(self, *args, **kwargs): # noqa: ANN001, ANN002, ANN003 + ranges_via_get.append(kwargs.get("options", {}).get("range")) + return await original_get(self, *args, **kwargs) + + async def spy_get_range(self, *args, **kwargs): # noqa: ANN001, ANN002, ANN003 + bounds_via_get_range.append((kwargs["start"], kwargs["end"])) + return await original_get_range(self, *args, **kwargs) + + monkeypatch.setattr(ObjectStoreMethods, "get_async", spy_get) + monkeypatch.setattr(ObjectStoreMethods, "get_range_async", spy_get_range) + + # The whole object, with and without an explicit zero start. + assert fs.cat_file(path) == data + assert fs.cat_file(path, start=0) == data + assert ranges_via_get == [None, None] + assert bounds_via_get_range == [] + + # Both bounds given go through `get_range` instead of `get`. + ranges_via_get.clear() + bounds_via_get_range.clear() + assert fs.cat_file(path, start=10, end=20) == data[10:20] + assert ranges_via_get == [] + assert bounds_via_get_range == [(10, 20)] + + # Either bound on its own. + ranges_via_get.clear() + bounds_via_get_range.clear() + assert fs.cat_file(path, start=10) == data[10:] + assert fs.cat_file(path, end=20) == data[:20] + assert ranges_via_get == [{"offset": 10}] + assert bounds_via_get_range == [(0, 20)] # An `end` alone is a bounded read. + + # Bounds counted back from the end of the object. + ranges_via_get.clear() + bounds_via_get_range.clear() + assert fs.cat_file(path, start=-10) == data[-10:] + assert fs.cat_file(path, start=-20000) == data # The store clamps the suffix. + assert fs.cat_file(path, start=10, end=-10) == data[10:-10] + assert fs.cat_file(path, start=-20, end=-10) == data[-20:-10] + assert fs.cat_file(path, start=-20, end=9995) == data[-20:9995] + assert ranges_via_get == [{"suffix": 10}, {"suffix": 20000}] + # Resolved ends become bounded reads. + assert bounds_via_get_range == [(10, 9990), (9980, 9990), (9980, 9995)] + + +def test_suffix_fallback(fs: FsspecStore, monkeypatch: pytest.MonkeyPatch): + """Test that a suffix request falls back to a bounded read, as Azure needs.""" + data = os.urandom(10000) + path = f"{TEST_BUCKET_NAME}/data1" + empty = f"{TEST_BUCKET_NAME}/empty" + fs.pipe_file(path, data) + fs.pipe_file(empty, b"") + + ranges_via_get: list[dict[str, int] | None] = [] + original_get = ObjectStoreMethods.get_async + + async def refuse_ranges(self, *args, **kwargs): # noqa: ANN001, ANN002, ANN003 + byte_range = kwargs.get("options", {}).get("range") + ranges_via_get.append(byte_range) + if byte_range is not None: + raise NotSupportedError("Azure does not support suffix range requests") + return await original_get(self, *args, **kwargs) + + monkeypatch.setattr(ObjectStoreMethods, "get_async", refuse_ranges) + + assert fs.cat_file(path, start=-10) == data[-10:] + assert fs.cat_ranges([path], [-10], [None]) == [data[-10:]] + # Each call tries the suffix once; the fallback read bypasses `get`. + assert ranges_via_get == [{"suffix": 10}, {"suffix": 10}] + + # A suffix covering the whole object clamps, exactly as the native path does: + # the refused suffix request becomes a plain `get`. + ranges_via_get.clear() + assert fs.cat_file(path, start=-20000) == data + assert fs.cat_file(empty, start=-10) == b"" + assert ranges_via_get == [{"suffix": 20000}, None, {"suffix": 10}, None] + + # Only a suffix is retried, so an offset still surfaces the error. + with pytest.raises(NotSupportedError): + fs.cat_file(path, start=10) + + +def test_cat_ranges_max_gap(fs: FsspecStore, monkeypatch: pytest.MonkeyPatch): + """Test that `max_gap` is forwarded to obstore as `coalesce`.""" + data = os.urandom(10000) + path = f"{TEST_BUCKET_NAME}/data1" + fs.pipe_file(path, data) + + forwarded_coalesce: list[int | None] = [] + original_get_ranges = ObjectStoreMethods.get_ranges_async + + async def spy_get_ranges(self, *args, **kwargs): # noqa: ANN001, ANN002, ANN003 + forwarded_coalesce.append(kwargs.get("coalesce")) + return await original_get_ranges(self, *args, **kwargs) + + monkeypatch.setattr(ObjectStoreMethods, "get_ranges_async", spy_get_ranges) + + # No max_gap: obstore uses its own default. + assert fs.cat_ranges([path, path], [0, 20], [10, 30]) == [ + data[0:10], + data[20:30], + ] + assert forwarded_coalesce == [None] + + # With max_gap: passed straight through, including 0 to turn coalescing off. + for max_gap in (0, 1000): + forwarded_coalesce.clear() + assert fs.cat_ranges([path, path], [0, 20], [10, 30], max_gap=max_gap) == [ + data[0:10], + data[20:30], + ] + assert forwarded_coalesce == [max_gap] + + +def test_cat_ranges_routing(fs: FsspecStore, monkeypatch: pytest.MonkeyPatch): + """Test that bounded ranges use `get_ranges` and open-ended ones use `get`.""" + data = os.urandom(10000) + path = f"{TEST_BUCKET_NAME}/data1" + fs.pipe_file(path, data) + + bounds_via_get_ranges: list[tuple[list[int], list[int]]] = [] + ranges_via_get: list[dict[str, int] | None] = [] + original_get_ranges = ObjectStoreMethods.get_ranges_async + original_get = ObjectStoreMethods.get_async + + async def spy_get_ranges(self, *args, **kwargs): # noqa: ANN001, ANN002, ANN003 + bounds_via_get_ranges.append((list(kwargs["starts"]), list(kwargs["ends"]))) + return await original_get_ranges(self, *args, **kwargs) + + async def spy_get(self, *args, **kwargs): # noqa: ANN001, ANN002, ANN003 + ranges_via_get.append(kwargs.get("options", {}).get("range")) + return await original_get(self, *args, **kwargs) + + monkeypatch.setattr(ObjectStoreMethods, "get_ranges_async", spy_get_ranges) + monkeypatch.setattr(ObjectStoreMethods, "get_async", spy_get) + + out = fs.cat_ranges([path] * 4, [0, 20, -5, 9000], [10, None, None, -10]) + assert out == [data[0:10], data[20:], data[-5:], data[9000:-10]] + + # The bounded range [0:10] and the range with negative end [9000:-10] both + # reach `get_ranges`. The latter is resolved beforehand against the object + # size, so the two share a single call. + assert bounds_via_get_ranges == [([0, 9000], [10, 9990])] + + # The other two, [20:] and [-5:], go through `get` instead. The latter stays + # a suffix request, even if another range on the same object causes the size + # to be fetched. + assert len(ranges_via_get) == 2 + assert {"offset": 20} in ranges_via_get + assert {"suffix": 5} in ranges_via_get + + +def test_cat_ranges_batch_size(fs: FsspecStore, monkeypatch: pytest.MonkeyPatch): + """Test that `batch_size` is forwarded to fsspec's request batching.""" + data = os.urandom(10000) + path = f"{TEST_BUCKET_NAME}/data1" + fs.pipe_file(path, data) + + forwarded_batch_sizes: list[int | None] = [] + original_run_coros = fsspec.asyn._run_coros_in_chunks + + async def spy_run_coros(coros, **kwargs): # noqa: ANN001, ANN003 + forwarded_batch_sizes.append(kwargs.get("batch_size")) + return await original_run_coros(coros, **kwargs) + + monkeypatch.setattr(fsspec.asyn, "_run_coros_in_chunks", spy_run_coros) + + # Unset: forwarded as None, so fsspec infers its own default. + assert fs.cat_ranges([path], [0], [10]) == [data[0:10]] + assert forwarded_batch_sizes == [None] + + forwarded_batch_sizes.clear() + assert fs.cat_ranges([path], [0], [10], batch_size=4) == [data[0:10]] + assert forwarded_batch_sizes == [4] + + # The size lookup for a negative end is batched with the same batch size. + forwarded_batch_sizes.clear() + assert fs.cat_ranges([path], [0], [-10], batch_size=4) == [data[0:-10]] + assert forwarded_batch_sizes == [4, 4] + + +def test_cat_ranges_on_error(fs: FsspecStore): + """Test that `on_error` controls whether a failure is returned or raised.""" + data = os.urandom(10000) + path = f"{TEST_BUCKET_NAME}/data1" + fs.pipe_file(path, data) + missing1 = f"{TEST_BUCKET_NAME}/missing1" + missing2 = f"{TEST_BUCKET_NAME}/missing2" + + # Default returns the error in place. + out = fs.cat_ranges( + [path, missing1, missing1, missing1], + [0, 0, 9000, -5], + [10, 10, 9990, None], + ) + assert out[0] == data[0:10] + assert isinstance(out[1], FileNotFoundError) + assert isinstance(out[3], FileNotFoundError) + + # A single `get_ranges` call covers both bounded ranges of `missing1`, so they + # fail together with the same error. + assert out[1] is out[2] + + # "raise" reports the earliest failing range rather than the first request to + # finish: [-5:] is read by a later request than [0:10], but comes first here. + with pytest.raises(FileNotFoundError, match="missing1"): + fs.cat_ranges([missing1, missing2], [-5, 0], [None, 10], on_error="raise") + + def test_cat_ranges_one(fs: FsspecStore): data1 = os.urandom(10000) fs.pipe_file(f"{TEST_BUCKET_NAME}/data1", data1) @@ -631,16 +854,35 @@ def test_cat_ranges_two(fs: FsspecStore): assert out == [data1[10:20], data2[10:20]] -@pytest.mark.xfail(reason="negative and mixed ranges not implemented") def test_cat_ranges_mixed(fs: FsspecStore): data1 = os.urandom(10000) data2 = os.urandom(10000) - fs.pipe({"data1": data1, "data2": data2}) + path1 = f"{TEST_BUCKET_NAME}/data1" + path2 = f"{TEST_BUCKET_NAME}/data2" + empty = f"{TEST_BUCKET_NAME}/empty" + fs.pipe({path1: data1, path2: data2, empty: b""}) - # single range in each file - out = fs.cat_ranges(["data1", "data1", "data2"], [-10, None, 10], [None, -10, -10]) + # negative and None bounds mixed across two files + out = fs.cat_ranges([path1, path1, path2], [-10, None, 10], [None, -10, -10]) assert out == [data1[-10:], data1[:-10], data2[10:-10]] + # an unbounded range reads the whole object, even a zero-length one + out = fs.cat_ranges([empty, empty], [0, None], [None, None]) + assert out == [b"", b""] + + +def test_cat_ranges_broadcast(fs: FsspecStore): + """Test that a scalar or `None` start and end broadcast across all paths.""" + data1 = os.urandom(10000) + data2 = os.urandom(10000) + path1 = f"{TEST_BUCKET_NAME}/data1" + path2 = f"{TEST_BUCKET_NAME}/data2" + fs.pipe({path1: data1, path2: data2}) + + assert fs.cat_ranges([path1, path2], 10, 20) == [data1[10:20], data2[10:20]] + assert fs.cat_ranges([path1, path2], 10, None) == [data1[10:], data2[10:]] + assert fs.cat_ranges([path1, path2], None, 20) == [data1[:20], data2[:20]] + @pytest.mark.xfail(reason="atomic writes not working on moto") def test_atomic_write(fs: FsspecStore):