Fixes to _get_partial_values

Author: kylebarron

src/zarr/storage/object_store.py

@@ -3,7 +3,7 @@
 import asyncio
 from collections import defaultdict
 from collections.abc import Iterable
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, TypedDict
 
 import obstore as obs
 
@@ -15,8 +15,7 @@
     from collections.abc import AsyncGenerator, Coroutine, Iterable
     from typing import Any
 
-    from obstore import Buffer as ObjectStoreBuffer
-    from obstore import ListStream, ObjectMeta
+    from obstore import ListStream, ObjectMeta, OffsetRange, SuffixRange
     from obstore.store import ObjectStore as _ObjectStore
 
     from zarr.core.buffer import Buffer, BufferPrototype
@@ -62,33 +61,7 @@ async def get_partial_values(
         prototype: BufferPrototype,
         key_ranges: Iterable[tuple[str, ByteRangeRequest]],
     ) -> list[Buffer | None]:
-        # TODO: this is a bit hacky and untested. ObjectStore has a `get_ranges` method
-        # that will additionally merge nearby ranges, but it's _per_ file. So we need to
-        # split these key_ranges into **per-file** key ranges, and then reassemble the
-        # results in the original order.
-        key_ranges = list(key_ranges)
-
-        per_file_requests: dict[str, list[tuple[int | None, int | None, int]]] = defaultdict(list)
-        for idx, (path, range_) in enumerate(key_ranges):
-            per_file_requests[path].append((range_[0], range_[1], idx))
-
-        futs: list[Coroutine[Any, Any, list[ObjectStoreBuffer]]] = []
-        for path, ranges in per_file_requests.items():
-            starts = [r[0] for r in ranges]
-            ends = [r[1] for r in ranges]
-            fut = obs.get_ranges_async(self.store, path, starts=starts, ends=ends)
-            futs.append(fut)
-
-        result = await asyncio.gather(*futs)
-
-        output_buffers: list[type[BufferPrototype]] = [b""] * len(key_ranges)
-        for per_file_request, buffers in zip(per_file_requests.items(), result, strict=True):
-            path, ranges = per_file_request
-            for buffer, ranges_ in zip(buffers, ranges, strict=True):
-                initial_index = ranges_[2]
-                output_buffers[initial_index] = prototype.buffer.from_buffer(memoryview(buffer))
-
-        return output_buffers
+        return await _get_partial_values(self.store, prototype=prototype, key_ranges=key_ranges)
 
     async def exists(self, key: str) -> bool:
         try:
