Mark chunked arrays in encode_decode_data_array

Author: juntyr

src/numcodecs_combinators/abc.py

@@ -15,6 +15,8 @@ class CodecCombinatorMixin(ABC):
     Mixin class for combinators over [`Codec`][numcodecs.abc.Codec]s.
     """
 
+    __slots__ = ()
+
     @abstractmethod
     def map(self, mapper: Callable[[Codec], Codec]) -> Codec:
         """

src/numcodecs_combinators/stack.py

@@ -140,6 +140,8 @@ def encode_decode(self, buf: Buffer) -> Buffer:
             buffer protocol.
         """
 
+        chunked = getattr(buf, "chunked", False)
+
         encoded = np.asarray(
             numcodecs.compat.ensure_contiguous_ndarray_like(buf, flatten=False)
         )
@@ -149,16 +151,23 @@ def encode_decode(self, buf: Buffer) -> Buffer:
             silhouettes.append((encoded.shape, encoded.dtype))
             encoded = np.asarray(
                 numcodecs.compat.ensure_contiguous_ndarray_like(
-                    codec.encode((encoded)), flatten=False
+                    codec.encode(_MaybeChunkedNdArray(encoded) if chunked else encoded),
+                    flatten=False,
                 )
             )
 
-        decoded = encoded
+        decoded = encoded.view(np.ndarray)
 
         for codec in reversed(self):
             shape, dtype = silhouettes.pop()
             out = np.empty(shape=shape, dtype=dtype)
-            decoded = codec.decode(decoded, out).view(dtype).reshape(shape)
+            decoded = (
+                codec.decode(decoded, _MaybeChunkedNdArray(out) if chunked else out)
+                .view(dtype)
+                .reshape(shape)
+            )
+
+        decoded = decoded.view(np.ndarray)
 
         if isinstance(decoded, type(buf)):
             return decoded
@@ -167,7 +176,8 @@ def encode_decode(self, buf: Buffer) -> Buffer:
 
     def encode_decode_data_array(self, da: "xr.DataArray") -> "xr.DataArray":
         """
-        Encode, then decode each chunk (independently) in the data array `da`.
+        Encode, then decode the data array `da`. If `da` is chunked, each chunk
+        is encoded and decoded *independently*.
 
         Since each chunk is encoded *independently*, this method may cause
         chunk boundary artifacts. Do *not* use this method if the codec
@@ -195,6 +205,8 @@ def encode_decode_data_array(self, da: "xr.DataArray") -> "xr.DataArray":
 
         import xarray as xr
 
+        chunked = da.chunks is not None
+
         def encode_decode_data_array_single_chunk(
             da: xr.DataArray,
         ) -> xr.DataArray:
@@ -205,9 +217,11 @@ def encode_decode_data_array_single_chunk(
                 return da.copy(deep=False).chunk(single_chunk)
 
             # eagerly compute the input chunk and encode and decode it
-            decoded = self.encode_decode(da.values)  # type: ignore
+            decoded = self.encode_decode(_MaybeChunkedNdArray(da.values, chunked))  # type: ignore
 
-            return da.copy(deep=False, data=decoded).chunk(single_chunk)
+            return da.copy(deep=False, data=np.array(decoded).view(np.ndarray)).chunk(
+                single_chunk
+            )
 
         return xr.map_blocks(encode_decode_data_array_single_chunk, da)
 
@@ -293,3 +307,22 @@ def __rmul__(self, other) -> "CodecStack":
 
 
 numcodecs.registry.register_codec(CodecStack)
+
+
+class _MaybeChunkedNdArray(np.ndarray):
+    __slots__ = ("_chunked",)
+    _chunked: bool
+
+    def __new__(cls, array, chunked: bool = True):
+        obj = np.asarray(array).view(cls)
+        obj._chunked = chunked
+        return obj
+
+    def __array_finalize__(self, obj):
+        if obj is None:
+            return
+        self._chunked = getattr(obj, "chunked", True)
+
+    @property
+    def chunked(self) -> bool:
+        return self._chunked

tests/test_stack.py

@@ -1,6 +1,7 @@
 import numcodecs
 import numpy as np
 import xarray as xr
+from numcodecs.abc import Codec
 
 import numcodecs_combinators
 from numcodecs_combinators.stack import CodecStack
@@ -51,6 +52,39 @@ def test_encode_decode():
     assert encoded_decoded.equals(xr.DataArray([1.0, 2.0, 3.0]))
 
 
+def test_chunked_encode_decode():
+    class CheckChunkedCodec(Codec):
+        __slots__ = ("is_chunked",)
+        is_chunked: bool
+
+        def __init__(self, is_chunked: bool):
+            self.is_chunked = is_chunked
+
+        def encode(self, buf):
+            assert getattr(buf, "chunked", False) == self.is_chunked
+            return buf
+
+        def decode(self, buf, out=None):
+            assert getattr(buf, "chunked", False) is False
+            assert getattr(out, "chunked", False) == self.is_chunked
+            return numcodecs.compat.ndarray_copy(buf, out)
+
+    stack = CodecStack(CheckChunkedCodec(False))
+
+    encoded_decoded = stack.encode_decode(np.array([1.0, 2.0, 3.0]))
+    assert np.all(encoded_decoded == np.array([1.0, 2.0, 3.0]))
+
+    encoded_decoded = stack.encode_decode_data_array(xr.DataArray([1.0, 2.0, 3.0]))
+    assert encoded_decoded.equals(xr.DataArray([1.0, 2.0, 3.0]))
+
+    stack = CodecStack(CheckChunkedCodec(True))
+
+    encoded_decoded = stack.encode_decode_data_array(
+        xr.DataArray([1.0, 2.0, 3.0]).chunk(1)
+    )
+    assert encoded_decoded.equals(xr.DataArray([1.0, 2.0, 3.0]))
+
+
 def test_map():
     stack = CodecStack(numcodecs.Zlib(level=9), numcodecs.CRC32())