[v3] First step to generalizes ndarray and bytes

src/zarr/abc/codec.py

@@ -3,16 +3,16 @@
 from abc import abstractmethod
 from typing import TYPE_CHECKING, Optional
 
-import numpy as np
 from zarr.abc.metadata import Metadata
 
+from zarr.buffer import Buffer, NDBuffer
 from zarr.common import ArraySpec
 from zarr.store import StorePath
 
 
 if TYPE_CHECKING:
     from typing_extensions import Self
-    from zarr.common import BytesLike, SliceSelection
+    from zarr.common import SliceSelection
     from zarr.metadata import ArrayMetadata
 
 
@@ -37,35 +37,35 @@ class ArrayArrayCodec(Codec):
     @abstractmethod
     async def decode(
         self,
-        chunk_array: np.ndarray,
+        chunk_array: NDBuffer,
         chunk_spec: ArraySpec,
-    ) -> np.ndarray:
+    ) -> NDBuffer:
         pass
 
     @abstractmethod
     async def encode(
         self,
-        chunk_array: np.ndarray,
+        chunk_array: NDBuffer,
         chunk_spec: ArraySpec,
-    ) -> Optional[np.ndarray]:
+    ) -> Optional[NDBuffer]:
         pass
 
 
 class ArrayBytesCodec(Codec):
     @abstractmethod
     async def decode(
         self,
-        chunk_array: BytesLike,
+        chunk_array: Buffer,
         chunk_spec: ArraySpec,
-    ) -> np.ndarray:
+    ) -> NDBuffer:
         pass
 
     @abstractmethod
     async def encode(
         self,
-        chunk_array: np.ndarray,
+        chunk_array: NDBuffer,
         chunk_spec: ArraySpec,
-    ) -> Optional[BytesLike]:
+    ) -> Optional[Buffer]:
         pass
 
 
@@ -76,7 +76,7 @@ async def decode_partial(
         store_path: StorePath,
         selection: SliceSelection,
         chunk_spec: ArraySpec,
-    ) -> Optional[np.ndarray]:
+    ) -> Optional[NDBuffer]:
         pass
 
 
@@ -85,7 +85,7 @@ class ArrayBytesCodecPartialEncodeMixin:
     async def encode_partial(
         self,
         store_path: StorePath,
-        chunk_array: np.ndarray,
+        chunk_array: NDBuffer,
         selection: SliceSelection,
         chunk_spec: ArraySpec,
     ) -> None:
@@ -96,15 +96,15 @@ class BytesBytesCodec(Codec):
     @abstractmethod
     async def decode(
         self,
-        chunk_array: BytesLike,
+        chunk_array: Buffer,
         chunk_spec: ArraySpec,
-    ) -> BytesLike:
+    ) -> Buffer:
         pass
 
     @abstractmethod
     async def encode(
         self,
-        chunk_array: BytesLike,
+        chunk_array: Buffer,
         chunk_spec: ArraySpec,
-    ) -> Optional[BytesLike]:
+    ) -> Optional[Buffer]:
         pass

src/zarr/abc/store.py

@@ -3,12 +3,14 @@
 from collections.abc import AsyncGenerator
 from typing import List, Tuple, Optional
 
+from zarr.buffer import Buffer
+
 
 class Store(ABC):
     @abstractmethod
     async def get(
         self, key: str, byte_range: Optional[Tuple[int, Optional[int]]] = None
-    ) -> Optional[bytes]:
+    ) -> Optional[Buffer]:
         """Retrieve the value associated with a given key.
 
         Parameters
@@ -18,14 +20,14 @@ async def get(
 
         Returns
         -------
-        bytes
+        Buffer
         """
         ...
 
     @abstractmethod
     async def get_partial_values(
         self, key_ranges: List[Tuple[str, Tuple[int, int]]]
-    ) -> List[Optional[bytes]]:
+    ) -> List[Optional[Buffer]]:
         """Retrieve possibly partial values from given key_ranges.
 
         Parameters
@@ -35,8 +37,7 @@ async def get_partial_values(
 
         Returns
         -------
-        list[bytes]
-            list of values, in the order of the key_ranges, may contain null/none for missing keys
+        list of values, in the order of the key_ranges, may contain null/none for missing keys
         """
         ...
 
