feat:get_ranges

[d-v-b]

this PR adds a get_ranges protocol for stores. the protocol defines the shape of a function that fetches multiple byte ranges within the same stored object. The purpose is to define a method stores can opt into if they offer an efficient way to fetch multiple byte ranges from the same object, which would be immediately useful for the sharding codec. The protocol looks like this:

async def get_ranges(
    key: str,
    byte_ranges: Iterable[ByteRequest | None],
    *,
    prototype: BufferPrototype,
) -> AsyncIterator[Sequence[tuple[int, Buffer | None]]]:

the return type is an async iterator over sequences, where each sequence is the result of an IO operation the store implemented. this provides some observability to the caller about the actual coalescing, if any, that occurred. Results are returned in computed order, so the inner result type is tuple[int, Buffer | None], where the int is the index into the input byte_ranges for that result.

Only byte range requests that declare an explicit interval (RangeByteRequest) are coalesced. Any other byte range, or None, results in no coalescing and so the ranges will be fetched separately. I assume here that we do not care about coalescing overlapping suffix or prefix range requests, but we could add support for that if we need to.

In addition to this protocol, there's a freestanding function that takes:

• a f(byte range) -> Awaitable[Buffer] function (which we would generate by combining Store.get with functools.partial)
• an iterable of byte range requests
• options (max gap bytes, max total bytes, etc)

this function contains basic byte range coalescing logic, and it can be re-used for multiple stores. This is a non-abstract-base-class alternative to a default implementation on an abc.

That freestanding function is used to implement get_ranges on the FsspecStore. This is probably not useful for local- or memory-backed storage, but is useful for remote storage. The actual implementation is lightweight:

    async def get_ranges(
        self,
        key: str,
        byte_ranges: Iterable[ByteRequest | None],
        *,
        prototype: BufferPrototype,
    ) -> AsyncIterator[Sequence[tuple[int, Buffer | None]]]:
        """Read many byte ranges from ``key``, coalescing nearby ranges and fetching concurrently.

        See :class:`zarr.storage._protocols.SupportsGetRanges` for the contract and
        :func:`zarr.core._coalesce.coalesced_get` for the full semantics.
        """
        fetch = partial(self.get, key, prototype)
        async for group in coalesced_get(fetch, byte_ranges, options=self.coalesce_options):
            yield group

cc @aldenks, the idea here is to build a basis for your range coalescing work for the sharding codec
cc @kylebarron, would love your feedback on this design.

related issues/PRs:

#1758
#3004

[d-v-b]

for context, we do already have a get_partial_values method, but it returns eagerly (no streaming), and it handles fetching from multiple keys. get_ranges only targets a single key, and it supports streaming. This is IMO a more targeted and performance-friendly design, and it's narrowly what the sharding codec needs for IO against a single shard.

[codecov]

Codecov Report

❌ Patch coverage is 96.29630% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 93.30%. Comparing base (7c2f372) to head (5eb25bc).

Files with missing lines	Patch %	Lines
src/zarr/core/_coalesce.py	95.78%	4 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3925      +/-   ##
==========================================
+ Coverage   93.26%   93.30%   +0.03%     
==========================================
  Files          87       89       +2     
  Lines       11721    11829     +108     
==========================================
+ Hits        10932    11037     +105     
- Misses        789      792       +3     

Files with missing lines	Coverage Δ
src/zarr/storage/_fsspec.py	91.71% <100.00%> (+0.38%)	⬆️
src/zarr/storage/_protocols.py	100.00% <100.00%> (ø)
src/zarr/core/_coalesce.py	95.78% <95.78%> (ø)

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[kylebarron]

I think making it iterable adds complexity and adds confusion to whether request coalescing is expected to be applied here or not.

If you have requests:

• 0-1000
• 1000-2000
• 2000-3000
• 100_000 - 100_005

Then of course we should coalesce all the first 3 requests into one. But then the async iterable implies that we might want to use one of the first responses before the last one arrives... but the last request would've arrived first.

Either, if you really want the response type to be async iterable, the responses should probably have an index for which of the input requests it's associated to.

