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(array): thread order parameter through to Array.__init__ #2417

Merged
merged 9 commits into from
Oct 21, 2024
10 changes: 3 additions & 7 deletions src/zarr/api/asynchronous.py
Original file line number Diff line number Diff line change
Expand Up @@ -712,7 +712,7 @@ async def create(
dtype: npt.DTypeLike | None = None,
compressor: dict[str, JSON] | None = None, # TODO: default and type change
fill_value: Any | None = 0, # TODO: need type
order: MemoryOrder | None = None, # TODO: default change
order: MemoryOrder | None = None,
store: str | StoreLike | None = None,
synchronizer: Any | None = None,
overwrite: bool = False,
Expand Down Expand Up @@ -761,6 +761,7 @@ async def create(
Default value to use for uninitialized portions of the array.
order : {'C', 'F'}, optional
Memory layout to be used within each chunk.
Default is set in Zarr's config (`array.order`).
store : Store or str
Store or path to directory in file system or name of zip file.
synchronizer : object, optional
Expand Down Expand Up @@ -834,12 +835,6 @@ async def create(
else:
chunk_shape = shape

if order is not None:
warnings.warn(
"order is deprecated, use config `array.order` instead",
DeprecationWarning,
stacklevel=2,
)
if synchronizer is not None:
warnings.warn("synchronizer is not yet implemented", RuntimeWarning, stacklevel=2)
if chunk_store is not None:
Expand Down Expand Up @@ -889,6 +884,7 @@ async def create(
codecs=codecs,
dimension_names=dimension_names,
attributes=attributes,
order=order,
**kwargs,
)

Expand Down
35 changes: 17 additions & 18 deletions src/zarr/core/array.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
ZARRAY_JSON,
ZATTRS_JSON,
ChunkCoords,
MemoryOrder,
ShapeLike,
ZarrFormat,
concurrent_map,
Expand Down Expand Up @@ -203,29 +204,29 @@ class AsyncArray(Generic[T_ArrayMetadata]):
metadata: T_ArrayMetadata
store_path: StorePath
codec_pipeline: CodecPipeline = field(init=False)
order: Literal["C", "F"]
order: MemoryOrder

@overload
def __init__(
self: AsyncArray[ArrayV2Metadata],
metadata: ArrayV2Metadata | ArrayV2MetadataDict,
store_path: StorePath,
order: Literal["C", "F"] | None = None,
order: MemoryOrder | None = None,
) -> None: ...

@overload
def __init__(
self: AsyncArray[ArrayV3Metadata],
metadata: ArrayV3Metadata | ArrayV3MetadataDict,
store_path: StorePath,
order: Literal["C", "F"] | None = None,
order: MemoryOrder | None = None,
) -> None: ...

def __init__(
self,
metadata: ArrayMetadata | ArrayMetadataDict,
store_path: StorePath,
order: Literal["C", "F"] | None = None,
order: MemoryOrder | None = None,
) -> None:
if isinstance(metadata, dict):
zarr_format = metadata["zarr_format"]
Expand Down Expand Up @@ -261,7 +262,7 @@ async def create(
attributes: dict[str, JSON] | None = None,
chunks: ShapeLike | None = None,
dimension_separator: Literal[".", "/"] | None = None,
order: Literal["C", "F"] | None = None,
order: MemoryOrder | None = None,
filters: list[dict[str, JSON]] | None = None,
compressor: dict[str, JSON] | None = None,
# runtime
Expand Down Expand Up @@ -350,7 +351,7 @@ async def create(
# v2 only
chunks: ShapeLike | None = None,
dimension_separator: Literal[".", "/"] | None = None,
order: Literal["C", "F"] | None = None,
order: MemoryOrder | None = None,
filters: list[dict[str, JSON]] | None = None,
compressor: dict[str, JSON] | None = None,
# runtime
Expand Down Expand Up @@ -382,7 +383,7 @@ async def create(
# v2 only
chunks: ShapeLike | None = None,
dimension_separator: Literal[".", "/"] | None = None,
order: Literal["C", "F"] | None = None,
order: MemoryOrder | None = None,
filters: list[dict[str, JSON]] | None = None,
compressor: dict[str, JSON] | None = None,
# runtime
Expand Down Expand Up @@ -422,7 +423,6 @@ async def create(
V2 only. V3 arrays cannot have a dimension separator.
order : Literal["C", "F"], optional
The order of the array (default is None).
V2 only. V3 arrays should not have 'order' parameter.
filters : list[dict[str, JSON]], optional
The filters used to compress the data (default is None).
V2 only. V3 arrays should not have 'filters' parameter.
Expand Down Expand Up @@ -471,10 +471,6 @@ async def create(
raise ValueError(
"dimension_separator cannot be used for arrays with version 3. Use chunk_key_encoding instead."
)
if order is not None:
raise ValueError(
"order cannot be used for arrays with version 3. Use a transpose codec instead."
)
if filters is not None:
raise ValueError(
"filters cannot be used for arrays with version 3. Use array-to-array codecs instead."
Expand All @@ -494,6 +490,7 @@ async def create(
dimension_names=dimension_names,
attributes=attributes,
exists_ok=exists_ok,
order=order,
)
elif zarr_format == 2:
if dtype is str or dtype == "str":
Expand Down Expand Up @@ -545,6 +542,7 @@ async def _create_v3(
dtype: npt.DTypeLike,
chunk_shape: ChunkCoords,
fill_value: Any | None = None,
order: MemoryOrder | None = None,
chunk_key_encoding: (
ChunkKeyEncoding
| tuple[Literal["default"], Literal[".", "/"]]
Expand Down Expand Up @@ -588,7 +586,7 @@ async def _create_v3(
attributes=attributes or {},
)

array = cls(metadata=metadata, store_path=store_path)
array = cls(metadata=metadata, store_path=store_path, order=order)
await array._save_metadata(metadata, ensure_parents=True)
return array

Expand All @@ -602,16 +600,17 @@ async def _create_v2(
chunks: ChunkCoords,
dimension_separator: Literal[".", "/"] | None = None,
fill_value: None | float = None,
order: Literal["C", "F"] | None = None,
order: MemoryOrder | None = None,
filters: list[dict[str, JSON]] | None = None,
compressor: dict[str, JSON] | None = None,
attributes: dict[str, JSON] | None = None,
exists_ok: bool = False,
) -> AsyncArray[ArrayV2Metadata]:
if not exists_ok:
await ensure_no_existing_node(store_path, zarr_format=2)

if order is None:
order = "C"
order = parse_indexing_order(config.get("array.order"))

if dimension_separator is None:
dimension_separator = "."
Expand All @@ -627,7 +626,7 @@ async def _create_v2(
filters=filters,
attributes=attributes,
)
array = cls(metadata=metadata, store_path=store_path)
array = cls(metadata=metadata, store_path=store_path, order=order)
await array._save_metadata(metadata, ensure_parents=True)
return array

Expand Down Expand Up @@ -1179,7 +1178,7 @@ def create(
# v2 only
chunks: ChunkCoords | None = None,
dimension_separator: Literal[".", "/"] | None = None,
order: Literal["C", "F"] | None = None,
order: MemoryOrder | None = None,
filters: list[dict[str, JSON]] | None = None,
compressor: dict[str, JSON] | None = None,
# runtime
Expand Down Expand Up @@ -1370,7 +1369,7 @@ def store_path(self) -> StorePath:
return self._async_array.store_path

@property
def order(self) -> Literal["C", "F"]:
def order(self) -> MemoryOrder:
return self._async_array.order

@property
Expand Down
8 changes: 4 additions & 4 deletions src/zarr/core/array_spec.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,11 @@
from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Literal
from typing import TYPE_CHECKING, Any

import numpy as np

from zarr.core.common import parse_fill_value, parse_order, parse_shapelike
from zarr.core.common import MemoryOrder, parse_fill_value, parse_order, parse_shapelike

if TYPE_CHECKING:
from zarr.core.buffer import BufferPrototype
Expand All @@ -17,15 +17,15 @@ class ArraySpec:
shape: ChunkCoords
dtype: np.dtype[Any]
fill_value: Any
order: Literal["C", "F"]
order: MemoryOrder
prototype: BufferPrototype

def __init__(
self,
shape: ChunkCoords,
dtype: np.dtype[Any],
fill_value: Any,
order: Literal["C", "F"],
order: MemoryOrder,
prototype: BufferPrototype,
) -> None:
shape_parsed = parse_shapelike(shape)
Expand Down
8 changes: 4 additions & 4 deletions src/zarr/core/metadata/v2.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@
from zarr.core.array_spec import ArraySpec
from zarr.core.chunk_grids import RegularChunkGrid
from zarr.core.chunk_key_encodings import parse_separator
from zarr.core.common import ZARRAY_JSON, ZATTRS_JSON, parse_shapelike
from zarr.core.common import ZARRAY_JSON, ZATTRS_JSON, MemoryOrder, parse_shapelike
from zarr.core.config import config, parse_indexing_order
from zarr.core.metadata.common import parse_attributes

Expand All @@ -45,7 +45,7 @@ class ArrayV2Metadata(Metadata):
chunks: tuple[int, ...]
dtype: np.dtype[Any]
fill_value: None | int | float | str | bytes = 0
order: Literal["C", "F"] = "C"
order: MemoryOrder = "C"
filters: tuple[numcodecs.abc.Codec, ...] | None = None
dimension_separator: Literal[".", "/"] = "."
compressor: numcodecs.abc.Codec | None = None
Expand All @@ -59,7 +59,7 @@ def __init__(
dtype: npt.DTypeLike,
chunks: ChunkCoords,
fill_value: Any,
order: Literal["C", "F"],
order: MemoryOrder,
dimension_separator: Literal[".", "/"] = ".",
compressor: numcodecs.abc.Codec | dict[str, JSON] | None = None,
filters: Iterable[numcodecs.abc.Codec | dict[str, JSON]] | None = None,
Expand Down Expand Up @@ -185,7 +185,7 @@ def to_dict(self) -> dict[str, JSON]:
return zarray_dict

def get_chunk_spec(
self, _chunk_coords: ChunkCoords, order: Literal["C", "F"], prototype: BufferPrototype
self, _chunk_coords: ChunkCoords, order: MemoryOrder, prototype: BufferPrototype
) -> ArraySpec:
return ArraySpec(
shape=self.chunks,
Expand Down
3 changes: 2 additions & 1 deletion src/zarr/core/metadata/v3.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
JSON,
ZARR_JSON,
ChunkCoords,
MemoryOrder,
parse_named_configuration,
parse_shapelike,
)
Expand Down Expand Up @@ -289,7 +290,7 @@ def ndim(self) -> int:
return len(self.shape)

def get_chunk_spec(
self, _chunk_coords: ChunkCoords, order: Literal["C", "F"], prototype: BufferPrototype
self, _chunk_coords: ChunkCoords, order: MemoryOrder, prototype: BufferPrototype
) -> ArraySpec:
assert isinstance(
self.chunk_grid, RegularChunkGrid
Expand Down
18 changes: 17 additions & 1 deletion tests/test_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
save_array,
save_group,
)
from zarr.core.common import ZarrFormat
from zarr.core.common import MemoryOrder, ZarrFormat
from zarr.errors import MetadataValidationError
from zarr.storage._utils import normalize_path
from zarr.storage.memory import MemoryStore
Expand Down Expand Up @@ -206,6 +206,22 @@ def test_open_with_mode_w_minus(tmp_path: pathlib.Path) -> None:
zarr.open(store=tmp_path, mode="w-")


@pytest.mark.parametrize("order", ["C", "F", None])
@pytest.mark.parametrize("zarr_format", [2, 3])
def test_array_order(order: MemoryOrder | None, zarr_format: ZarrFormat) -> None:
arr = zarr.ones(shape=(2, 2), order=order, zarr_format=zarr_format)
expected = order or zarr.config.get("array.order")
assert arr.order == expected

vals = np.asarray(arr)
if expected == "C":
assert vals.flags.c_contiguous
elif expected == "F":
assert vals.flags.f_contiguous
else:
raise AssertionError


# def test_lazy_loader():
# foo = np.arange(100)
# bar = np.arange(100, 0, -1)
Expand Down
21 changes: 20 additions & 1 deletion tests/test_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
from zarr.codecs import BytesCodec, VLenBytesCodec
from zarr.core.array import chunks_initialized
from zarr.core.buffer.cpu import NDBuffer
from zarr.core.common import JSON, ZarrFormat
from zarr.core.common import JSON, MemoryOrder, ZarrFormat
from zarr.core.group import AsyncGroup
from zarr.core.indexing import ceildiv
from zarr.core.sync import sync
Expand Down Expand Up @@ -417,3 +417,22 @@ def test_update_attrs(zarr_format: int) -> None:

arr2 = zarr.open_array(store=store, zarr_format=zarr_format)
assert arr2.attrs["foo"] == "bar"


@pytest.mark.parametrize("order", ["C", "F", None])
@pytest.mark.parametrize("zarr_format", [2, 3])
@pytest.mark.parametrize("store", ["memory"], indirect=True)
def test_array_create_order(
order: MemoryOrder | None, zarr_format: int, store: MemoryStore
) -> None:
arr = Array.create(store=store, shape=(2, 2), order=order, zarr_format=zarr_format, dtype="i4")
expected = order or zarr.config.get("array.order")
assert arr.order == expected

vals = np.asarray(arr)
if expected == "C":
assert vals.flags.c_contiguous
elif expected == "F":
assert vals.flags.f_contiguous
else:
raise AssertionError