Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix: order for v2 arrays #2679

Merged
merged 2 commits into from
Jan 10, 2025
Merged
Show file tree
Hide file tree
Changes from 1 commit
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
18 changes: 13 additions & 5 deletions src/zarr/core/array.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
import warnings
from asyncio import gather
from collections.abc import Iterable
from dataclasses import dataclass, field
from dataclasses import dataclass, field, replace
from itertools import starmap
from logging import getLogger
from typing import (
Expand Down Expand Up @@ -1226,14 +1226,17 @@ async def _get_selection(
fill_value=self.metadata.fill_value,
)
if product(indexer.shape) > 0:
# need to use the order from the metadata for v2
_config = self._config
if self.metadata.zarr_format == 2:
_config = replace(_config, order=self.metadata.order)

# reading chunks and decoding them
await self.codec_pipeline.read(
[
(
self.store_path / self.metadata.encode_chunk_key(chunk_coords),
self.metadata.get_chunk_spec(
chunk_coords, self._config, prototype=prototype
),
self.metadata.get_chunk_spec(chunk_coords, _config, prototype=prototype),
chunk_selection,
out_selection,
)
Expand Down Expand Up @@ -1350,12 +1353,17 @@ async def _set_selection(
# Buffer and NDBuffer between components.
value_buffer = prototype.nd_buffer.from_ndarray_like(value)

# need to use the order from the metadata for v2
_config = self._config
if self.metadata.zarr_format == 2:
_config = replace(_config, order=self.metadata.order)

# merging with existing data and encoding chunks
await self.codec_pipeline.write(
[
(
self.store_path / self.metadata.encode_chunk_key(chunk_coords),
self.metadata.get_chunk_spec(chunk_coords, self._config, prototype),
self.metadata.get_chunk_spec(chunk_coords, _config, prototype),
chunk_selection,
out_selection,
)
Expand Down
38 changes: 29 additions & 9 deletions tests/test_v2.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@
import zarr.core.buffer
import zarr.storage
from zarr import config
from zarr.core.buffer.core import default_buffer_prototype
from zarr.core.sync import sync
from zarr.storage import MemoryStore, StorePath


Expand Down Expand Up @@ -166,36 +168,54 @@ def test_v2_filters_codecs(filters: Any, order: Literal["C", "F"]) -> None:

@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:
@pytest.mark.parametrize("memory_order", ["C", "F"])
def test_v2_non_contiguous(
array_order: Literal["C", "F"], data_order: Literal["C", "F"], memory_order: Literal["C", "F"]
) -> None:
store = MemoryStore()
arr = zarr.create_array(
MemoryStore({}),
store,
shape=(10, 8),
chunks=(3, 3),
fill_value=np.nan,
dtype="float64",
zarr_format=2,
filters=None,
compressors=None,
overwrite=True,
order=array_order,
config={"order": memory_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
arr[6:9, 3:6] = a[6:9, 3:6] # The slice on the RHS is important
np.testing.assert_array_equal(arr[6:9, 3:6], a[6:9, 3:6])

np.testing.assert_array_equal(
arr[slice(6, 9, None), slice(3, 6, None)], a[slice(6, 9, None), slice(3, 6, None)]
a[6:9, 3:6],
np.frombuffer(
sync(store.get("2.1", default_buffer_prototype())).to_bytes(), dtype="float64"
).reshape((3, 3), order=array_order),
)
if memory_order == "F":
assert (arr[6:9, 3:6]).flags.f_contiguous
else:
assert (arr[6:9, 3:6]).flags.c_contiguous

store = MemoryStore()
arr = zarr.create_array(
MemoryStore({}),
store,
shape=(10, 8),
chunks=(3, 3),
fill_value=np.nan,
dtype="float64",
zarr_format=2,
compressors=None,
filters=None,
overwrite=True,
order=array_order,
config={"order": memory_order},
)

# Contiguous write
Expand All @@ -204,8 +224,8 @@ def test_v2_non_contiguous(array_order: Literal["C", "F"], data_order: Literal["
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)
arr[6:9, 3:6] = a
np.testing.assert_array_equal(arr[6:9, 3:6], a)


def test_default_compressor_deprecation_warning():
Expand Down
Loading