Support all indexing variants (#1917)

zarr-developers · Jun 3, 2024 · 24e855c · 24e855c
1 parent 56d36e3
commit 24e855c
Show file tree

Hide file tree

Showing 10 changed files with 3,200 additions and 169 deletions.
diff --git a/src/zarr/abc/codec.py b/src/zarr/abc/codec.py
@@ -13,7 +13,8 @@
 if TYPE_CHECKING:
     from typing_extensions import Self
 
-    from zarr.common import ArraySpec, SliceSelection
+    from zarr.common import ArraySpec
+    from zarr.indexing import SelectorTuple
     from zarr.metadata import ArrayMetadata
 
 
@@ -155,21 +156,21 @@ class ArrayBytesCodecPartialDecodeMixin:
     """Mixin for array-to-bytes codecs that implement partial decoding."""
 
     async def _decode_partial_single(
-        self, byte_getter: ByteGetter, selection: SliceSelection, chunk_spec: ArraySpec
+        self, byte_getter: ByteGetter, selection: SelectorTuple, chunk_spec: ArraySpec
     ) -> NDBuffer | None:
         raise NotImplementedError
 
     async def decode_partial(
         self,
-        batch_info: Iterable[tuple[ByteGetter, SliceSelection, ArraySpec]],
+        batch_info: Iterable[tuple[ByteGetter, SelectorTuple, ArraySpec]],
     ) -> Iterable[NDBuffer | None]:
         """Partially decodes a batch of chunks.
         This method determines parts of a chunk from the slice selection,
         fetches these parts from the store (via ByteGetter) and decodes them.
 
         Parameters
         ----------
-        batch_info : Iterable[tuple[ByteGetter, SliceSelection, ArraySpec]]
+        batch_info : Iterable[tuple[ByteGetter, SelectorTuple, ArraySpec]]
             Ordered set of information about slices of encoded chunks.
             The slice selection determines which parts of the chunk will be fetched.
             The ByteGetter is used to fetch the necessary bytes.
@@ -196,14 +197,14 @@ async def _encode_partial_single(
         self,
         byte_setter: ByteSetter,
         chunk_array: NDBuffer,
-        selection: SliceSelection,
+        selection: SelectorTuple,
         chunk_spec: ArraySpec,
     ) -> None:
         raise NotImplementedError
 
     async def encode_partial(
         self,
-        batch_info: Iterable[tuple[ByteSetter, NDBuffer, SliceSelection, ArraySpec]],
+        batch_info: Iterable[tuple[ByteSetter, NDBuffer, SelectorTuple, ArraySpec]],
     ) -> None:
         """Partially encodes a batch of chunks.
         This method determines parts of a chunk from the slice selection, encodes them and
@@ -213,7 +214,7 @@ async def encode_partial(
 
         Parameters
         ----------
-        batch_info : Iterable[tuple[ByteSetter, NDBuffer, SliceSelection, ArraySpec]]
+        batch_info : Iterable[tuple[ByteSetter, NDBuffer, SelectorTuple, ArraySpec]]
             Ordered set of information about slices of to-be-encoded chunks.
             The slice selection determines which parts of the chunk will be encoded.
             The ByteSetter is used to write the necessary bytes and fetch bytes for existing chunk data.
@@ -342,15 +343,16 @@ async def encode(
     @abstractmethod
     async def read(
         self,
-        batch_info: Iterable[tuple[ByteGetter, ArraySpec, SliceSelection, SliceSelection]],
+        batch_info: Iterable[tuple[ByteGetter, ArraySpec, SelectorTuple, SelectorTuple]],
         out: NDBuffer,
+        drop_axes: tuple[int, ...] = (),
     ) -> None:
         """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, SliceSelection, SliceSelection]]
+        batch_info : Iterable[tuple[ByteGetter, ArraySpec, SelectorTuple, SelectorTuple]]
             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.
@@ -363,16 +365,17 @@ async def read(
     @abstractmethod
     async def write(
         self,
-        batch_info: Iterable[tuple[ByteSetter, ArraySpec, SliceSelection, SliceSelection]],
+        batch_info: Iterable[tuple[ByteSetter, ArraySpec, SelectorTuple, SelectorTuple]],
         value: NDBuffer,
+        drop_axes: tuple[int, ...] = (),
     ) -> None:
         """Encodes chunk data and writes it to the store.
         Merges with existing chunk data by reading first, if necessary.
         Partial encoding may be utilized if the codecs and stores support it.
 
         Parameters
         ----------
-        batch_info : Iterable[tuple[ByteSetter, ArraySpec, SliceSelection, SliceSelection]]
+        batch_info : Iterable[tuple[ByteSetter, ArraySpec, SelectorTuple, SelectorTuple]]
             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.