Repository: zarr-developers/zarr-python

Issue:
bug in v2 codec pipeline
### Zarr version

main


### Description

There is a bug in these lines:

`cdata` is not of type bytes in the `else` clause:

https://github.com/zarr-developers/zarr-python/blob/eb85ee7ba32df4983b80ae9e60686e7d2f35349e/src/zarr/codecs/_v2.py#L80-L86

cc @normanrz 

### Steps to reproduce


This is triggering failures in xarray's test suite. I haven't succeeded at making a minimal reproducer.

### Additional output

_No response_

Discussion:

Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 6506f734..6068d000 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -7,7 +7,7 @@ default_language_version:
   python: python3
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.8.0
+    rev: v0.8.1
     hooks:
     - id: ruff
       args: ["--fix", "--show-fixes"]
diff --git a/pyproject.toml b/pyproject.toml
index 63ecdd85..c4c3a7a4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -294,11 +294,7 @@ extend-select = [
     "W",    # pycodestyle warnings
 ]
 ignore = [
-    "ANN101", # deprecated
-    "ANN102", # deprecated
     "ANN401",
-    "PT004",  # deprecated
-    "PT005",  # deprecated
     "PT011",  # TODO: apply this rule
     "PT012",  # TODO: apply this rule
     "RET505",
@@ -306,7 +302,6 @@ ignore = [
     "RUF005",
     "SIM108",
     "TRY003",
-    "UP027",  # deprecated
     "UP038",  # https://github.com/astral-sh/ruff/issues/7871
     # https://docs.astral.sh/ruff/formatter/#conflicting-lint-rules
     "W191",
diff --git a/src/zarr/codecs/_v2.py b/src/zarr/codecs/_v2.py
index 30504ad2..df0d8ecb 100644
--- a/src/zarr/codecs/_v2.py
+++ b/src/zarr/codecs/_v2.py
@@ -5,7 +5,7 @@ from dataclasses import dataclass
 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 @@ class V2Codec(ArrayBytesCodec):
     ) -> 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 @@ class V2Codec(ArrayBytesCodec):
         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:
diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py
index 1e815d4d..71a6f9d3 100644
--- a/src/zarr/core/array.py
+++ b/src/zarr/core/array.py
@@ -1346,18 +1346,53 @@ class AsyncArray(Generic[T_ArrayMetadata]):
         AsyncArray.info_complete
             All information about a group, including dynamic information
             like the number of bytes and chunks written.
+
+        Examples
+        --------
+
+        >>> arr = await zarr.api.asynchronous.create(
+        ...     path="array", shape=(3, 4, 5), chunks=(2, 2, 2))
+        ... )
+        >>> arr.info
+        Type               : Array
+        Zarr format        : 3
+        Data type          : DataType.float64
+        Shape              : (3, 4, 5)
+        Chunk shape        : (2, 2, 2)
+        Order              : C
+        Read-only          : False
+        Store type         : MemoryStore
+        Codecs             : [{'endian': <Endian.little: 'little'>}]
+        No. bytes          : 480
         """
         return self._info()
 
     async def info_complete(self) -> Any:
-        # TODO: get the size of the object from the store.
-        extra = {
-            "count_chunks_initialized": await self.nchunks_initialized(),
-            # count_bytes_stored isn't yet implemented.
-        }
-        return self._info(extra=extra)
-
-    def _info(self, extra: dict[str, int] | None = None) -> Any:
+        """
+        Return all the information for an array, including dynamic information like a storage size.
+
+        In addition to the static information, this provides
+
+        - The count of chunks initialized
+        - The sum of the bytes written
+
+        Returns
+        -------
+        ArrayInfo
+
+        See Also
+        --------
+        AsyncArray.info
+            A property giving just the statically known information about an array.
+        """
+        return self._info(
+            await self.nchunks_initialized(),
+            await self.store_path.store.getsize_prefix(self.store_path.path),
+        )
+
+    def _info(
+        self, count_chunks_initialized: int | None = None, count_bytes_stored: int | None = None
+    ) -> Any:
         kwargs: dict[str, Any] = {}
         if self.metadata.zarr_format == 2:
             assert isinstance(self.metadata, ArrayV2Metadata)
@@ -1386,6 +1421,8 @@ class AsyncArray(Generic[T_ArrayMetadata]):
             _read_only=self.read_only,
             _store_type=type(self.store_path.store).__name__,
             _count_bytes=self.dtype.itemsize * self.size,
+            _count_bytes_stored=count_bytes_stored,
+            _count_chunks_initialized=count_chunks_initialized,
             **kwargs,
         )
 
diff --git a/src/zarr/core/buffer/core.py b/src/zarr/core/buffer/core.py
index 9a07583c..7ddedfe0 100644
--- a/src/zarr/core/buffer/core.py
+++ b/src/zarr/core/buffer/core.py
@@ -80,7 +80,13 @@ class NDArrayLike(Protocol):
 
     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: ...
 
diff --git a/src/zarr/core/group.py b/src/zarr/core/group.py
index 13a8c720..2ca6e209 100644
--- a/src/zarr/core/group.py
+++ b/src/zarr/core/group.py
@@ -434,7 +434,7 @@ class AsyncGroup:
     async def open(
         cls,
         store: StoreLike,
-        zarr_format: Literal[2, 3, None] = 3,
+        zarr_format: Literal[2, 3] | None = 3,
         use_consolidated: bool | str | None = None,
     ) -> AsyncGroup:
         """Open a new AsyncGroup