@@ -163,3 +136,156 @@ async def _transform_list_dir(
             # Yield this item if "/" does not exist after the prefix.
             if "/" not in item["path"][prefix_len:]:
                 yield item["path"]
+
+
+class BoundedRequest(TypedDict):
+    """Range request with a known start and end byte.
+
+    These requests can be multiplexed natively on the Rust side with
+    `obstore.get_ranges_async`.
+    """
+
+    original_request_index: int
+    """The positional index in the original key_ranges input"""
+
+    start: int
+    """Start byte offset."""
+
+    end: int
+    """End byte offset."""
+
+
+class OtherRequest(TypedDict):
+    """Offset or suffix range requests.
+
+    These requests cannot be concurrent on the Rust side, and each need their own call
+    to `obstore.get_async`, passing in the `range` parameter.
+    """
+
+    original_request_index: int
+    """The positional index in the original key_ranges input"""
+
+    path: str
+    """The path to request from."""
+
+    range: OffsetRange | SuffixRange
+    """The range request type."""
+
+
+class Response(TypedDict):
+    """A response buffer associated with the original index that it should be restored to."""
+
+    original_request_index: int
+    """The positional index in the original key_ranges input"""
+
+    buffer: Buffer
+    """The buffer returned from obstore's range request."""
+
+
+async def _make_bounded_requests(
+    store: obs.store.ObjectStore,
+    path: str,
+    requests: list[BoundedRequest],
+    prototype: BufferPrototype,
+) -> list[Response]:
+    """Make all bounded requests for a specific file.
+
+    `obstore.get_ranges_async` allows for making concurrent requests for multiple ranges
+    within a single file, and will e.g. merge concurrent requests. This only uses one
+    single Python coroutine.
+    """
+
+    starts = [r["start"] for r in requests]
+    ends = [r["end"] for r in requests]
+    responses = await obs.get_ranges_async(store, path=path, starts=starts, ends=ends)
+
+    buffer_responses: list[Response] = []
+    for request, response in zip(requests, responses, strict=True):
+        buffer_responses.append(
+            {
+                "original_request_index": request["original_request_index"],
+                "buffer": prototype.buffer.from_bytes(memoryview(response)),
+            }
+        )
+
+    return buffer_responses
+
+
+async def _make_other_request(
+    store: obs.store.ObjectStore,
+    request: OtherRequest,
+    prototype: BufferPrototype,
+) -> list[Response]:
+    """Make suffix or offset requests.
+
+    We return a `list[Response]` for symmetry with `_make_bounded_requests` so that all
+    futures can be gathered together.
+    """
+    resp = await obs.get_async(store, request["path"], options={"range": request["range"]})
+    buffer = await resp.bytes_async()
+    return [
+        {
+            "original_request_index": request["original_request_index"],
+            "buffer": prototype.buffer.from_bytes(buffer),
+        }
+    ]
+
+
+async def _get_partial_values(
+    store: obs.store.ObjectStore,
+    prototype: BufferPrototype,
+    key_ranges: Iterable[tuple[str, ByteRangeRequest]],
+) -> list[Buffer | None]:
+    """Make multiple range requests.
+
+    ObjectStore has a `get_ranges` method that will additionally merge nearby ranges,
+    but it's _per_ file. So we need to split these key_ranges into **per-file** key
+    ranges, and then reassemble the results in the original order.
+
+    We separate into different requests:
+
+    - One call to `obstore.get_ranges_async` **per target file**
+    - One call to `obstore.get_async` for each other request.
+    """
+    key_ranges = list(key_ranges)
+    per_file_bounded_requests: dict[str, list[BoundedRequest]] = defaultdict(list)
+    other_requests: list[OtherRequest] = []
+
+    for idx, (path, (start, end)) in enumerate(key_ranges):
+        if start is None:
+            raise ValueError("Cannot pass `None` for the start of the range request.")
+
+        if end is not None:
+            # This is a bounded request with known start and end byte.
+            per_file_bounded_requests[path].append(
+                {"original_request_index": idx, "start": start, "end": end}
+            )
+        elif end is None and start < 0:
+            # Suffix request from the end
+            other_requests.append(
+                {"original_request_index": idx, "path": path, "range": {"suffix": abs(start)}}
+            )
+        elif end is None and start > 0:
+            # Offset request to the end
+            other_requests.append(
+                {"original_request_index": idx, "path": path, "range": {"offset": start}}
+            )
+        else:
+            raise ValueError(f"Unsupported range input: {start=}, {end=}")
+
+    futs: list[Coroutine[Any, Any, list[Response]]] = []
+    for path, bounded_ranges in per_file_bounded_requests.items():
+        futs.append(_make_bounded_requests(store, path, bounded_ranges, prototype))
+
+    for request in other_requests:
+        futs.append(_make_other_request(store, request, prototype))  # noqa: PERF401
+
+    buffers: list[Buffer | None] = [None] * len(key_ranges)
+
+    # TODO: this gather a list of list of Response; not sure if there's a way to
+    # unpack these lists inside of an `asyncio.gather`?
+    for responses in await asyncio.gather(*futs):
+        for resp in responses:
+            buffers[resp["original_request_index"]] = resp["buffer"]
+
+    return buffers