But I think it would be much simpler to take in a sequence of byte ranges and return a sequence of results. Just like object-store/obstore/obspec.

[d-v-b]

If you have requests:

* 0-1000

* 1000-2000

* 2000-3000

* 100_000 - 100_005

In the design in this PR we are iterating over IO calls the reader actually did, which is less than or equal to the number of byte ranges requested. So, assuming the first three are fused, we would either see:

yields ((0, <payload>), (1, <payload>), (2, <payload>))
yields ((3, <payload>),)

or

yields ((3, <payload>),)
yields ((0, <payload>), (1, <payload>), (2, <payload>))

Which, for sharding is useful -- you can start decoding chunks immediately while you wait for the rest of the sub-chunks to come in. Does that make sense? or am I misunderstanding something.

[ilan-gold]

Nice!

At an architecture level, it's not immediately clear to me how this will be exposed to users. But maybe that is out of the scope of this PR? I think we don't want to turn this on by default without some rough estimates.

[d-v-b]

@ilan-gold have another look, I cleaned things up here a bit. one notable change that I plan on re-using: define subfunctions that explicitly depend on a shared state, instead of relying implicitly on variables defined in the outer scope (like the queue and the semaphore)

[chuckwondo]

Is there any reason you chose not to implement a range coalescing function independent of fetching? Perhaps I'm missing something, particularly since I'm still very much getting my feet wet with this codebase, but it seems to me that having an independent range coalescing function would allow for far simpler and easier testing.

As currently written, your tests must include fetching since you have coupled coalescing and fetching, which seems to have added quite a bit of verbosity/complexity to your tests that you might otherwise be able to avoid, since existing tests should already be covering the fetching.

[d-v-b]

No good reason! I think factoring the range coalescing into a separate routine is a great idea for the reasons you mention. I'll implement that

[chuckwondo]

I would also argue that OffsetByteRequests are coalescable.

Although you may not know how many total bytes are available, you certainly know, by definition, where the range starts (the offset), so as you coalesce earlier RangeByteRequests, when you bump into your first OffsetByteRequest, and you still haven't hit your max_coalesced_bytes limit, you can pull as many bytes from the front of the OffsetByteRequest as you need to reach the limit, and shift the offset by the corresponding number of bytes.

You can also coalesce beyond that point. For example, all RangeByteRequests and OffsetByteRequests starting after the first OffsetByteRequest are subsumed by that first OffsetByteRequest.

This, of course, would seem to add a bit of complexity to the coalescing logic, so I suspect you could exclude such logic from this PR, and separately consider whether or not it would be worth the additional complexity at a later date (or never).

[d-v-b]

Agree, but the only consumer of this method is the sharding codec which exclusively requests explicit byte ranges. So if we implemented coalescing for other byte range requests, it would never be used.

[chuckwondo]

I'm a bit confused then. Why do you accept all types of byte requests as input if you expect only byte range requests?

[d-v-b]

The coalescing is a performance optimization, and we don't want to push branching on the request type to the caller. A non-coalesced request falls through to a regular fetch, which is submitted in the same batch as the coalesced fetches. The alternative requires two batches.

And we can't predict the requirements of future codecs. We know today what sharding needs, so we optimize for that. If a future caller needs to coalesce offset range requests, we can implement it inside this function without changing the signature.

[chuckwondo]

That clarifies a bit for me, but I suspect I might need more time getting familiar with the library before that will be much more clear to me.

After looking at the PR code a bit more, I see that you group the range requests, but don't actually collapse them (except when you actually fetch the bytes) so that once the bytes are fetched, you can slice them up according to the original range requests that you preserved.

[ilan-gold]

Why does this need its own maximal concurrency? How will this interact with the global setting?

[d-v-b]

this needs a max concurrency because implementations may launch concurrent range reads. Right now this setting doesn't interact at all with the global setting, and that's hard to support given our current concurrency limit design -- codecs like sharding call concurrent_map recursively, so the global concurrency limit is alrady on shaky ground for sharding. We should do some testing with this parameter to see if there are any pathologies possible with too much concurrency.