@@ -1691,7 +1691,7 @@ class Group(SyncMixin):
     def open(
         cls,
         store: StoreLike,
-        zarr_format: Literal[2, 3, None] = 3,
+        zarr_format: Literal[2, 3] | None = 3,
     ) -> Group:
         """Open a group from an initialized store.
 
diff --git a/src/zarr/core/indexing.py b/src/zarr/core/indexing.py
index 3d47f5f1..ca227be0 100644
--- a/src/zarr/core/indexing.py
+++ b/src/zarr/core/indexing.py
@@ -1346,8 +1346,15 @@ def decode_morton(z: int, chunk_shape: ChunkCoords) -> ChunkCoords:
 
 
 def morton_order_iter(chunk_shape: ChunkCoords) -> Iterator[ChunkCoords]:
-    for i in range(product(chunk_shape)):
-        yield decode_morton(i, chunk_shape)
+    i = 0
+    order: list[ChunkCoords] = []
+    while len(order) < product(chunk_shape):
+        m = decode_morton(i, chunk_shape)
+        if m not in order and all(x < y for x, y in zip(m, chunk_shape, strict=False)):
+            order.append(m)
+        i += 1
+    for j in range(product(chunk_shape)):
+        yield order[j]
 
 
 def c_order_iter(chunks_per_shard: ChunkCoords) -> Iterator[ChunkCoords]:


Test patch:
diff --git a/tests/test_array.py b/tests/test_array.py
index f0f36cf7..86da801d 100644
--- a/tests/test_array.py
+++ b/tests/test_array.py
@@ -1,3 +1,4 @@
+import dataclasses
 import json
 import math
 import pickle
@@ -474,6 +475,87 @@ class TestInfo:
         )
         assert result == expected
 
+    def test_info_complete(self) -> None:
+        arr = zarr.create(shape=(4, 4), chunks=(2, 2), zarr_format=3)
+        result = arr.info_complete()
+        expected = ArrayInfo(
+            _zarr_format=3,
+            _data_type=DataType.parse("float64"),
+            _shape=(4, 4),
+            _chunk_shape=(2, 2),
+            _order="C",
+            _read_only=False,
+            _store_type="MemoryStore",
+            _codecs=[BytesCodec()],
+            _count_bytes=128,
+            _count_chunks_initialized=0,
+            _count_bytes_stored=373,  # the metadata?
+        )
+        assert result == expected
+
+        arr[:2, :2] = 10
+        result = arr.info_complete()
+        expected = dataclasses.replace(
+            expected, _count_chunks_initialized=1, _count_bytes_stored=405
+        )
+        assert result == expected
+
+    async def test_info_v2_async(self) -> None:
+        arr = await zarr.api.asynchronous.create(shape=(4, 4), chunks=(2, 2), zarr_format=2)
+        result = arr.info
+        expected = ArrayInfo(
+            _zarr_format=2,
+            _data_type=np.dtype("float64"),
+            _shape=(4, 4),
+            _chunk_shape=(2, 2),
+            _order="C",
+            _read_only=False,
+            _store_type="MemoryStore",
+            _count_bytes=128,
+        )
+        assert result == expected
+
+    async def test_info_v3_async(self) -> None:
+        arr = await zarr.api.asynchronous.create(shape=(4, 4), chunks=(2, 2), zarr_format=3)
+        result = arr.info
+        expected = ArrayInfo(
+            _zarr_format=3,
+            _data_type=DataType.parse("float64"),
+            _shape=(4, 4),
+            _chunk_shape=(2, 2),
+            _order="C",
+            _read_only=False,
+            _store_type="MemoryStore",
+            _codecs=[BytesCodec()],
+            _count_bytes=128,
+        )
+        assert result == expected
+
+    async def test_info_complete_async(self) -> None:
+        arr = await zarr.api.asynchronous.create(shape=(4, 4), chunks=(2, 2), zarr_format=3)
+        result = await arr.info_complete()
+        expected = ArrayInfo(
+            _zarr_format=3,
+            _data_type=DataType.parse("float64"),
+            _shape=(4, 4),
+            _chunk_shape=(2, 2),
+            _order="C",
+            _read_only=False,
+            _store_type="MemoryStore",
+            _codecs=[BytesCodec()],
+            _count_bytes=128,
+            _count_chunks_initialized=0,
+            _count_bytes_stored=373,  # the metadata?
+        )
+        assert result == expected
+
+        await arr.setitem((slice(2), slice(2)), 10)
+        result = await arr.info_complete()
+        expected = dataclasses.replace(
+            expected, _count_chunks_initialized=1, _count_bytes_stored=405
+        )
+        assert result == expected
+
 
 @pytest.mark.parametrize("store", ["memory"], indirect=True)
 @pytest.mark.parametrize("zarr_format", [2, 3])
diff --git a/tests/test_codecs/test_codecs.py b/tests/test_codecs/test_codecs.py
index dfb8e1c5..2025e729 100644
--- a/tests/test_codecs/test_codecs.py
+++ b/tests/test_codecs/test_codecs.py
@@ -204,6 +204,26 @@ def test_morton() -> None:
     ]
 
 
+@pytest.mark.parametrize(
+    "shape",
+    [
+        [2, 2, 2],
+        [5, 2],
+        [2, 5],
+        [2, 9, 2],
+        [3, 2, 12],
+        [2, 5, 1],
+        [4, 3, 6, 2, 7],
+        [3, 2, 1, 6, 4, 5, 2],
+    ],
+)
+def test_morton2(shape) -> None:
+    order = list(morton_order_iter(shape))
+    for i, x in enumerate(order):
+        assert x not in order[:i]  # no duplicates
+        assert all(x[j] < shape[j] for j in range(len(shape)))  # all indices are within bounds
+
+
 @pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"])
 def test_write_partial_chunks(store: Store) -> None:
     data = np.arange(0, 256, dtype="uint16").reshape((16, 16))
diff --git a/tests/test_codecs/test_sharding.py b/tests/test_codecs/test_sharding.py
index 78f32fef..51c82067 100644
--- a/tests/test_codecs/test_sharding.py
+++ b/tests/test_codecs/test_sharding.py
@@ -393,3 +393,32 @@ async def test_sharding_with_empty_inner_chunk(
     print("read data")
     data_read = await a.getitem(...)
     assert np.array_equal(data_read, data)
+
+
+@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"])
+@pytest.mark.parametrize(
+    "index_location",
+    [ShardingCodecIndexLocation.start, ShardingCodecIndexLocation.end],
+)
+@pytest.mark.parametrize("chunks_per_shard", [(5, 2), (2, 5), (5, 5)])
+async def test_sharding_with_chunks_per_shard(
+    store: Store, index_location: ShardingCodecIndexLocation, chunks_per_shard: tuple[int]
+) -> None:
+    chunk_shape = (2, 1)
+    shape = [x * y for x, y in zip(chunks_per_shard, chunk_shape, strict=False)]
+    data = np.ones(np.prod(shape), dtype="int32").reshape(shape)
+    fill_value = 42
+
+    path = f"test_sharding_with_chunks_per_shard_{index_location}"
+    spath = StorePath(store, path)
+    a = Array.create(
+        spath,
+        shape=shape,
+        chunk_shape=shape,
+        dtype="int32",
+        fill_value=fill_value,
+        codecs=[ShardingCodec(chunk_shape=chunk_shape, index_location=index_location)],
+    )
+    a[...] = data
+    data_read = a[...]
+    assert np.array_equal(data_read, data)
diff --git a/tests/test_indexing.py b/tests/test_indexing.py
index 479ce201..3dc93ba4 100644
--- a/tests/test_indexing.py
+++ b/tests/test_indexing.py
@@ -12,7 +12,7 @@ from numpy.testing import assert_array_equal
 
 import zarr
 from zarr import Array
-from zarr.core.buffer import BufferPrototype, default_buffer_prototype
+from zarr.core.buffer import default_buffer_prototype
 from zarr.core.indexing import (
     BasicSelection,
     CoordinateSelection,
diff --git a/tests/test_store/test_logging.py b/tests/test_store/test_logging.py
index aed1f06f..c0630dff 100644
--- a/tests/test_store/test_logging.py
+++ b/tests/test_store/test_logging.py
@@ -5,7 +5,6 @@ from typing import TYPE_CHECKING
 import pytest
 
 import zarr
-import zarr.storage
 from zarr.core.buffer import default_buffer_prototype
 from zarr.storage.logging import LoggingStore
 
diff --git a/tests/test_v2.py b/tests/test_v2.py
index e6b50ab2..68c07e20 100644
--- a/tests/test_v2.py
+++ b/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)