@@ -61,7 +62,7 @@ def supports_writes(self) -> bool:
         ...
 
     @abstractmethod
-    async def set(self, key: str, value: bytes) -> None:
+    async def set(self, key: str, value: Buffer) -> None:
         """Store a (key, value) pair.
 
         Parameters

src/zarr/array.py

@@ -20,6 +20,7 @@
 
 
 # from zarr.array_v2 import ArrayV2
+from zarr.buffer import Buffer, Factory, NDArrayLike, NDBuffer
 from zarr.codecs import BytesCodec
 from zarr.codecs.pipeline import CodecPipeline
 from zarr.common import (
@@ -145,7 +146,7 @@ async def open(
         assert zarr_json_bytes is not None
         return cls.from_dict(
             store_path,
-            json.loads(zarr_json_bytes),
+            json.loads(zarr_json_bytes.to_bytes()),
         )
 
     @classmethod
@@ -158,7 +159,7 @@ async def open_auto(
         if v3_metadata_bytes is not None:
             return cls.from_dict(
                 store_path,
-                json.loads(v3_metadata_bytes),
+                json.loads(v3_metadata_bytes.to_bytes()),
             )
         else:
             raise ValueError("no v2 support yet")
@@ -184,7 +185,9 @@ def dtype(self) -> np.dtype[Any]:
     def attrs(self) -> dict[str, Any]:
         return self.metadata.attributes
 
-    async def getitem(self, selection: Selection) -> npt.NDArray[Any]:
+    async def getitem(
+        self, selection: Selection, *, factory: Factory.Create = NDBuffer.create
+    ) -> NDArrayLike:
         assert isinstance(self.metadata.chunk_grid, RegularChunkGrid)
         indexer = BasicIndexer(
             selection,
@@ -193,10 +196,8 @@ async def getitem(self, selection: Selection) -> npt.NDArray[Any]:
         )
 
         # setup output array
-        out = np.zeros(
-            indexer.shape,
-            dtype=self.metadata.dtype,
-            order=self.order,
+        out = factory(
+            shape=indexer.shape, dtype=self.metadata.dtype, order=self.order, fill_value=0
         )
 
         # reading chunks and decoding them
@@ -208,21 +209,17 @@ async def getitem(self, selection: Selection) -> npt.NDArray[Any]:
             self._read_chunk,
             config.get("async.concurrency"),
         )
-
-        if out.shape:
-            return out
-        else:
-            return out[()]
+        return out.as_ndarray_like()
 
     async def _save_metadata(self) -> None:
-        await (self.store_path / ZARR_JSON).set(self.metadata.to_bytes())
+        await (self.store_path / ZARR_JSON).set(Buffer.from_bytes(self.metadata.to_bytes()))
 
     async def _read_chunk(
         self,
         chunk_coords: ChunkCoords,
         chunk_selection: SliceSelection,
         out_selection: SliceSelection,
-        out: npt.NDArray[Any],
+        out: NDBuffer,
     ) -> None:
         chunk_spec = self.metadata.get_chunk_spec(chunk_coords, self.order)
         chunk_key_encoding = self.metadata.chunk_key_encoding
@@ -244,7 +241,12 @@ async def _read_chunk(
             else:
                 out[out_selection] = self.metadata.fill_value
 
-    async def setitem(self, selection: Selection, value: npt.NDArray[Any]) -> None:
+    async def setitem(
+        self,
+        selection: Selection,
+        value: NDArrayLike,
+        factory: Factory.NDArrayLike = NDBuffer.from_ndarray_like,
+    ) -> None:
         assert isinstance(self.metadata.chunk_grid, RegularChunkGrid)
         chunk_shape = self.metadata.chunk_grid.chunk_shape
         indexer = BasicIndexer(
@@ -257,15 +259,19 @@ async def setitem(self, selection: Selection, value: npt.NDArray[Any]) -> None:
 
         # check value shape
         if np.isscalar(value):
-            # setting a scalar value
-            pass
+            value = np.asanyarray(value)
         else:
             if not hasattr(value, "shape"):
                 value = np.asarray(value, self.metadata.dtype)
             assert value.shape == sel_shape
             if value.dtype.name != self.metadata.dtype.name:
                 value = value.astype(self.metadata.dtype, order="A")
 
+        # We accept any ndarray like object from the user and convert it
+        # to a NDBuffer (or subclass). From this point onwards, we only pass
+        # Buffer and NDBuffer between components.
+        value = factory(value)
+
         # merging with existing data and encoding chunks
         await concurrent_map(
             [
@@ -284,7 +290,7 @@ async def setitem(self, selection: Selection, value: npt.NDArray[Any]) -> None:
 
     async def _write_chunk(
         self,
-        value: npt.NDArray[Any],
+        value: NDBuffer,
         chunk_shape: ChunkCoords,
         chunk_coords: ChunkCoords,
         chunk_selection: SliceSelection,
@@ -298,11 +304,9 @@ async def _write_chunk(
         if is_total_slice(chunk_selection, chunk_shape):
             # write entire chunks
             if np.isscalar(value):
-                chunk_array = np.empty(
-                    chunk_shape,
-                    dtype=self.metadata.dtype,
+                chunk_array = NDBuffer.create(
+                    shape=chunk_shape, dtype=self.metadata.dtype, fill_value=value
                 )
-                chunk_array.fill(value)
             else:
                 chunk_array = value[out_selection]
             await self._write_chunk_to_store(store_path, chunk_array, chunk_spec)
@@ -322,11 +326,11 @@ async def _write_chunk(
 
             # merge new value
             if chunk_bytes is None:
-                chunk_array = np.empty(
-                    chunk_shape,
+                chunk_array = NDBuffer.create(
+                    shape=chunk_shape,
                     dtype=self.metadata.dtype,
+                    fill_value=self.metadata.fill_value,
                 )
-                chunk_array.fill(self.metadata.fill_value)
             else:
                 chunk_array = (
                     await self.codecs.decode(chunk_bytes, chunk_spec)
@@ -336,9 +340,9 @@ async def _write_chunk(
             await self._write_chunk_to_store(store_path, chunk_array, chunk_spec)
 
     async def _write_chunk_to_store(
-        self, store_path: StorePath, chunk_array: npt.NDArray[Any], chunk_spec: ArraySpec
+        self, store_path: StorePath, chunk_array: NDBuffer, chunk_spec: ArraySpec
     ) -> None:
-        if np.all(chunk_array == self.metadata.fill_value):
+        if chunk_array.all_equal(self.metadata.fill_value):
             # chunks that only contain fill_value will be removed
             await store_path.delete()
         else:
@@ -377,14 +381,14 @@ async def _delete_key(key: str) -> None:
         )
 
         # Write new metadata
-        await (self.store_path / ZARR_JSON).set(new_metadata.to_bytes())
+        await (self.store_path / ZARR_JSON).set(Buffer.from_bytes(new_metadata.to_bytes()))
         return replace(self, metadata=new_metadata)
 
     async def update_attributes(self, new_attributes: Dict[str, Any]) -> AsyncArray:
         new_metadata = replace(self.metadata, attributes=new_attributes)
 
         # Write new metadata
-        await (self.store_path / ZARR_JSON).set(new_metadata.to_bytes())
+        await (self.store_path / ZARR_JSON).set(Buffer.from_bytes(new_metadata.to_bytes()))
         return replace(self, metadata=new_metadata)
 
     def __repr__(self) -> str: