Ensure chunk markings are preserved by other combinators

Author: juntyr

mkdocs.yml

@@ -24,6 +24,7 @@ plugins:
       source_dirs:
         - nav_heading: [Documentation]
           base: src
+          ignore: ["_chunked.py"]
   - exclude:
       glob:
         - requirements.txt

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "numcodecs-combinators"
-version = "0.2.9"
+version = "0.2.10"
 description = "Combinator codecs for the `numcodecs` buffer compression API"
 readme = "README.md"
 license = { file = "LICENSE" }

src/numcodecs_combinators/_chunked.py

@@ -0,0 +1,12 @@
+import numpy as np
+
+
+class ChunkedNdArray(np.ndarray):
+    __slots__ = ()
+
+    def __new__(cls, array):
+        return np.asarray(array).view(cls)
+
+    @property
+    def chunked(self) -> bool:
+        return True

src/numcodecs_combinators/best.py

@@ -70,7 +70,9 @@ def encode(self, buf: Buffer) -> bytes:
         if len(self) == 0:
             return buf
 
-        data = numcodecs.compat.ensure_ndarray(buf)
+        data = (
+            buf if isinstance(buf, np.ndarray) else numcodecs.compat.ensure_ndarray(buf)
+        )
 
         best_size = np.inf
         best_index = None

src/numcodecs_combinators/framed.py

@@ -16,6 +16,7 @@
 from numcodecs.abc import Codec
 from typing_extensions import Buffer, Self  # MSPV 3.12
 
+from ._chunked import ChunkedNdArray
 from .abc import CodecCombinatorMixin
 
 
@@ -74,6 +75,8 @@ def encode(self, buf: Buffer) -> bytes:
             Encoded and framed data as a bytestring.
         """
 
+        chunked = getattr(buf, "chunked", False)
+
         encoded = buf
         encoded_ndarray = np.asarray(
             numcodecs.compat.ensure_contiguous_ndarray_like(encoded, flatten=False)
@@ -82,7 +85,9 @@ def encode(self, buf: Buffer) -> bytes:
         frames = [(encoded_ndarray.dtype, encoded_ndarray.shape)]
 
         for codec in self:
-            encoded = codec.encode(encoded_ndarray)
+            encoded = codec.encode(
+                ChunkedNdArray(encoded_ndarray) if chunked else encoded_ndarray
+            )
             encoded_ndarray = np.asarray(
                 numcodecs.compat.ensure_contiguous_ndarray_like(encoded, flatten=False)
             )
@@ -132,6 +137,8 @@ def decode(self, buf: Buffer, out: Optional[Buffer] = None) -> Buffer:
             buffer protocol.
         """
 
+        chunked = getattr(out, "chunked", False)
+
         b = numcodecs.compat.ensure_bytes(buf)
 
         b_io = BytesIO(b)
@@ -165,10 +172,11 @@ def decode(self, buf: Buffer, out: Optional[Buffer] = None) -> Buffer:
             decoded = decoded.byteswap()
 
         for codec, (dtype, shape) in zip(reversed(self), frames[:-1][::-1]):
+            empty = np.empty(shape, dtype)
             decoded = (
                 codec.decode(
                     decoded,
-                    out=np.empty(shape, dtype),
+                    out=ChunkedNdArray(empty) if chunked else empty,
                 )
                 .view(dtype)
                 .reshape(shape)

src/numcodecs_combinators/stack.py

@@ -19,6 +19,7 @@
 
 from numcodecs.abc import Codec
 
+from ._chunked import ChunkedNdArray
 from .abc import CodecCombinatorMixin
 
 
@@ -89,11 +90,22 @@ def encode(self, buf: Buffer) -> Buffer:
             protocol.
         """
 
+        if len(self) == 0:
+            return buf
+
+        chunked = getattr(buf, "chunked", False)
+
         encoded = buf
         for codec in self:
-            encoded = codec.encode(
+            encoded_ndarray = np.asarray(
                 numcodecs.compat.ensure_contiguous_ndarray_like(encoded, flatten=False)
             )
+            encoded = codec.encode(
+                ChunkedNdArray(encoded_ndarray) if chunked else encoded_ndarray
+            )
+
+        if getattr(encoded, "chunked", False):
+            return np.array(encoded).view(np.ndarray)  # type: ignore
         return encoded
 
     def decode(self, buf: Buffer, out: Optional[Buffer] = None) -> Buffer:
@@ -151,7 +163,7 @@ def encode_decode(self, buf: Buffer) -> Buffer:
             silhouettes.append((encoded.shape, encoded.dtype))
             encoded = np.asarray(
                 numcodecs.compat.ensure_contiguous_ndarray_like(
-                    codec.encode(_MaybeChunkedNdArray(encoded) if chunked else encoded),
+                    codec.encode(ChunkedNdArray(encoded) if chunked else encoded),
                     flatten=False,
                 )
             )
@@ -162,12 +174,13 @@ def encode_decode(self, buf: Buffer) -> Buffer:
             shape, dtype = silhouettes.pop()
             out = np.empty(shape=shape, dtype=dtype)
             decoded = (
-                codec.decode(decoded, _MaybeChunkedNdArray(out) if chunked else out)
+                codec.decode(decoded, ChunkedNdArray(out) if chunked else out)
                 .view(dtype)
                 .reshape(shape)
             )
 
-        decoded = decoded.view(np.ndarray)
+        if getattr(decoded, "chunked", False):
+            decoded = decoded.view(np.ndarray)
 
         if isinstance(decoded, type(buf)):
             return decoded
@@ -205,7 +218,8 @@ def encode_decode_data_array(self, da: "xr.DataArray") -> "xr.DataArray":
 
         import xarray as xr
 
-        chunked = da.chunks is not None
+        if da.chunks is None:
+            return da.copy(data=self.encode_decode(da.values))  # type: ignore
 
         def encode_decode_data_array_single_chunk(
             da: xr.DataArray,
@@ -217,7 +231,7 @@ 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(_MaybeChunkedNdArray(da.values, chunked))  # type: ignore
+            decoded = self.encode_decode(ChunkedNdArray(da.values))  # type: ignore
 
             return da.copy(deep=False, data=np.array(decoded).view(np.ndarray)).chunk(
                 single_chunk
@@ -307,22 +321,3 @@ 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