From 4e2c2fc0f47cab8e2471a8db0708223b12a7916a Mon Sep 17 00:00:00 2001 From: Gunnar Schulze Date: Wed, 12 Aug 2026 22:01:21 +0200 Subject: [PATCH 01/22] feat: Handle all forms of range requests --- obstore/python/obstore/_get.pyi | 70 +++++++++++++---- obstore/python/obstore/fsspec.py | 131 +++++++++++++++++++++++++------ obstore/python/obstore/store.py | 8 +- obstore/src/get.rs | 122 +++++++++++++++++++--------- tests/test_fsspec.py | 25 +++++- tests/test_get.py | 46 +++++++++++ 6 files changed, 320 insertions(+), 82 deletions(-) diff --git a/obstore/python/obstore/_get.pyi b/obstore/python/obstore/_get.pyi index ff0cbf37..b109b58e 100644 --- a/obstore/python/obstore/_get.pyi +++ b/obstore/python/obstore/_get.pyi @@ -328,19 +328,33 @@ def get_range( ) -> Bytes: """Return the bytes that are stored at the specified location in the given byte range. - If the given range is zero-length or starts after the end of the object, an error - will be returned. Additionally, if the range ends after the end of the object, the - entire remainder of the object will be returned. Otherwise, the exact requested - range will be returned. + The requested range may be bounded, open-ended, or relative to the end of the + object: + + - `start` plus either `end` or `length` requests a specific range of bytes. + If the given range is zero-length or starts after the end of the object, an + error will be returned. Additionally, if the range ends after the end of the + object, the entire remainder of the object will be returned. Otherwise, the + exact requested range will be returned. + - A non-negative `start` on its own requests all bytes from `start` onwards. + This is equivalent to `bytes={start}-` as an HTTP header. + - A negative `start` on its own requests the last `abs(start)` bytes. Note that + here, `abs(start)` is _the size of the request_, not a byte offset. This is + equivalent to `bytes=-{abs(start)}` as an HTTP header. Args: store: The ObjectStore instance to use. path: The path within ObjectStore to retrieve. Keyword Args: - start: The start of the byte range. - end: The end of the byte range (exclusive). Either `end` or `length` must be non-None. - length: The number of bytes of the byte range. Either `end` or `length` must be non-None. + start: The start of the byte range. If negative, the last `abs(start)` bytes + of the object are requested, and `end` and `length` must both be None. + end: The end of the byte range (exclusive). Mutually exclusive with `length`. + Defaults to None, in which case the range continues to the end of the + object. + length: The number of bytes of the byte range. Mutually exclusive with `end`. + Defaults to None, in which case the range continues to the end of the + object. Returns: A `Bytes` object implementing the Python buffer protocol, allowing @@ -366,27 +380,53 @@ def get_ranges( path: str, *, starts: Sequence[int], - ends: Sequence[int] | None = None, - lengths: Sequence[int] | None = None, + ends: Sequence[int | None] | None = None, + lengths: Sequence[int | None] | None = None, coalesce: int = 1024 * 1024, ) -> list[Bytes]: """Return the bytes stored at the specified location in the given byte ranges. + Each range is described by one element of `starts` and the element at the same + position in `ends` or `lengths`, and follows the same semantics as + [get_range][obstore.get_range]. In particular, a negative element of `starts` + requests the last `abs(start)` bytes, and a range whose end is left unspecified + continues to the end of the object. Because ranges are matched up by position, + every sequence given must have the same length as `starts`; a mismatch raises + `ValueError`. + + `ends` and `lengths` may both be given at once, so long as at most one of the two + is non-None for any individual range — for example `ends=[10, None], + lengths=[None, 5]` bounds the first range by its end offset and the second by its + length. Omitting both reads every range to the end of the object. + To improve performance this will: - Transparently combine ranges less than `coalesce` bytes apart into a single underlying request (defaults to 1MB) - Make multiple `fetch` requests in parallel (up to maximum of 10) + !!! note + + Combining ranges requires knowing where each one ends, so if any requested + range is open-ended or relative to the end of the object, one additional + `head` request is made to resolve the size of the object. Requests made up + entirely of bounded ranges do not incur this cost. + Args: store: The ObjectStore instance to use. path: The path within ObjectStore to retrieve. Other Args: - starts: A sequence of `int` where each offset starts. - ends: A sequence of `int` where each offset ends (exclusive). Either `ends` or `lengths` must be non-None. - lengths: A sequence of `int` with the number of bytes of each byte range. Either `ends` or `lengths` must be non-None. - coalesce: Maximum distance in bytes between ranges that will be coalesced into a single request. Defaults to 1MiB. Set to `0` to disable coalescing. + starts: A sequence of `int` where each offset starts. A negative value + requests the last `abs(start)` bytes for that range, in which case the + corresponding elements of `ends` and `lengths` must be None. + ends: A sequence of `int` where each offset ends (exclusive). An element may + be None to continue that range to the end of the object. + lengths: A sequence of `int` with the number of bytes of each byte range. An + element may be None to continue that range to the end of the object. + coalesce: Maximum distance in bytes between ranges that will be coalesced + into a single request. Defaults to 1MiB. Set to `0` to disable + coalescing. Returns: A sequence of `Bytes`, one for each range. This `Bytes` object implements the @@ -400,8 +440,8 @@ async def get_ranges_async( path: str, *, starts: Sequence[int], - ends: Sequence[int] | None = None, - lengths: Sequence[int] | None = None, + ends: Sequence[int | None] | None = None, + lengths: Sequence[int | None] | None = None, coalesce: int = 1024 * 1024, ) -> list[Bytes]: """Call `get_ranges` asynchronously. diff --git a/obstore/python/obstore/fsspec.py b/obstore/python/obstore/fsspec.py index d6c76a92..016704a5 100644 --- a/obstore/python/obstore/fsspec.py +++ b/obstore/python/obstore/fsspec.py @@ -35,9 +35,10 @@ 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 @@ -49,7 +50,7 @@ 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 @@ -113,6 +114,35 @@ """A type hint for all supported protocols.""" +def _needs_object_size(start: int | None, end: int | None) -> bool: + """Whether a range must be resolved against the object size first. + + A negative `start` on its own is a suffix request, which obstore supports + directly. Anything else counting back from the end of the object is not. + """ + return end is not None and (end < 0 or (start is not None and start < 0)) + + +def _apply_object_size( + start: int | None, + end: int | None, + size: int | None, +) -> tuple[int, int | None]: + """Rewrite the bounds of a range that count back from the end of an object. + + A `size` of None leaves the range alone, for objects whose size was never + needed. + """ + if start is None: + start = 0 + if size is None: + return start, end + return ( + max(0, size + start) if start < 0 else start, + max(0, size + end) if end is not None and end < 0 else end, + ) + + class FsspecStore(fsspec.asyn.AsyncFileSystem): """An fsspec implementation based on a obstore Store. @@ -389,12 +419,19 @@ async def _cat_file( resp = await store.get_async(path) return (await resp.bytes_async()).to_bytes() - if start is None or end is None: - raise NotImplementedError( - "cat_file not implemented for start=None xor end=None", - ) - - range_bytes = await store.get_range_async(path, start=start, end=end) + if _needs_object_size(start, end): + # See `_resolve_ranges` for why this needs the object size. + size = (await store.head_async(path))["size"] + if start is not None and start < 0: + start = max(0, size + start) + if end is not None and end < 0: + end = max(0, size + end) + + range_bytes = await store.get_range_async( + path, + start=0 if start is None else start, + end=end, + ) return range_bytes.to_bytes() async def _cat( # type: ignore (fsspec has bad typing) @@ -426,23 +463,30 @@ 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, # noqa: ARG002 + batch_size: int | None = None, # noqa: ARG002 + on_error: str = "return", # noqa: ARG002 **_kwargs: Any, ) -> list[bytes]: - if isinstance(starts, int): + # A non-iterable start or end applies to every path, per fsspec's + # AsyncFileSystem._cat_ranges. + 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)): + resolved = await self._resolve_ranges(paths, starts, ends) + + per_file_requests: dict[str, list[tuple[int, int | None, int]]] = defaultdict( + list, + ) + for idx, (path, (start, end)) in enumerate( + zip(paths, resolved, strict=True), + ): per_file_requests[path].append((start, end, idx)) futs: list[Coroutine[Any, Any, list[Bytes]]] = [] @@ -451,23 +495,62 @@ async def _cat_ranges( # noqa: PLR0913, PLR0917 # type: ignore (fsspec has bad store = self._construct_store(bucket) 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) + file_ends = [r[1] for r in ranges] + fut = store.get_ranges_async( + path_no_bucket, + starts=offsets, + ends=file_ends, + ) futs.append(fut) result = await asyncio.gather(*futs) 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): + for per_file_request, buffers in zip( + per_file_requests.items(), + result, + strict=True, + ): path, ranges = per_file_request - # When upgrading to Python 3.10, use strict=True - for buffer, ranges_ in zip(buffers, ranges): + for buffer, ranges_ in zip(buffers, ranges, strict=True): initial_index = ranges_[2] output_buffers[initial_index] = buffer.to_bytes() return output_buffers + async def _resolve_ranges( + self, + paths: list[str], + starts: Sequence[int | None], + ends: Sequence[int | None], + ) -> list[tuple[int, int | None]]: + """Rewrite fsspec's range vocabulary into the subset obstore accepts. + + `start` is normalized from None to 0. Ranges counting back from the end of + the object are resolved against its size, at the cost of one `head` request + per object; a lone negative `start` is left alone, since obstore expresses + that directly as a suffix request. + """ + sized_paths = sorted( + { + path + for path, start, end in zip(paths, starts, ends, strict=True) + if _needs_object_size(start, end) + }, + ) + sizes = dict( + zip( + sized_paths, + # fsspec types `_sizes` as `list[None]`; the values are really sizes. + cast("list[int]", await self._sizes(sized_paths)), + strict=True, + ), + ) + return [ + _apply_object_size(start, end, sizes.get(path)) + for path, start, end in zip(paths, starts, ends, strict=True) + ] + async def _put_file( self, lpath: str, diff --git a/obstore/python/obstore/store.py b/obstore/python/obstore/store.py index 72fe1b80..bbd95fb2 100644 --- a/obstore/python/obstore/store.py +++ b/obstore/python/obstore/store.py @@ -244,8 +244,8 @@ def get_ranges( path: str, *, starts: Sequence[int], - ends: Sequence[int] | None = None, - lengths: Sequence[int] | None = None, + ends: Sequence[int | None] | None = None, + lengths: Sequence[int | None] | None = None, coalesce: int = 1024 * 1024, ) -> list[Bytes]: """Return the bytes stored at the specified location in the given byte ranges. @@ -266,8 +266,8 @@ async def get_ranges_async( path: str, *, starts: Sequence[int], - ends: Sequence[int] | None = None, - lengths: Sequence[int] | None = None, + ends: Sequence[int | None] | None = None, + lengths: Sequence[int | None] | None = None, coalesce: int = 1024 * 1024, ) -> list[Bytes]: """Call `get_ranges` asynchronously. diff --git a/obstore/src/get.rs b/obstore/src/get.rs index afdcc978..c723d847 100644 --- a/obstore/src/get.rs +++ b/obstore/src/get.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use bytes::Bytes; use chrono::{DateTime, Utc}; use futures::stream::{BoxStream, Fuse}; -use futures::StreamExt; +use futures::{StreamExt, TryFutureExt}; use object_store::{ coalesce_ranges, Attributes, GetOptions, GetRange, GetResult, ObjectMeta, ObjectStore, ObjectStoreExt, OBJECT_STORE_COALESCE_DEFAULT, @@ -117,7 +117,7 @@ impl<'py> FromPyObject<'_, 'py> for PyGetRange { } else if let Ok(suffix_range) = obj.extract::() { Ok(Self(suffix_range.into())) } else { - Err(PyValueError::new_err("Unexpected input for byte range.\nExpected two-integer tuple or list, or dict with 'offset' or 'suffix' key." )) + Err(PyValueError::new_err("Unexpected input for byte range.\nExpected two-integer tuple or list, or dict with 'offset' or 'suffix' key.")) } } } @@ -382,15 +382,20 @@ pub(crate) fn get_range( py: Python, store: PyObjectStore, path: PyPath, - start: u64, + start: i64, end: Option, length: Option, -) -> PyObjectStoreResult { +) -> PyObjectStoreResult { let runtime = get_runtime(); let range = params_to_range(start, end, length)?; py.detach(|| { - let out = runtime.block_on(store.as_ref().get_range(path.as_ref(), range))?; - Ok::<_, PyObjectStoreError>(pyo3_bytes::PyBytes::new(out)) + let out = runtime.block_on( + store + .as_ref() + .get_opts(path.as_ref(), GetOptions::new().with_range(range.into())) + .and_then(GetResult::bytes), + )?; + Ok::<_, PyObjectStoreError>(PyBytes::new(out)) }) } @@ -400,7 +405,7 @@ pub(crate) fn get_range_async( py: Python, store: PyObjectStore, path: PyPath, - start: u64, + start: i64, end: Option, length: Option, ) -> PyResult> { @@ -408,36 +413,64 @@ pub(crate) fn get_range_async( pyo3_async_runtimes::tokio::future_into_py(py, async move { let out = store .as_ref() - .get_range(path.as_ref(), range) + .get_opts(path.as_ref(), GetOptions::new().with_range(range.into())) + .and_then(GetResult::bytes) .await .map_err(PyObjectStoreError::ObjectStoreError)?; - Ok(pyo3_bytes::PyBytes::new(out)) + Ok(PyBytes::new(out)) }) } fn params_to_range( - start: u64, + start: i64, end: Option, length: Option, -) -> PyObjectStoreResult> { +) -> PyObjectStoreResult { + if start < 0 { + if end.is_some() || length.is_some() { + return Err( + PyValueError::new_err("end and length must be None if start is negative.").into(), + ); + } + return Ok(GetRange::Suffix(start.unsigned_abs())); + } + + let start = start as u64; match (end, length) { (Some(_), Some(_)) => { Err(PyValueError::new_err("end and length cannot both be non-None.").into()) } - (None, None) => Err(PyValueError::new_err("Either end or length must be non-None.").into()), - (Some(end), None) => validate_range(start..end), - (None, Some(length)) => validate_range(start..start + length), + (None, None) => Ok(GetRange::Offset(start)), + (Some(end), None) => validate_range(start..end).map(GetRange::Bounded), + (None, Some(length)) => validate_range(start..start + length).map(GetRange::Bounded), } } async fn _get_ranges( store: PyObjectStore, path: PyPath, - ranges: &[Range], + ranges: &[GetRange], coalesce: u64, ) -> PyObjectStoreResult> { + let mut len: Option = None; + let mut resolved = Vec::with_capacity(ranges.len()); + for range in ranges { + resolved.push(match range { + GetRange::Bounded(r) => r.clone(), + other => { + let size = match len { + Some(len) => len, + None => *len.insert(store.as_ref().head(path.as_ref()).await?.size), + }; + other + .as_range(size) + .map_err(|err| PyValueError::new_err(err.to_string()))? + } + }); + } + let out = coalesce_ranges( - ranges, + &resolved, |range| store.as_ref().get_range(path.as_ref(), range), coalesce, ) @@ -451,9 +484,9 @@ pub(crate) fn get_ranges( py: Python, store: PyObjectStore, path: PyPath, - starts: Vec, - ends: Option>, - lengths: Option>, + starts: Vec, + ends: Option>>, + lengths: Option>>, coalesce: u64, ) -> PyObjectStoreResult> { let runtime = get_runtime(); @@ -467,9 +500,9 @@ pub(crate) fn get_ranges_async( py: Python, store: PyObjectStore, path: PyPath, - starts: Vec, - ends: Option>, - lengths: Option>, + starts: Vec, + ends: Option>>, + lengths: Option>>, coalesce: u64, ) -> PyResult> { let ranges = params_to_ranges(starts, ends, lengths)?; @@ -479,28 +512,45 @@ pub(crate) fn get_ranges_async( } fn params_to_ranges( - starts: Vec, - ends: Option>, - lengths: Option>, -) -> PyObjectStoreResult>> { - match (ends, lengths) { - (Some(_), Some(_)) => { - Err(PyValueError::new_err("ends and lengths cannot both be non-None.").into()) - } - (None, None) => { - Err(PyValueError::new_err("Either ends or lengths must be non-None.").into()) + starts: Vec, + ends: Option>>, + lengths: Option>>, +) -> PyObjectStoreResult> { + for (name, len) in [("ends", ends.as_ref()), ("lengths", lengths.as_ref())] + .into_iter() + .filter_map(|(name, seq)| seq.map(|seq| (name, seq.len()))) + { + if len != starts.len() { + return Err(PyValueError::new_err(format!( + "starts and {name} must have the same length, got {} and {len}.", + starts.len(), + )) + .into()); } + } + + match (ends, lengths) { + // Consistent with `get_range`: an unbounded range reads to the end of the + // object. + (None, None) => starts + .into_iter() + .map(|start| params_to_range(start, None, None)) + .collect(), + (Some(ends), Some(lengths)) => starts + .into_iter() + .zip(ends) + .zip(lengths) + .map(|((start, end), length)| params_to_range(start, end, length)) + .collect(), (Some(ends), None) => starts .into_iter() .zip(ends) - .map(|(start, end)| start..end) - .map(validate_range) + .map(|(start, end)| params_to_range(start, end, None)) .collect(), (None, Some(lengths)) => starts .into_iter() .zip(lengths) - .map(|(start, length)| start..start + length) - .map(validate_range) + .map(|(start, length)| params_to_range(start, None, length)) .collect(), } } diff --git a/tests/test_fsspec.py b/tests/test_fsspec.py index 3da2f5d3..3f673c65 100644 --- a/tests/test_fsspec.py +++ b/tests/test_fsspec.py @@ -572,6 +572,24 @@ def test_multi_file_ops(minio_bucket: tuple[S3Config, ClientConfig]): assert out == [f"{bucket}/afile"] +def test_cat_file(fs: FsspecStore): + data = os.urandom(10000) + path = f"{TEST_BUCKET_NAME}/data1" + fs.pipe_file(path, data) + + assert fs.cat_file(path) == data + assert fs.cat_file(path, start=100, end=200) == data[100:200] + + # Either bound on its own. + assert fs.cat_file(path, start=100) == data[100:] + assert fs.cat_file(path, end=200) == data[:200] + + # Bounds counted back from the end of the object. + assert fs.cat_file(path, start=-100) == data[-100:] + assert fs.cat_file(path, start=100, end=-100) == data[100:-100] + assert fs.cat_file(path, start=-200, end=-100) == data[-200:-100] + + def test_cat_ranges_one(fs: FsspecStore): data1 = os.urandom(10000) fs.pipe_file(f"{TEST_BUCKET_NAME}/data1", data1) @@ -631,14 +649,15 @@ 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" + fs.pipe({path1: data1, path2: data2}) # single range in each file - out = fs.cat_ranges(["data1", "data1", "data2"], [-10, None, 10], [None, -10, -10]) + out = fs.cat_ranges([path1, path1, path2], [-10, None, 10], [None, -10, -10]) assert out == [data1[-10:], data1[:-10], data2[10:-10]] diff --git a/tests/test_get.py b/tests/test_get.py index a301d32e..edfa9f1c 100644 --- a/tests/test_get.py +++ b/tests/test_get.py @@ -113,6 +113,19 @@ def test_get_range(): view = memoryview(buffer) assert view == data[5:15] + buffer = store.get_range(path, start=4300) + view = memoryview(buffer) + assert view == data[4300:4400] + + buffer = store.get_range(path, start=-100) + view = memoryview(buffer) + assert view == data[4300:4400] + + # A suffix longer than the object yields the whole object. + buffer = store.get_range(path, start=-999999) + view = memoryview(buffer) + assert view == data + def test_get_ranges(): store = MemoryStore() @@ -130,6 +143,27 @@ def test_get_ranges(): for start, end, buffer in zip(starts, ends, buffers): assert memoryview(buffer) == data[start:end] + # A `None` element leaves that one range open-ended; a negative start makes it + # a suffix request. Omitting `ends` and `lengths` reads every range to the end. + buffers = store.get_ranges(path, starts=[5, 4300, -100], ends=[15, None, None]) + assert [memoryview(b) for b in buffers] == [ + data[5:15], + data[4300:4400], + data[4300:4400], + ] + + buffers = store.get_ranges(path, starts=[4300, -100]) + assert [memoryview(b) for b in buffers] == [data[4300:4400], data[4300:4400]] + + # `ends` and `lengths` may be mixed, at most one per range. + buffers = store.get_ranges( + path, + starts=[5, 20], + ends=[15, None], + lengths=[None, 10], + ) + assert [memoryview(b) for b in buffers] == [data[5:15], data[20:30]] + lengths = [10, 10, 10, 10] buffers = store.get_ranges(path, starts=starts, lengths=lengths) @@ -207,6 +241,12 @@ def test_get_range_invalid_range(): with pytest.raises(ValueError, match="Invalid range"): store.get_range(path, start=10, length=0) + with pytest.raises(ValueError, match="end and length must be None"): + store.get_range(path, start=-10, end=10) + + with pytest.raises(ValueError, match="end and length must be None"): + store.get_range(path, start=-10, length=10) + def test_get_ranges_invalid_range(): store = MemoryStore() @@ -224,6 +264,12 @@ def test_get_ranges_invalid_range(): with pytest.raises(ValueError, match="Invalid range"): store.get_ranges(path, starts=[10, 20], lengths=[10, 0]) + with pytest.raises(ValueError, match="starts and ends must have the same length"): + store.get_ranges(path, starts=[10, 20], ends=[30]) + + with pytest.raises(ValueError, match="starts and lengths must have the same"): + store.get_ranges(path, starts=[10], lengths=[10, 20]) + def test_access_getresult_attributes_after_reading_stream(): store = MemoryStore() From 0c44f4ddaf577361d94b86f895d492049eac317f Mon Sep 17 00:00:00 2001 From: Gunnar Schulze Date: Wed, 12 Aug 2026 21:54:35 +0200 Subject: [PATCH 02/22] fix: Respect `max_gap` in `FsspecStore.cat_ranges` --- obstore/python/obstore/fsspec.py | 20 +++++++++++++++++-- tests/test_fsspec.py | 33 ++++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+), 2 deletions(-) diff --git a/obstore/python/obstore/fsspec.py b/obstore/python/obstore/fsspec.py index 016704a5..0fad2202 100644 --- a/obstore/python/obstore/fsspec.py +++ b/obstore/python/obstore/fsspec.py @@ -38,7 +38,7 @@ from collections.abc import Iterable from functools import cached_property, lru_cache from pathlib import Path -from typing import TYPE_CHECKING, Literal, cast, overload +from typing import TYPE_CHECKING, Literal, TypedDict, cast, overload from urllib.parse import urlparse import fsspec.asyn @@ -114,6 +114,16 @@ """A type hint for all supported protocols.""" +class _CoalesceKwarg(TypedDict, total=False): + """The optional `coalesce` argument of [obstore.get_ranges][]. + + Omitting the key entirely lets obstore apply its own default, so that default + does not have to be restated here. + """ + + coalesce: int + + def _needs_object_size(start: int | None, end: int | None) -> bool: """Whether a range must be resolved against the object size first. @@ -465,7 +475,7 @@ async def _cat_ranges( # noqa: PLR0913, PLR0917 # type: ignore (fsspec has bad paths: list[str], starts: Sequence[int | None] | int | None, ends: Sequence[int | None] | int | None, - max_gap: int | None = None, # noqa: ARG002 + max_gap: int | None = None, batch_size: int | None = None, # noqa: ARG002 on_error: str = "return", # noqa: ARG002 **_kwargs: Any, @@ -481,6 +491,11 @@ async def _cat_ranges( # noqa: PLR0913, PLR0917 # type: ignore (fsspec has bad resolved = await self._resolve_ranges(paths, starts, ends) + # fsspec's `max_gap` is obstore's `coalesce`: the largest gap between two + # ranges that may still be served by a single request. Left unset, defer to + # obstore's own default rather than restating it here. + coalesce: _CoalesceKwarg = {} if max_gap is None else {"coalesce": max_gap} + per_file_requests: dict[str, list[tuple[int, int | None, int]]] = defaultdict( list, ) @@ -500,6 +515,7 @@ async def _cat_ranges( # noqa: PLR0913, PLR0917 # type: ignore (fsspec has bad path_no_bucket, starts=offsets, ends=file_ends, + **coalesce, ) futs.append(fut) diff --git a/tests/test_fsspec.py b/tests/test_fsspec.py index 3f673c65..265346ab 100644 --- a/tests/test_fsspec.py +++ b/tests/test_fsspec.py @@ -13,6 +13,7 @@ from fsspec.registry import _registry from obstore.fsspec import FsspecStore, register +from obstore.store import ObjectStoreMethods from tests.conftest import TEST_BUCKET_NAME if TYPE_CHECKING: @@ -590,6 +591,38 @@ def test_cat_file(fs: FsspecStore): assert fs.cat_file(path, start=-200, end=-100) == data[-200:-100] +def test_cat_ranges_max_gap(fs: FsspecStore, monkeypatch: pytest.MonkeyPatch): + data = os.urandom(10000) + path = f"{TEST_BUCKET_NAME}/data1" + fs.pipe_file(path, data) + + seen: list[int | None] = [] + original = ObjectStoreMethods.get_ranges_async + + async def spy(self, *args, **kwargs): # noqa: ANN001, ANN002, ANN003 + seen.append(kwargs.get("coalesce")) + return await original(self, *args, **kwargs) + + monkeypatch.setattr(ObjectStoreMethods, "get_ranges_async", spy) + + # Unset: obstore's own default applies, so `coalesce` is not forwarded at all. + assert fs.cat_ranges([path, path], [0, 200], [100, 300]) == [ + data[0:100], + data[200:300], + ] + assert seen == [None] + + # Set: forwarded verbatim as obstore's `coalesce`, including 0 to disable + # coalescing entirely. + for max_gap in (0, 1000): + seen.clear() + assert fs.cat_ranges([path, path], [0, 200], [100, 300], max_gap=max_gap) == [ + data[0:100], + data[200:300], + ] + assert seen == [max_gap] + + def test_cat_ranges_one(fs: FsspecStore): data1 = os.urandom(10000) fs.pipe_file(f"{TEST_BUCKET_NAME}/data1", data1) From ebca42848ac79c98dfff942884a8d40c7b5f5806 Mon Sep 17 00:00:00 2001 From: Gunnar Schulze Date: Wed, 12 Aug 2026 23:05:23 +0200 Subject: [PATCH 03/22] refactor: Accept None for coalesce so the default lives in one place --- obstore/python/obstore/_get.pyi | 8 ++++---- obstore/python/obstore/fsspec.py | 21 ++++----------------- obstore/python/obstore/store.py | 4 ++-- obstore/src/get.rs | 10 ++++++---- tests/test_fsspec.py | 2 +- 5 files changed, 17 insertions(+), 28 deletions(-) diff --git a/obstore/python/obstore/_get.pyi b/obstore/python/obstore/_get.pyi index b109b58e..ed062d17 100644 --- a/obstore/python/obstore/_get.pyi +++ b/obstore/python/obstore/_get.pyi @@ -382,7 +382,7 @@ def get_ranges( starts: Sequence[int], ends: Sequence[int | None] | None = None, lengths: Sequence[int | None] | None = None, - coalesce: int = 1024 * 1024, + coalesce: int | None = None, ) -> list[Bytes]: """Return the bytes stored at the specified location in the given byte ranges. @@ -425,8 +425,8 @@ def get_ranges( lengths: A sequence of `int` with the number of bytes of each byte range. An element may be None to continue that range to the end of the object. coalesce: Maximum distance in bytes between ranges that will be coalesced - into a single request. Defaults to 1MiB. Set to `0` to disable - coalescing. + into a single request. Defaults to None, in which case obstore applies + its own default of 1MiB. Set to `0` to disable coalescing. Returns: A sequence of `Bytes`, one for each range. This `Bytes` object implements the @@ -442,7 +442,7 @@ async def get_ranges_async( starts: Sequence[int], ends: Sequence[int | None] | None = None, lengths: Sequence[int | None] | None = None, - coalesce: int = 1024 * 1024, + coalesce: int | None = None, ) -> list[Bytes]: """Call `get_ranges` asynchronously. diff --git a/obstore/python/obstore/fsspec.py b/obstore/python/obstore/fsspec.py index 0fad2202..3e9409bc 100644 --- a/obstore/python/obstore/fsspec.py +++ b/obstore/python/obstore/fsspec.py @@ -38,7 +38,7 @@ from collections.abc import Iterable from functools import cached_property, lru_cache from pathlib import Path -from typing import TYPE_CHECKING, Literal, TypedDict, cast, overload +from typing import TYPE_CHECKING, Literal, cast, overload from urllib.parse import urlparse import fsspec.asyn @@ -114,16 +114,6 @@ """A type hint for all supported protocols.""" -class _CoalesceKwarg(TypedDict, total=False): - """The optional `coalesce` argument of [obstore.get_ranges][]. - - Omitting the key entirely lets obstore apply its own default, so that default - does not have to be restated here. - """ - - coalesce: int - - def _needs_object_size(start: int | None, end: int | None) -> bool: """Whether a range must be resolved against the object size first. @@ -491,11 +481,6 @@ async def _cat_ranges( # noqa: PLR0913, PLR0917 # type: ignore (fsspec has bad resolved = await self._resolve_ranges(paths, starts, ends) - # fsspec's `max_gap` is obstore's `coalesce`: the largest gap between two - # ranges that may still be served by a single request. Left unset, defer to - # obstore's own default rather than restating it here. - coalesce: _CoalesceKwarg = {} if max_gap is None else {"coalesce": max_gap} - per_file_requests: dict[str, list[tuple[int, int | None, int]]] = defaultdict( list, ) @@ -511,11 +496,13 @@ async def _cat_ranges( # noqa: PLR0913, PLR0917 # type: ignore (fsspec has bad offsets = [r[0] for r in ranges] file_ends = [r[1] for r in ranges] + # fsspec's `max_gap` is obstore's `coalesce`: the largest gap between + # two ranges that may still be served by a single request. fut = store.get_ranges_async( path_no_bucket, starts=offsets, ends=file_ends, - **coalesce, + coalesce=max_gap, ) futs.append(fut) diff --git a/obstore/python/obstore/store.py b/obstore/python/obstore/store.py index bbd95fb2..764d2c42 100644 --- a/obstore/python/obstore/store.py +++ b/obstore/python/obstore/store.py @@ -246,7 +246,7 @@ def get_ranges( starts: Sequence[int], ends: Sequence[int | None] | None = None, lengths: Sequence[int | None] | None = None, - coalesce: int = 1024 * 1024, + coalesce: int | None = None, ) -> list[Bytes]: """Return the bytes stored at the specified location in the given byte ranges. @@ -268,7 +268,7 @@ async def get_ranges_async( starts: Sequence[int], ends: Sequence[int | None] | None = None, lengths: Sequence[int | None] | None = None, - coalesce: int = 1024 * 1024, + coalesce: int | None = None, ) -> list[Bytes]: """Call `get_ranges` asynchronously. diff --git a/obstore/src/get.rs b/obstore/src/get.rs index c723d847..7de648be 100644 --- a/obstore/src/get.rs +++ b/obstore/src/get.rs @@ -479,7 +479,7 @@ async fn _get_ranges( } #[pyfunction] -#[pyo3(signature = (store, path, *, starts, ends=None, lengths=None, coalesce=OBJECT_STORE_COALESCE_DEFAULT))] +#[pyo3(signature = (store, path, *, starts, ends=None, lengths=None, coalesce=None))] pub(crate) fn get_ranges( py: Python, store: PyObjectStore, @@ -487,15 +487,16 @@ pub(crate) fn get_ranges( starts: Vec, ends: Option>>, lengths: Option>>, - coalesce: u64, + coalesce: Option, ) -> PyObjectStoreResult> { let runtime = get_runtime(); let ranges = params_to_ranges(starts, ends, lengths)?; + let coalesce = coalesce.unwrap_or(OBJECT_STORE_COALESCE_DEFAULT); py.detach(|| runtime.block_on(_get_ranges(store, path, &ranges, coalesce))) } #[pyfunction] -#[pyo3(signature = (store, path, *, starts, ends=None, lengths=None, coalesce=OBJECT_STORE_COALESCE_DEFAULT))] +#[pyo3(signature = (store, path, *, starts, ends=None, lengths=None, coalesce=None))] pub(crate) fn get_ranges_async( py: Python, store: PyObjectStore, @@ -503,9 +504,10 @@ pub(crate) fn get_ranges_async( starts: Vec, ends: Option>>, lengths: Option>>, - coalesce: u64, + coalesce: Option, ) -> PyResult> { let ranges = params_to_ranges(starts, ends, lengths)?; + let coalesce = coalesce.unwrap_or(OBJECT_STORE_COALESCE_DEFAULT); pyo3_async_runtimes::tokio::future_into_py(py, async move { Ok(_get_ranges(store, path, &ranges, coalesce).await?) }) diff --git a/tests/test_fsspec.py b/tests/test_fsspec.py index 265346ab..beba5d36 100644 --- a/tests/test_fsspec.py +++ b/tests/test_fsspec.py @@ -605,7 +605,7 @@ async def spy(self, *args, **kwargs): # noqa: ANN001, ANN002, ANN003 monkeypatch.setattr(ObjectStoreMethods, "get_ranges_async", spy) - # Unset: obstore's own default applies, so `coalesce` is not forwarded at all. + # Unset: forwarded as None, which lets obstore apply its own default. assert fs.cat_ranges([path, path], [0, 200], [100, 300]) == [ data[0:100], data[200:300], From 275e3c68b2a4ba547fdd0b5da1edc58be346db89 Mon Sep 17 00:00:00 2001 From: Gunnar Schulze Date: Fri, 21 Aug 2026 07:09:33 +0200 Subject: [PATCH 04/22] revert: Undo range request changes --- obstore/python/obstore/_get.pyi | 74 ++++------------- obstore/python/obstore/fsspec.py | 134 ++++++------------------------- obstore/python/obstore/store.py | 12 +-- obstore/src/get.rs | 132 +++++++++--------------------- tests/test_fsspec.py | 58 +------------ tests/test_get.py | 46 ----------- 6 files changed, 90 insertions(+), 366 deletions(-) diff --git a/obstore/python/obstore/_get.pyi b/obstore/python/obstore/_get.pyi index ed062d17..ff0cbf37 100644 --- a/obstore/python/obstore/_get.pyi +++ b/obstore/python/obstore/_get.pyi @@ -328,33 +328,19 @@ def get_range( ) -> Bytes: """Return the bytes that are stored at the specified location in the given byte range. - The requested range may be bounded, open-ended, or relative to the end of the - object: - - - `start` plus either `end` or `length` requests a specific range of bytes. - If the given range is zero-length or starts after the end of the object, an - error will be returned. Additionally, if the range ends after the end of the - object, the entire remainder of the object will be returned. Otherwise, the - exact requested range will be returned. - - A non-negative `start` on its own requests all bytes from `start` onwards. - This is equivalent to `bytes={start}-` as an HTTP header. - - A negative `start` on its own requests the last `abs(start)` bytes. Note that - here, `abs(start)` is _the size of the request_, not a byte offset. This is - equivalent to `bytes=-{abs(start)}` as an HTTP header. + If the given range is zero-length or starts after the end of the object, an error + will be returned. Additionally, if the range ends after the end of the object, the + entire remainder of the object will be returned. Otherwise, the exact requested + range will be returned. Args: store: The ObjectStore instance to use. path: The path within ObjectStore to retrieve. Keyword Args: - start: The start of the byte range. If negative, the last `abs(start)` bytes - of the object are requested, and `end` and `length` must both be None. - end: The end of the byte range (exclusive). Mutually exclusive with `length`. - Defaults to None, in which case the range continues to the end of the - object. - length: The number of bytes of the byte range. Mutually exclusive with `end`. - Defaults to None, in which case the range continues to the end of the - object. + start: The start of the byte range. + end: The end of the byte range (exclusive). Either `end` or `length` must be non-None. + length: The number of bytes of the byte range. Either `end` or `length` must be non-None. Returns: A `Bytes` object implementing the Python buffer protocol, allowing @@ -380,53 +366,27 @@ def get_ranges( path: str, *, starts: Sequence[int], - ends: Sequence[int | None] | None = None, - lengths: Sequence[int | None] | None = None, - coalesce: int | None = None, + ends: Sequence[int] | None = None, + lengths: Sequence[int] | None = None, + coalesce: int = 1024 * 1024, ) -> list[Bytes]: """Return the bytes stored at the specified location in the given byte ranges. - Each range is described by one element of `starts` and the element at the same - position in `ends` or `lengths`, and follows the same semantics as - [get_range][obstore.get_range]. In particular, a negative element of `starts` - requests the last `abs(start)` bytes, and a range whose end is left unspecified - continues to the end of the object. Because ranges are matched up by position, - every sequence given must have the same length as `starts`; a mismatch raises - `ValueError`. - - `ends` and `lengths` may both be given at once, so long as at most one of the two - is non-None for any individual range — for example `ends=[10, None], - lengths=[None, 5]` bounds the first range by its end offset and the second by its - length. Omitting both reads every range to the end of the object. - To improve performance this will: - Transparently combine ranges less than `coalesce` bytes apart into a single underlying request (defaults to 1MB) - Make multiple `fetch` requests in parallel (up to maximum of 10) - !!! note - - Combining ranges requires knowing where each one ends, so if any requested - range is open-ended or relative to the end of the object, one additional - `head` request is made to resolve the size of the object. Requests made up - entirely of bounded ranges do not incur this cost. - Args: store: The ObjectStore instance to use. path: The path within ObjectStore to retrieve. Other Args: - starts: A sequence of `int` where each offset starts. A negative value - requests the last `abs(start)` bytes for that range, in which case the - corresponding elements of `ends` and `lengths` must be None. - ends: A sequence of `int` where each offset ends (exclusive). An element may - be None to continue that range to the end of the object. - lengths: A sequence of `int` with the number of bytes of each byte range. An - element may be None to continue that range to the end of the object. - coalesce: Maximum distance in bytes between ranges that will be coalesced - into a single request. Defaults to None, in which case obstore applies - its own default of 1MiB. Set to `0` to disable coalescing. + starts: A sequence of `int` where each offset starts. + ends: A sequence of `int` where each offset ends (exclusive). Either `ends` or `lengths` must be non-None. + lengths: A sequence of `int` with the number of bytes of each byte range. Either `ends` or `lengths` must be non-None. + coalesce: Maximum distance in bytes between ranges that will be coalesced into a single request. Defaults to 1MiB. Set to `0` to disable coalescing. Returns: A sequence of `Bytes`, one for each range. This `Bytes` object implements the @@ -440,9 +400,9 @@ async def get_ranges_async( path: str, *, starts: Sequence[int], - ends: Sequence[int | None] | None = None, - lengths: Sequence[int | None] | None = None, - coalesce: int | None = None, + ends: Sequence[int] | None = None, + lengths: Sequence[int] | None = None, + coalesce: int = 1024 * 1024, ) -> list[Bytes]: """Call `get_ranges` asynchronously. diff --git a/obstore/python/obstore/fsspec.py b/obstore/python/obstore/fsspec.py index 3e9409bc..d6c76a92 100644 --- a/obstore/python/obstore/fsspec.py +++ b/obstore/python/obstore/fsspec.py @@ -35,10 +35,9 @@ 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, cast, overload +from typing import TYPE_CHECKING, Literal, overload from urllib.parse import urlparse import fsspec.asyn @@ -50,7 +49,7 @@ if TYPE_CHECKING: import sys - from collections.abc import Coroutine, Sequence + from collections.abc import Coroutine, Iterable from datetime import datetime from typing import Any @@ -114,35 +113,6 @@ """A type hint for all supported protocols.""" -def _needs_object_size(start: int | None, end: int | None) -> bool: - """Whether a range must be resolved against the object size first. - - A negative `start` on its own is a suffix request, which obstore supports - directly. Anything else counting back from the end of the object is not. - """ - return end is not None and (end < 0 or (start is not None and start < 0)) - - -def _apply_object_size( - start: int | None, - end: int | None, - size: int | None, -) -> tuple[int, int | None]: - """Rewrite the bounds of a range that count back from the end of an object. - - A `size` of None leaves the range alone, for objects whose size was never - needed. - """ - if start is None: - start = 0 - if size is None: - return start, end - return ( - max(0, size + start) if start < 0 else start, - max(0, size + end) if end is not None and end < 0 else end, - ) - - class FsspecStore(fsspec.asyn.AsyncFileSystem): """An fsspec implementation based on a obstore Store. @@ -419,19 +389,12 @@ async def _cat_file( resp = await store.get_async(path) return (await resp.bytes_async()).to_bytes() - if _needs_object_size(start, end): - # See `_resolve_ranges` for why this needs the object size. - size = (await store.head_async(path))["size"] - if start is not None and start < 0: - start = max(0, size + start) - if end is not None and end < 0: - end = max(0, size + end) - - range_bytes = await store.get_range_async( - path, - start=0 if start is None else start, - end=end, - ) + if start is None or end is None: + raise NotImplementedError( + "cat_file not implemented for start=None xor end=None", + ) + + range_bytes = await store.get_range_async(path, start=start, end=end) return range_bytes.to_bytes() async def _cat( # type: ignore (fsspec has bad typing) @@ -463,30 +426,23 @@ 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: Sequence[int | None] | int | None, - ends: Sequence[int | None] | int | None, - max_gap: int | None = None, - batch_size: int | None = None, # noqa: ARG002 - on_error: str = "return", # noqa: ARG002 + 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 **_kwargs: Any, ) -> list[bytes]: - # A non-iterable start or end applies to every path, per fsspec's - # AsyncFileSystem._cat_ranges. - if not isinstance(starts, Iterable): + if isinstance(starts, int): starts = [starts] * len(paths) - if not isinstance(ends, Iterable): + if isinstance(ends, int): ends = [ends] * len(paths) if not len(paths) == len(starts) == len(ends): raise ValueError - resolved = await self._resolve_ranges(paths, starts, ends) - - per_file_requests: dict[str, list[tuple[int, int | None, int]]] = defaultdict( - list, - ) - for idx, (path, (start, end)) in enumerate( - zip(paths, resolved, strict=True), - ): + 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)) futs: list[Coroutine[Any, Any, list[Bytes]]] = [] @@ -495,65 +451,23 @@ async def _cat_ranges( # noqa: PLR0913, PLR0917 # type: ignore (fsspec has bad store = self._construct_store(bucket) offsets = [r[0] for r in ranges] - file_ends = [r[1] for r in ranges] - # fsspec's `max_gap` is obstore's `coalesce`: the largest gap between - # two ranges that may still be served by a single request. - fut = store.get_ranges_async( - path_no_bucket, - starts=offsets, - ends=file_ends, - coalesce=max_gap, - ) + ends = [r[1] for r in ranges] + fut = store.get_ranges_async(path_no_bucket, starts=offsets, ends=ends) futs.append(fut) result = await asyncio.gather(*futs) output_buffers: list[bytes] = [b""] * len(paths) - for per_file_request, buffers in zip( - per_file_requests.items(), - result, - strict=True, - ): + # 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 - for buffer, ranges_ in zip(buffers, ranges, strict=True): + # 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() return output_buffers - async def _resolve_ranges( - self, - paths: list[str], - starts: Sequence[int | None], - ends: Sequence[int | None], - ) -> list[tuple[int, int | None]]: - """Rewrite fsspec's range vocabulary into the subset obstore accepts. - - `start` is normalized from None to 0. Ranges counting back from the end of - the object are resolved against its size, at the cost of one `head` request - per object; a lone negative `start` is left alone, since obstore expresses - that directly as a suffix request. - """ - sized_paths = sorted( - { - path - for path, start, end in zip(paths, starts, ends, strict=True) - if _needs_object_size(start, end) - }, - ) - sizes = dict( - zip( - sized_paths, - # fsspec types `_sizes` as `list[None]`; the values are really sizes. - cast("list[int]", await self._sizes(sized_paths)), - strict=True, - ), - ) - return [ - _apply_object_size(start, end, sizes.get(path)) - for path, start, end in zip(paths, starts, ends, strict=True) - ] - async def _put_file( self, lpath: str, diff --git a/obstore/python/obstore/store.py b/obstore/python/obstore/store.py index 764d2c42..72fe1b80 100644 --- a/obstore/python/obstore/store.py +++ b/obstore/python/obstore/store.py @@ -244,9 +244,9 @@ def get_ranges( path: str, *, starts: Sequence[int], - ends: Sequence[int | None] | None = None, - lengths: Sequence[int | None] | None = None, - coalesce: int | None = None, + ends: Sequence[int] | None = None, + lengths: Sequence[int] | None = None, + coalesce: int = 1024 * 1024, ) -> list[Bytes]: """Return the bytes stored at the specified location in the given byte ranges. @@ -266,9 +266,9 @@ async def get_ranges_async( path: str, *, starts: Sequence[int], - ends: Sequence[int | None] | None = None, - lengths: Sequence[int | None] | None = None, - coalesce: int | None = None, + ends: Sequence[int] | None = None, + lengths: Sequence[int] | None = None, + coalesce: int = 1024 * 1024, ) -> list[Bytes]: """Call `get_ranges` asynchronously. diff --git a/obstore/src/get.rs b/obstore/src/get.rs index 7de648be..afdcc978 100644 --- a/obstore/src/get.rs +++ b/obstore/src/get.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use bytes::Bytes; use chrono::{DateTime, Utc}; use futures::stream::{BoxStream, Fuse}; -use futures::{StreamExt, TryFutureExt}; +use futures::StreamExt; use object_store::{ coalesce_ranges, Attributes, GetOptions, GetRange, GetResult, ObjectMeta, ObjectStore, ObjectStoreExt, OBJECT_STORE_COALESCE_DEFAULT, @@ -117,7 +117,7 @@ impl<'py> FromPyObject<'_, 'py> for PyGetRange { } else if let Ok(suffix_range) = obj.extract::() { Ok(Self(suffix_range.into())) } else { - Err(PyValueError::new_err("Unexpected input for byte range.\nExpected two-integer tuple or list, or dict with 'offset' or 'suffix' key.")) + Err(PyValueError::new_err("Unexpected input for byte range.\nExpected two-integer tuple or list, or dict with 'offset' or 'suffix' key." )) } } } @@ -382,20 +382,15 @@ pub(crate) fn get_range( py: Python, store: PyObjectStore, path: PyPath, - start: i64, + start: u64, end: Option, length: Option, -) -> PyObjectStoreResult { +) -> PyObjectStoreResult { let runtime = get_runtime(); let range = params_to_range(start, end, length)?; py.detach(|| { - let out = runtime.block_on( - store - .as_ref() - .get_opts(path.as_ref(), GetOptions::new().with_range(range.into())) - .and_then(GetResult::bytes), - )?; - Ok::<_, PyObjectStoreError>(PyBytes::new(out)) + let out = runtime.block_on(store.as_ref().get_range(path.as_ref(), range))?; + Ok::<_, PyObjectStoreError>(pyo3_bytes::PyBytes::new(out)) }) } @@ -405,7 +400,7 @@ pub(crate) fn get_range_async( py: Python, store: PyObjectStore, path: PyPath, - start: i64, + start: u64, end: Option, length: Option, ) -> PyResult> { @@ -413,64 +408,36 @@ pub(crate) fn get_range_async( pyo3_async_runtimes::tokio::future_into_py(py, async move { let out = store .as_ref() - .get_opts(path.as_ref(), GetOptions::new().with_range(range.into())) - .and_then(GetResult::bytes) + .get_range(path.as_ref(), range) .await .map_err(PyObjectStoreError::ObjectStoreError)?; - Ok(PyBytes::new(out)) + Ok(pyo3_bytes::PyBytes::new(out)) }) } fn params_to_range( - start: i64, + start: u64, end: Option, length: Option, -) -> PyObjectStoreResult { - if start < 0 { - if end.is_some() || length.is_some() { - return Err( - PyValueError::new_err("end and length must be None if start is negative.").into(), - ); - } - return Ok(GetRange::Suffix(start.unsigned_abs())); - } - - let start = start as u64; +) -> PyObjectStoreResult> { match (end, length) { (Some(_), Some(_)) => { Err(PyValueError::new_err("end and length cannot both be non-None.").into()) } - (None, None) => Ok(GetRange::Offset(start)), - (Some(end), None) => validate_range(start..end).map(GetRange::Bounded), - (None, Some(length)) => validate_range(start..start + length).map(GetRange::Bounded), + (None, None) => Err(PyValueError::new_err("Either end or length must be non-None.").into()), + (Some(end), None) => validate_range(start..end), + (None, Some(length)) => validate_range(start..start + length), } } async fn _get_ranges( store: PyObjectStore, path: PyPath, - ranges: &[GetRange], + ranges: &[Range], coalesce: u64, ) -> PyObjectStoreResult> { - let mut len: Option = None; - let mut resolved = Vec::with_capacity(ranges.len()); - for range in ranges { - resolved.push(match range { - GetRange::Bounded(r) => r.clone(), - other => { - let size = match len { - Some(len) => len, - None => *len.insert(store.as_ref().head(path.as_ref()).await?.size), - }; - other - .as_range(size) - .map_err(|err| PyValueError::new_err(err.to_string()))? - } - }); - } - let out = coalesce_ranges( - &resolved, + ranges, |range| store.as_ref().get_range(path.as_ref(), range), coalesce, ) @@ -479,80 +446,61 @@ async fn _get_ranges( } #[pyfunction] -#[pyo3(signature = (store, path, *, starts, ends=None, lengths=None, coalesce=None))] +#[pyo3(signature = (store, path, *, starts, ends=None, lengths=None, coalesce=OBJECT_STORE_COALESCE_DEFAULT))] pub(crate) fn get_ranges( py: Python, store: PyObjectStore, path: PyPath, - starts: Vec, - ends: Option>>, - lengths: Option>>, - coalesce: Option, + starts: Vec, + ends: Option>, + lengths: Option>, + coalesce: u64, ) -> PyObjectStoreResult> { let runtime = get_runtime(); let ranges = params_to_ranges(starts, ends, lengths)?; - let coalesce = coalesce.unwrap_or(OBJECT_STORE_COALESCE_DEFAULT); py.detach(|| runtime.block_on(_get_ranges(store, path, &ranges, coalesce))) } #[pyfunction] -#[pyo3(signature = (store, path, *, starts, ends=None, lengths=None, coalesce=None))] +#[pyo3(signature = (store, path, *, starts, ends=None, lengths=None, coalesce=OBJECT_STORE_COALESCE_DEFAULT))] pub(crate) fn get_ranges_async( py: Python, store: PyObjectStore, path: PyPath, - starts: Vec, - ends: Option>>, - lengths: Option>>, - coalesce: Option, + starts: Vec, + ends: Option>, + lengths: Option>, + coalesce: u64, ) -> PyResult> { let ranges = params_to_ranges(starts, ends, lengths)?; - let coalesce = coalesce.unwrap_or(OBJECT_STORE_COALESCE_DEFAULT); pyo3_async_runtimes::tokio::future_into_py(py, async move { Ok(_get_ranges(store, path, &ranges, coalesce).await?) }) } fn params_to_ranges( - starts: Vec, - ends: Option>>, - lengths: Option>>, -) -> PyObjectStoreResult> { - for (name, len) in [("ends", ends.as_ref()), ("lengths", lengths.as_ref())] - .into_iter() - .filter_map(|(name, seq)| seq.map(|seq| (name, seq.len()))) - { - if len != starts.len() { - return Err(PyValueError::new_err(format!( - "starts and {name} must have the same length, got {} and {len}.", - starts.len(), - )) - .into()); - } - } - + starts: Vec, + ends: Option>, + lengths: Option>, +) -> PyObjectStoreResult>> { match (ends, lengths) { - // Consistent with `get_range`: an unbounded range reads to the end of the - // object. - (None, None) => starts - .into_iter() - .map(|start| params_to_range(start, None, None)) - .collect(), - (Some(ends), Some(lengths)) => starts - .into_iter() - .zip(ends) - .zip(lengths) - .map(|((start, end), length)| params_to_range(start, end, length)) - .collect(), + (Some(_), Some(_)) => { + Err(PyValueError::new_err("ends and lengths cannot both be non-None.").into()) + } + (None, None) => { + Err(PyValueError::new_err("Either ends or lengths must be non-None.").into()) + } (Some(ends), None) => starts .into_iter() .zip(ends) - .map(|(start, end)| params_to_range(start, end, None)) + .map(|(start, end)| start..end) + .map(validate_range) .collect(), (None, Some(lengths)) => starts .into_iter() .zip(lengths) - .map(|(start, length)| params_to_range(start, None, length)) + .map(|(start, length)| start..start + length) + .map(validate_range) .collect(), } } diff --git a/tests/test_fsspec.py b/tests/test_fsspec.py index beba5d36..3da2f5d3 100644 --- a/tests/test_fsspec.py +++ b/tests/test_fsspec.py @@ -13,7 +13,6 @@ from fsspec.registry import _registry from obstore.fsspec import FsspecStore, register -from obstore.store import ObjectStoreMethods from tests.conftest import TEST_BUCKET_NAME if TYPE_CHECKING: @@ -573,56 +572,6 @@ def test_multi_file_ops(minio_bucket: tuple[S3Config, ClientConfig]): assert out == [f"{bucket}/afile"] -def test_cat_file(fs: FsspecStore): - data = os.urandom(10000) - path = f"{TEST_BUCKET_NAME}/data1" - fs.pipe_file(path, data) - - assert fs.cat_file(path) == data - assert fs.cat_file(path, start=100, end=200) == data[100:200] - - # Either bound on its own. - assert fs.cat_file(path, start=100) == data[100:] - assert fs.cat_file(path, end=200) == data[:200] - - # Bounds counted back from the end of the object. - assert fs.cat_file(path, start=-100) == data[-100:] - assert fs.cat_file(path, start=100, end=-100) == data[100:-100] - assert fs.cat_file(path, start=-200, end=-100) == data[-200:-100] - - -def test_cat_ranges_max_gap(fs: FsspecStore, monkeypatch: pytest.MonkeyPatch): - data = os.urandom(10000) - path = f"{TEST_BUCKET_NAME}/data1" - fs.pipe_file(path, data) - - seen: list[int | None] = [] - original = ObjectStoreMethods.get_ranges_async - - async def spy(self, *args, **kwargs): # noqa: ANN001, ANN002, ANN003 - seen.append(kwargs.get("coalesce")) - return await original(self, *args, **kwargs) - - monkeypatch.setattr(ObjectStoreMethods, "get_ranges_async", spy) - - # Unset: forwarded as None, which lets obstore apply its own default. - assert fs.cat_ranges([path, path], [0, 200], [100, 300]) == [ - data[0:100], - data[200:300], - ] - assert seen == [None] - - # Set: forwarded verbatim as obstore's `coalesce`, including 0 to disable - # coalescing entirely. - for max_gap in (0, 1000): - seen.clear() - assert fs.cat_ranges([path, path], [0, 200], [100, 300], max_gap=max_gap) == [ - data[0:100], - data[200:300], - ] - assert seen == [max_gap] - - def test_cat_ranges_one(fs: FsspecStore): data1 = os.urandom(10000) fs.pipe_file(f"{TEST_BUCKET_NAME}/data1", data1) @@ -682,15 +631,14 @@ 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) - path1 = f"{TEST_BUCKET_NAME}/data1" - path2 = f"{TEST_BUCKET_NAME}/data2" - fs.pipe({path1: data1, path2: data2}) + fs.pipe({"data1": data1, "data2": data2}) # single range in each file - out = fs.cat_ranges([path1, path1, path2], [-10, None, 10], [None, -10, -10]) + out = fs.cat_ranges(["data1", "data1", "data2"], [-10, None, 10], [None, -10, -10]) assert out == [data1[-10:], data1[:-10], data2[10:-10]] diff --git a/tests/test_get.py b/tests/test_get.py index edfa9f1c..a301d32e 100644 --- a/tests/test_get.py +++ b/tests/test_get.py @@ -113,19 +113,6 @@ def test_get_range(): view = memoryview(buffer) assert view == data[5:15] - buffer = store.get_range(path, start=4300) - view = memoryview(buffer) - assert view == data[4300:4400] - - buffer = store.get_range(path, start=-100) - view = memoryview(buffer) - assert view == data[4300:4400] - - # A suffix longer than the object yields the whole object. - buffer = store.get_range(path, start=-999999) - view = memoryview(buffer) - assert view == data - def test_get_ranges(): store = MemoryStore() @@ -143,27 +130,6 @@ def test_get_ranges(): for start, end, buffer in zip(starts, ends, buffers): assert memoryview(buffer) == data[start:end] - # A `None` element leaves that one range open-ended; a negative start makes it - # a suffix request. Omitting `ends` and `lengths` reads every range to the end. - buffers = store.get_ranges(path, starts=[5, 4300, -100], ends=[15, None, None]) - assert [memoryview(b) for b in buffers] == [ - data[5:15], - data[4300:4400], - data[4300:4400], - ] - - buffers = store.get_ranges(path, starts=[4300, -100]) - assert [memoryview(b) for b in buffers] == [data[4300:4400], data[4300:4400]] - - # `ends` and `lengths` may be mixed, at most one per range. - buffers = store.get_ranges( - path, - starts=[5, 20], - ends=[15, None], - lengths=[None, 10], - ) - assert [memoryview(b) for b in buffers] == [data[5:15], data[20:30]] - lengths = [10, 10, 10, 10] buffers = store.get_ranges(path, starts=starts, lengths=lengths) @@ -241,12 +207,6 @@ def test_get_range_invalid_range(): with pytest.raises(ValueError, match="Invalid range"): store.get_range(path, start=10, length=0) - with pytest.raises(ValueError, match="end and length must be None"): - store.get_range(path, start=-10, end=10) - - with pytest.raises(ValueError, match="end and length must be None"): - store.get_range(path, start=-10, length=10) - def test_get_ranges_invalid_range(): store = MemoryStore() @@ -264,12 +224,6 @@ def test_get_ranges_invalid_range(): with pytest.raises(ValueError, match="Invalid range"): store.get_ranges(path, starts=[10, 20], lengths=[10, 0]) - with pytest.raises(ValueError, match="starts and ends must have the same length"): - store.get_ranges(path, starts=[10, 20], ends=[30]) - - with pytest.raises(ValueError, match="starts and lengths must have the same"): - store.get_ranges(path, starts=[10], lengths=[10, 20]) - def test_access_getresult_attributes_after_reading_stream(): store = MemoryStore() From 671296651258e454aa5f6a5c151027da859cdfbe Mon Sep 17 00:00:00 2001 From: Gunnar Schulze Date: Fri, 21 Aug 2026 14:59:29 +0200 Subject: [PATCH 05/22] feat: Handle all forms of range requests in fsspec --- obstore/python/obstore/fsspec.py | 217 ++++++++++++++++++++++++++----- tests/test_fsspec.py | 138 +++++++++++++++++++- 2 files changed, 316 insertions(+), 39 deletions(-) diff --git a/obstore/python/obstore/fsspec.py b/obstore/python/obstore/fsspec.py index d6c76a92..1e20fdb6 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, TypedDict, cast, overload from urllib.parse import urlparse import fsspec.asyn @@ -49,11 +49,18 @@ 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, + GetOptions, + OffsetRange, + ReadableFile, + SuffixRange, + WritableFile, + ) from obstore.store import ( AzureConfig, AzureCredentialProvider, @@ -113,6 +120,31 @@ """A type hint for all supported protocols.""" +class _CoalesceKwarg(TypedDict, total=False): + """The optional `coalesce` argument of [obstore.get_ranges][].""" + + coalesce: int + + +def _needs_object_size(start: int | None, end: int | None) -> bool: + """Whether resolving a range requires knowing the size of the object.""" + return end is not None and (end < 0 or (start is not None and start < 0)) + + +def _apply_object_size( + start: int, + end: int | None, + size: int | None, +) -> tuple[int, int | None]: + """Resolve bounds that count back from the end of an object of `size`.""" + if size is None: + return start, end + return ( + max(0, size + start) if start < 0 else start, + max(0, size + end) if end is not None and end < 0 else end, + ) + + class FsspecStore(fsspec.asyn.AsyncFileSystem): """An fsspec implementation based on a obstore Store. @@ -382,17 +414,34 @@ async def _cat_file( end: int | None = None, **_kwargs: Any, ) -> bytes: + # `_info` splits the path itself, so keep the original before rebinding. + full_path = path bucket, path = self._split_path(path) store = self._construct_store(bucket) - if start is None and end is None: + # A zero start with no end is the whole object, so skip the range request. + if not start and end is None: resp = await store.get_async(path) return (await resp.bytes_async()).to_bytes() - if start is None or end is None: - raise NotImplementedError( - "cat_file not implemented for start=None xor end=None", - ) + start = 0 if start is None else start + + if _needs_object_size(start, end): + # A negative `end` has no `GetOptions` equivalent, so resolve it against + # the object size. + size = (await self._info(full_path))["size"] + start, end = _apply_object_size(start, end, size) + + if end is None: + # `get_range` mirrors the Rust API and takes bounded ranges only. `get` + # covers the other two forms. + options: GetOptions + if start < 0: + options = {"range": {"suffix": -start}} + else: + options = {"range": {"offset": start}} + resp = await store.get_async(path, options=options) + return (await resp.bytes_async()).to_bytes() range_bytes = await store.get_range_async(path, start=start, end=end) return range_bytes.to_bytes() @@ -426,48 +475,144 @@ 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", # noqa: ARG002 **_kwargs: Any, ) -> list[bytes]: - if isinstance(starts, int): + # A non-iterable start or end applies to every path, per fsspec's + # AsyncFileSystem._cat_ranges. + 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)) + resolved_starts, resolved_ends = await self._resolve_inexpressible_bounds( + paths, + starts, + ends, + ) - 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) + # fsspec's `max_gap` is obstore's `coalesce`. + coalesce = ( + _CoalesceKwarg() if max_gap is None else _CoalesceKwarg(coalesce=max_gap) + ) - 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) + # `get_ranges` only merges bounded ranges, and 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, OffsetRange | SuffixRange]] = [] + for idx, (path, start, end) in enumerate( + zip(paths, resolved_starts, resolved_ends, strict=True), + ): + if end is not None: + per_file_bounded_requests[path].append((idx, start, end)) + elif start < 0: + open_ended_requests.append((idx, path, {"suffix": -start})) + else: + open_ended_requests.append((idx, path, {"offset": start})) + + futs: list[Coroutine[Any, Any, list[tuple[int, bytes]]]] = [ + self._cat_bounded_ranges(path, ranges, coalesce) + for path, ranges in per_file_bounded_requests.items() + ] + futs += [ + self._cat_open_ended_range(idx, path, byte_range) + for idx, path, byte_range in open_ended_requests + ] - result = await asyncio.gather(*futs) + # Batched, to limit how many requests are in flight at once. Like fsspec's + # own `_cat_ranges`, `on_error` is ignored, so failures propagate. + results = cast( + "list[list[tuple[int, bytes]]]", + await fsspec.asyn._run_coros_in_chunks( # noqa: SLF001 + futs, + batch_size=batch_size or self.batch_size, + nofiles=True, + ), + ) 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() + for responses in results: + for idx, buffer in responses: + output_buffers[idx] = buffer return output_buffers + async def _cat_bounded_ranges( + self, + path: str, + ranges: list[tuple[int, int, int]], + coalesce: _CoalesceKwarg, + ) -> list[tuple[int, bytes]]: + """Read several bounded ranges of one object, merging nearby ones.""" + bucket, path = self._split_path(path) + store = self._construct_store(bucket) + buffers = await store.get_ranges_async( + path, + starts=[start for _, start, _ in ranges], + ends=[end for _, _, end in ranges], + **coalesce, + ) + return [ + (idx, buffer.to_bytes()) + for (idx, _, _), buffer in zip(ranges, buffers, strict=True) + ] + + async def _cat_open_ended_range( + self, + idx: int, + path: str, + byte_range: OffsetRange | SuffixRange, + ) -> list[tuple[int, bytes]]: + """Read 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. + """ + bucket, path = self._split_path(path) + store = self._construct_store(bucket) + resp = await store.get_async(path, options={"range": byte_range}) + return [(idx, (await resp.bytes_async()).to_bytes())] + + async def _resolve_inexpressible_bounds( + self, + paths: list[str], + starts: Sequence[int | None], + ends: Sequence[int | None], + ) -> tuple[list[int], list[int | None]]: + """Resolve the bounds that obstore cannot express.""" + paths_needing_size = sorted( + { + path + for path, start, end in zip(paths, 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. + # fsspec types the result as `list[None]`, but the values are sizes. + sizes = dict( + zip( + paths_needing_size, + cast("list[int]", await self._sizes(paths_needing_size)), + strict=True, + ), + ) + resolved = [ + _apply_object_size(0 if start is None else start, end, sizes.get(path)) + for path, start, end in zip(paths, starts, ends, strict=True) + ] + return [start for start, _ in resolved], [end for _, end in resolved] + async def _put_file( self, lpath: str, diff --git a/tests/test_fsspec.py b/tests/test_fsspec.py index 3da2f5d3..5c0da59f 100644 --- a/tests/test_fsspec.py +++ b/tests/test_fsspec.py @@ -8,11 +8,13 @@ from unittest.mock import patch import fsspec +import fsspec.asyn import pyarrow.parquet as pq import pytest from fsspec.registry import _registry from obstore.fsspec import FsspecStore, register +from obstore.store import ObjectStoreMethods from tests.conftest import TEST_BUCKET_NAME if TYPE_CHECKING: @@ -572,6 +574,135 @@ def test_multi_file_ops(minio_bucket: tuple[S3Config, ClientConfig]): assert out == [f"{bucket}/afile"] +def test_cat_file(fs: FsspecStore): + """Test that `cat_file` accepts every range form fsspec documents.""" + data = os.urandom(10000) + path = f"{TEST_BUCKET_NAME}/data1" + fs.pipe_file(path, data) + + assert fs.cat_file(path) == data + assert fs.cat_file(path, start=10, end=20) == data[10:20] + + # Either bound on its own. + assert fs.cat_file(path, start=10) == data[10:] + assert fs.cat_file(path, end=20) == data[:20] + + # Bounds counted back from the end of the object. + assert fs.cat_file(path, start=-10) == data[-10:] + assert fs.cat_file(path, start=10, end=-10) == data[10:-10] + assert fs.cat_file(path, start=-20, end=-10) == data[-20:-10] + + +def test_cat_file_zero_start(fs: FsspecStore, monkeypatch: pytest.MonkeyPatch): + """Test that a zero start with no end is not sent as a range request.""" + data = os.urandom(10000) + path = f"{TEST_BUCKET_NAME}/data1" + fs.pipe_file(path, data) + + ranges: list[dict[str, int] | None] = [] + original = ObjectStoreMethods.get_async + + async def spy(self, *args, **kwargs): # noqa: ANN001, ANN002, ANN003 + ranges.append((kwargs.get("options") or {}).get("range")) + return await original(self, *args, **kwargs) + + monkeypatch.setattr(ObjectStoreMethods, "get_async", spy) + + assert fs.cat_file(path, start=0) == data + assert fs.cat_file(path, start=10) == data[10:] + assert ranges == [None, {"offset": 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) + + seen: list[int | None] = [] + original = ObjectStoreMethods.get_ranges_async + + async def spy(self, *args, **kwargs): # noqa: ANN001, ANN002, ANN003 + seen.append(kwargs.get("coalesce")) + return await original(self, *args, **kwargs) + + monkeypatch.setattr(ObjectStoreMethods, "get_ranges_async", spy) + + # 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 seen == [None] + + # With max_gap: passed straight through, including 0 to turn coalescing off. + for max_gap in (0, 1000): + seen.clear() + assert fs.cat_ranges([path, path], [0, 20], [10, 30], max_gap=max_gap) == [ + data[0:10], + data[20:30], + ] + assert seen == [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) + + batched: list[tuple[list[int], list[int]]] = [] + individual: list[dict[str, int]] = [] + original_get_ranges = ObjectStoreMethods.get_ranges_async + original_get = ObjectStoreMethods.get_async + + async def spy_get_ranges(self, *args, **kwargs): # noqa: ANN001, ANN002, ANN003 + batched.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 + individual.append(kwargs["options"]["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] * 3, [0, 20, -5], [10, None, None]) + assert out == [data[0:10], data[20:], data[-5:]] + + # Only the bounded range reaches `get_ranges`, so only it gets coalesced. + assert batched == [([0], [10])] + + # The other two go through `get` instead, and may complete in either order. + assert len(individual) == 2 + assert {"offset": 20} in individual + assert {"suffix": 5} in individual + + +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) + + seen: list[int | None] = [] + original = fsspec.asyn._run_coros_in_chunks + + async def spy(coros, **kwargs): # noqa: ANN001, ANN003 + seen.append(kwargs.get("batch_size")) + return await original(coros, **kwargs) + + monkeypatch.setattr(fsspec.asyn, "_run_coros_in_chunks", spy) + + # Unset: forwarded as None, so fsspec infers its own default. + assert fs.cat_ranges([path], [0], [10]) == [data[0:10]] + assert seen == [None] + + seen.clear() + assert fs.cat_ranges([path], [0], [10], batch_size=4) == [data[0:10]] + assert seen == [4] + + def test_cat_ranges_one(fs: FsspecStore): data1 = os.urandom(10000) fs.pipe_file(f"{TEST_BUCKET_NAME}/data1", data1) @@ -631,14 +762,15 @@ 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" + fs.pipe({path1: data1, path2: data2}) # single range in each file - out = fs.cat_ranges(["data1", "data1", "data2"], [-10, None, 10], [None, -10, -10]) + out = fs.cat_ranges([path1, path1, path2], [-10, None, 10], [None, -10, -10]) assert out == [data1[-10:], data1[:-10], data2[10:-10]] From b89eb43ed2ede0a22001000fddb7be4acf14882e Mon Sep 17 00:00:00 2001 From: Gunnar Schulze Date: Tue, 25 Aug 2026 17:40:02 +0200 Subject: [PATCH 06/22] refactor: Remove `_CoalesceKwarg` --- obstore/python/obstore/fsspec.py | 30 ++++++++++++------------------ 1 file changed, 12 insertions(+), 18 deletions(-) diff --git a/obstore/python/obstore/fsspec.py b/obstore/python/obstore/fsspec.py index 1e20fdb6..a18d426a 100644 --- a/obstore/python/obstore/fsspec.py +++ b/obstore/python/obstore/fsspec.py @@ -37,7 +37,7 @@ from collections.abc import Iterable from functools import cached_property, lru_cache from pathlib import Path -from typing import TYPE_CHECKING, Literal, TypedDict, cast, overload +from typing import TYPE_CHECKING, Literal, cast, overload from urllib.parse import urlparse import fsspec.asyn @@ -120,12 +120,6 @@ """A type hint for all supported protocols.""" -class _CoalesceKwarg(TypedDict, total=False): - """The optional `coalesce` argument of [obstore.get_ranges][].""" - - coalesce: int - - def _needs_object_size(start: int | None, end: int | None) -> bool: """Whether resolving a range requires knowing the size of the object.""" return end is not None and (end < 0 or (start is not None and start < 0)) @@ -497,11 +491,6 @@ async def _cat_ranges( # noqa: PLR0913, PLR0917 # type: ignore (fsspec has bad ends, ) - # fsspec's `max_gap` is obstore's `coalesce`. - coalesce = ( - _CoalesceKwarg() if max_gap is None else _CoalesceKwarg(coalesce=max_gap) - ) - # `get_ranges` only merges bounded ranges, and 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 @@ -520,7 +509,7 @@ async def _cat_ranges( # noqa: PLR0913, PLR0917 # type: ignore (fsspec has bad open_ended_requests.append((idx, path, {"offset": start})) futs: list[Coroutine[Any, Any, list[tuple[int, bytes]]]] = [ - self._cat_bounded_ranges(path, ranges, coalesce) + self._cat_bounded_ranges(path, ranges, max_gap) for path, ranges in per_file_bounded_requests.items() ] futs += [ @@ -549,21 +538,26 @@ async def _cat_ranges( # noqa: PLR0913, PLR0917 # type: ignore (fsspec has bad async def _cat_bounded_ranges( self, path: str, - ranges: list[tuple[int, int, int]], - coalesce: _CoalesceKwarg, + ranges: list[tuple[int, int, int]], # (output index, start, end) + max_gap: int | None, ) -> list[tuple[int, bytes]]: """Read several bounded ranges of one object, merging nearby ones.""" bucket, path = self._split_path(path) store = self._construct_store(bucket) + 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} buffers = await store.get_ranges_async( path, - starts=[start for _, start, _ in ranges], - ends=[end for _, _, end in ranges], + starts=starts, + ends=ends, + lengths=None, **coalesce, ) return [ (idx, buffer.to_bytes()) - for (idx, _, _), buffer in zip(ranges, buffers, strict=True) + for idx, buffer in zip(indices, buffers, strict=True) ] async def _cat_open_ended_range( From 33eed4692f99242d7fc4a8e21b412c2729216e16 Mon Sep 17 00:00:00 2001 From: Gunnar Schulze Date: Tue, 25 Aug 2026 17:46:43 +0200 Subject: [PATCH 07/22] fix: Apply object size only to ranges that need it --- obstore/python/obstore/fsspec.py | 10 +++++----- tests/test_fsspec.py | 14 +++++++++----- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/obstore/python/obstore/fsspec.py b/obstore/python/obstore/fsspec.py index a18d426a..0d70d8f2 100644 --- a/obstore/python/obstore/fsspec.py +++ b/obstore/python/obstore/fsspec.py @@ -128,11 +128,9 @@ def _needs_object_size(start: int | None, end: int | None) -> bool: def _apply_object_size( start: int, end: int | None, - size: int | None, + size: int, ) -> tuple[int, int | None]: """Resolve bounds that count back from the end of an object of `size`.""" - if size is None: - return start, end return ( max(0, size + start) if start < 0 else start, max(0, size + end) if end is not None and end < 0 else end, @@ -593,7 +591,7 @@ async def _resolve_inexpressible_bounds( sizes: dict[str, int] = {} if paths_needing_size: # `_sizes` goes through `_info`, so the dircache may serve these. - # fsspec types the result as `list[None]`, but the values are sizes. + # Type checkers infer `list[None]` here, but the values are sizes. sizes = dict( zip( paths_needing_size, @@ -602,7 +600,9 @@ async def _resolve_inexpressible_bounds( ), ) resolved = [ - _apply_object_size(0 if start is None else start, end, sizes.get(path)) + _apply_object_size(0 if start is None else start, end, sizes[path]) + if _needs_object_size(start, end) + else (0 if start is None else start, end) for path, start, end in zip(paths, starts, ends, strict=True) ] return [start for start, _ in resolved], [end for _, end in resolved] diff --git a/tests/test_fsspec.py b/tests/test_fsspec.py index 5c0da59f..5c6396fa 100644 --- a/tests/test_fsspec.py +++ b/tests/test_fsspec.py @@ -667,13 +667,17 @@ async def spy_get(self, *args, **kwargs): # noqa: ANN001, ANN002, ANN003 monkeypatch.setattr(ObjectStoreMethods, "get_ranges_async", spy_get_ranges) monkeypatch.setattr(ObjectStoreMethods, "get_async", spy_get) - out = fs.cat_ranges([path] * 3, [0, 20, -5], [10, None, None]) - assert out == [data[0:10], data[20:], data[-5:]] + 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]] - # Only the bounded range reaches `get_ranges`, so only it gets coalesced. - assert batched == [([0], [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 batched == [([0, 9000], [10, 9990])] - # The other two go through `get` instead, and may complete in either order. + # 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(individual) == 2 assert {"offset": 20} in individual assert {"suffix": 5} in individual From a964a5e244b4fb655ee13a1550de2b98726d68bf Mon Sep 17 00:00:00 2001 From: Gunnar Schulze Date: Tue, 25 Aug 2026 19:58:19 +0200 Subject: [PATCH 08/22] docs: Add reference that supports typing of `_cat_ranges` --- obstore/python/obstore/fsspec.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/obstore/python/obstore/fsspec.py b/obstore/python/obstore/fsspec.py index 0d70d8f2..2c1094f9 100644 --- a/obstore/python/obstore/fsspec.py +++ b/obstore/python/obstore/fsspec.py @@ -474,8 +474,13 @@ async def _cat_ranges( # noqa: PLR0913, PLR0917 # type: ignore (fsspec has bad on_error: str = "return", # noqa: ARG002 **_kwargs: Any, ) -> list[bytes]: - # A non-iterable start or end applies to every path, per fsspec's - # AsyncFileSystem._cat_ranges. + # 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 not isinstance(ends, Iterable): From a148dacb99933ec5b90bdacb59f4f5d63aa46d1b Mon Sep 17 00:00:00 2001 From: Gunnar Schulze Date: Thu, 27 Aug 2026 06:53:39 +0200 Subject: [PATCH 09/22] refactor: Extract splitting of requests from `_cat_ranges` --- obstore/python/obstore/fsspec.py | 52 +++++++++++++++++--------------- 1 file changed, 27 insertions(+), 25 deletions(-) diff --git a/obstore/python/obstore/fsspec.py b/obstore/python/obstore/fsspec.py index 2c1094f9..a593eb03 100644 --- a/obstore/python/obstore/fsspec.py +++ b/obstore/python/obstore/fsspec.py @@ -137,6 +137,29 @@ def _apply_object_size( ) +def _split_requests( + paths: list[str], + bounds: Sequence[tuple[int, int | None]], +) -> tuple[ + dict[str, list[tuple[int, int, int]]], + list[tuple[int, str, OffsetRange | SuffixRange]], +]: + """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, OffsetRange | SuffixRange]] = [] + 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)) + elif start < 0: + open_ended_requests.append((idx, path, {"suffix": -start})) + else: + open_ended_requests.append((idx, path, {"offset": start})) + return per_file_bounded_requests, open_ended_requests + + class FsspecStore(fsspec.asyn.AsyncFileSystem): """An fsspec implementation based on a obstore Store. @@ -488,28 +511,8 @@ async def _cat_ranges( # noqa: PLR0913, PLR0917 # type: ignore (fsspec has bad if not len(paths) == len(starts) == len(ends): raise ValueError - resolved_starts, resolved_ends = await self._resolve_inexpressible_bounds( - paths, - starts, - ends, - ) - - # `get_ranges` only merges bounded ranges, and 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, OffsetRange | SuffixRange]] = [] - for idx, (path, start, end) in enumerate( - zip(paths, resolved_starts, resolved_ends, strict=True), - ): - if end is not None: - per_file_bounded_requests[path].append((idx, start, end)) - elif start < 0: - open_ended_requests.append((idx, path, {"suffix": -start})) - else: - open_ended_requests.append((idx, path, {"offset": start})) + bounds = await self._resolve_inexpressible_bounds(paths, starts, ends) + per_file_bounded_requests, open_ended_requests = _split_requests(paths, bounds) futs: list[Coroutine[Any, Any, list[tuple[int, bytes]]]] = [ self._cat_bounded_ranges(path, ranges, max_gap) @@ -584,7 +587,7 @@ async def _resolve_inexpressible_bounds( paths: list[str], starts: Sequence[int | None], ends: Sequence[int | None], - ) -> tuple[list[int], list[int | None]]: + ) -> list[tuple[int, int | None]]: """Resolve the bounds that obstore cannot express.""" paths_needing_size = sorted( { @@ -604,13 +607,12 @@ async def _resolve_inexpressible_bounds( strict=True, ), ) - resolved = [ + return [ _apply_object_size(0 if start is None else start, end, sizes[path]) if _needs_object_size(start, end) else (0 if start is None else start, end) for path, start, end in zip(paths, starts, ends, strict=True) ] - return [start for start, _ in resolved], [end for _, end in resolved] async def _put_file( self, From 361e2d3dffc7d238fd3eb426fbcdde2dc2041b91 Mon Sep 17 00:00:00 2001 From: Gunnar Schulze Date: Thu, 27 Aug 2026 06:55:29 +0200 Subject: [PATCH 10/22] fix: Handle errors correctly in `_cat_ranges` --- obstore/python/obstore/fsspec.py | 57 ++++++++++++++++++++------------ tests/test_fsspec.py | 28 ++++++++++++++++ 2 files changed, 63 insertions(+), 22 deletions(-) diff --git a/obstore/python/obstore/fsspec.py b/obstore/python/obstore/fsspec.py index a593eb03..8a7a4210 100644 --- a/obstore/python/obstore/fsspec.py +++ b/obstore/python/obstore/fsspec.py @@ -494,9 +494,9 @@ async def _cat_ranges( # noqa: PLR0913, PLR0917 # type: ignore (fsspec has bad ends: Sequence[int | None] | int | None, max_gap: int | None = None, batch_size: int | None = None, - on_error: str = "return", # noqa: ARG002 + on_error: str = "return", **_kwargs: Any, - ) -> list[bytes]: + ) -> list[bytes | BaseException]: # The base class implementation `AsyncFileSystem._cat_ranges` forwards each # element to `_cat_file`, which documents negative bounds and `None` for # either end. @@ -514,7 +514,7 @@ async def _cat_ranges( # noqa: PLR0913, PLR0917 # type: ignore (fsspec has bad bounds = await self._resolve_inexpressible_bounds(paths, starts, ends) per_file_bounded_requests, open_ended_requests = _split_requests(paths, bounds) - futs: list[Coroutine[Any, Any, list[tuple[int, bytes]]]] = [ + 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() ] @@ -523,10 +523,9 @@ async def _cat_ranges( # noqa: PLR0913, PLR0917 # type: ignore (fsspec has bad for idx, path, byte_range in open_ended_requests ] - # Batched, to limit how many requests are in flight at once. Like fsspec's - # own `_cat_ranges`, `on_error` is ignored, so failures propagate. + # Batched, to limit how many requests are in flight at once. results = cast( - "list[list[tuple[int, bytes]]]", + "list[list[tuple[int, bytes | BaseException]]]", await fsspec.asyn._run_coros_in_chunks( # noqa: SLF001 futs, batch_size=batch_size or self.batch_size, @@ -534,11 +533,17 @@ async def _cat_ranges( # noqa: PLR0913, PLR0917 # type: ignore (fsspec has bad ), ) - output_buffers: list[bytes] = [b""] * len(paths) + output_buffers: list[bytes | BaseException] = [b""] * len(paths) for responses in results: for idx, buffer in responses: output_buffers[idx] = buffer + # 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( @@ -546,21 +551,25 @@ async def _cat_bounded_ranges( path: str, ranges: list[tuple[int, int, int]], # (output index, start, end) max_gap: int | None, - ) -> list[tuple[int, bytes]]: + ) -> list[tuple[int, bytes | BaseException]]: """Read several bounded ranges of one object, merging nearby ones.""" - bucket, path = self._split_path(path) - store = self._construct_store(bucket) 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} - buffers = await store.get_ranges_async( - path, - starts=starts, - ends=ends, - lengths=None, - **coalesce, - ) + 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) @@ -571,16 +580,20 @@ async def _cat_open_ended_range( idx: int, path: str, byte_range: OffsetRange | SuffixRange, - ) -> list[tuple[int, bytes]]: + ) -> list[tuple[int, bytes | BaseException]]: """Read 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. """ - bucket, path = self._split_path(path) - store = self._construct_store(bucket) - resp = await store.get_async(path, options={"range": byte_range}) - return [(idx, (await resp.bytes_async()).to_bytes())] + try: + bucket, path_in_bucket = self._split_path(path) + store = self._construct_store(bucket) + resp = await store.get_async(path_in_bucket, options={"range": byte_range}) + buffer = await resp.bytes_async() + except Exception as exc: # noqa: BLE001 + return [(idx, exc)] + return [(idx, buffer.to_bytes())] async def _resolve_inexpressible_bounds( self, diff --git a/tests/test_fsspec.py b/tests/test_fsspec.py index 5c6396fa..ec93f4ec 100644 --- a/tests/test_fsspec.py +++ b/tests/test_fsspec.py @@ -707,6 +707,34 @@ async def spy(coros, **kwargs): # noqa: ANN001, ANN003 assert seen == [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) From ad5d2f138d46e39809aa8f146653aa45ce897924 Mon Sep 17 00:00:00 2001 From: Gunnar Schulze Date: Thu, 27 Aug 2026 21:39:54 +0200 Subject: [PATCH 11/22] fix: Read the whole object when a range in `_cat_ranges` is unbounded --- obstore/python/obstore/fsspec.py | 66 ++++++++++++++++---------------- tests/test_fsspec.py | 11 ++++-- 2 files changed, 41 insertions(+), 36 deletions(-) diff --git a/obstore/python/obstore/fsspec.py b/obstore/python/obstore/fsspec.py index 8a7a4210..f4faf312 100644 --- a/obstore/python/obstore/fsspec.py +++ b/obstore/python/obstore/fsspec.py @@ -55,10 +55,9 @@ from obstore import ( Attributes, + Bytes, GetOptions, - OffsetRange, ReadableFile, - SuffixRange, WritableFile, ) from obstore.store import ( @@ -142,24 +141,40 @@ def _split_requests( bounds: Sequence[tuple[int, int | None]], ) -> tuple[ dict[str, list[tuple[int, int, int]]], - list[tuple[int, str, OffsetRange | SuffixRange]], + 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, OffsetRange | SuffixRange]] = [] + 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)) - elif start < 0: - open_ended_requests.append((idx, path, {"suffix": -start})) else: - open_ended_requests.append((idx, path, {"offset": start})) + 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.""" + 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}} + resp = await store.get_async(path, options=options) + return await resp.bytes_async() + + class FsspecStore(fsspec.asyn.AsyncFileSystem): """An fsspec implementation based on a obstore Store. @@ -429,36 +444,22 @@ async def _cat_file( end: int | None = None, **_kwargs: Any, ) -> bytes: - # `_info` splits the path itself, so keep the original before rebinding. - full_path = path - bucket, path = self._split_path(path) + bucket, path_in_bucket = self._split_path(path) store = self._construct_store(bucket) - # A zero start with no end is the whole object, so skip the range request. - if not start 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 _needs_object_size(start, end): - # A negative `end` has no `GetOptions` equivalent, so resolve it against - # the object size. - size = (await self._info(full_path))["size"] + # 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: - # `get_range` mirrors the Rust API and takes bounded ranges only. `get` - # covers the other two forms. - options: GetOptions - if start < 0: - options = {"range": {"suffix": -start}} - else: - options = {"range": {"offset": start}} - resp = await store.get_async(path, options=options) - return (await resp.bytes_async()).to_bytes() + 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) @@ -519,8 +520,8 @@ async def _cat_ranges( # noqa: PLR0913, PLR0917 # type: ignore (fsspec has bad for path, ranges in per_file_bounded_requests.items() ] futs += [ - self._cat_open_ended_range(idx, path, byte_range) - for idx, path, byte_range in open_ended_requests + self._cat_open_ended_range(idx, path, start) + for idx, path, start in open_ended_requests ] # Batched, to limit how many requests are in flight at once. @@ -579,7 +580,7 @@ async def _cat_open_ended_range( self, idx: int, path: str, - byte_range: OffsetRange | SuffixRange, + start: int, ) -> list[tuple[int, bytes | BaseException]]: """Read one open-ended range, which `get_ranges` cannot express. @@ -589,8 +590,7 @@ async def _cat_open_ended_range( try: bucket, path_in_bucket = self._split_path(path) store = self._construct_store(bucket) - resp = await store.get_async(path_in_bucket, options={"range": byte_range}) - buffer = await resp.bytes_async() + buffer = await _get_open_ended(store, path_in_bucket, start) except Exception as exc: # noqa: BLE001 return [(idx, exc)] return [(idx, buffer.to_bytes())] diff --git a/tests/test_fsspec.py b/tests/test_fsspec.py index ec93f4ec..b2119fb9 100644 --- a/tests/test_fsspec.py +++ b/tests/test_fsspec.py @@ -652,7 +652,7 @@ def test_cat_ranges_routing(fs: FsspecStore, monkeypatch: pytest.MonkeyPatch): fs.pipe_file(path, data) batched: list[tuple[list[int], list[int]]] = [] - individual: list[dict[str, int]] = [] + individual: list[dict[str, int] | None] = [] original_get_ranges = ObjectStoreMethods.get_ranges_async original_get = ObjectStoreMethods.get_async @@ -661,7 +661,7 @@ async def spy_get_ranges(self, *args, **kwargs): # noqa: ANN001, ANN002, ANN003 return await original_get_ranges(self, *args, **kwargs) async def spy_get(self, *args, **kwargs): # noqa: ANN001, ANN002, ANN003 - individual.append(kwargs["options"]["range"]) + individual.append(kwargs.get("options", {}).get("range")) return await original_get(self, *args, **kwargs) monkeypatch.setattr(ObjectStoreMethods, "get_ranges_async", spy_get_ranges) @@ -799,12 +799,17 @@ def test_cat_ranges_mixed(fs: FsspecStore): data2 = os.urandom(10000) path1 = f"{TEST_BUCKET_NAME}/data1" path2 = f"{TEST_BUCKET_NAME}/data2" - fs.pipe({path1: data1, path2: data2}) + empty = f"{TEST_BUCKET_NAME}/empty" + fs.pipe({path1: data1, path2: data2, empty: b""}) # single range in each file 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""] + @pytest.mark.xfail(reason="atomic writes not working on moto") def test_atomic_write(fs: FsspecStore): From 3782cd221bd25252604ca91d9baf097dcae6a7f6 Mon Sep 17 00:00:00 2001 From: Gunnar Schulze Date: Fri, 28 Aug 2026 08:18:50 +0200 Subject: [PATCH 12/22] docs: Add docstrings for `_cat_file` and `_cat_ranges` --- obstore/python/obstore/fsspec.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/obstore/python/obstore/fsspec.py b/obstore/python/obstore/fsspec.py index f4faf312..4dbd9aff 100644 --- a/obstore/python/obstore/fsspec.py +++ b/obstore/python/obstore/fsspec.py @@ -444,6 +444,11 @@ async def _cat_file( end: int | None = None, **_kwargs: Any, ) -> bytes: + """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) @@ -498,6 +503,11 @@ async def _cat_ranges( # noqa: PLR0913, PLR0917 # type: ignore (fsspec has bad on_error: str = "return", **_kwargs: Any, ) -> 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. @@ -553,7 +563,7 @@ async def _cat_bounded_ranges( ranges: list[tuple[int, int, int]], # (output index, start, end) max_gap: int | None, ) -> list[tuple[int, bytes | BaseException]]: - """Read several bounded ranges of one object, merging nearby ones.""" + """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. @@ -582,7 +592,7 @@ async def _cat_open_ended_range( path: str, start: int, ) -> list[tuple[int, bytes | BaseException]]: - """Read one open-ended range, which `get_ranges` cannot express. + """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. From 1a33600f1a09abc2d07ccf2abd46aa99aef13f86 Mon Sep 17 00:00:00 2001 From: Gunnar Schulze Date: Fri, 28 Aug 2026 16:32:16 +0200 Subject: [PATCH 13/22] fix: Fall back to a bounded read where suffix ranges are unsupported --- obstore/python/obstore/fsspec.py | 22 +++++++++++++++--- tests/test_fsspec.py | 38 ++++++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 3 deletions(-) diff --git a/obstore/python/obstore/fsspec.py b/obstore/python/obstore/fsspec.py index 4dbd9aff..868a37ba 100644 --- a/obstore/python/obstore/fsspec.py +++ b/obstore/python/obstore/fsspec.py @@ -45,6 +45,7 @@ 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: @@ -162,7 +163,7 @@ async def _get_open_ended( path: str, start: int, ) -> Bytes: - """Get an open-ended range.""" + """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. @@ -171,8 +172,23 @@ async def _get_open_ended( options = {"range": {"suffix": -start}} else: options = {"range": {"offset": start}} - resp = await store.get_async(path, options=options) - return await resp.bytes_async() + 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): diff --git a/tests/test_fsspec.py b/tests/test_fsspec.py index b2119fb9..2715f59a 100644 --- a/tests/test_fsspec.py +++ b/tests/test_fsspec.py @@ -13,6 +13,7 @@ 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 @@ -613,6 +614,43 @@ async def spy(self, *args, **kwargs): # noqa: ANN001, ANN002, ANN003 assert ranges == [None, {"offset": 10}] +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) From 33ed3e2e357f1b6b01f4c4292007a1c36922d14d Mon Sep 17 00:00:00 2001 From: Gunnar Schulze Date: Tue, 1 Sep 2026 20:48:46 +0200 Subject: [PATCH 14/22] fix: Pass `batch_size` to `_sizes` --- obstore/python/obstore/fsspec.py | 14 ++++++++++++-- tests/test_fsspec.py | 5 +++++ 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/obstore/python/obstore/fsspec.py b/obstore/python/obstore/fsspec.py index 868a37ba..de9df8a5 100644 --- a/obstore/python/obstore/fsspec.py +++ b/obstore/python/obstore/fsspec.py @@ -538,7 +538,12 @@ async def _cat_ranges( # noqa: PLR0913, PLR0917 # type: ignore (fsspec has bad if not len(paths) == len(starts) == len(ends): raise ValueError - bounds = await self._resolve_inexpressible_bounds(paths, starts, ends) + 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[tuple[int, bytes | BaseException]]]] = [ @@ -626,6 +631,8 @@ async def _resolve_inexpressible_bounds( 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.""" paths_needing_size = sorted( @@ -642,7 +649,10 @@ async def _resolve_inexpressible_bounds( sizes = dict( zip( paths_needing_size, - cast("list[int]", await self._sizes(paths_needing_size)), + cast( + "list[int]", + await self._sizes(paths_needing_size, batch_size=batch_size), + ), strict=True, ), ) diff --git a/tests/test_fsspec.py b/tests/test_fsspec.py index 2715f59a..073c0041 100644 --- a/tests/test_fsspec.py +++ b/tests/test_fsspec.py @@ -744,6 +744,11 @@ async def spy(coros, **kwargs): # noqa: ANN001, ANN003 assert fs.cat_ranges([path], [0], [10], batch_size=4) == [data[0:10]] assert seen == [4] + # The size lookup for a negative end is batched with the same batch size. + seen.clear() + assert fs.cat_ranges([path], [0], [-10], batch_size=4) == [data[0:-10]] + assert seen == [4, 4] + def test_cat_ranges_on_error(fs: FsspecStore): """Test that `on_error` controls whether a failure is returned or raised.""" From 833a6bab9f73d839aef67448adf210ddd6f1b9e1 Mon Sep 17 00:00:00 2001 From: Gunnar Schulze Date: Wed, 2 Sep 2026 22:38:25 +0200 Subject: [PATCH 15/22] refactor: Normalize starts once in `_resolve_inexpressible_bounds` --- obstore/python/obstore/fsspec.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/obstore/python/obstore/fsspec.py b/obstore/python/obstore/fsspec.py index de9df8a5..0cbc3fb7 100644 --- a/obstore/python/obstore/fsspec.py +++ b/obstore/python/obstore/fsspec.py @@ -120,9 +120,13 @@ """A type hint for all supported protocols.""" -def _needs_object_size(start: int | None, end: int | None) -> bool: - """Whether resolving a range requires knowing the size of the object.""" - return end is not None and (end < 0 or (start is not None and start < 0)) +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( @@ -635,10 +639,11 @@ async def _resolve_inexpressible_bounds( 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 = sorted( { path - for path, start, end in zip(paths, starts, ends, strict=True) + for path, start, end in zip(paths, normalized_starts, ends, strict=True) if _needs_object_size(start, end) }, ) @@ -657,10 +662,10 @@ async def _resolve_inexpressible_bounds( ), ) return [ - _apply_object_size(0 if start is None else start, end, sizes[path]) + _apply_object_size(start, end, sizes[path]) if _needs_object_size(start, end) - else (0 if start is None else start, end) - for path, start, end in zip(paths, starts, ends, strict=True) + else (start, end) + for path, start, end in zip(paths, normalized_starts, ends, strict=True) ] async def _put_file( From a2e3e73e6ed905d10951549b16b0775ef7cf6fc1 Mon Sep 17 00:00:00 2001 From: Gunnar Schulze Date: Wed, 2 Sep 2026 23:12:00 +0200 Subject: [PATCH 16/22] test: Fold the zero-start check into test_cat_file --- tests/test_fsspec.py | 44 ++++++++++++++++++++++---------------------- 1 file changed, 22 insertions(+), 22 deletions(-) diff --git a/tests/test_fsspec.py b/tests/test_fsspec.py index 073c0041..c6fba2c8 100644 --- a/tests/test_fsspec.py +++ b/tests/test_fsspec.py @@ -575,43 +575,43 @@ def test_multi_file_ops(minio_bucket: tuple[S3Config, ClientConfig]): assert out == [f"{bucket}/afile"] -def test_cat_file(fs: FsspecStore): - """Test that `cat_file` accepts every range form fsspec documents.""" +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] = [] + original = ObjectStoreMethods.get_async + + async def spy(self, *args, **kwargs): # noqa: ANN001, ANN002, ANN003 + ranges_via_get.append(kwargs.get("options", {}).get("range")) + return await original(self, *args, **kwargs) + + monkeypatch.setattr(ObjectStoreMethods, "get_async", spy) + + # 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] + + # Both bounds given go through `get_range` instead of `get`. + ranges_via_get.clear() assert fs.cat_file(path, start=10, end=20) == data[10:20] + assert ranges_via_get == [] # Either bound on its own. + ranges_via_get.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}] # An `end` alone is a bounded read. # Bounds counted back from the end of the object. + ranges_via_get.clear() assert fs.cat_file(path, start=-10) == data[-10:] assert fs.cat_file(path, start=10, end=-10) == data[10:-10] assert fs.cat_file(path, start=-20, end=-10) == data[-20:-10] - - -def test_cat_file_zero_start(fs: FsspecStore, monkeypatch: pytest.MonkeyPatch): - """Test that a zero start with no end is not sent as a range request.""" - data = os.urandom(10000) - path = f"{TEST_BUCKET_NAME}/data1" - fs.pipe_file(path, data) - - ranges: list[dict[str, int] | None] = [] - original = ObjectStoreMethods.get_async - - async def spy(self, *args, **kwargs): # noqa: ANN001, ANN002, ANN003 - ranges.append((kwargs.get("options") or {}).get("range")) - return await original(self, *args, **kwargs) - - monkeypatch.setattr(ObjectStoreMethods, "get_async", spy) - - assert fs.cat_file(path, start=0) == data - assert fs.cat_file(path, start=10) == data[10:] - assert ranges == [None, {"offset": 10}] + assert ranges_via_get == [{"suffix": 10}] # Resolved ends become bounded reads. def test_suffix_fallback(fs: FsspecStore, monkeypatch: pytest.MonkeyPatch): From de6c0ec93cdc5d5c925bc483301c535c5dec9e13 Mon Sep 17 00:00:00 2001 From: Gunnar Schulze Date: Sat, 5 Sep 2026 11:29:39 +0200 Subject: [PATCH 17/22] test: Name spies and captures after what they observe --- tests/test_fsspec.py | 60 ++++++++++++++++++++++---------------------- 1 file changed, 30 insertions(+), 30 deletions(-) diff --git a/tests/test_fsspec.py b/tests/test_fsspec.py index c6fba2c8..c233d5fb 100644 --- a/tests/test_fsspec.py +++ b/tests/test_fsspec.py @@ -584,11 +584,11 @@ def test_cat_file(fs: FsspecStore, monkeypatch: pytest.MonkeyPatch): ranges_via_get: list[dict[str, int] | None] = [] original = ObjectStoreMethods.get_async - async def spy(self, *args, **kwargs): # noqa: ANN001, ANN002, ANN003 + async def spy_get(self, *args, **kwargs): # noqa: ANN001, ANN002, ANN003 ranges_via_get.append(kwargs.get("options", {}).get("range")) return await original(self, *args, **kwargs) - monkeypatch.setattr(ObjectStoreMethods, "get_async", spy) + monkeypatch.setattr(ObjectStoreMethods, "get_async", spy_get) # The whole object, with and without an explicit zero start. assert fs.cat_file(path) == data @@ -657,30 +657,30 @@ def test_cat_ranges_max_gap(fs: FsspecStore, monkeypatch: pytest.MonkeyPatch): path = f"{TEST_BUCKET_NAME}/data1" fs.pipe_file(path, data) - seen: list[int | None] = [] - original = ObjectStoreMethods.get_ranges_async + forwarded_coalesce: list[int | None] = [] + original_get_ranges = ObjectStoreMethods.get_ranges_async - async def spy(self, *args, **kwargs): # noqa: ANN001, ANN002, ANN003 - seen.append(kwargs.get("coalesce")) - return await original(self, *args, **kwargs) + 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) + 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 seen == [None] + assert forwarded_coalesce == [None] # With max_gap: passed straight through, including 0 to turn coalescing off. for max_gap in (0, 1000): - seen.clear() + forwarded_coalesce.clear() assert fs.cat_ranges([path, path], [0, 20], [10, 30], max_gap=max_gap) == [ data[0:10], data[20:30], ] - assert seen == [max_gap] + assert forwarded_coalesce == [max_gap] def test_cat_ranges_routing(fs: FsspecStore, monkeypatch: pytest.MonkeyPatch): @@ -689,17 +689,17 @@ def test_cat_ranges_routing(fs: FsspecStore, monkeypatch: pytest.MonkeyPatch): path = f"{TEST_BUCKET_NAME}/data1" fs.pipe_file(path, data) - batched: list[tuple[list[int], list[int]]] = [] - individual: list[dict[str, int] | None] = [] + 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 - batched.append((list(kwargs["starts"]), list(kwargs["ends"]))) + 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 - individual.append(kwargs.get("options", {}).get("range")) + 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) @@ -711,14 +711,14 @@ async def spy_get(self, *args, **kwargs): # noqa: ANN001, ANN002, ANN003 # 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 batched == [([0, 9000], [10, 9990])] + 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(individual) == 2 - assert {"offset": 20} in individual - assert {"suffix": 5} in individual + 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): @@ -727,27 +727,27 @@ def test_cat_ranges_batch_size(fs: FsspecStore, monkeypatch: pytest.MonkeyPatch) path = f"{TEST_BUCKET_NAME}/data1" fs.pipe_file(path, data) - seen: list[int | None] = [] - original = fsspec.asyn._run_coros_in_chunks + forwarded_batch_sizes: list[int | None] = [] + original_run_coros = fsspec.asyn._run_coros_in_chunks - async def spy(coros, **kwargs): # noqa: ANN001, ANN003 - seen.append(kwargs.get("batch_size")) - return await original(coros, **kwargs) + 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) + 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 seen == [None] + assert forwarded_batch_sizes == [None] - seen.clear() + forwarded_batch_sizes.clear() assert fs.cat_ranges([path], [0], [10], batch_size=4) == [data[0:10]] - assert seen == [4] + assert forwarded_batch_sizes == [4] # The size lookup for a negative end is batched with the same batch size. - seen.clear() + forwarded_batch_sizes.clear() assert fs.cat_ranges([path], [0], [-10], batch_size=4) == [data[0:-10]] - assert seen == [4, 4] + assert forwarded_batch_sizes == [4, 4] def test_cat_ranges_on_error(fs: FsspecStore): From 5f2dac24c90f7634003dd0968eb62b16f1843f5e Mon Sep 17 00:00:00 2001 From: Gunnar Schulze Date: Sat, 5 Sep 2026 11:54:38 +0200 Subject: [PATCH 18/22] test: Assert that bounded reads reach `get_range` --- tests/test_fsspec.py | 23 +++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/tests/test_fsspec.py b/tests/test_fsspec.py index c233d5fb..2b9d8e8e 100644 --- a/tests/test_fsspec.py +++ b/tests/test_fsspec.py @@ -582,36 +582,51 @@ def test_cat_file(fs: FsspecStore, monkeypatch: pytest.MonkeyPatch): fs.pipe_file(path, data) ranges_via_get: list[dict[str, int] | None] = [] - original = ObjectStoreMethods.get_async + 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(self, *args, **kwargs) + 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}] # An `end` alone is a bounded read. + 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=10, end=-10) == data[10:-10] assert fs.cat_file(path, start=-20, end=-10) == data[-20:-10] - assert ranges_via_get == [{"suffix": 10}] # Resolved ends become bounded reads. + assert ranges_via_get == [{"suffix": 10}] + # Resolved ends become bounded reads. + assert bounds_via_get_range == [(10, 9990), (9980, 9990)] def test_suffix_fallback(fs: FsspecStore, monkeypatch: pytest.MonkeyPatch): From 14c097bdec6142c85cafc5ac447f87803f14a0f1 Mon Sep 17 00:00:00 2001 From: Gunnar Schulze Date: Sat, 5 Sep 2026 13:30:04 +0200 Subject: [PATCH 19/22] refactor: Drop unnecessary sort in `_resolve_inexpressible_bounds` --- obstore/python/obstore/fsspec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/obstore/python/obstore/fsspec.py b/obstore/python/obstore/fsspec.py index 0cbc3fb7..a936e238 100644 --- a/obstore/python/obstore/fsspec.py +++ b/obstore/python/obstore/fsspec.py @@ -640,7 +640,7 @@ async def _resolve_inexpressible_bounds( ) -> 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 = sorted( + paths_needing_size = list( { path for path, start, end in zip(paths, normalized_starts, ends, strict=True) From 86ebbdee4e23a0b9e230334bf3283332f56a66b9 Mon Sep 17 00:00:00 2001 From: Gunnar Schulze Date: Sat, 5 Sep 2026 13:00:21 +0200 Subject: [PATCH 20/22] test: Fix an inaccurate comment in `test_cat_ranges_mixed` --- tests/test_fsspec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_fsspec.py b/tests/test_fsspec.py index 2b9d8e8e..ce2b9036 100644 --- a/tests/test_fsspec.py +++ b/tests/test_fsspec.py @@ -860,7 +860,7 @@ def test_cat_ranges_mixed(fs: FsspecStore): empty = f"{TEST_BUCKET_NAME}/empty" fs.pipe({path1: data1, path2: data2, empty: b""}) - # single range in each file + # 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]] From f719d0ab0ac333fc8729497ecc359de347d7624b Mon Sep 17 00:00:00 2001 From: Gunnar Schulze Date: Sat, 5 Sep 2026 13:43:19 +0200 Subject: [PATCH 21/22] test: Cover oversized suffixes and mixed-sign bounds --- tests/test_fsspec.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/test_fsspec.py b/tests/test_fsspec.py index ce2b9036..66dbcbf8 100644 --- a/tests/test_fsspec.py +++ b/tests/test_fsspec.py @@ -622,11 +622,13 @@ async def spy_get_range(self, *args, **kwargs): # noqa: ANN001, ANN002, ANN003 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 ranges_via_get == [{"suffix": 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)] + assert bounds_via_get_range == [(10, 9990), (9980, 9990), (9980, 9995)] def test_suffix_fallback(fs: FsspecStore, monkeypatch: pytest.MonkeyPatch): From a0e58fe065d3dcbbc1fec6dfe9d4406e8dbe2d80 Mon Sep 17 00:00:00 2001 From: Gunnar Schulze Date: Sat, 5 Sep 2026 13:44:30 +0200 Subject: [PATCH 22/22] test: Cover scalar broadcasts in `cat_ranges` --- tests/test_fsspec.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/test_fsspec.py b/tests/test_fsspec.py index 66dbcbf8..022eba53 100644 --- a/tests/test_fsspec.py +++ b/tests/test_fsspec.py @@ -871,6 +871,19 @@ def test_cat_ranges_mixed(fs: FsspecStore): 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): fs.pipe_file("data1", b"data1")