refactor: write paths use generic merge_and_encode with ShardIndex.leaf_transform

Author: d-v-b

src/zarr/core/codec_pipeline.py

@@ -877,6 +877,152 @@ def decode_chunks_from_index(
     return out
 
 
+def merge_and_encode_from_index(
+    existing_raw: dict[tuple[int, ...], Buffer | None],
+    index: ShardIndex,
+    value: NDBuffer,
+    chunk_spec: ArraySpec,
+    chunk_selection: SelectorTuple,
+    out_selection: SelectorTuple,
+    drop_axes: tuple[int, ...],
+) -> dict[tuple[int, ...], Buffer | None]:
+    """Merge new data into existing chunk(s) and encode, using index.leaf_transform.
+
+    For non-sharded layouts (``index.is_sharded`` is False): decode the single
+    existing chunk (or create from fill value), merge *value* at the given
+    selection, and encode. Returns ``{(0,...): encoded}``.
+
+    For sharded layouts (``index.is_sharded`` is True): start with existing raw
+    chunks, fill missing coords with None, then iterate over affected inner
+    chunks using ``get_indexer``. Decode/merge/encode each. Returns the full
+    chunk dict for subsequent packing into a shard blob.
+    """
+    from zarr.core.indexing import get_indexer
+
+    assert index.leaf_transform is not None
+    transform = index.leaf_transform
+
+    if not index.is_sharded:
+        # --- Simple (non-sharded) path ---
+        coord = next(iter(existing_raw)) if existing_raw else (0,) * len(chunk_spec.shape)
+
+        existing_bytes = existing_raw.get(coord)
+        if existing_bytes is not None:
+            chunk_array = transform.decode_chunk(existing_bytes, chunk_shape=chunk_spec.shape)
+            if not chunk_array.as_ndarray_like().flags.writeable:  # type: ignore[attr-defined]
+                chunk_array = chunk_spec.prototype.nd_buffer.from_ndarray_like(
+                    chunk_array.as_ndarray_like().copy()
+                )
+        else:
+            chunk_array = chunk_spec.prototype.nd_buffer.create(
+                shape=chunk_spec.shape,
+                dtype=chunk_spec.dtype.to_native_dtype(),
+                fill_value=fill_value_or_default(chunk_spec),
+            )
+
+        # Merge value
+        if chunk_selection == () or is_scalar(
+            value.as_ndarray_like(), chunk_spec.dtype.to_native_dtype()
+        ):
+            chunk_value = value
+        else:
+            chunk_value = value[out_selection]
+            if drop_axes:
+                item = tuple(
+                    None if idx in drop_axes else slice(None) for idx in range(chunk_spec.ndim)
+                )
+                chunk_value = chunk_value[item]
+        chunk_array[chunk_selection] = chunk_value
+
+        # Check write_empty_chunks
+        if not chunk_spec.config.write_empty_chunks and chunk_array.all_equal(
+            chunk_spec.fill_value
+        ):
+            return {coord: None}
+
+        chunk_shape = chunk_spec.shape if chunk_spec.shape != transform.array_spec.shape else None
+        encoded = transform.encode_chunk(chunk_array, chunk_shape=chunk_shape)
+        return {coord: encoded}
+
+    # --- Sharded path ---
+    inner_shape = transform.array_spec.shape
+    chunks_per_shard = tuple(
+        s // cs for s, cs in zip(chunk_spec.shape, inner_shape, strict=True)
+    )
+
+    chunk_dict: dict[tuple[int, ...], Buffer | None] = dict(existing_raw)
+
+    # Fill missing coords with None
+    for coord in np.ndindex(chunks_per_shard):
+        if coord not in chunk_dict:
+            chunk_dict[coord] = None
+
+    inner_spec = ArraySpec(
+        shape=inner_shape,
+        dtype=chunk_spec.dtype,
+        fill_value=chunk_spec.fill_value,
+        config=chunk_spec.config,
+        prototype=chunk_spec.prototype,
+    )
+
+    # Extract the shard's portion of the write value
+    if is_scalar(value.as_ndarray_like(), chunk_spec.dtype.to_native_dtype()):
+        shard_value = value
+    else:
+        shard_value = value[out_selection]
+        if drop_axes:
+            item = tuple(
+                None if idx in drop_axes else slice(None)
+                for idx in range(len(chunk_spec.shape))
+            )
+            shard_value = shard_value[item]
+
+    # Determine which inner chunks are affected
+    from zarr.core.chunk_grids import ChunkGrid as _ChunkGrid
+
+    indexer = get_indexer(
+        chunk_selection,
+        shape=chunk_spec.shape,
+        chunk_grid=_ChunkGrid.from_sizes(chunk_spec.shape, inner_shape),
+    )
+
+    for inner_coords, inner_sel, value_sel, _ in indexer:
+        existing_bytes = chunk_dict.get(inner_coords)
+
+        # Decode just this inner chunk
+        if existing_bytes is not None:
+            inner_array = transform.decode_chunk(existing_bytes)
+            if not inner_array.as_ndarray_like().flags.writeable:  # type: ignore[attr-defined]
+                inner_array = inner_spec.prototype.nd_buffer.from_ndarray_like(
+                    inner_array.as_ndarray_like().copy()
+                )
+        else:
+            inner_array = inner_spec.prototype.nd_buffer.create(
+                shape=inner_spec.shape,
+                dtype=inner_spec.dtype.to_native_dtype(),
+                fill_value=fill_value_or_default(inner_spec),
+            )
+
+        # Merge new data
+        if inner_sel == () or is_scalar(
+            shard_value.as_ndarray_like(), inner_spec.dtype.to_native_dtype()
+        ):
+            inner_value = shard_value
+        else:
+            inner_value = shard_value[value_sel]
+        inner_array[inner_sel] = inner_value
+
+        # Re-encode
+        if not chunk_spec.config.write_empty_chunks and inner_array.all_equal(
+            chunk_spec.fill_value
+        ):
+            chunk_dict[inner_coords] = None
+        else:
+            chunk_dict[inner_coords] = transform.encode_chunk(inner_array)
+
+    return chunk_dict
+
+
 class ChunkLayout:
     """Describes how a stored blob maps to one or more inner chunks.
 
@@ -934,6 +1080,12 @@ def store_chunks_sync(self, byte_setter: Any, encoded_chunks: dict[tuple[int, ..
     async def store_chunks_async(self, byte_setter: Any, encoded_chunks: dict[tuple[int, ...], Buffer | None], chunk_spec: ArraySpec) -> None:
         raise NotImplementedError
 
+    def pack_and_store_sync(self, byte_setter: Any, encoded_chunks: dict[tuple[int, ...], Buffer | None]) -> None:
+        raise NotImplementedError
+
+    async def pack_and_store_async(self, byte_setter: Any, encoded_chunks: dict[tuple[int, ...], Buffer | None]) -> None:
+        raise NotImplementedError
+
     # -- Low-level helpers --
 
     def unpack_blob(self, blob: Buffer) -> dict[tuple[int, ...], Buffer | None]:
@@ -1054,6 +1206,22 @@ async def store_chunks_async(self, byte_setter: Any, encoded_chunks: dict[tuple[
         else:
             await byte_setter.set(blob)
 
+    def pack_and_store_sync(self, byte_setter: Any, encoded_chunks: dict[tuple[int, ...], Buffer | None]) -> None:
+        coord = (0,) * len(self.chunks_per_shard)
+        blob = encoded_chunks.get(coord)
+        if blob is None:
+            byte_setter.delete_sync()  # type: ignore[attr-defined]
+        else:
+            byte_setter.set_sync(blob)  # type: ignore[attr-defined]
+
+    async def pack_and_store_async(self, byte_setter: Any, encoded_chunks: dict[tuple[int, ...], Buffer | None]) -> None:
+        coord = (0,) * len(self.chunks_per_shard)
+        blob = encoded_chunks.get(coord)
+        if blob is None:
+            await byte_setter.delete()
+        else:
+            await byte_setter.set(blob)
+
     # -- Low-level --
 
     def unpack_blob(self, blob: Buffer) -> dict[tuple[int, ...], Buffer | None]:
@@ -1347,6 +1515,30 @@ async def store_chunks_async(self, byte_setter: Any, encoded_chunks: dict[tuple[
             else:
                 await byte_setter.set(blob)
 
+    def pack_and_store_sync(self, byte_setter: Any, encoded_chunks: dict[tuple[int, ...], Buffer | None]) -> None:
+        from zarr.core.buffer import default_buffer_prototype
+
+        if all(v is None for v in encoded_chunks.values()):
+            byte_setter.delete_sync()  # type: ignore[attr-defined]
+            return
+        blob = self.pack_blob(encoded_chunks, default_buffer_prototype())
+        if blob is None:
+            byte_setter.delete_sync()  # type: ignore[attr-defined]
+        else:
+            byte_setter.set_sync(blob)  # type: ignore[attr-defined]
+
+    async def pack_and_store_async(self, byte_setter: Any, encoded_chunks: dict[tuple[int, ...], Buffer | None]) -> None:
+        from zarr.core.buffer import default_buffer_prototype
+
+        if all(v is None for v in encoded_chunks.values()):
+            await byte_setter.delete()
+            return
+        blob = self.pack_blob(encoded_chunks, default_buffer_prototype())
+        if blob is None:
+            await byte_setter.delete()
+        else:
+            await byte_setter.set(blob)
+
     def _decode_per_chunk(
         self,
         chunk_dict: dict[tuple[int, ...], Buffer | None],
@@ -1760,7 +1952,7 @@ async def _process_chunk(
 
             # Phase 1: resolve index (IO)
             if is_complete:
-                index = ShardIndex(key=key)
+                index = ShardIndex(key=key, leaf_transform=layout.inner_transform, is_sharded=layout.is_sharded)
             elif layout.is_sharded:
                 async with sem:
                     index = await layout.resolve_index_async(byte_setter, key, chunk_selection=None)  # ALL coords
@@ -1771,25 +1963,26 @@ async def _process_chunk(
             # Phase 2: fetch existing chunks (IO)
             if index.chunks:
                 async with sem:
-                    existing_chunks = await layout.fetch_chunks_async(byte_setter, index, prototype=chunk_spec.prototype)
+                    existing = await fetch_chunks_async(byte_setter, index, prototype=chunk_spec.prototype)
             else:
-                existing_chunks = {}
+                existing = {}
 
             # Phase 3: merge and encode (compute)
-            encoded_chunks = await loop.run_in_executor(
+            encoded = await loop.run_in_executor(
                 pool,
-                layout.merge_and_encode,
-                existing_chunks,
+                merge_and_encode_from_index,
+                existing,
+                index,
                 value,
                 chunk_spec,
                 chunk_selection,
                 out_selection,
                 drop_axes,
             )
 
-            # Phase 4: store (IO)
+            # Phase 4: pack + store (IO)
             async with sem:
-                await layout.store_chunks_async(byte_setter, encoded_chunks, chunk_spec)
+                await layout.pack_and_store_async(byte_setter, encoded)
 
         await asyncio.gather(
             *[
@@ -1868,24 +2061,33 @@ def write_sync(
         if not batch:
             return
 
+        assert self.layout is not None
+        default_layout = self.layout
+
         for bs, chunk_spec, chunk_selection, out_selection, is_complete in batch:
             layout = (
-                self.layout
-                if self.layout is not None and chunk_spec.shape == self.layout.chunk_shape
+                default_layout
+                if chunk_spec.shape == default_layout.chunk_shape
                 else self._get_layout(chunk_spec)
             )
             key = bs.path if hasattr(bs, "path") else ""
 
+            # Phase 1: resolve index
             if is_complete:
-                index = ShardIndex(key=key)
+                index = ShardIndex(key=key, leaf_transform=layout.inner_transform, is_sharded=layout.is_sharded)
             elif layout.is_sharded:
                 index = layout.resolve_index(bs, key, chunk_selection=None)  # ALL coords
             else:
                 index = layout.resolve_index(bs, key, chunk_selection=chunk_selection)
 
-            existing_chunks = layout.fetch_chunks(bs, index, prototype=chunk_spec.prototype) if index.chunks else {}
-            encoded_chunks = layout.merge_and_encode(existing_chunks, value, chunk_spec, chunk_selection, out_selection, drop_axes)
-            layout.store_chunks_sync(bs, encoded_chunks, chunk_spec)
+            # Phase 2: fetch existing
+            existing = fetch_chunks_sync(bs, index, prototype=chunk_spec.prototype) if index.chunks else {}
+
+            # Phase 3: merge + encode (compute)
+            encoded = merge_and_encode_from_index(existing, index, value, chunk_spec, chunk_selection, out_selection, drop_axes)
+
+            # Phase 4: pack + store
+            layout.pack_and_store_sync(bs, encoded)
 
 
 register_pipeline(PhasedCodecPipeline)