[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
     from zarr.config import RuntimeConfiguration
 
@@ -38,39 +38,39 @@ class ArrayArrayCodec(Codec):
     @abstractmethod
     async def decode(
         self,
-        chunk_array: np.ndarray,
+        chunk_array: NDBuffer,
         chunk_spec: ArraySpec,
         runtime_configuration: RuntimeConfiguration,
-    ) -> np.ndarray:
+    ) -> NDBuffer:
         pass
 
     @abstractmethod
     async def encode(
         self,
-        chunk_array: np.ndarray,
+        chunk_array: NDBuffer,
         chunk_spec: ArraySpec,
         runtime_configuration: RuntimeConfiguration,
-    ) -> Optional[np.ndarray]:
+    ) -> Optional[NDBuffer]:
         pass
 
 
 class ArrayBytesCodec(Codec):
     @abstractmethod
     async def decode(
         self,
-        chunk_array: BytesLike,
+        chunk_array: Buffer,
         chunk_spec: ArraySpec,
         runtime_configuration: RuntimeConfiguration,
-    ) -> np.ndarray:
+    ) -> NDBuffer:
         pass
 
     @abstractmethod
     async def encode(
         self,
-        chunk_array: np.ndarray,
+        chunk_array: NDBuffer,
         chunk_spec: ArraySpec,
         runtime_configuration: RuntimeConfiguration,
-    ) -> Optional[BytesLike]:
+    ) -> Optional[Buffer]:
         pass
 
 
@@ -82,7 +82,7 @@ async def decode_partial(
         selection: SliceSelection,
         chunk_spec: ArraySpec,
         runtime_configuration: RuntimeConfiguration,
-    ) -> Optional[np.ndarray]:
+    ) -> Optional[NDBuffer]:
         pass
 
 
@@ -91,7 +91,7 @@ class ArrayBytesCodecPartialEncodeMixin:
     async def encode_partial(
         self,
         store_path: StorePath,
-        chunk_array: np.ndarray,
+        chunk_array: NDBuffer,
         selection: SliceSelection,
         chunk_spec: ArraySpec,
         runtime_configuration: RuntimeConfiguration,
@@ -103,17 +103,17 @@ class BytesBytesCodec(Codec):
     @abstractmethod
     async def decode(
         self,
-        chunk_array: BytesLike,
+        chunk_array: Buffer,
         chunk_spec: ArraySpec,
         runtime_configuration: RuntimeConfiguration,
-    ) -> BytesLike:
+    ) -> Buffer:
         pass
 
     @abstractmethod
     async def encode(
         self,
-        chunk_array: BytesLike,
+        chunk_array: Buffer,
         chunk_spec: ArraySpec,
         runtime_configuration: RuntimeConfiguration,
-    ) -> 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 NDBuffer, as_buffer, as_nd_buffer
 from zarr.codecs import BytesCodec
 from zarr.common import (
     ZARR_JSON,
@@ -150,7 +151,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()),
             runtime_configuration=runtime_configuration,
         )
 
@@ -165,7 +166,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()),
                 runtime_configuration=runtime_configuration or RuntimeConfiguration(),
             )
         else:
@@ -201,8 +202,8 @@ async def getitem(self, selection: Selection) -> np.ndarray:
         )
 
         # setup output array
-        out = np.zeros(
-            indexer.shape,
+        out = NDBuffer.create_zeros(
+            shape=indexer.shape,
             dtype=self.metadata.dtype,
             order=self.runtime_configuration.order,
         )
@@ -217,20 +218,21 @@ async def getitem(self, selection: Selection) -> np.ndarray:
             self.runtime_configuration.concurrency,
         )
 
+        # We always return a numpy array to the user
         if out.shape:
-            return out
+            return out.as_numpy_array()
         else:
-            return out[()]
+            return out.as_numpy_array()[()]
 
     async def _save_metadata(self) -> None:
-        await (self.store_path / ZARR_JSON).set(self.metadata.to_bytes())
+        await (self.store_path / ZARR_JSON).set(as_buffer(self.metadata.to_bytes()))
 
     async def _read_chunk(
         self,
         chunk_coords: ChunkCoords,
         chunk_selection: SliceSelection,
         out_selection: SliceSelection,
-        out: np.ndarray,
+        out: NDBuffer,
     ) -> None:
         chunk_spec = self.metadata.get_chunk_spec(chunk_coords)
         chunk_key_encoding = self.metadata.chunk_key_encoding
@@ -257,6 +259,7 @@ async def _read_chunk(
                 out[out_selection] = self.metadata.fill_value
 
     async def setitem(self, selection: Selection, value: np.ndarray) -> None:
+        assert isinstance(value, np.ndarray)
         assert isinstance(self.metadata.chunk_grid, RegularChunkGrid)
         chunk_shape = self.metadata.chunk_grid.chunk_shape
         indexer = BasicIndexer(
@@ -278,6 +281,10 @@ async def setitem(self, selection: Selection, value: np.ndarray) -> None:
             if value.dtype.name != self.metadata.dtype.name:
                 value = value.astype(self.metadata.dtype, order="A")
 
+        # We accept a numpy array as input from the user and convert it to a NDBuffer.
+        # From this point onwards, we only pass Buffer and NDBuffer between components.
+        value = as_nd_buffer(value)
+
         # merging with existing data and encoding chunks
         await concurrent_map(
             [
@@ -296,12 +303,13 @@ async def setitem(self, selection: Selection, value: np.ndarray) -> None:
 
     async def _write_chunk(
         self,
-        value: np.ndarray,
+        value: NDBuffer,
         chunk_shape: ChunkCoords,
         chunk_coords: ChunkCoords,
         chunk_selection: SliceSelection,
         out_selection: SliceSelection,
     ) -> None:
+        assert isinstance(value, NDBuffer)
         chunk_spec = self.metadata.get_chunk_spec(chunk_coords)
         chunk_key_encoding = self.metadata.chunk_key_encoding
         chunk_key = chunk_key_encoding.encode_chunk_key(chunk_coords)
@@ -310,8 +318,8 @@ 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,
+                chunk_array = NDBuffer.create_empty(
+                    shape=chunk_shape,
                     dtype=self.metadata.dtype,
                 )
                 chunk_array.fill(value)
@@ -335,8 +343,8 @@ async def _write_chunk(
 
             # merge new value
             if chunk_bytes is None:
-                chunk_array = np.empty(
-                    chunk_shape,
+                chunk_array = NDBuffer.create_empty(
+                    shape=chunk_shape,
                     dtype=self.metadata.dtype,
                 )
                 chunk_array.fill(self.metadata.fill_value)
@@ -349,9 +357,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: np.ndarray, chunk_spec: ArraySpec
+        self, store_path: StorePath, chunk_array: NDBuffer, chunk_spec: ArraySpec
     ) -> None:
-        if np.all(chunk_array == self.metadata.fill_value):
+        if np.all(chunk_array.as_numpy_array() == self.metadata.fill_value):
             # chunks that only contain fill_value will be removed
             await store_path.delete()
         else:
@@ -392,14 +400,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(as_buffer(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(as_buffer(new_metadata.to_bytes()))
         return replace(self, metadata=new_metadata)
 
     def __repr__(self):