Merge branch 'main' of https://github.com/zarr-developers/zarr-python into feat/serializable-protocol

Author: d-v-b

.github/workflows/needs_release_notes.yml

@@ -1,8 +1,8 @@
 name: "Pull Request Labeler"
 
 on:
-  - pull_request_target:
-      types: [opened, reopened, synchronize]
+  pull_request_target:
+    types: [opened, reopened, synchronize]
 
 jobs:
   labeler:

changes/3797.bugfix.md

@@ -0,0 +1,2 @@
+Fix an issue that prevents the correct parsing of special NumPy ``uint32`` dtypes resulting e.g.
+from bit wise operations on ``uint32`` arrays on Windows.

changes/3828.misc.md

@@ -0,0 +1,2 @@
+`CodecPipeline.read` and `CodecPipeline.read_batch` now return a tuple of typeddict objects
+that each carry information about the request for a chunk from storage.

changes/3836.doc.md

@@ -0,0 +1,2 @@
+Corrects the type annotation reported for the `batch_info` parameter in the `CodecPipeline.write`
+method docstring.

mkdocs.yml

@@ -216,7 +216,6 @@ plugins:
         'developers/index.html.md': 'contributing.md'
         'developers/roadmap.html.md': 'https://zarr.readthedocs.io/en/v3.0.8/developers/roadmap.html'
         'api/zarr/creation.md': 'api/zarr/deprecated/creation.md'
-        'api/zarr/codecs/numcodecs.md': 'api/zarr/deprecated/creation.md'
         'api.md': 'api/zarr/index.md'
         'api/zarr/metadata/migrate_v3.md': 'api/zarr/metadata.md'
 

src/zarr/abc/codec.py

@@ -2,7 +2,7 @@
 
 from abc import abstractmethod
 from collections.abc import Mapping
-from typing import TYPE_CHECKING, Generic, Protocol, TypeGuard, TypeVar, runtime_checkable
+from typing import TYPE_CHECKING, Generic, Literal, Protocol, TypeGuard, TypeVar, runtime_checkable
 
 from typing_extensions import ReadOnly, TypedDict
 
@@ -32,9 +32,17 @@
     "CodecInput",
     "CodecOutput",
     "CodecPipeline",
+    "GetResult",
     "SupportsSyncCodec",
 ]
 
+
+class GetResult(TypedDict):
+    """Metadata about a store get operation."""
+
+    status: Literal["present", "missing"]
+
+
 CodecInput = TypeVar("CodecInput", bound=NDBuffer | Buffer)
 CodecOutput = TypeVar("CodecOutput", bound=NDBuffer | Buffer)
 
@@ -433,13 +441,13 @@ async def read(
         batch_info: Iterable[tuple[ByteGetter, ArraySpec, SelectorTuple, SelectorTuple, bool]],
         out: NDBuffer,
         drop_axes: tuple[int, ...] = (),
-    ) -> None:
+    ) -> tuple[GetResult, ...]:
         """Reads chunk data from the store, decodes it and writes it into an output array.
         Partial decoding may be utilized if the codecs and stores support it.
 
         Parameters
         ----------
-        batch_info : Iterable[tuple[ByteGetter, ArraySpec, SelectorTuple, SelectorTuple]]
+        batch_info : Iterable[tuple[ByteGetter, ArraySpec, SelectorTuple, SelectorTuple, bool]]
             Ordered set of information about the chunks.
             The first slice selection determines which parts of the chunk will be fetched.
             The second slice selection determines where in the output array the chunk data will be written.
@@ -451,6 +459,11 @@ async def read(
             ``out``) to the fill value for the array.
 
         out : NDBuffer
+
+        Returns
+        -------
+        tuple[GetResult, ...]
+            One result per chunk in ``batch_info``.
         """
         ...
 
