Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion src/zarr/codecs/_v2.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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:
Expand Down
8 changes: 7 additions & 1 deletion src/zarr/core/buffer/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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: ...

Expand Down
51 changes: 48 additions & 3 deletions tests/test_v2.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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)