Sharding: do not merge new/old if old is empty + avoid sequential buffer concatenation

Author: balbasty

src/zarr/codecs/sharding.py

@@ -251,7 +251,7 @@ def create_empty(
         if buffer_prototype is None:
             buffer_prototype = default_buffer_prototype()
         obj = cls()
-        obj.buf = buffer_prototype.buffer.create_zero_length()
+        obj.buf = buffer_prototype.buffer.Delayed.create_zero_length()
         obj.index = _ShardIndex.create_empty(chunks_per_shard)
         return obj
 
@@ -585,15 +585,16 @@ async def _encode_partial_single(
         chunks_per_shard = self._get_chunks_per_shard(shard_spec)
         chunk_spec = self._get_chunk_spec(shard_spec)
 
-        shard_dict = _MergingShardBuilder(
-            await self._load_full_shard_maybe(
-                byte_getter=byte_setter,
-                prototype=chunk_spec.prototype,
-                chunks_per_shard=chunks_per_shard,
-            )
-            or _ShardReader.create_empty(chunks_per_shard),
-            _ShardBuilder.create_empty(chunks_per_shard),
+        shard_read = await self._load_full_shard_maybe(
+            byte_getter=byte_setter,
+            prototype=chunk_spec.prototype,
+            chunks_per_shard=chunks_per_shard,
         )
+        shard_build = _ShardBuilder.create_empty(chunks_per_shard)
+        if shard_read:
+            shard_dict = _MergingShardBuilder(shard_read, shard_build)
+        else:
+            shard_dict = shard_build
 
         indexer = list(
             get_indexer(

src/zarr/core/buffer/cpu.py

@@ -185,6 +185,48 @@ def __setitem__(self, key: Any, value: Any) -> None:
         self._data.__setitem__(key, value)
 
 
+class DelayedBuffer(Buffer):
+    """
+    A Buffer that is the virtual concatenation of other buffers.
+    """
+
+    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))
+
+
+Buffer.Delayed = DelayedBuffer
+
+
 def as_numpy_array_wrapper(
     func: Callable[[npt.NDArray[Any]], bytes], buf: core.Buffer, prototype: core.BufferPrototype
 ) -> core.Buffer:

src/zarr/core/buffer/gpu.py

@@ -218,6 +218,49 @@ def __setitem__(self, key: Any, value: Any) -> None:
         self._data.__setitem__(key, value)
 
 
+class DelayedBuffer(Buffer):
+    """
+    A Buffer that is the virtual concatenation of other buffers.
+    """
+
+    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")
+        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
+
+
 buffer_prototype = BufferPrototype(buffer=Buffer, nd_buffer=NDBuffer)
 
 register_buffer(Buffer)