Refactor DelayedBuffer + implement efficient __getitem__

Author: balbasty

src/zarr/core/buffer/core.py

@@ -502,6 +502,89 @@ class BufferPrototype(NamedTuple):
     nd_buffer: type[NDBuffer]
 
 
+class DelayedBuffer(Buffer):
+    """
+    A Buffer that is the virtual concatenation of other buffers.
+    """
+    _BufferImpl: type
+    _concatenate: callable
+
+    def __init__(self, array: NDArrayLike | list[NDArrayLike] | None) -> None:
+        if array is None:
+            self._data_list = []
+        elif isinstance(array, list):
+            self._data_list = list(array)
+        else:
+            self._data_list = [array]
+        for array in self._data_list:
+            if array.ndim != 1:
+                raise ValueError("array: only 1-dim allowed")
+            if array.dtype != np.dtype("b"):
+                raise ValueError("array: only byte dtype allowed")
+
+    @property
+    def _data(self) -> npt.NDArray[Any]:
+        return type(self)._concatenate(self._data_list)
+
+    @classmethod
+    def from_buffer(cls, buffer: Buffer) -> Self:
+        if isinstance(buffer, cls):
+            return cls(buffer._data_list)
+        else:
+            return cls(buffer._data)
+
+    def __add__(self, other: Buffer) -> Self:
+        if isinstance(other, self.__class__):
+            return self.__class__(self._data_list + other._data_list)
+        else:
+            return self.__class__(self._data_list + [other._data])
+
+    def __radd__(self, other: Buffer) -> Self:
+        if isinstance(other, self.__class__):
+            return self.__class__(other._data_list + self._data_list)
+        else:
+            return self.__class__([other._data] + self._data_list)
+
+    def __len__(self) -> int:
+        return sum(map(len, self._data_list))
+
+    def __getitem__(self, key: slice) -> Self:
+        check_item_key_is_1d_contiguous(key)
+        start, stop = key.start, key.stop
+        if start is None:
+            start = 0
+        if stop is None:
+            stop = len(self)
+        new_list = []
+        offset = 0
+        found_last = False
+        for chunk in self._data_list:
+            chunk_size = len(chunk)
+            skip = False
+            if offset <= start < offset + chunk_size:
+                # first chunk
+                if stop <= offset + chunk_size:
+                    # also last chunk
+                    chunk = chunk[start-offset:stop-offset]
+                    found_last = True
+                else:
+                    chunk = chunk[start-offset:]
+            elif offset <= stop <= offset + chunk_size:
+                # last chunk
+                chunk = chunk[:stop-offset]
+                found_last = True
+            elif offset + chunk_size <= start:
+                skip = True
+
+            if not skip:
+                new_list.append(chunk)
+            if found_last:
+                break
+            offset += chunk_size
+        assert sum(map(len, new_list)) == stop - start
+        return self.__class__(new_list)
+
+
 # The default buffer prototype used throughout the Zarr codebase.
 def default_buffer_prototype() -> BufferPrototype:
     from zarr.registry import (

src/zarr/core/buffer/cpu.py

@@ -185,43 +185,16 @@ def __setitem__(self, key: Any, value: Any) -> None:
         self._data.__setitem__(key, value)
 
 
-class DelayedBuffer(Buffer):
+class DelayedBuffer(core.DelayedBuffer, Buffer):
     """
     A Buffer that is the virtual concatenation of other buffers.
     """
+    _BufferImpl = Buffer
+    _concatenate = np.concatenate
 
     def __init__(self, array: NDArrayLike | list[NDArrayLike] | None) -> None:
-        if array is None:
-            self._data_list = []
-        elif isinstance(array, list):
-            self._data_list = list(array)
-        else:
-            self._data_list = [array]
-        for array in self._data_list:
-            if array.ndim != 1:
-                raise ValueError("array: only 1-dim allowed")
-            if array.dtype != np.dtype("b"):
-                raise ValueError("array: only byte dtype allowed")
-
-    @property
-    def _data(self) -> npt.NDArray[Any]:
-        return np.concatenate(self._data_list)
-
-    @classmethod
-    def from_buffer(cls, buffer: core.Buffer) -> Self:
-        if isinstance(buffer, cls):
-            return cls(buffer._data_list)
-        else:
-            return cls(buffer._data)
-
-    def __add__(self, other: core.Buffer) -> Self:
-        if isinstance(other, self.__class__):
-            return self.__class__(self._data_list + other._data_list)
-        else:
-            return self.__class__(self._data_list + [other._data])
-
-    def __len__(self) -> int:
-        return sum(map(len, self._data_list))
+        core.DelayedBuffer.__init__(self, array)
+        self._data_list = list(map(np.asanyarray, self._data_list))
 
 
 Buffer.Delayed = DelayedBuffer

src/zarr/core/buffer/gpu.py

@@ -218,45 +218,17 @@ def __setitem__(self, key: Any, value: Any) -> None:
         self._data.__setitem__(key, value)
 
 
-class DelayedBuffer(Buffer):
+class DelayedBuffer(core.DelayedBuffer, Buffer):
     """
     A Buffer that is the virtual concatenation of other buffers.
     """
+    _BufferImpl = Buffer
+    _concatenate = getattr(cp, 'concatenate', None)
 
     def __init__(self, array: NDArrayLike | list[NDArrayLike] | None) -> None:
-        if array is None:
-            self._data_list = []
-        elif isinstance(array, list):
-            self._data_list = list(array)
-        else:
-            self._data_list = [array]
-        for array in self._data_list:
-            if array.ndim != 1:
-                raise ValueError("array: only 1-dim allowed")
-            if array.dtype != np.dtype("b"):
-                raise ValueError("array: only byte dtype allowed")
+        core.DelayedBuffer.__init__(self, array)
         self._data_list = list(map(cp.asarray, self._data_list))
 
-    @property
-    def _data(self) -> npt.NDArray[Any]:
-        return cp.concatenate(self._data_list)
-
-    @classmethod
-    def from_buffer(cls, buffer: core.Buffer) -> Self:
-        if isinstance(buffer, cls):
-            return cls(buffer._data_list)
-        else:
-            return cls(buffer._data)
-
-    def __add__(self, other: core.Buffer) -> Self:
-        if isinstance(other, self.__class__):
-            return self.__class__(self._data_list + other._data_list)
-        else:
-            return self.__class__(self._data_list + [other._data])
-
-    def __len__(self) -> int:
-        return sum(map(len, self._data_list))
-
 
 Buffer.Delayed = DelayedBuffer