@@ -467,7 +480,7 @@ async def write(
 
         Parameters
         ----------
-        batch_info : Iterable[tuple[ByteSetter, ArraySpec, SelectorTuple, SelectorTuple]]
+        batch_info : Iterable[tuple[ByteSetter, ArraySpec, SelectorTuple, SelectorTuple, bool]]
             Ordered set of information about the chunks.
             The first slice selection determines which parts of the chunk will be encoded.
             The second slice selection determines where in the value array the chunk data is located.

src/zarr/codecs/sharding.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from collections.abc import Iterable, Mapping, MutableMapping
+from collections.abc import Iterable, Mapping, MutableMapping, Sequence
 from dataclasses import dataclass, replace
 from enum import Enum
 from functools import lru_cache
@@ -45,6 +45,7 @@
 from zarr.core.dtype.npy.int import UInt64
 from zarr.core.indexing import (
     BasicIndexer,
+    ChunkProjection,
     SelectorTuple,
     _morton_order,
     _morton_order_keys,
@@ -574,21 +575,26 @@ async def _encode_partial_single(
         chunks_per_shard = self._get_chunks_per_shard(shard_spec)
         chunk_spec = self._get_chunk_spec(shard_spec)
 
-        shard_reader = await self._load_full_shard_maybe(
-            byte_getter=byte_setter,
-            prototype=chunk_spec.prototype,
-            chunks_per_shard=chunks_per_shard,
-        )
-        shard_reader = shard_reader or _ShardReader.create_empty(chunks_per_shard)
-        # Use vectorized lookup for better performance
-        shard_dict = shard_reader.to_dict_vectorized(np.asarray(_morton_order(chunks_per_shard)))
-
         indexer = list(
             get_indexer(
                 selection, shape=shard_shape, chunk_grid=RegularChunkGrid(chunk_shape=chunk_shape)
             )
         )
 
+        if self._is_complete_shard_write(indexer, chunks_per_shard):
+            shard_dict = dict.fromkeys(morton_order_iter(chunks_per_shard))
+        else:
+            shard_reader = await self._load_full_shard_maybe(
+                byte_getter=byte_setter,
+                prototype=chunk_spec.prototype,
+                chunks_per_shard=chunks_per_shard,
+            )
+            shard_reader = shard_reader or _ShardReader.create_empty(chunks_per_shard)
+            # Use vectorized lookup for better performance
+            shard_dict = shard_reader.to_dict_vectorized(
+                np.asarray(_morton_order(chunks_per_shard))
+            )
+
         await self.codec_pipeline.write(
             [
                 (
@@ -661,6 +667,16 @@ def _is_total_shard(
             chunk_coords in all_chunk_coords for chunk_coords in c_order_iter(chunks_per_shard)
         )
 
+    def _is_complete_shard_write(
+        self,
+        indexed_chunks: Sequence[ChunkProjection],
+        chunks_per_shard: tuple[int, ...],
+    ) -> bool:
+        all_chunk_coords = {chunk_coords for chunk_coords, *_ in indexed_chunks}
+        return self._is_total_shard(all_chunk_coords, chunks_per_shard) and all(
+            is_complete_chunk for *_, is_complete_chunk in indexed_chunks
+        )
+
     async def _decode_shard_index(
         self, index_bytes: Buffer, chunks_per_shard: tuple[int, ...]
     ) -> _ShardIndex:

src/zarr/core/codec_pipeline.py

@@ -13,6 +13,7 @@
     BytesBytesCodec,
     Codec,
     CodecPipeline,
+    GetResult,
 )
 from zarr.core.common import concurrent_map
 from zarr.core.config import config
@@ -251,47 +252,58 @@ async def read_batch(
         batch_info: Iterable[tuple[ByteGetter, ArraySpec, SelectorTuple, SelectorTuple, bool]],
         out: NDBuffer,
         drop_axes: tuple[int, ...] = (),
-    ) -> None:
+    ) -> tuple[GetResult, ...]:
+        results: list[GetResult] = []
         if self.supports_partial_decode:
+            batch_info_list = list(batch_info)
             chunk_array_batch = await self.decode_partial_batch(
                 [
                     (byte_getter, chunk_selection, chunk_spec)
-                    for byte_getter, chunk_spec, chunk_selection, *_ in batch_info
+                    for byte_getter, chunk_spec, chunk_selection, *_ in batch_info_list
                 ]
             )
             for chunk_array, (_, chunk_spec, _, out_selection, _) in zip(
-                chunk_array_batch, batch_info, strict=False
+                chunk_array_batch, batch_info_list, strict=False
             ):
                 if chunk_array is not None:
                     if drop_axes:
                         chunk_array = chunk_array.squeeze(axis=drop_axes)
                     out[out_selection] = chunk_array
+                    results.append(GetResult(status="present"))
                 else:
                     out[out_selection] = fill_value_or_default(chunk_spec)
+                    results.append(GetResult(status="missing"))
         else:
+            batch_info_list = list(batch_info)
             chunk_bytes_batch = await concurrent_map(
-                [(byte_getter, array_spec.prototype) for byte_getter, array_spec, *_ in batch_info],
+                [
+                    (byte_getter, array_spec.prototype)
+                    for byte_getter, array_spec, *_ in batch_info_list
+                ],
                 lambda byte_getter, prototype: byte_getter.get(prototype),
                 config.get("async.concurrency"),
             )
             chunk_array_batch = await self.decode_batch(
                 [
                     (chunk_bytes, chunk_spec)
                     for chunk_bytes, (_, chunk_spec, *_) in zip(
-                        chunk_bytes_batch, batch_info, strict=False
+                        chunk_bytes_batch, batch_info_list, strict=False
                     )
                 ],
             )
             for chunk_array, (_, chunk_spec, chunk_selection, out_selection, _) in zip(
-                chunk_array_batch, batch_info, strict=False
+                chunk_array_batch, batch_info_list, strict=False
             ):
                 if chunk_array is not None:
                     tmp = chunk_array[chunk_selection]
                     if drop_axes:
                         tmp = tmp.squeeze(axis=drop_axes)
                     out[out_selection] = tmp
+                    results.append(GetResult(status="present"))
                 else:
                     out[out_selection] = fill_value_or_default(chunk_spec)
+                    results.append(GetResult(status="missing"))
+        return tuple(results)
 
     def _merge_chunk_array(
         self,
@@ -471,15 +483,19 @@ async def read(
         batch_info: Iterable[tuple[ByteGetter, ArraySpec, SelectorTuple, SelectorTuple, bool]],
         out: NDBuffer,
         drop_axes: tuple[int, ...] = (),
-    ) -> None:
-        await concurrent_map(
+    ) -> tuple[GetResult, ...]:
+        batch_results = await concurrent_map(
             [
                 (single_batch_info, out, drop_axes)
                 for single_batch_info in batched(batch_info, self.batch_size)
             ],
             self.read_batch,
             config.get("async.concurrency"),
         )
+        results: list[GetResult] = []
+        for batch in batch_results:
+            results.extend(batch)
+        return tuple(results)
 
     async def write(
         self,

src/zarr/core/common.py

@@ -49,7 +49,7 @@
 AccessModeLiteral = Literal["r", "r+", "a", "w", "w-"]
 ANY_ACCESS_MODE: Final = "r", "r+", "a", "w", "w-"
 DimensionNamesLike = Iterable[str | None] | None
-DimensionNames = DimensionNamesLike  # backwards compatibility
+DimensionNames = DimensionNamesLike  # for backwards compatibility
 
 TName = TypeVar("TName", bound=str)
 TConfig = TypeVar("TConfig", bound=Mapping[str, object])

src/zarr/core/dtype/npy/int.py

@@ -1070,6 +1070,28 @@ class UInt32(BaseInt[np.dtypes.UInt32DType, np.uint32], HasEndianness):
     _zarr_v3_name: ClassVar[Literal["uint32"]] = "uint32"
     _zarr_v2_names: ClassVar[tuple[Literal[">u4"], Literal["<u4"]]] = (">u4", "<u4")
 
+    @classmethod
+    def _check_native_dtype(cls: type[Self], dtype: TBaseDType) -> TypeGuard[np.dtypes.UInt32DType]:
+        """
+        A type guard that checks if the input is assignable to the type of ``cls.dtype_class``
+
+        This method is overridden for this particular data type because of a Windows-specific issue
+        where ``np.array([1], dtype=np.uint32) & 1`` creates an instance of ``np.dtypes.UIntDType``,
+        rather than an instance of ``np.dtypes.UInt32DType``, even though both represent 32-bit
+        unsigned integers. (In contrast to ``np.dtype('i')``, ``np.dtype('u')`` raises an error.)
+
+        Parameters
+        ----------
+        dtype : TDType
+            The dtype to check.
+
+        Returns
+        -------
+        Bool
+            True if the dtype matches, False otherwise.
+        """
+        return super()._check_native_dtype(dtype) or dtype == np.dtypes.UInt32DType()
+
     @classmethod
     def from_native_dtype(cls, dtype: TBaseDType) -> Self:
         """