Merge branch 'main' into release-notes

Author: d-v-b

src/zarr/codecs/bytes.py

@@ -114,7 +114,7 @@ async def _encode_single(
 
         nd_array = chunk_array.as_ndarray_like()
         # Flatten the nd-array (only copy if needed) and reinterpret as bytes
-        nd_array = nd_array.ravel().view(dtype="b")
+        nd_array = nd_array.ravel().view(dtype="B")
         return chunk_spec.prototype.buffer.from_array_like(nd_array)
 
     def compute_encoded_size(self, input_byte_length: int, _chunk_spec: ArraySpec) -> int:

src/zarr/codecs/crc32c_.py

@@ -57,7 +57,7 @@ async def _encode_single(
         # Calculate the checksum and "cast" it to a numpy array
         checksum = np.array([crc32c(cast(typing_extensions.Buffer, data))], dtype=np.uint32)
         # Append the checksum (as bytes) to the data
-        return chunk_spec.prototype.buffer.from_array_like(np.append(data, checksum.view("b")))
+        return chunk_spec.prototype.buffer.from_array_like(np.append(data, checksum.view("B")))
 
     def compute_encoded_size(self, input_byte_length: int, _chunk_spec: ArraySpec) -> int:
         return input_byte_length + 4

src/zarr/core/buffer/core.py

@@ -143,7 +143,7 @@ class Buffer(ABC):
     def __init__(self, array_like: ArrayLike) -> None:
         if array_like.ndim != 1:
             raise ValueError("array_like: only 1-dim allowed")
-        if array_like.dtype != np.dtype("b"):
+        if array_like.dtype != np.dtype("B"):
             raise ValueError("array_like: only byte dtype allowed")
         self._data = array_like
 
@@ -306,7 +306,7 @@ class NDBuffer:
     Notes
     -----
     The two buffer classes Buffer and NDBuffer are very similar. In fact, Buffer
-    is a special case of NDBuffer where dim=1, stride=1, and dtype="b". However,
+    is a special case of NDBuffer where dim=1, stride=1, and dtype="B". However,
     in order to use Python's type system to differentiate between the contiguous
     Buffer and the n-dim (non-contiguous) NDBuffer, we keep the definition of the
     two classes separate.

src/zarr/core/buffer/cpu.py

@@ -49,7 +49,7 @@ def __init__(self, array_like: ArrayLike) -> None:
 
     @classmethod
     def create_zero_length(cls) -> Self:
-        return cls(np.array([], dtype="b"))
+        return cls(np.array([], dtype="B"))
 
     @classmethod
     def from_buffer(cls, buffer: core.Buffer) -> Self:
@@ -92,7 +92,7 @@ def from_bytes(cls, bytes_like: BytesLike) -> Self:
         -------
             New buffer representing `bytes_like`
         """
-        return cls.from_array_like(np.frombuffer(bytes_like, dtype="b"))
+        return cls.from_array_like(np.frombuffer(bytes_like, dtype="B"))
 
     def as_numpy_array(self) -> npt.NDArray[Any]:
         """Returns the buffer as a NumPy array (host memory).
@@ -111,7 +111,7 @@ def __add__(self, other: core.Buffer) -> Self:
         """Concatenate two buffers"""
 
         other_array = other.as_array_like()
-        assert other_array.dtype == np.dtype("b")
+        assert other_array.dtype == np.dtype("B")
         return self.__class__(
             np.concatenate((np.asanyarray(self._data), np.asanyarray(other_array)))
         )
@@ -131,7 +131,7 @@ class NDBuffer(core.NDBuffer):
     Notes
     -----
     The two buffer classes Buffer and NDBuffer are very similar. In fact, Buffer
-    is a special case of NDBuffer where dim=1, stride=1, and dtype="b". However,
+    is a special case of NDBuffer where dim=1, stride=1, and dtype="B". However,
     in order to use Python's type system to differentiate between the contiguous
     Buffer and the n-dim (non-contiguous) NDBuffer, we keep the definition of the
     two classes separate.

src/zarr/core/buffer/gpu.py

@@ -59,7 +59,7 @@ def __init__(self, array_like: ArrayLike) -> None:
 
         if array_like.ndim != 1:
             raise ValueError("array_like: only 1-dim allowed")
-        if array_like.dtype != np.dtype("b"):
+        if array_like.dtype != np.dtype("B"):
             raise ValueError("array_like: only byte dtype allowed")
 
         if not hasattr(array_like, "__cuda_array_interface__"):
@@ -84,7 +84,7 @@ def create_zero_length(cls) -> Self:
         -------
             New empty 0-length buffer
         """
-        return cls(cp.array([], dtype="b"))
+        return cls(cp.array([], dtype="B"))
 
     @classmethod
     def from_buffer(cls, buffer: core.Buffer) -> Self:
@@ -100,14 +100,14 @@ def from_buffer(cls, buffer: core.Buffer) -> Self:
 
     @classmethod
     def from_bytes(cls, bytes_like: BytesLike) -> Self:
-        return cls.from_array_like(cp.frombuffer(bytes_like, dtype="b"))
+        return cls.from_array_like(cp.frombuffer(bytes_like, dtype="B"))
 
     def as_numpy_array(self) -> npt.NDArray[Any]:
         return cast(npt.NDArray[Any], cp.asnumpy(self._data))
 
     def __add__(self, other: core.Buffer) -> Self:
         other_array = other.as_array_like()
-        assert other_array.dtype == np.dtype("b")
+        assert other_array.dtype == np.dtype("B")
         gpu_other = Buffer(other_array)
         gpu_other_array = gpu_other.as_array_like()
         return self.__class__(
@@ -129,7 +129,7 @@ class NDBuffer(core.NDBuffer):
     Notes
     -----
     The two buffer classes Buffer and NDBuffer are very similar. In fact, Buffer
-    is a special case of NDBuffer where dim=1, stride=1, and dtype="b". However,
+    is a special case of NDBuffer where dim=1, stride=1, and dtype="B". However,
     in order to use Python's type system to differentiate between the contiguous
     Buffer and the n-dim (non-contiguous) NDBuffer, we keep the definition of the
     two classes separate.

src/zarr/storage/_fsspec.py

@@ -32,7 +32,7 @@
 
 class FsspecStore(Store):
     """
-    A remote Store based on FSSpec
+    Store for remote data based on FSSpec.
 
     Parameters
     ----------

src/zarr/storage/_local.py

@@ -67,7 +67,7 @@ def _put(
 
 class LocalStore(Store):
     """
-    Local file system store.
+    Store for the local file system.
 
     Parameters
     ----------

src/zarr/storage/_logging.py

@@ -24,7 +24,7 @@
 
 class LoggingStore(WrapperStore[T_Store]):
     """
-    Store wrapper that logs all calls to the wrapped store.
+    Store that logs all calls to another wrapped store.
 
     Parameters
     ----------

src/zarr/storage/_memory.py

@@ -19,7 +19,7 @@
 
 class MemoryStore(Store):
     """
-    In-memory store.
+    Store for local memory.
 
     Parameters
     ----------
@@ -173,8 +173,10 @@ async def list_dir(self, prefix: str) -> AsyncIterator[str]:
 
 
 class GpuMemoryStore(MemoryStore):
-    """A GPU only memory store that stores every chunk in GPU memory irrespective
-    of the original location.
+    """
+    Store for GPU memory.
+
+    Stores every chunk in GPU memory irrespective of the original location.
 
     The dictionary of buffers to initialize this memory store with *must* be
     GPU Buffers.

src/zarr/storage/_obstore.py

@@ -37,7 +37,8 @@
 
 
 class ObjectStore(Store):
-    """A Zarr store that uses obstore for fast read/write from AWS, GCP, Azure.
+    """
+    Store that uses obstore for fast read/write from AWS, GCP, Azure.
 
     Parameters
     ----------
@@ -106,10 +107,25 @@ async def get(
                 )
                 return prototype.buffer.from_bytes(await resp.bytes_async())  # type: ignore[arg-type]
             elif isinstance(byte_range, SuffixByteRequest):
-                resp = await obs.get_async(
-                    self.store, key, options={"range": {"suffix": byte_range.suffix}}
-                )
-                return prototype.buffer.from_bytes(await resp.bytes_async())  # type: ignore[arg-type]
+                # some object stores (Azure) don't support suffix requests. In this
+                # case, our workaround is to first get the length of the object and then
+                # manually request the byte range at the end.
+                try:
+                    resp = await obs.get_async(
+                        self.store, key, options={"range": {"suffix": byte_range.suffix}}
+                    )
+                    return prototype.buffer.from_bytes(await resp.bytes_async())  # type: ignore[arg-type]
+                except obs.exceptions.NotSupportedError:
+                    head_resp = await obs.head_async(self.store, key)
+                    file_size = head_resp["size"]
+                    suffix_len = byte_range.suffix
+                    buffer = await obs.get_range_async(
+                        self.store,
+                        key,
+                        start=file_size - suffix_len,
+                        length=suffix_len,
+                    )
+                    return prototype.buffer.from_bytes(buffer)  # type: ignore[arg-type]
             else:
                 raise ValueError(f"Unexpected byte_range, got {byte_range}")
         except _ALLOWED_EXCEPTIONS:
@@ -265,10 +281,29 @@ class _OtherRequest(TypedDict):
     path: str
     """The path to request from."""
 
-    range: OffsetRange | SuffixRange | None
+    range: OffsetRange | None
+    # Note: suffix requests are handled separately because some object stores (Azure)
+    # don't support them
     """The range request type."""
 
 
+class _SuffixRequest(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: SuffixRange
+    """The suffix range."""
+
+
 class _Response(TypedDict):
     """A response buffer associated with the original index that it should be restored to."""
 
@@ -317,7 +352,7 @@ async def _make_other_request(
     prototype: BufferPrototype,
     semaphore: asyncio.Semaphore,
 ) -> list[_Response]:
-    """Make suffix or offset requests.
+    """Make offset or full-file requests.
 
     We return a `list[_Response]` for symmetry with `_make_bounded_requests` so that all
     futures can be gathered together.
@@ -339,6 +374,46 @@ async def _make_other_request(
     ]
 
 
+async def _make_suffix_request(
+    store: _UpstreamObjectStore,
+    request: _SuffixRequest,
+    prototype: BufferPrototype,
+    semaphore: asyncio.Semaphore,
+) -> list[_Response]:
+    """Make suffix requests.
+
+    This is separated out from `_make_other_request` because some object stores (Azure)
+    don't support suffix requests. In this case, our workaround is to first get the
+    length of the object and then manually request the byte range at the end.
+
+    We return a `list[_Response]` for symmetry with `_make_bounded_requests` so that all
+    futures can be gathered together.
+    """
+    import obstore as obs
+
+    async with semaphore:
+        try:
+            resp = await obs.get_async(store, request["path"], options={"range": request["range"]})
+            buffer = await resp.bytes_async()
+        except obs.exceptions.NotSupportedError:
+            head_resp = await obs.head_async(store, request["path"])
+            file_size = head_resp["size"]
+            suffix_len = request["range"]["suffix"]
+            buffer = await obs.get_range_async(
+                store,
+                request["path"],
+                start=file_size - suffix_len,
+                length=suffix_len,
+            )
+
+    return [
+        {
+            "original_request_index": request["original_request_index"],
+            "buffer": prototype.buffer.from_bytes(buffer),  # type: ignore[arg-type]
+        }
+    ]
+
+
 async def _get_partial_values(
     store: _UpstreamObjectStore,
     prototype: BufferPrototype,
@@ -358,6 +433,7 @@ async def _get_partial_values(
     key_ranges = list(key_ranges)
     per_file_bounded_requests: dict[str, list[_BoundedRequest]] = defaultdict(list)
     other_requests: list[_OtherRequest] = []
+    suffix_requests: list[_SuffixRequest] = []
 
     for idx, (path, byte_range) in enumerate(key_ranges):
         if byte_range is None:
@@ -381,7 +457,7 @@ async def _get_partial_values(
                 }
             )
         elif isinstance(byte_range, SuffixByteRequest):
-            other_requests.append(
+            suffix_requests.append(
                 {
                     "original_request_index": idx,
                     "path": path,
@@ -402,6 +478,9 @@ async def _get_partial_values(
     for request in other_requests:
         futs.append(_make_other_request(store, request, prototype, semaphore=semaphore))  # noqa: PERF401
 
+    for suffix_request in suffix_requests:
+        futs.append(_make_suffix_request(store, suffix_request, prototype, semaphore=semaphore))  # noqa: PERF401
+
     buffers: list[Buffer | None] = [None] * len(key_ranges)
 
     for responses in await asyncio.gather(*futs):