[v3] Makes data contiguous in v2 codec (zarr-developers#2515)

* fixes zarr-developers#2501

* typing

* only use c-contiguous

* more tests

* typing

* tests

* astype with copy=False

---------

Co-authored-by: Deepak Cherian <[email protected]>

Author: normanrz

src/zarr/codecs/_v2.py

@@ -5,7 +5,7 @@
 from typing import TYPE_CHECKING
 
 import numcodecs
-from numcodecs.compat import ensure_ndarray_like
+from numcodecs.compat import ensure_bytes, ensure_ndarray_like
 
 from zarr.abc.codec import ArrayBytesCodec
 from zarr.registry import get_ndbuffer_class
@@ -68,6 +68,9 @@ async def _encode_single(
     ) -> Buffer | None:
         chunk = chunk_array.as_ndarray_like()
 
+        # ensure contiguous and correct order
+        chunk = chunk.astype(chunk_spec.dtype, order=chunk_spec.order, copy=False)
+
         # apply filters
         if self.filters:
             for f in self.filters:
@@ -83,6 +86,7 @@ async def _encode_single(
         else:
             cdata = chunk
 
+        cdata = ensure_bytes(cdata)
         return chunk_spec.prototype.buffer.from_bytes(cdata)
 
     def compute_encoded_size(self, _input_byte_length: int, _chunk_spec: ArraySpec) -> int:

src/zarr/core/buffer/core.py

@@ -80,7 +80,13 @@ def reshape(
 
     def view(self, dtype: npt.DTypeLike) -> Self: ...
 
-    def astype(self, dtype: npt.DTypeLike, order: Literal["K", "A", "C", "F"] = ...) -> Self: ...
+    def astype(
+        self,
+        dtype: npt.DTypeLike,
+        order: Literal["K", "A", "C", "F"] = ...,
+        *,
+        copy: bool = ...,
+    ) -> Self: ...
 
     def fill(self, value: Any) -> None: ...
 

tests/test_v2.py

@@ -1,6 +1,6 @@
 import json
 from collections.abc import Iterator
-from typing import Any
+from typing import Any, Literal
 
 import numcodecs.vlen
 import numpy as np
@@ -126,9 +126,54 @@ async def test_create_dtype_str(dtype: Any) -> None:
 
 
 @pytest.mark.parametrize("filters", [[], [numcodecs.Delta(dtype="<i4")], [numcodecs.Zlib(level=2)]])
-def test_v2_filters_codecs(filters: Any) -> None:
+@pytest.mark.parametrize("order", ["C", "F"])
+def test_v2_filters_codecs(filters: Any, order: Literal["C", "F"]) -> None:
     array_fixture = [42]
-    arr = zarr.create(shape=1, dtype="<i4", zarr_format=2, filters=filters)
+    arr = zarr.create(shape=1, dtype="<i4", zarr_format=2, filters=filters, order=order)
     arr[:] = array_fixture
     result = arr[:]
     np.testing.assert_array_equal(result, array_fixture)
+
+
+@pytest.mark.parametrize("array_order", ["C", "F"])
+@pytest.mark.parametrize("data_order", ["C", "F"])
+def test_v2_non_contiguous(array_order: Literal["C", "F"], data_order: Literal["C", "F"]) -> None:
+    arr = zarr.Array.create(
+        MemoryStore({}),
+        shape=(10, 8),
+        chunks=(3, 3),
+        fill_value=np.nan,
+        dtype="float64",
+        zarr_format=2,
+        exists_ok=True,
+        order=array_order,
+    )
+
+    # Non-contiguous write
+    a = np.arange(arr.shape[0] * arr.shape[1]).reshape(arr.shape, order=data_order)
+    arr[slice(6, 9, None), slice(3, 6, None)] = a[
+        slice(6, 9, None), slice(3, 6, None)
+    ]  # The slice on the RHS is important
+    np.testing.assert_array_equal(
+        arr[slice(6, 9, None), slice(3, 6, None)], a[slice(6, 9, None), slice(3, 6, None)]
+    )
+
+    arr = zarr.Array.create(
+        MemoryStore({}),
+        shape=(10, 8),
+        chunks=(3, 3),
+        fill_value=np.nan,
+        dtype="float64",
+        zarr_format=2,
+        exists_ok=True,
+        order=array_order,
+    )
+
+    # Contiguous write
+    a = np.arange(9).reshape((3, 3), order=data_order)
+    if data_order == "F":
+        assert a.flags.f_contiguous
+    else:
+        assert a.flags.c_contiguous
+    arr[slice(6, 9, None), slice(3, 6, None)] = a
+    np.testing.assert_array_equal(arr[slice(6, 9, None), slice(3, 6, None)], a)