diff --git a/.github/workflows/engine.yml b/.github/workflows/engine.yml
new file mode 100644
index 0000000000..d02544fac8
--- /dev/null
+++ b/.github/workflows/engine.yml
@@ -0,0 +1,49 @@
+name: Array engine
+
+on:
+ push:
+ branches: [ main ]
+ paths:
+ - 'src/zarr/**'
+ - 'tests/engine/**'
+ - 'tests/zarrista/**'
+ - 'pyproject.toml'
+ - '.github/workflows/engine.yml'
+ pull_request:
+ branches: [ main ]
+ paths:
+ - 'src/zarr/**'
+ - 'tests/engine/**'
+ - 'tests/zarrista/**'
+ - 'pyproject.toml'
+ - '.github/workflows/engine.yml'
+ workflow_dispatch:
+
+permissions: {}
+
+concurrency:
+ group: ${{ github.workflow }}-${{ github.ref }}
+ cancel-in-progress: true
+
+jobs:
+ test:
+ runs-on: ubuntu-latest
+ permissions:
+ contents: read
+ steps:
+ - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3
+ with:
+ fetch-depth: 0 # hatch-vcs needs tags to compute zarr's version
+ persist-credentials: false
+ - name: Install Rust
+ uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable
+ - name: Install uv
+ uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # v8.2.0
+ with:
+ python-version: '3.12'
+ - name: Sync zarrista group
+ run: uv sync --group zarrista
+ - name: Run engine tests
+ # the ubuntu runner image ships a Rust toolchain; the maturin build
+ # backend for zarrista is fetched by uv on demand
+ run: uv run --group zarrista pytest tests/engine tests/zarrista -v
diff --git a/changes/4181.feature.md b/changes/4181.feature.md
new file mode 100644
index 0000000000..2d5cb135ed
--- /dev/null
+++ b/changes/4181.feature.md
@@ -0,0 +1,21 @@
+`Array` and `AsyncArray` now route data I/O through pluggable *array engines*
+(`zarr.abc.engine.ArrayEngine` / `AsyncArrayEngine`). An engine is selected
+with `engine=` on `zarr.create_array`, `zarr.open_array`, `Group.create_array`
+and the other array entry points; `zarr.list_engines()` lists the built-in
+names.
+
+The default engine is the existing codec-pipeline path, so behavior and
+performance are unchanged when `engine=` is omitted.
+
+Pass `engine="zarrista"` (requires the `zarrista` package) to serve Zarr v3
+arrays through the Rust `zarrs` implementation, on local-filesystem, obstore,
+or icechunk storage. It supports basic, orthogonal, coordinate, mask and block
+selections, for both reads and writes.
+
+An engine receives the user's selection as a `SelectionRequest`, which carries
+the selection unresolved along with the dialect it is written in. A backend
+that reads chunks takes the `Indexer` the request builds on demand; a backend
+that reads boxes — an FFI binding, an HTTP range endpoint — resolves the
+selection through [zarr-indexing](https://zarr-indexing.readthedocs.io/), which
+grafts the full NumPy indexing dialect onto a source offering only step-1
+slices. The zarrista engine takes the latter route.
diff --git a/docs/api/zarr/abc/engine.md b/docs/api/zarr/abc/engine.md
new file mode 100644
index 0000000000..fe38c62f40
--- /dev/null
+++ b/docs/api/zarr/abc/engine.md
@@ -0,0 +1,5 @@
+---
+title: engine
+---
+
+::: zarr.abc.engine
diff --git a/docs/api/zarr/zarrista.md b/docs/api/zarr/zarrista.md
new file mode 100644
index 0000000000..faf749495f
--- /dev/null
+++ b/docs/api/zarr/zarrista.md
@@ -0,0 +1,5 @@
+---
+title: zarrista
+---
+
+::: zarr.zarrista
diff --git a/docs/user-guide/examples/open_with_engine.md b/docs/user-guide/examples/open_with_engine.md
new file mode 100644
index 0000000000..30848612c9
--- /dev/null
+++ b/docs/user-guide/examples/open_with_engine.md
@@ -0,0 +1,7 @@
+--8<-- "examples/open_with_engine/README.md"
+
+## Source Code
+
+```python exec="false" reason="pymdownx snippet include directive, not python source"
+--8<-- "examples/open_with_engine/open_with_engine.py"
+```
diff --git a/examples/open_with_engine/README.md b/examples/open_with_engine/README.md
new file mode 100644
index 0000000000..de691ca4f1
--- /dev/null
+++ b/examples/open_with_engine/README.md
@@ -0,0 +1,51 @@
+# Open With a Different Engine Example
+
+This example demonstrates how to open the **same array data with a different
+backend** -- what zarr-python calls an *engine*.
+
+An engine is purely an *execution* setting: it selects which compute backend
+reads and writes an array's chunks. The bytes on disk are identical regardless
+of the engine, so the same Zarr array can be driven by a different backend
+without rewriting any data.
+
+The example shows how to:
+
+- Discover the available engines with `zarr.list_engines()`.
+- Inspect the engine an array is using via its public `array.engine` property.
+- Write one array once and read it back through several engines, asserting the
+ results are byte-for-byte identical (`numpy.testing.assert_array_equal`).
+- Select an engine **per call** via the `engine=` kwarg on `zarr.open_array`
+ and `zarr.create_array`.
+- Use the `"default"` engine, which works on any store and format.
+- Use the `"zarrista"` (Rust-backed) engine, which serves Zarr v3 arrays on a
+ `LocalStore` or an obstore-backed `ObjectStore` (the example guards this so it
+ still runs if the package is absent).
+
+## Available engines
+
+| Engine | Backend | Stores | Notes |
+| ------------ | ----------------- | ----------------------------------- | ------------------------------- |
+| `"default"` | built-in Python | any | used when `engine=` is omitted |
+| `"zarrista"` | Rust (`zarrista`) | `LocalStore`, obstore `ObjectStore` | requires the `zarrista` package |
+
+## Running the Example
+
+The `"zarrista"` engine is a Rust extension, so it is deliberately left out of
+this example's PEP 723 inline-dependency header — resolving it would put a
+toolchain build in the path of every run. The zarrista section is guarded and
+skips when the package is absent.
+
+Run it from a checkout of this branch:
+
+```bash
+uv run python examples/open_with_engine/open_with_engine.py
+```
+
+To exercise the Rust-backed engine, sync the `zarrista` dependency group:
+
+```bash
+uv run --group zarrista python examples/open_with_engine/open_with_engine.py
+```
+
+If `zarrista` is not installed, the example still runs and demonstrates the
+default engine; the `zarrista` portion is skipped.
diff --git a/examples/open_with_engine/open_with_engine.py b/examples/open_with_engine/open_with_engine.py
new file mode 100644
index 0000000000..6b53305522
--- /dev/null
+++ b/examples/open_with_engine/open_with_engine.py
@@ -0,0 +1,120 @@
+# /// script
+# requires-python = ">=3.12"
+# dependencies = [
+# "zarr @ git+https://github.com/zarr-developers/zarr-python.git@main",
+# ]
+# ///
+#
+# `zarrista` is deliberately *not* pinned here: it is a Rust extension, so
+# resolving it would put a toolchain build in the path of every run of this
+# example. The zarrista section below is guarded and skips when the package is
+# absent, so the example still runs and demonstrates the default engine. To see
+# the Rust-backed engine, run it from a checkout with that dependency group:
+#
+# uv run --group zarrista python examples/open_with_engine/open_with_engine.py
+
+"""
+Open the same array data with a different backend ("engine").
+
+zarr-python routes an array's data I/O through a selectable *engine*. The engine
+is purely an *execution* setting: it chooses which compute backend reads and
+writes the array's chunks. The bytes on disk are identical regardless of the
+engine -- the same Zarr array, just a different machine doing the work.
+
+Engines demonstrated here:
+
+- `"default"` -- the built-in engine (used when `engine=` is omitted). Works on
+ every store and format.
+- `"zarrista"` -- a Rust-backed engine (via the `zarrista` package) that serves
+ Zarr v3 arrays on a `LocalStore` or an obstore-backed `ObjectStore`. Guarded:
+ skipped if the package is not installed.
+
+We write one array once, then open and read it back through each engine and
+assert the results are byte-for-byte identical.
+"""
+
+import tempfile
+from pathlib import Path
+
+import numpy as np
+
+import zarr
+from zarr.storage import LocalStore
+
+
+def zarrista_available() -> bool:
+ """Report whether the optional Rust-backed `"zarrista"` engine can be used.
+
+ The `"zarrista"` engine imports the `zarrista` package lazily, so we probe
+ that package directly: importing `zarr.zarrista` alone succeeds even when the
+ package is absent (the actual `ImportError` would only surface on first I/O).
+ """
+ try:
+ import zarrista # noqa: F401
+ except ImportError:
+ return False
+ return True
+
+
+def main() -> None:
+ # Discover which engines exist. `zarr.list_engines()` returns the built-in
+ # engine names; `"zarrista"` additionally requires the `zarrista` package.
+ print("available engines:", zarr.list_engines())
+
+ # zarrista ingests a LocalStore (used here, on a temp directory) or an
+ # obstore-backed ObjectStore. We keep a single store so every engine reads
+ # the exact same bytes off the same disk.
+ with tempfile.TemporaryDirectory() as tmp:
+ store = LocalStore(Path(tmp) / "store")
+
+ # The data we will write once and read back through several engines.
+ data = np.arange(8 * 8, dtype="uint16").reshape(8, 8)
+
+ # --- Write the array once, with the default engine. -----------------
+ # No engine= here, so this uses the "default" engine.
+ source = zarr.create_array(
+ store=store, name="a", shape=(8, 8), chunks=(4, 4), dtype="uint16"
+ )
+ source[:] = data
+
+ # --- 1. Default engine (the baseline). ------------------------------
+ # Either omit engine= or pass engine="default"; both mean the built-in.
+ default = zarr.open_array(store=store, path="a", engine="default")
+ # `array.engine` is the resolved engine instance backing this array.
+ print("default engine:", type(default.engine).__name__)
+ np.testing.assert_array_equal(default[:], data)
+ # Basic indexing (ints and slices) is what engines route; check a slice.
+ np.testing.assert_array_equal(default[2:6, 1:5], data[2:6, 1:5])
+ print("default (engine='default') : read back OK")
+
+ # --- 2. zarrista engine (Rust), guarded behind availability. --------
+ if zarrista_available():
+ zst = zarr.open_array(store=store, path="a", engine="zarrista")
+ print("zarrista engine:", type(zst.engine).__name__)
+ np.testing.assert_array_equal(zst[:], data)
+ np.testing.assert_array_equal(zst[2:6, 1:5], data[2:6, 1:5])
+ # Same bytes on disk, different compute backend: identical to default.
+ np.testing.assert_array_equal(zst[:], default[:])
+ print("zarrista (engine='zarrista') : read back OK, equals default")
+
+ # Writes also route through the engine. Create + write + read back
+ # entirely through zarrista, then confirm a *default* reader agrees.
+ written = zarr.create_array(
+ store=store,
+ name="b",
+ shape=(8, 8),
+ chunks=(4, 4),
+ dtype="uint16",
+ engine="zarrista",
+ )
+ written[:] = data
+ np.testing.assert_array_equal(zarr.open_array(store=store, path="b")[:], data)
+ print("zarrista (write path) : round-trip OK, default reader agrees")
+ else:
+ print("zarrista : SKIPPED (package not installed)")
+
+ print("\nAll engines returned identical data. Same bytes on disk, different backend.")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/mkdocs.yml b/mkdocs.yml
index ca8165af4c..25110f9f20 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -30,6 +30,7 @@ nav:
- user-guide/glossary.md
- Examples:
- user-guide/examples/custom_dtype.md
+ - user-guide/examples/open_with_engine.md
- user-guide/examples/rectilinear_chunks.md
- user-guide/examples/codec_pipeline_performance.md
- user-guide/examples/sharding_coalescing.md
@@ -40,6 +41,7 @@ nav:
- api/zarr/abc/index.md
- 'zarr.abc.buffer': api/zarr/abc/buffer.md
- 'zarr.abc.codec': api/zarr/abc/codec.md
+ - 'zarr.abc.engine': api/zarr/abc/engine.md
- 'zarr.abc.metadata': api/zarr/abc/metadata.md
- 'zarr.abc.numcodec': api/zarr/abc/numcodec.md
- 'zarr.abc.store': api/zarr/abc/store.md
@@ -93,6 +95,7 @@ nav:
- 'zarr.testing.store': api/zarr/testing/store.md
- 'zarr.testing.strategies': api/zarr/testing/strategies.md
- 'zarr.testing.utils': api/zarr/testing/utils.md
+ - 'zarr.zarrista': api/zarr/zarrista.md
- 'zarr.zeros': api/zarr/functions/zeros.md
- 'zarr.zeros_like': api/zarr/functions/zeros_like.md
# The companion packages are Read the Docs subprojects of this one; link
diff --git a/pyproject.toml b/pyproject.toml
index 6626f8f0bc..29e648caa4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -145,6 +145,20 @@ dev = [
"universal-pathlib",
"mypy==2.3.0",
]
+zarrista = [
+ {include-group = "test"},
+ # Tracked from git rather than PyPI: the engine uses `store_array_subset`
+ # and the `Tensor` family, and `uv.lock` pins the exact commit.
+ "zarrista @ git+https://github.com/developmentseed/zarrista",
+ # The zarrista engine resolves selections through zarr-indexing's
+ # `LazyArray`; zarr itself does not depend on it.
+ "zarr-indexing>=0.2.1",
+ # zarrista's `store` submodule unconditionally imports both of these
+ # (for its `AsyncStore` type alias), even though neither is declared in
+ # zarrista's own package metadata as a runtime dependency.
+ "icechunk>=1.1.21",
+ "obstore>=0.10.1",
+]
[tool.coverage.report]
exclude_also = [
diff --git a/src/zarr/__init__.py b/src/zarr/__init__.py
index cdf3840c3b..261fb6cfdf 100644
--- a/src/zarr/__init__.py
+++ b/src/zarr/__init__.py
@@ -36,6 +36,7 @@
)
from zarr.core.array import Array, AsyncArray
from zarr.core.config import config
+from zarr.core.engine import list_engines
from zarr.core.group import AsyncGroup, Group
# in case setuptools scm screw up and find version to be 0.0.0
@@ -164,6 +165,7 @@ def set_format(log_format: str) -> None:
"full",
"full_like",
"group",
+ "list_engines",
"load",
"ones",
"ones_like",
diff --git a/src/zarr/abc/engine.py b/src/zarr/abc/engine.py
new file mode 100644
index 0000000000..3aa61288d0
--- /dev/null
+++ b/src/zarr/abc/engine.py
@@ -0,0 +1,171 @@
+"""Array engine protocols.
+
+An *array engine* owns the data path of one open array: reading and writing
+decoded data for a selection. `AsyncArray` wraps an object satisfying
+`AsyncArrayEngine`; `Array` wraps an object satisfying `ArrayEngine`. A
+*hierarchy engine* is bound to a store and mints array engines that share
+resources.
+
+What crosses the boundary is a `SelectionRequest`: the user's selection,
+unresolved, tagged with the dialect it is written in. Both descriptions an
+engine might want are derived from it, losslessly:
+
+- `SelectionRequest.indexer` is zarr-python's native `Indexer`, a chunk-gather
+ plan. The built-in engine consumes this and so drives the existing codec
+ pipeline with no round-trip.
+- The raw `selection` and `kind` are what a backend that reads *boxes* rather
+ than chunks (an FFI binding, an HTTP range endpoint) needs, typically to
+ drive a `zarr_indexing.LazyArray`.
+
+The request deliberately carries the *unresolved* selection. An `Indexer` has
+already thrown away what a box-reading backend needs — `CoordinateIndexer`
+stores its coordinates in chunk-sorted order, not the user's — so an engine
+that reconstructed a selection from an indexer would silently reorder results.
+Deriving both views from the same raw selection avoids that entirely, and
+keeps `zarr_indexing` out of this module: only a backend that wants a
+transform imports it.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from functools import cached_property
+from typing import TYPE_CHECKING, Any, Literal, Protocol, runtime_checkable
+
+if TYPE_CHECKING:
+ import numpy.typing as npt
+
+ from zarr.core.array_spec import ArrayConfig
+ from zarr.core.buffer import BufferPrototype, NDArrayLikeOrScalar, NDBuffer
+ from zarr.core.chunk_grids import ChunkGrid
+ from zarr.core.indexing import Fields, Indexer
+ from zarr.core.metadata import ArrayMetadata
+
+__all__ = [
+ "ArrayEngine",
+ "AsyncArrayEngine",
+ "AsyncHierarchyEngine",
+ "HierarchyEngine",
+ "SelectionKind",
+ "SelectionRequest",
+]
+
+SelectionKind = Literal["basic", "orthogonal", "coordinate", "mask", "block"]
+"""Which indexing dialect a `SelectionRequest.selection` is written in."""
+
+
+@dataclass(frozen=True)
+class SelectionRequest:
+ """One read or write request: a selection, its dialect, and its context.
+
+ See the module docstring for why this carries the *unresolved* selection.
+ """
+
+ kind: SelectionKind
+ selection: Any
+ shape: tuple[int, ...]
+ chunk_grid: ChunkGrid
+
+ @cached_property
+ def indexer(self) -> Indexer:
+ """This request as a zarr-python `Indexer`.
+
+ Resolved on first access and then cached, so an engine that never asks
+ for it does not pay to build it.
+ """
+ from zarr.core.indexing import (
+ BasicIndexer,
+ BlockIndexer,
+ CoordinateIndexer,
+ MaskIndexer,
+ OrthogonalIndexer,
+ )
+
+ indexer_for: dict[str, Any] = {
+ "basic": BasicIndexer,
+ "orthogonal": OrthogonalIndexer,
+ "coordinate": CoordinateIndexer,
+ "mask": MaskIndexer,
+ "block": BlockIndexer,
+ }
+ return indexer_for[self.kind](self.selection, self.shape, self.chunk_grid) # type: ignore[no-any-return]
+
+
+@runtime_checkable
+class AsyncArrayEngine(Protocol):
+ """The asynchronous data path of one open array.
+
+ Bound to `(store, path, metadata, config)` at construction. The signatures
+ mirror zarr-python's own `_get_selection` / `_set_selection` — including
+ the scalar return for a rank-0 basic selection — so that the built-in
+ engine is a pure delegation.
+
+ Note: `runtime_checkable` isinstance checks only verify method names;
+ mypy is the authoritative conformance check.
+ """
+
+ async def read_selection(
+ self,
+ request: SelectionRequest,
+ *,
+ prototype: BufferPrototype,
+ out: NDBuffer | None = None,
+ fields: Fields | None = None,
+ ) -> NDArrayLikeOrScalar: ...
+
+ async def write_selection(
+ self,
+ request: SelectionRequest,
+ value: npt.ArrayLike,
+ *,
+ prototype: BufferPrototype,
+ fields: Fields | None = None,
+ ) -> None: ...
+
+ def with_metadata(self, metadata: ArrayMetadata) -> AsyncArrayEngine: ...
+
+
+@runtime_checkable
+class ArrayEngine(Protocol):
+ """The synchronous data path of one open array.
+
+ Methods must not require a running event loop. See `AsyncArrayEngine`.
+ """
+
+ def read_selection(
+ self,
+ request: SelectionRequest,
+ *,
+ prototype: BufferPrototype,
+ out: NDBuffer | None = None,
+ fields: Fields | None = None,
+ ) -> NDArrayLikeOrScalar: ...
+
+ def write_selection(
+ self,
+ request: SelectionRequest,
+ value: npt.ArrayLike,
+ *,
+ prototype: BufferPrototype,
+ fields: Fields | None = None,
+ ) -> None: ...
+
+ def with_metadata(self, metadata: ArrayMetadata) -> ArrayEngine: ...
+
+
+@runtime_checkable
+class AsyncHierarchyEngine(Protocol):
+ """A store-bound factory for asynchronous array engines."""
+
+ def array_engine(
+ self, path: str, metadata: ArrayMetadata, config: ArrayConfig | None = None
+ ) -> AsyncArrayEngine: ...
+
+
+@runtime_checkable
+class HierarchyEngine(Protocol):
+ """A store-bound factory for synchronous array engines."""
+
+ def array_engine(
+ self, path: str, metadata: ArrayMetadata, config: ArrayConfig | None = None
+ ) -> ArrayEngine: ...
diff --git a/src/zarr/api/asynchronous.py b/src/zarr/api/asynchronous.py
index 3bdc254ea5..897741f0ec 100644
--- a/src/zarr/api/asynchronous.py
+++ b/src/zarr/api/asynchronous.py
@@ -53,9 +53,11 @@
from collections.abc import Iterable
from zarr.abc.codec import Codec
+ from zarr.abc.engine import ArrayEngine, AsyncArrayEngine
from zarr.abc.numcodec import Numcodec
from zarr.core.buffer import NDArrayLikeOrScalar
from zarr.core.chunk_key_encodings import ChunkKeyEncoding
+ from zarr.core.engine import EngineName
from zarr.core.metadata.v2 import CompressorLikev2
from zarr.storage import StoreLike
from zarr.types import AnyArray, AnyAsyncArray
@@ -915,6 +917,7 @@ async def create(
dimension_names: DimensionNamesLike = None,
storage_options: dict[str, Any] | None = None,
config: ArrayConfigLike | None = None,
+ engine: ArrayEngine | AsyncArrayEngine | EngineName | None = None,
**kwargs: Any,
) -> AnyAsyncArray:
"""Create an array.
@@ -1035,6 +1038,12 @@ async def create(
config : ArrayConfigLike, optional
Runtime configuration of the array. If provided, will override the
default values from `zarr.config.array`.
+ engine : ArrayEngine | AsyncArrayEngine | Literal["default", "zarrista"] | None, optional
+ The data-path engine backing the created array: a name (`"default"`,
+ `"zarrista"`) or a pre-built engine instance. A synchronous
+ `ArrayEngine` instance only makes sense from the sync API; an
+ `AsyncArrayEngine` instance only from the async API; a name works
+ from either. When omitted, the `"default"` behavior is unchanged.
Returns
-------
@@ -1095,6 +1104,10 @@ async def create(
dimension_names=dimension_names,
attributes=attributes,
config=config_parsed,
+ # `AsyncArray._create` accepts only `AsyncArrayEngine`; a wrong-kind
+ # (sync) instance is rejected downstream, when `AsyncArray.__init__`
+ # resolves it via `resolve_async_engine`.
+ engine=cast("AsyncArrayEngine | EngineName | None", engine),
**kwargs,
)
@@ -1234,6 +1247,7 @@ async def open_array(
zarr_format: ZarrFormat | None = None,
path: PathLike = "",
storage_options: dict[str, Any] | None = None,
+ engine: ArrayEngine | AsyncArrayEngine | EngineName | None = None,
**kwargs: Any, # TODO: type kwargs as valid args to save
) -> AnyAsyncArray:
"""Open an array using file-mode-like semantics.
@@ -1251,6 +1265,13 @@ async def open_array(
storage_options : dict
If using an fsspec URL to create the store, these will be passed to
the backend implementation. Ignored otherwise.
+ engine : ArrayEngine | AsyncArrayEngine | Literal["default", "zarrista"] | None, optional
+ The data-path engine backing the opened (or, if missing, created)
+ array: a name (`"default"`, `"zarrista"`) or a pre-built engine
+ instance. A synchronous `ArrayEngine` instance only makes sense from
+ the sync API; an `AsyncArrayEngine` instance only from the async API;
+ a name works from either. When omitted, the `"default"` behavior is
+ unchanged.
**kwargs
Any keyword arguments to pass to [`create`][zarr.api.asynchronous.create].
@@ -1267,7 +1288,14 @@ async def open_array(
_warn_write_empty_chunks_kwarg()
try:
- return await AsyncArray.open(store_path, zarr_format=zarr_format)
+ # `AsyncArray.open` accepts only `AsyncArrayEngine`; a wrong-kind
+ # (sync) instance is rejected downstream, when `AsyncArray.__init__`
+ # resolves it via `resolve_async_engine`.
+ return await AsyncArray.open(
+ store_path,
+ zarr_format=zarr_format,
+ engine=cast("AsyncArrayEngine | EngineName | None", engine),
+ )
except FileNotFoundError as err:
if not store_path.read_only and mode in _CREATE_MODES:
overwrite = _infer_overwrite(mode)
@@ -1276,6 +1304,7 @@ async def open_array(
store=store_path,
zarr_format=_zarr_format,
overwrite=overwrite,
+ engine=engine,
**kwargs,
)
msg = f"No array found in store {store_path.store} at path {store_path.path}"
diff --git a/src/zarr/api/synchronous.py b/src/zarr/api/synchronous.py
index 30ddf9ce6a..319f5d60b1 100644
--- a/src/zarr/api/synchronous.py
+++ b/src/zarr/api/synchronous.py
@@ -7,6 +7,7 @@
import zarr.api.asynchronous as async_api
import zarr.core.array
from zarr.core.array import DEFAULT_FILL_VALUE, Array, AsyncArray, CompressorLike
+from zarr.core.engine import route_sync_engine_arg
from zarr.core.group import Group
from zarr.core.sync import sync
from zarr.core.sync_group import create_hierarchy
@@ -19,6 +20,7 @@
import numpy.typing as npt
from zarr.abc.codec import Codec
+ from zarr.abc.engine import ArrayEngine, AsyncArrayEngine
from zarr.abc.numcodec import Numcodec
from zarr.api.asynchronous import ArrayLike, PathLike
from zarr.core.array import (
@@ -40,6 +42,7 @@
ZarrFormat,
)
from zarr.core.dtype import ZDTypeLike
+ from zarr.core.engine import EngineName
from zarr.storage import StoreLike
from zarr.types import AnyArray
@@ -653,6 +656,7 @@ def create(
dimension_names: DimensionNamesLike = None,
storage_options: dict[str, Any] | None = None,
config: ArrayConfigLike | None = None,
+ engine: ArrayEngine | AsyncArrayEngine | EngineName | None = None,
**kwargs: Any,
) -> AnyArray:
"""Create an array.
@@ -773,12 +777,19 @@ def create(
config : ArrayConfigLike, optional
Runtime configuration of the array. If provided, will override the
default values from `zarr.config.array`.
+ engine : ArrayEngine | AsyncArrayEngine | Literal["default", "zarrista"] | None, optional
+ The data-path engine backing the created array: a name (`"default"`,
+ `"zarrista"`) or a pre-built engine instance. A synchronous
+ `ArrayEngine` instance only makes sense from the sync API; an
+ `AsyncArrayEngine` instance only from the async API; a name works
+ from either. When omitted, the `"default"` behavior is unchanged.
Returns
-------
z : Array
The array.
"""
+ engine_for_async, engine_for_array = route_sync_engine_arg(engine)
return Array(
sync(
async_api.create(
@@ -809,9 +820,11 @@ def create(
dimension_names=dimension_names,
storage_options=storage_options,
config=config,
+ engine=engine_for_async,
**kwargs,
)
- )
+ ),
+ engine_spec=engine_for_array,
)
@@ -837,6 +850,7 @@ def create_array(
overwrite: bool = False,
config: ArrayConfigLike | None = None,
write_data: bool = True,
+ engine: ArrayEngine | AsyncArrayEngine | EngineName | None = None,
) -> AnyArray:
"""Create an array.
@@ -943,6 +957,12 @@ def create_array(
then `write_data` determines whether the values in that array-like object should be
written to the Zarr array created by this function. If `write_data` is `False`, then the
array will be left empty.
+ engine : ArrayEngine | AsyncArrayEngine | Literal["default", "zarrista"] | None, optional
+ The data-path engine backing the created array: a name (`"default"`,
+ `"zarrista"`) or a pre-built engine instance. A synchronous
+ `ArrayEngine` instance only makes sense from the sync API; an
+ `AsyncArrayEngine` instance only from the async API; a name works
+ from either. When omitted, the `"default"` behavior is unchanged.
Returns
-------
@@ -963,6 +983,7 @@ def create_array(
#
```
"""
+ engine_for_async, engine_for_array = route_sync_engine_arg(engine)
return Array(
sync(
zarr.core.array.create_array(
@@ -986,8 +1007,10 @@ def create_array(
overwrite=overwrite,
config=config,
write_data=write_data,
+ engine=engine_for_async,
)
- )
+ ),
+ engine_spec=engine_for_array,
)
@@ -1349,6 +1372,7 @@ def open_array(
zarr_format: ZarrFormat | None = None,
path: PathLike = "",
storage_options: dict[str, Any] | None = None,
+ engine: ArrayEngine | AsyncArrayEngine | EngineName | None = None,
**kwargs: Any,
) -> AnyArray:
"""Open an array using file-mode-like semantics.
@@ -1366,6 +1390,13 @@ def open_array(
storage_options : dict
If using an fsspec URL to create the store, these will be passed to
the backend implementation. Ignored otherwise.
+ engine : ArrayEngine | AsyncArrayEngine | Literal["default", "zarrista"] | None, optional
+ The data-path engine backing the opened (or, if missing, created)
+ array: a name (`"default"`, `"zarrista"`) or a pre-built engine
+ instance. A synchronous `ArrayEngine` instance only makes sense from
+ the sync API; an `AsyncArrayEngine` instance only from the async API;
+ a name works from either. When omitted, the `"default"` behavior is
+ unchanged.
**kwargs
Any keyword arguments to pass to [`create`][zarr.api.asynchronous.create].
@@ -1375,6 +1406,7 @@ def open_array(
AsyncArray
The opened array.
"""
+ engine_for_async, engine_for_array = route_sync_engine_arg(engine)
return Array(
sync(
async_api.open_array(
@@ -1382,9 +1414,11 @@ def open_array(
zarr_format=zarr_format,
path=path,
storage_options=storage_options,
+ engine=engine_for_async,
**kwargs,
)
- )
+ ),
+ engine_spec=engine_for_array,
)
diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py
index 9cb66f339e..564ccd7d3f 100644
--- a/src/zarr/core/array.py
+++ b/src/zarr/core/array.py
@@ -4,7 +4,7 @@
import warnings
from asyncio import gather
from collections.abc import Iterable, Mapping, Sequence
-from dataclasses import dataclass, field, replace
+from dataclasses import InitVar, dataclass, field, replace
from itertools import starmap
from logging import getLogger
from typing import (
@@ -22,6 +22,7 @@
import zarr
from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec, Codec
+from zarr.abc.engine import SelectionKind, SelectionRequest
from zarr.abc.numcodec import Numcodec, _is_numcodec
from zarr.codecs._v2 import V2Codec
from zarr.codecs.bytes import BytesCodec
@@ -82,21 +83,23 @@
parse_dtype,
)
from zarr.core.dtype.common import HasEndianness, HasItemSize, HasObjectCodec
+from zarr.core.engine import (
+ adopt_codec_pipeline,
+ resolve_async_engine,
+ resolve_sync_engine,
+ route_sync_engine_arg,
+)
from zarr.core.indexing import (
AsyncOIndex,
AsyncVIndex,
BasicIndexer,
BasicSelection,
BlockIndex,
- BlockIndexer,
- CoordinateIndexer,
CoordinateSelection,
Fields,
Indexer,
- MaskIndexer,
MaskSelection,
OIndex,
- OrthogonalIndexer,
OrthogonalSelection,
Selection,
VIndex,
@@ -153,9 +156,12 @@
import numpy.typing as npt
from zarr.abc.codec import CodecPipeline
+ from zarr.abc.engine import ArrayEngine, AsyncArrayEngine
from zarr.abc.store import Store
from zarr.codecs.sharding import IndexLocation
from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar
+ from zarr.core.engine import EngineName
+ from zarr.core.indexing import CoordinateIndexer
from zarr.storage import StoreLike
from zarr.types import AnyArray, AnyAsyncArray, ArrayV2, ArrayV3, AsyncArrayV2, AsyncArrayV3
@@ -370,6 +376,8 @@ class AsyncArray[T_ArrayMetadata: (ArrayV2Metadata, ArrayV3Metadata)]:
The codec pipeline used for encoding and decoding chunks.
config : ArrayConfig
The runtime configuration of the array.
+ engine : AsyncArrayEngine
+ The engine backing this array's data path.
"""
metadata: T_ArrayMetadata
@@ -377,6 +385,13 @@ class AsyncArray[T_ArrayMetadata: (ArrayV2Metadata, ArrayV3Metadata)]:
codec_pipeline: CodecPipeline = field(init=False)
_chunk_grid: ChunkGrid = field(init=False)
config: ArrayConfig
+ engine: AsyncArrayEngine = field(init=False, compare=False, repr=False)
+ # The `engine=` argument as given, kept so that a derived array (`with_config`,
+ # `update_attributes`) re-resolves it rather than inheriting an engine bound to
+ # the old metadata or config.
+ _engine_spec: AsyncArrayEngine | EngineName | None = field(
+ init=False, compare=False, repr=False
+ )
@overload
def __init__(
@@ -384,6 +399,7 @@ def __init__(
metadata: ArrayV2Metadata | ArrayV2MetadataDict,
store_path: StorePath,
config: ArrayConfigLike | None = None,
+ engine: AsyncArrayEngine | EngineName | None = None,
) -> None: ...
@overload
@@ -392,6 +408,7 @@ def __init__(
metadata: ArrayV3Metadata | ArrayMetadataJSON_V3,
store_path: StorePath,
config: ArrayConfigLike | None = None,
+ engine: AsyncArrayEngine | EngineName | None = None,
) -> None: ...
def __init__(
@@ -399,6 +416,7 @@ def __init__(
metadata: ArrayMetadata | ArrayMetadataDict,
store_path: StorePath,
config: ArrayConfigLike | None = None,
+ engine: AsyncArrayEngine | EngineName | None = None,
) -> None:
metadata_parsed = parse_array_metadata(metadata)
config_parsed = parse_array_config(config)
@@ -412,6 +430,16 @@ def __init__(
"codec_pipeline",
create_codec_pipeline(metadata=metadata_parsed, store=store_path.store),
)
+ object.__setattr__(self, "_engine_spec", engine)
+ resolved = resolve_async_engine(
+ engine,
+ store=store_path.store,
+ path=store_path.path,
+ metadata=metadata_parsed,
+ config=config_parsed,
+ )
+ adopt_codec_pipeline(resolved, self.codec_pipeline)
+ object.__setattr__(self, "engine", resolved)
@classmethod
async def _create(
@@ -444,6 +472,7 @@ async def _create(
overwrite: bool = False,
data: npt.ArrayLike | None = None,
config: ArrayConfigLike | None = None,
+ engine: AsyncArrayEngine | EngineName | None = None,
) -> AnyAsyncArray:
"""Method to create a new asynchronous array instance.
Deprecated in favor of [`zarr.api.asynchronous.create_array`][].
@@ -500,6 +529,7 @@ async def _create(
overwrite=overwrite,
config=config_parsed,
chunk_grid=chunk_grid,
+ engine=engine,
)
elif zarr_format == 2:
if codecs is not None:
@@ -544,6 +574,7 @@ async def _create(
compressor=compressor,
attributes=attributes,
overwrite=overwrite,
+ engine=engine,
)
else:
raise ValueError(f"zarr_format must be 2 or 3, got {zarr_format}") # pragma: no cover
@@ -624,6 +655,7 @@ async def _create_v3(
dimension_names: DimensionNamesLike = None,
attributes: dict[str, JSON] | None = None,
overwrite: bool = False,
+ engine: AsyncArrayEngine | EngineName | None = None,
) -> AsyncArrayV3:
await _prepare_overwrite(store_path, zarr_format=3, overwrite=overwrite)
@@ -645,7 +677,7 @@ async def _create_v3(
attributes=attributes,
)
- array = cls(metadata=metadata, store_path=store_path, config=config)
+ array = cls(metadata=metadata, store_path=store_path, config=config, engine=engine)
await array._save_metadata(metadata, ensure_parents=True)
return array
@@ -699,6 +731,7 @@ async def _create_v2(
compressor: CompressorLike = "auto",
attributes: dict[str, JSON] | None = None,
overwrite: bool = False,
+ engine: AsyncArrayEngine | EngineName | None = None,
) -> AsyncArrayV2:
await _prepare_overwrite(store_path, zarr_format=2, overwrite=overwrite)
@@ -728,7 +761,7 @@ async def _create_v2(
attributes=attributes,
)
- array = cls(metadata=metadata, store_path=store_path, config=config)
+ array = cls(metadata=metadata, store_path=store_path, config=config, engine=engine)
await array._save_metadata(metadata, ensure_parents=True)
return array
@@ -737,6 +770,7 @@ def from_dict(
cls,
store_path: StorePath,
data: dict[str, JSON],
+ engine: AsyncArrayEngine | EngineName | None = None,
) -> AnyAsyncArray:
"""
Create a Zarr array from a dictionary, with support for both Zarr format 2 and 3 metadata.
@@ -762,13 +796,14 @@ def from_dict(
If the dictionary data is invalid or incompatible with either Zarr format 2 or 3 array creation.
"""
metadata = parse_array_metadata(data)
- return cls(metadata=metadata, store_path=store_path)
+ return cls(metadata=metadata, store_path=store_path, engine=engine)
@classmethod
async def open(
cls,
store: StoreLike,
zarr_format: ZarrFormat | None = 3,
+ engine: AsyncArrayEngine | EngineName | None = None,
) -> AnyAsyncArray:
"""
Async method to open an existing Zarr array from a given store.
@@ -812,7 +847,7 @@ async def example():
metadata_dict = await get_array_metadata(store_path, zarr_format=zarr_format)
# TODO: remove this cast when we have better type hints
_metadata_dict = cast("ArrayMetadataJSON_V3", metadata_dict)
- return cls(store_path=store_path, metadata=_metadata_dict)
+ return cls(store_path=store_path, metadata=_metadata_dict, engine=engine)
@property
def store(self) -> Store:
@@ -1206,7 +1241,14 @@ def with_config(self, config: ArrayConfigLike) -> Self:
# Merge new config with existing config, so missing keys are inherited
# from the current array rather than from global defaults
new_config = ArrayConfig(**{**self.config.to_dict(), **config}) # type: ignore[arg-type]
- return type(self)(metadata=self.metadata, store_path=self.store_path, config=new_config)
+ # `_engine_spec`, not `engine`: passing the already-resolved engine would
+ # freeze the *old* config into the copy.
+ return type(self)(
+ metadata=self.metadata,
+ store_path=self.store_path,
+ config=new_config,
+ engine=self._engine_spec,
+ )
async def nchunks_initialized(self) -> int:
"""
@@ -1445,24 +1487,25 @@ def nbytes(self) -> int:
"""
return self.size * self.dtype.itemsize
+ def _request(self, kind: SelectionKind, selection: Any) -> SelectionRequest:
+ """Package a selection for this array's engine."""
+ return SelectionRequest(
+ kind=kind,
+ selection=selection,
+ shape=self.metadata.shape,
+ chunk_grid=self._chunk_grid,
+ )
+
async def _get_selection(
self,
- indexer: Indexer,
+ request: SelectionRequest,
*,
prototype: BufferPrototype,
out: NDBuffer | None = None,
fields: Fields | None = None,
) -> NDArrayLikeOrScalar:
- return await _get_selection(
- self.store_path,
- self.metadata,
- self.codec_pipeline,
- self.config,
- self._chunk_grid,
- indexer,
- prototype=prototype,
- out=out,
- fields=fields,
+ return await self.engine.read_selection(
+ request, prototype=prototype, out=out, fields=fields
)
async def getitem(
@@ -1507,15 +1550,9 @@ async def getitem(
np.int32(0)
"""
- return await _getitem(
- self.store_path,
- self.metadata,
- self.codec_pipeline,
- self.config,
- self._chunk_grid,
- selection,
- prototype=prototype,
- )
+ if prototype is None:
+ prototype = default_buffer_prototype()
+ return await self._get_selection(self._request("basic", selection), prototype=prototype)
async def get_orthogonal_selection(
self,
@@ -1527,14 +1564,8 @@ async def get_orthogonal_selection(
) -> NDArrayLikeOrScalar:
if prototype is None:
prototype = default_buffer_prototype()
- indexer = OrthogonalIndexer(selection, self.metadata.shape, self._chunk_grid)
- return await _get_selection(
- self.store_path,
- self.metadata,
- self.codec_pipeline,
- self.config,
- self._chunk_grid,
- indexer=indexer,
+ return await self._get_selection(
+ self._request("orthogonal", selection),
out=out,
fields=fields,
prototype=prototype,
@@ -1550,14 +1581,8 @@ async def get_mask_selection(
) -> NDArrayLikeOrScalar:
if prototype is None:
prototype = default_buffer_prototype()
- indexer = MaskIndexer(mask, self.metadata.shape, self._chunk_grid)
- return await _get_selection(
- self.store_path,
- self.metadata,
- self.codec_pipeline,
- self.config,
- self._chunk_grid,
- indexer=indexer,
+ return await self._get_selection(
+ self._request("mask", mask),
out=out,
fields=fields,
prototype=prototype,
@@ -1573,21 +1598,19 @@ async def get_coordinate_selection(
) -> NDArrayLikeOrScalar:
if prototype is None:
prototype = default_buffer_prototype()
- indexer = CoordinateIndexer(selection, self.metadata.shape, self._chunk_grid)
- out_array = await _get_selection(
- self.store_path,
- self.metadata,
- self.codec_pipeline,
- self.config,
- self._chunk_grid,
- indexer=indexer,
+ request = self._request("coordinate", selection)
+ out_array = await self._get_selection(
+ request,
out=out,
fields=fields,
prototype=prototype,
)
if hasattr(out_array, "shape"):
# restore shape
- out_array = cast("NDArrayLikeOrScalar", np.array(out_array).reshape(indexer.sel_shape))
+ out_array = cast(
+ "NDArrayLikeOrScalar",
+ np.array(out_array).reshape(cast("CoordinateIndexer", request.indexer).sel_shape),
+ )
return out_array
async def _save_metadata(self, metadata: ArrayMetadata, ensure_parents: bool = False) -> None:
@@ -1598,23 +1621,13 @@ async def _save_metadata(self, metadata: ArrayMetadata, ensure_parents: bool = F
async def _set_selection(
self,
- indexer: Indexer,
+ request: SelectionRequest,
value: npt.ArrayLike,
*,
prototype: BufferPrototype,
fields: Fields | None = None,
) -> None:
- return await _set_selection(
- self.store_path,
- self.metadata,
- self.codec_pipeline,
- self.config,
- self._chunk_grid,
- indexer,
- value,
- prototype=prototype,
- fields=fields,
- )
+ return await self.engine.write_selection(request, value, prototype=prototype, fields=fields)
async def setitem(
self,
@@ -1655,15 +1668,10 @@ async def setitem(
- This method is asynchronous and should be awaited.
- Supports basic indexing, where the selection is contiguous and does not involve advanced indexing.
"""
- return await _setitem(
- self.store_path,
- self.metadata,
- self.codec_pipeline,
- self.config,
- self._chunk_grid,
- selection,
- value,
- prototype=prototype,
+ if prototype is None:
+ prototype = default_buffer_prototype()
+ return await self._set_selection(
+ self._request("basic", selection), value, prototype=prototype
)
@property
@@ -1850,6 +1858,69 @@ class Array[T_ArrayMetadata: (ArrayV2Metadata, ArrayV3Metadata)]:
"""
_async_array: AsyncArray[T_ArrayMetadata]
+ engine_spec: InitVar[ArrayEngine | EngineName | None] = None
+ _engine: ArrayEngine | None = field(default=None, init=False, repr=False, compare=False)
+
+ def __post_init__(self, engine_spec: ArrayEngine | EngineName | None) -> None:
+ self._engine_spec = engine_spec
+
+ @property
+ def engine(self) -> ArrayEngine:
+ """The synchronous engine backing this array's data path.
+
+ Resolved on first access rather than at construction: the sync data
+ path must never require a running event loop, and an engine that
+ cannot serve this store should only fail when it is actually used.
+ """
+ if self._engine is None:
+ aa = self._async_array
+ self._engine = resolve_sync_engine(
+ self._engine_spec,
+ store=aa.store_path.store,
+ path=aa.store_path.path,
+ metadata=aa.metadata,
+ config=aa.config,
+ )
+ adopt_codec_pipeline(self._engine, aa.codec_pipeline)
+ return self._engine
+
+ def _rebind_engine(self) -> None:
+ """Point the resolved engine at the async array's current metadata.
+
+ Called after an in-place metadata change (`resize`, `append`). A no-op
+ if the engine has never been resolved.
+ """
+ if self._engine is not None:
+ self._engine = self._engine.with_metadata(self._async_array.metadata)
+
+ def _request(self, kind: SelectionKind, selection: Any) -> SelectionRequest:
+ """Package a selection for this array's engine."""
+ return SelectionRequest(
+ kind=kind,
+ selection=selection,
+ shape=self._async_array.metadata.shape,
+ chunk_grid=self._async_array._chunk_grid,
+ )
+
+ def _get_selection(
+ self,
+ request: SelectionRequest,
+ *,
+ prototype: BufferPrototype,
+ out: NDBuffer | None = None,
+ fields: Fields | None = None,
+ ) -> NDArrayLikeOrScalar:
+ return self.engine.read_selection(request, prototype=prototype, out=out, fields=fields)
+
+ def _set_selection(
+ self,
+ request: SelectionRequest,
+ value: npt.ArrayLike,
+ *,
+ prototype: BufferPrototype,
+ fields: Fields | None = None,
+ ) -> None:
+ self.engine.write_selection(request, value, prototype=prototype, fields=fields)
@property
def async_array(self) -> AsyncArray[T_ArrayMetadata]:
@@ -1909,10 +1980,12 @@ def _create(
# runtime
overwrite: bool = False,
config: ArrayConfigLike | None = None,
+ engine: ArrayEngine | AsyncArrayEngine | EngineName | None = None,
) -> Self:
"""Creates a new Array instance from an initialized store.
Deprecated in favor of [`zarr.create_array`][].
"""
+ engine_for_async, engine_for_array = route_sync_engine_arg(engine)
async_array = sync(
AsyncArray._create(
store=store,
@@ -1932,9 +2005,10 @@ def _create(
compressor=compressor,
overwrite=overwrite,
config=config,
+ engine=engine_for_async,
),
)
- return cls(async_array)
+ return cls(async_array, engine_spec=engine_for_array)
@classmethod
def from_dict(
@@ -1971,6 +2045,7 @@ def from_dict(
def open(
cls,
store: StoreLike,
+ engine: ArrayEngine | AsyncArrayEngine | EngineName | None = None,
) -> Self:
"""Opens an existing Array from a store.
@@ -1986,8 +2061,9 @@ def open(
Array
Array opened from the store.
"""
- async_array = sync(AsyncArray.open(store))
- return cls(async_array)
+ engine_for_async, engine_for_array = route_sync_engine_arg(engine)
+ async_array = sync(AsyncArray.open(store, engine=engine_for_async))
+ return cls(async_array, engine_spec=engine_for_array)
@property
def store(self) -> Store:
@@ -2275,7 +2351,7 @@ def with_config(self, config: ArrayConfigLike) -> Self:
-------
A new Array
"""
- return type(self)(self._async_array.with_config(config))
+ return type(self)(self._async_array.with_config(config), engine_spec=self._engine_spec)
@property
def nbytes(self) -> int:
@@ -2875,13 +2951,8 @@ def get_basic_selection(
if prototype is None:
prototype = default_buffer_prototype()
- return sync(
- self.async_array._get_selection(
- BasicIndexer(selection, self.shape, self._chunk_grid),
- out=out,
- fields=fields,
- prototype=prototype,
- )
+ return self._get_selection(
+ self._request("basic", selection), out=out, fields=fields, prototype=prototype
)
def set_basic_selection(
@@ -2984,8 +3055,9 @@ def set_basic_selection(
"""
if prototype is None:
prototype = default_buffer_prototype()
- indexer = BasicIndexer(selection, self.shape, self._chunk_grid)
- sync(self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype))
+ self._set_selection(
+ self._request("basic", selection), value, fields=fields, prototype=prototype
+ )
def get_orthogonal_selection(
self,
@@ -3112,11 +3184,8 @@ def get_orthogonal_selection(
"""
if prototype is None:
prototype = default_buffer_prototype()
- indexer = OrthogonalIndexer(selection, self.shape, self._chunk_grid)
- return sync(
- self.async_array._get_selection(
- indexer=indexer, out=out, fields=fields, prototype=prototype
- )
+ return self._get_selection(
+ self._request("orthogonal", selection), out=out, fields=fields, prototype=prototype
)
def set_orthogonal_selection(
@@ -3230,9 +3299,8 @@ def set_orthogonal_selection(
"""
if prototype is None:
prototype = default_buffer_prototype()
- indexer = OrthogonalIndexer(selection, self.shape, self._chunk_grid)
- return sync(
- self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype)
+ return self._set_selection(
+ self._request("orthogonal", selection), value, fields=fields, prototype=prototype
)
def get_mask_selection(
@@ -3318,11 +3386,8 @@ def get_mask_selection(
if prototype is None:
prototype = default_buffer_prototype()
- indexer = MaskIndexer(mask, self.shape, self._chunk_grid)
- return sync(
- self.async_array._get_selection(
- indexer=indexer, out=out, fields=fields, prototype=prototype
- )
+ return self._get_selection(
+ self._request("mask", mask), out=out, fields=fields, prototype=prototype
)
def set_mask_selection(
@@ -3407,8 +3472,7 @@ def set_mask_selection(
"""
if prototype is None:
prototype = default_buffer_prototype()
- indexer = MaskIndexer(mask, self.shape, self._chunk_grid)
- sync(self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype))
+ self._set_selection(self._request("mask", mask), value, fields=fields, prototype=prototype)
def get_coordinate_selection(
self,
@@ -3495,16 +3559,14 @@ def get_coordinate_selection(
"""
if prototype is None:
prototype = default_buffer_prototype()
- indexer = CoordinateIndexer(selection, self.shape, self._chunk_grid)
- out_array = sync(
- self.async_array._get_selection(
- indexer=indexer, out=out, fields=fields, prototype=prototype
- )
- )
+ request = self._request("coordinate", selection)
+ out_array = self._get_selection(request, out=out, fields=fields, prototype=prototype)
if hasattr(out_array, "shape"):
# restore shape
- out_array = np.array(out_array).reshape(indexer.sel_shape)
+ out_array = np.array(out_array).reshape(
+ cast("CoordinateIndexer", request.indexer).sel_shape
+ )
return out_array
def set_coordinate_selection(
@@ -3587,7 +3649,8 @@ def set_coordinate_selection(
if prototype is None:
prototype = default_buffer_prototype()
# setup indexer
- indexer = CoordinateIndexer(selection, self.shape, self._chunk_grid)
+ request = self._request("coordinate", selection)
+ indexer = cast("CoordinateIndexer", request.indexer)
# handle value - need ndarray-like flatten value
if not is_scalar(value, self.dtype):
@@ -3609,7 +3672,7 @@ def set_coordinate_selection(
f"elements with an array of {value.shape[0]} elements."
)
- sync(self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype))
+ self._set_selection(request, value, fields=fields, prototype=prototype)
def get_block_selection(
self,
@@ -3708,11 +3771,8 @@ def get_block_selection(
"""
if prototype is None:
prototype = default_buffer_prototype()
- indexer = BlockIndexer(selection, self.shape, self._chunk_grid)
- return sync(
- self.async_array._get_selection(
- indexer=indexer, out=out, fields=fields, prototype=prototype
- )
+ return self._get_selection(
+ self._request("block", selection), out=out, fields=fields, prototype=prototype
)
def set_block_selection(
@@ -3808,8 +3868,9 @@ def set_block_selection(
"""
if prototype is None:
prototype = default_buffer_prototype()
- indexer = BlockIndexer(selection, self.shape, self._chunk_grid)
- sync(self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype))
+ self._set_selection(
+ self._request("block", selection), value, fields=fields, prototype=prototype
+ )
@property
def vindex(self) -> VIndex:
@@ -3874,6 +3935,7 @@ def resize(self, new_shape: ShapeLike) -> None:
```
"""
sync(self.async_array.resize(new_shape))
+ self._rebind_engine()
def append(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]:
"""Append `data` to `axis`.
@@ -3909,7 +3971,9 @@ def append(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]:
>>> z.shape
(20000, 2000)
"""
- return sync(self.async_array.append(data, axis=axis))
+ result = sync(self.async_array.append(data, axis=axis))
+ self._rebind_engine()
+ return result
def update_attributes(self, new_attributes: dict[str, JSON]) -> Self:
"""
@@ -3937,7 +4001,7 @@ def update_attributes(self, new_attributes: dict[str, JSON]) -> Self:
overwritten by the new values.
"""
new_array = sync(self.async_array.update_attributes(new_attributes))
- return type(self)(new_array)
+ return type(self)(new_array, engine_spec=self._engine_spec)
def __repr__(self) -> str:
return f""
@@ -4085,6 +4149,7 @@ async def from_array(
storage_options: dict[str, Any] | None = None,
overwrite: bool = False,
config: ArrayConfigLike | None = None,
+ engine: AsyncArrayEngine | EngineName | None = None,
) -> AnyAsyncArray:
"""Create an array from an existing array or array-like.
@@ -4309,6 +4374,7 @@ async def from_array(
dimension_names=dimension_names,
overwrite=overwrite,
config=config_parsed,
+ engine=engine,
)
if write_data:
@@ -4358,6 +4424,7 @@ async def init_array(
dimension_names: DimensionNamesLike = None,
overwrite: bool = False,
config: ArrayConfigLike | None = None,
+ engine: AsyncArrayEngine | EngineName | None = None,
) -> AnyAsyncArray:
"""Create and persist an array metadata document.
@@ -4558,7 +4625,7 @@ async def init_array(
attributes=attributes,
)
- arr = AsyncArray(metadata=meta, store_path=store_path, config=config)
+ arr = AsyncArray(metadata=meta, store_path=store_path, config=config, engine=engine)
await arr._save_metadata(meta, ensure_parents=True)
return arr
@@ -4585,6 +4652,7 @@ async def create_array(
overwrite: bool = False,
config: ArrayConfigLike | None = None,
write_data: bool = True,
+ engine: ArrayEngine | AsyncArrayEngine | EngineName | None = None,
) -> AnyAsyncArray:
"""Create an array.
@@ -4733,6 +4801,7 @@ async def create_array(
storage_options=storage_options,
overwrite=overwrite,
config=config,
+ engine=cast("AsyncArrayEngine | EngineName | None", engine),
)
else:
mode: Literal["a"] = "a"
@@ -4757,6 +4826,7 @@ async def create_array(
dimension_names=dimension_names,
overwrite=overwrite,
config=config,
+ engine=cast("AsyncArrayEngine | EngineName | None", engine),
)
@@ -5773,6 +5843,7 @@ async def _delete_key(key: str) -> None:
# Update metadata and chunk_grid (in place)
object.__setattr__(array, "metadata", new_metadata)
+ object.__setattr__(array, "engine", array.engine.with_metadata(new_metadata))
object.__setattr__(array, "_chunk_grid", new_chunk_grid)
diff --git a/src/zarr/core/engine/__init__.py b/src/zarr/core/engine/__init__.py
new file mode 100644
index 0000000000..25c96fcdab
--- /dev/null
+++ b/src/zarr/core/engine/__init__.py
@@ -0,0 +1,34 @@
+"""Array engine resolution and the built-in engine.
+
+The protocols themselves live in `zarr.abc.engine`.
+"""
+
+from zarr.core.engine._default import (
+ DefaultArrayEngine,
+ DefaultAsyncArrayEngine,
+ DefaultAsyncHierarchyEngine,
+ DefaultHierarchyEngine,
+ adopt_codec_pipeline,
+)
+from zarr.core.engine._resolve import (
+ EngineName,
+ classify_engine_arg,
+ list_engines,
+ resolve_async_engine,
+ resolve_sync_engine,
+ route_sync_engine_arg,
+)
+
+__all__ = [
+ "DefaultArrayEngine",
+ "DefaultAsyncArrayEngine",
+ "DefaultAsyncHierarchyEngine",
+ "DefaultHierarchyEngine",
+ "EngineName",
+ "adopt_codec_pipeline",
+ "classify_engine_arg",
+ "list_engines",
+ "resolve_async_engine",
+ "resolve_sync_engine",
+ "route_sync_engine_arg",
+]
diff --git a/src/zarr/core/engine/_default.py b/src/zarr/core/engine/_default.py
new file mode 100644
index 0000000000..9586774796
--- /dev/null
+++ b/src/zarr/core/engine/_default.py
@@ -0,0 +1,203 @@
+"""The built-in engine: zarr-python's own data path, behind the protocol.
+
+This engine owns no I/O logic. It forwards to the module-level
+`_get_selection` / `_set_selection` in `zarr.core.array` — the same functions
+the non-engine path calls — so the built-in engine is the identity engine by
+construction, and improvements to that path (codec-pipeline changes, chunk-spec
+caching, dtype handling) reach it without being mirrored here.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from zarr.core.array_spec import parse_array_config
+from zarr.core.chunk_grids import ChunkGrid
+from zarr.core.sync import sync
+from zarr.storage._common import StorePath
+
+if TYPE_CHECKING:
+ import numpy.typing as npt
+
+ from zarr.abc.codec import CodecPipeline
+ from zarr.abc.engine import SelectionRequest
+ from zarr.abc.store import Store
+ from zarr.core.array_spec import ArrayConfig
+ from zarr.core.buffer import BufferPrototype, NDArrayLikeOrScalar, NDBuffer
+ from zarr.core.indexing import Fields
+ from zarr.core.metadata import ArrayMetadata
+
+__all__ = [
+ "DefaultArrayEngine",
+ "DefaultAsyncArrayEngine",
+ "DefaultAsyncHierarchyEngine",
+ "DefaultHierarchyEngine",
+ "adopt_codec_pipeline",
+]
+
+
+def adopt_codec_pipeline(engine: object, codec_pipeline: CodecPipeline) -> None:
+ """Give a built-in engine a codec pipeline the caller has already built.
+
+ `Array` and `AsyncArray` each build one, and each resolves its own engine.
+ Letting the engine build another would duplicate that work on every array
+ and would re-emit the codec chain's advisory warnings (e.g. sharding's
+ "disables partial reads"), which are meant to fire once per user-facing
+ chain rather than once per construction.
+
+ A no-op for any other engine, and for a built-in engine that already has a
+ pipeline: an engine that does not use zarr-python's codec pipeline has no
+ use for one.
+ """
+ if isinstance(engine, DefaultArrayEngine):
+ engine = engine._async
+ if isinstance(engine, DefaultAsyncArrayEngine) and engine._pipeline is None:
+ engine._pipeline = codec_pipeline
+
+
+class DefaultAsyncArrayEngine:
+ """Codec-pipeline-backed engine. Any store, Zarr v2 and v3."""
+
+ def __init__(
+ self,
+ store_path: StorePath,
+ metadata: ArrayMetadata,
+ config: ArrayConfig,
+ codec_pipeline: CodecPipeline | None = None,
+ ) -> None:
+ self.store_path = store_path
+ self.metadata = metadata
+ self.config = config
+ self._chunk_grid = ChunkGrid.from_metadata(metadata)
+ self._pipeline = codec_pipeline
+
+ @property
+ def codec_pipeline(self) -> CodecPipeline:
+ """The codec pipeline, built on first use if one was not supplied.
+
+ `AsyncArray` passes in the pipeline it has already built. Building a
+ second one here would not just duplicate the work — `create_codec_pipeline`
+ emits the codec chain's advisory warnings (e.g. sharding's "disables
+ partial reads"), which are meant to fire once per user-facing chain, so
+ a second construction would also double the warnings.
+ """
+ if self._pipeline is None:
+ # Imported lazily: `zarr.core.array` imports the engine package to
+ # resolve `engine=`, while this engine calls back into it.
+ from zarr.core.array import create_codec_pipeline
+
+ self._pipeline = create_codec_pipeline(
+ metadata=self.metadata, store=self.store_path.store
+ )
+ return self._pipeline
+
+ def with_metadata(self, metadata: ArrayMetadata) -> DefaultAsyncArrayEngine:
+ # The pipeline is metadata-derived, so a new one is built for the new
+ # metadata rather than carried over.
+ return DefaultAsyncArrayEngine(
+ store_path=self.store_path, metadata=metadata, config=self.config
+ )
+
+ async def read_selection(
+ self,
+ request: SelectionRequest,
+ *,
+ prototype: BufferPrototype,
+ out: NDBuffer | None = None,
+ fields: Fields | None = None,
+ ) -> NDArrayLikeOrScalar:
+ from zarr.core.array import _get_selection
+
+ return await _get_selection(
+ self.store_path,
+ self.metadata,
+ self.codec_pipeline,
+ self.config,
+ self._chunk_grid,
+ request.indexer,
+ prototype=prototype,
+ out=out,
+ fields=fields,
+ )
+
+ async def write_selection(
+ self,
+ request: SelectionRequest,
+ value: npt.ArrayLike,
+ *,
+ prototype: BufferPrototype,
+ fields: Fields | None = None,
+ ) -> None:
+ from zarr.core.array import _set_selection
+
+ await _set_selection(
+ self.store_path,
+ self.metadata,
+ self.codec_pipeline,
+ self.config,
+ self._chunk_grid,
+ request.indexer,
+ value,
+ prototype=prototype,
+ fields=fields,
+ )
+
+
+class DefaultArrayEngine:
+ """Sync adapter over `DefaultAsyncArrayEngine` via `sync()`."""
+
+ def __init__(self, async_engine: DefaultAsyncArrayEngine) -> None:
+ self._async = async_engine
+
+ def read_selection(
+ self,
+ request: SelectionRequest,
+ *,
+ prototype: BufferPrototype,
+ out: NDBuffer | None = None,
+ fields: Fields | None = None,
+ ) -> NDArrayLikeOrScalar:
+ return sync(
+ self._async.read_selection(request, prototype=prototype, out=out, fields=fields)
+ )
+
+ def write_selection(
+ self,
+ request: SelectionRequest,
+ value: npt.ArrayLike,
+ *,
+ prototype: BufferPrototype,
+ fields: Fields | None = None,
+ ) -> None:
+ sync(self._async.write_selection(request, value, prototype=prototype, fields=fields))
+
+ def with_metadata(self, metadata: ArrayMetadata) -> DefaultArrayEngine:
+ return DefaultArrayEngine(self._async.with_metadata(metadata))
+
+
+class DefaultAsyncHierarchyEngine:
+ """Store-bound factory for built-in async engines."""
+
+ def __init__(self, store: Store) -> None:
+ self.store = store
+
+ def array_engine(
+ self, path: str, metadata: ArrayMetadata, config: ArrayConfig | None = None
+ ) -> DefaultAsyncArrayEngine:
+ return DefaultAsyncArrayEngine(
+ store_path=StorePath(self.store, path),
+ metadata=metadata,
+ config=config if config is not None else parse_array_config(None),
+ )
+
+
+class DefaultHierarchyEngine:
+ """Store-bound factory for built-in sync engines."""
+
+ def __init__(self, store: Store) -> None:
+ self._async = DefaultAsyncHierarchyEngine(store)
+
+ def array_engine(
+ self, path: str, metadata: ArrayMetadata, config: ArrayConfig | None = None
+ ) -> DefaultArrayEngine:
+ return DefaultArrayEngine(self._async.array_engine(path, metadata, config))
diff --git a/src/zarr/core/engine/_resolve.py b/src/zarr/core/engine/_resolve.py
new file mode 100644
index 0000000000..7b2f3bbe0c
--- /dev/null
+++ b/src/zarr/core/engine/_resolve.py
@@ -0,0 +1,223 @@
+"""Resolve an `engine=` argument to a bound array engine."""
+
+from __future__ import annotations
+
+import contextlib
+import inspect
+import weakref
+from typing import TYPE_CHECKING, Literal, get_args
+
+from zarr.core.engine._default import (
+ DefaultAsyncHierarchyEngine,
+ DefaultHierarchyEngine,
+)
+
+if TYPE_CHECKING:
+ from collections.abc import Callable
+
+ from zarr.abc.engine import (
+ ArrayEngine,
+ AsyncArrayEngine,
+ AsyncHierarchyEngine,
+ HierarchyEngine,
+ )
+ from zarr.abc.store import Store
+ from zarr.core.array_spec import ArrayConfig
+ from zarr.core.metadata import ArrayMetadata
+
+__all__ = [
+ "EngineName",
+ "classify_engine_arg",
+ "list_engines",
+ "resolve_async_engine",
+ "resolve_sync_engine",
+ "route_sync_engine_arg",
+]
+
+EngineName = Literal["default", "zarrista"]
+
+
+def list_engines() -> list[str]:
+ """Return the sorted names of the built-in array engines.
+
+ Any of these names can be passed as the `engine=` argument to
+ `zarr.open_array`, `zarr.create_array`, and the other array entry points to
+ select the data-path engine backing the array.
+
+ `"zarrista"` additionally requires the optional `zarrista` package to be
+ installed; without it, resolving that engine raises an `ImportError`.
+ """
+ # `EngineName` is the single source of truth for known engine names --
+ # `_hierarchy_factory` dispatches on exactly these literals.
+ return sorted(get_args(EngineName))
+
+
+def classify_engine_arg(engine: object) -> Literal["name", "sync", "async"]:
+ """Classify an `engine=` argument as a name, a sync instance, or an async instance.
+
+ `None` and `str` values classify as `"name"` -- valid wherever an `engine=`
+ argument is accepted. Any other value must implement `read_selection`: a
+ coroutine function classifies as `"async"` (an `AsyncArrayEngine`), anything
+ else as `"sync"` (an `ArrayEngine`). Objects with no `read_selection` at all
+ raise `TypeError` naming the two protocols.
+ """
+ if engine is None or isinstance(engine, str):
+ return "name"
+ read_selection = getattr(engine, "read_selection", None)
+ if read_selection is None:
+ raise TypeError(
+ f"{engine!r} does not implement the ArrayEngine or AsyncArrayEngine protocol "
+ "(missing a `read_selection` method)"
+ )
+ return "async" if inspect.iscoroutinefunction(read_selection) else "sync"
+
+
+def route_sync_engine_arg(
+ engine: ArrayEngine | AsyncArrayEngine | EngineName | None,
+) -> tuple[AsyncArrayEngine | EngineName | None, ArrayEngine | EngineName | None]:
+ """Route a sync-entry-point `engine=` argument to its two consumers.
+
+ The public sync entry points (`zarr.create_array`, `zarr.open_array`, ...)
+ accept the same broad `engine` type as their async counterparts so their
+ signatures and docstrings match; this function is where the sync-specific
+ rules actually get enforced.
+
+ Returns `(engine_for_async_array, engine_for_array)`. A name (or `None`) is
+ valid for both layers and is returned unchanged in both slots, so sync and
+ async access to the same object use the same engine family. A sync
+ `ArrayEngine` instance is returned only in the second slot -- the wrapped
+ `AsyncArray` keeps its default engine. An `AsyncArrayEngine` instance cannot
+ serve a sync entry point and raises `TypeError`.
+ """
+ kind = classify_engine_arg(engine)
+ if kind == "async":
+ # Fail fast at the API boundary rather than lazily when `Array`
+ # resolves its engine; message kept identical to
+ # `resolve_sync_engine`'s so the error looks the same regardless of
+ # where the wrong-kind instance was actually caught.
+ raise TypeError(
+ "Array requires a synchronous engine (ArrayEngine); got an "
+ f"async engine of type `{type(engine).__name__}`"
+ )
+ if kind == "name":
+ return engine, engine # type: ignore[return-value]
+ # kind == "sync": the instance only serves the sync Array; the inner
+ # AsyncArray keeps its default engine.
+ return None, engine # type: ignore[return-value]
+
+
+# (name, kind, id(store)) -> hierarchy engine; entries evicted automatically once
+# nothing keeps the hierarchy engine itself alive (see `_keepalive` below).
+#
+# Note: a hierarchy engine holds its `store` strongly (it must, to do I/O), so a
+# plain dict keyed by a `weakref.ref`/`weakref.finalize` on the *store* cannot
+# work here -- as long as the hierarchy sits in such a cache it keeps the store
+# alive, so the store's refcount never reaches zero and the finalizer never
+# fires. Using a `WeakValueDictionary` for the hierarchy itself sidesteps that:
+# the cache entry disappears as soon as nothing external holds the hierarchy.
+# `_keepalive` ties the hierarchy's lifetime to the array engines minted from
+# it, so engines resolved for the same store while at least one is still alive
+# share a hierarchy; once all of them (and the store) are unreferenced, both
+# the hierarchy and the cache entry are collected.
+_hierarchy_cache: weakref.WeakValueDictionary[tuple[str, str, int], object] = (
+ weakref.WeakValueDictionary()
+)
+
+
+def _cached_hierarchy(
+ name: str, kind: str, store: Store, factory: Callable[[Store], object]
+) -> object:
+ key = (name, kind, id(store))
+ hierarchy = _hierarchy_cache.get(key)
+ if hierarchy is None:
+ hierarchy = factory(store)
+ _hierarchy_cache[key] = hierarchy
+ return hierarchy
+
+
+def _keepalive(engine: object, hierarchy: object) -> object:
+ """Attach `hierarchy` to `engine` so the hierarchy (and thus the cache entry
+ tracking it) stays alive for as long as `engine` does. Best-effort: engines
+ that forbid arbitrary attributes (e.g. via `__slots__`) simply won't share
+ a cached hierarchy across calls.
+ """
+ with contextlib.suppress(AttributeError):
+ engine._resolve_hierarchy_keepalive = hierarchy # type: ignore[attr-defined]
+ return engine
+
+
+def _hierarchy_factory(name: str, *, sync: bool) -> Callable[[Store], object]:
+ if name == "default":
+ return DefaultHierarchyEngine if sync else DefaultAsyncHierarchyEngine
+ if name == "zarrista":
+ try:
+ from zarr.zarrista import (
+ ZarristaAsyncHierarchyEngine,
+ ZarristaHierarchyEngine,
+ )
+ except ImportError as e:
+ raise ImportError(
+ "engine='zarrista' requires the `zarrista` package; "
+ "install zarr with the `zarrista` extra"
+ ) from e
+ return ZarristaHierarchyEngine if sync else ZarristaAsyncHierarchyEngine
+ raise ValueError(f"unknown engine name {name!r}; expected 'default' or 'zarrista'")
+
+
+def resolve_async_engine(
+ engine: AsyncArrayEngine | EngineName | None,
+ *,
+ store: Store,
+ path: str,
+ metadata: ArrayMetadata,
+ config: ArrayConfig | None = None,
+) -> AsyncArrayEngine:
+ """Resolve an `engine=` argument to a bound `AsyncArrayEngine`.
+
+ `None` and `"default"` produce the built-in codec-pipeline engine;
+ `"zarrista"` lazily imports the `zarrista` package (raising a clear
+ `ImportError` if it is not installed); an existing engine instance is
+ returned unchanged. `config`, when given, is threaded to the engine so it
+ honours the owning array's runtime configuration (e.g. `order`,
+ `read_missing_chunks`); the hierarchy cache is keyed only by store, so
+ engines for arrays with differing configs still share resources.
+ """
+ if engine is None:
+ engine = "default"
+ if isinstance(engine, str):
+ factory = _hierarchy_factory(engine, sync=False)
+ hierarchy: AsyncHierarchyEngine = _cached_hierarchy( # type: ignore[assignment]
+ engine, "async", store, factory
+ )
+ return _keepalive(hierarchy.array_engine(path, metadata, config), hierarchy) # type: ignore[return-value]
+ if classify_engine_arg(engine) == "sync":
+ raise TypeError(
+ "AsyncArray requires an async engine (AsyncArrayEngine); got a "
+ f"synchronous engine of type `{type(engine).__name__}`"
+ )
+ return engine
+
+
+def resolve_sync_engine(
+ engine: ArrayEngine | EngineName | None,
+ *,
+ store: Store,
+ path: str,
+ metadata: ArrayMetadata,
+ config: ArrayConfig | None = None,
+) -> ArrayEngine:
+ """Resolve an `engine=` argument to a bound `ArrayEngine`. See `resolve_async_engine`."""
+ if engine is None:
+ engine = "default"
+ if isinstance(engine, str):
+ factory = _hierarchy_factory(engine, sync=True)
+ hierarchy: HierarchyEngine = _cached_hierarchy( # type: ignore[assignment]
+ engine, "sync", store, factory
+ )
+ return _keepalive(hierarchy.array_engine(path, metadata, config), hierarchy) # type: ignore[return-value]
+ if classify_engine_arg(engine) == "async":
+ raise TypeError(
+ "Array requires a synchronous engine (ArrayEngine); got an "
+ f"async engine of type `{type(engine).__name__}`"
+ )
+ return engine
diff --git a/src/zarr/core/group.py b/src/zarr/core/group.py
index 65f7767a29..1061f41a14 100644
--- a/src/zarr/core/group.py
+++ b/src/zarr/core/group.py
@@ -46,6 +46,7 @@
parse_shapelike,
)
from zarr.core.config import config
+from zarr.core.engine import route_sync_engine_arg
from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata
from zarr.core.metadata.io import save_metadata
from zarr.core.sync import SyncMixin, sync
@@ -73,11 +74,13 @@
)
from typing import Any
+ from zarr.abc.engine import ArrayEngine, AsyncArrayEngine
from zarr.core.array_spec import ArrayConfigLike
from zarr.core.buffer import Buffer, BufferPrototype
from zarr.core.chunk_key_encodings import ChunkKeyEncodingLike
from zarr.core.common import MemoryOrder
from zarr.core.dtype import ZDTypeLike
+ from zarr.core.engine import EngineName
from zarr.types import AnyArray, AnyAsyncArray, ArrayV2, ArrayV3, AsyncArrayV2, AsyncArrayV3
logger = logging.getLogger("zarr.group")
@@ -1095,6 +1098,7 @@ async def create_array(
overwrite: bool = False,
config: ArrayConfigLike | None = None,
write_data: bool = True,
+ engine: ArrayEngine | AsyncArrayEngine | EngineName | None = None,
) -> AnyAsyncArray:
"""Create an array within this group.
@@ -1188,6 +1192,12 @@ async def create_array(
then ``write_data`` determines whether the values in that array-like object should be
written to the Zarr array created by this function. If ``write_data`` is ``False``, then the
array will be left empty.
+ engine : ArrayEngine | AsyncArrayEngine | Literal["default", "zarrista"] | None, optional
+ The data-path engine backing the created array: a name (`"default"`,
+ `"zarrista"`) or a pre-built engine instance. A synchronous
+ `ArrayEngine` instance only makes sense from the sync API; an
+ `AsyncArrayEngine` instance only from the async API; a name works
+ from either. When omitted, the `"default"` behavior is unchanged.
Returns
-------
@@ -1218,6 +1228,7 @@ async def create_array(
overwrite=overwrite,
config=config,
write_data=write_data,
+ engine=engine,
)
async def require_array(
@@ -2677,6 +2688,7 @@ def create_array(
overwrite: bool = False,
config: ArrayConfigLike | None = None,
write_data: bool = True,
+ engine: ArrayEngine | AsyncArrayEngine | EngineName | None = None,
) -> AnyArray:
"""Create an array within this group.
@@ -2772,6 +2784,12 @@ def create_array(
then ``write_data`` determines whether the values in that array-like object should be
written to the Zarr array created by this function. If ``write_data`` is ``False``, then the
array will be left empty.
+ engine : ArrayEngine | AsyncArrayEngine | Literal["default", "zarrista"] | None, optional
+ The data-path engine backing the created array: a name (`"default"`,
+ `"zarrista"`) or a pre-built engine instance. A synchronous
+ `ArrayEngine` instance only makes sense from the sync API; an
+ `AsyncArrayEngine` instance only from the async API; a name works
+ from either. When omitted, the `"default"` behavior is unchanged.
Returns
-------
@@ -2780,6 +2798,7 @@ def create_array(
compressors = _parse_deprecated_compressor(
compressor, compressors, zarr_format=self.metadata.zarr_format
)
+ engine_for_async, engine_for_array = route_sync_engine_arg(engine)
return Array(
self._sync(
self._async_group.create_array(
@@ -2801,8 +2820,10 @@ def create_array(
storage_options=storage_options,
config=config,
write_data=write_data,
+ engine=engine_for_async,
)
- )
+ ),
+ engine_spec=engine_for_array,
)
def require_array(self, name: str, *, shape: ShapeLike, **kwargs: Any) -> AnyArray:
diff --git a/src/zarr/errors.py b/src/zarr/errors.py
index 781bebe534..d662a6849a 100644
--- a/src/zarr/errors.py
+++ b/src/zarr/errors.py
@@ -13,6 +13,7 @@
"NegativeStepError",
"NodeTypeValidationError",
"UnstableSpecificationWarning",
+ "UnsupportedEngineError",
"VindexInvalidSelectionError",
"ZarrDeprecationWarning",
"ZarrFutureWarning",
@@ -155,3 +156,12 @@ class ChunkNotFoundError(BaseZarrError):
"""
Raised when a chunk that was expected to exist in storage was not retrieved successfully.
"""
+
+
+class UnsupportedEngineError(ValueError):
+ """Raised when an array engine cannot serve the requested store or array.
+
+ Examples: a store the engine cannot translate, metadata (e.g. Zarr v2) the
+ engine does not support, or a config setting whose semantics the engine
+ cannot honour.
+ """
diff --git a/src/zarr/zarrista/__init__.py b/src/zarr/zarrista/__init__.py
new file mode 100644
index 0000000000..4873414bd6
--- /dev/null
+++ b/src/zarr/zarrista/__init__.py
@@ -0,0 +1,19 @@
+"""Zarrista-backed array engines for zarr-python.
+
+Requires the `zarrista` package (`pip install zarr[zarrista]` once released;
+currently the git-pinned `zarrista` dependency group).
+"""
+
+from zarr.zarrista._engine import (
+ ZarristaAsyncEngine,
+ ZarristaAsyncHierarchyEngine,
+ ZarristaEngine,
+ ZarristaHierarchyEngine,
+)
+
+__all__ = [
+ "ZarristaAsyncEngine",
+ "ZarristaAsyncHierarchyEngine",
+ "ZarristaEngine",
+ "ZarristaHierarchyEngine",
+]
diff --git a/src/zarr/zarrista/_engine.py b/src/zarr/zarrista/_engine.py
new file mode 100644
index 0000000000..bcf283db39
--- /dev/null
+++ b/src/zarr/zarrista/_engine.py
@@ -0,0 +1,378 @@
+"""Zarrista-backed array engines."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING, Any, cast
+
+import numpy as np
+
+from zarr.errors import UnsupportedEngineError
+from zarr.zarrista._resolve import dense_forward_box, lazy_view, scatter_points
+from zarr.zarrista._source import (
+ AsyncZarristaSource,
+ ZarristaSource,
+ await_on,
+ run_off_loop,
+ tensor_to_numpy,
+)
+from zarr.zarrista._translate import translate_store_async, translate_store_sync
+
+if TYPE_CHECKING:
+ import numpy.typing as npt
+ from zarr_metadata import ZarrV3ArrayMetadataJSON
+
+ from zarr.abc.engine import SelectionRequest
+ from zarr.abc.store import Store
+ from zarr.core.array_spec import ArrayConfig
+ from zarr.core.buffer import BufferPrototype, NDArrayLikeOrScalar, NDBuffer
+ from zarr.core.indexing import Fields
+ from zarr.core.metadata import ArrayMetadata
+
+__all__ = [
+ "ZarristaAsyncEngine",
+ "ZarristaAsyncHierarchyEngine",
+ "ZarristaEngine",
+ "ZarristaHierarchyEngine",
+]
+
+
+def _require_v3(metadata: ArrayMetadata) -> ZarrV3ArrayMetadataJSON:
+ if metadata.zarr_format != 3:
+ raise UnsupportedEngineError(
+ "the zarrista engine supports Zarr v3 only; this array is "
+ f"format v{metadata.zarr_format}"
+ )
+ # `metadata.to_dict()` is a plain `dict[str, JSON]`; zarrista's stubs type
+ # `from_metadata`'s argument as the `ZarrV3ArrayMetadataJSON` TypedDict,
+ # which the dict's runtime shape matches by construction.
+ return cast("ZarrV3ArrayMetadataJSON", metadata.to_dict())
+
+
+def _reject_unenforceable_config(config: ArrayConfig | None) -> None:
+ """Reject an `ArrayConfig` the zarrista engine cannot honour.
+
+ Most of `ArrayConfig` only affects the in-memory layout of the result,
+ which this engine normalizes anyway. Two fields change *semantics*, in
+ opposite directions, and per the project's fail-loud rule are refused
+ rather than silently downgraded:
+
+ - `read_missing_chunks=False` asks for a `ChunkNotFoundError` on a missing
+ chunk. zarrista fills it with the fill value and cannot be made to raise.
+ - `write_empty_chunks=True` asks for all-fill chunks to be written out.
+ zarrista elides them and cannot be made to keep them.
+
+ Both rejected values are the non-default, so an array using zarr's
+ defaults is served.
+ """
+ if config is None:
+ return
+ if not config.read_missing_chunks:
+ raise UnsupportedEngineError(
+ "the zarrista engine cannot enforce read_missing_chunks=False "
+ "(it fills missing chunks with the fill value instead of raising); "
+ "use the default engine to enforce this setting"
+ )
+ if config.write_empty_chunks:
+ raise UnsupportedEngineError(
+ "the zarrista engine cannot enforce write_empty_chunks=True "
+ "(it elides chunks that hold only the fill value); "
+ "use the default engine to enforce this setting"
+ )
+
+
+def _as_contiguous(array: npt.NDArray[Any]) -> npt.NDArray[Any]:
+ """Return `array` as memory zarrista will accept.
+
+ `np.ascontiguousarray` is not enough on its own: a broadcast array (which
+ is what a scalar write widens to) carries 0-strides, and NumPy still
+ reports it as C-contiguous when it holds at most one element per axis. The
+ buffer zarrista then sees is not the dense block it expects, so force a
+ real copy whenever a stride is 0.
+ """
+ if any(stride == 0 for stride in array.strides):
+ return np.ascontiguousarray(array.copy())
+ return np.ascontiguousarray(array)
+
+
+def _reject_fields(fields: Fields | None) -> None:
+ # Truthiness, not `is not None`: `pop_fields` returns an empty list, not
+ # `None`, when a tuple selection names no fields.
+ if fields:
+ raise UnsupportedEngineError(
+ "the zarrista engine does not support structured-dtype field selection; "
+ "use the default engine"
+ )
+
+
+def _finish_read(
+ result: npt.NDArray[Any],
+ request: SelectionRequest,
+ out: NDBuffer | None,
+) -> NDArrayLikeOrScalar:
+ """Match zarr-python's result conventions for this selection kind."""
+ # A coordinate or mask selection is flat at this boundary: zarr's own
+ # `_get_selection` returns the flattened points and the caller reshapes to
+ # `sel_shape`. `LazyArray` returns the un-flattened broadcast shape, whose
+ # C-order ravel is the same point order.
+ if request.kind in ("coordinate", "mask"):
+ result = result.reshape(-1)
+ if out is not None:
+ # `NDArrayLike` only types slice-key indexing, so the Ellipsis fill is
+ # written through an `Any`-typed local.
+ out_array: Any = out.as_ndarray_like()
+ out_array[...] = result
+ return cast("NDArrayLikeOrScalar", out_array)
+ # A rank-0 basic selection is a scalar, as in `_get_selection`.
+ if request.kind == "basic" and result.shape == ():
+ return cast("NDArrayLikeOrScalar", result[()])
+ return result
+
+
+class _ZarristaEngineBase:
+ """Shared selection resolution for the sync and async zarrista engines."""
+
+ _metadata: ArrayMetadata
+ _config: ArrayConfig | None
+
+ def _chunk_shape(self, chunk_grid: Any) -> tuple[int, ...] | None:
+ """The partition size for reads: the outer chunk shape, if regular."""
+ if chunk_grid.is_regular:
+ return cast("tuple[int, ...]", tuple(chunk_grid.chunk_shape))
+ return None
+
+ def _read(
+ self,
+ source: Any,
+ request: SelectionRequest,
+ *,
+ out: NDBuffer | None,
+ fields: Fields | None,
+ ) -> NDArrayLikeOrScalar:
+ _reject_fields(fields)
+ view = lazy_view(source, request, self._chunk_shape(request.chunk_grid))
+ return _finish_read(np.asarray(view.result()), request, out)
+
+ def _write(
+ self,
+ source: Any,
+ arr: Any,
+ request: SelectionRequest,
+ value: npt.ArrayLike,
+ *,
+ fields: Fields | None,
+ store_subset: Any,
+ read_subset: Any,
+ ) -> None:
+ _reject_fields(fields)
+ dtype = self._metadata.dtype.to_native_dtype()
+ view = lazy_view(source, request, None)
+ transform = view.transform
+ # `_set_selection` accepts scalars and broadcastable values, so widen
+ # to the selection's shape before laying anything out.
+ values = np.broadcast_to(np.asarray(value, dtype=dtype), transform.domain.shape)
+
+ box = dense_forward_box(transform)
+ if box is not None:
+ key, ndim_shape = box
+ store_subset(key, values.reshape(ndim_shape))
+ return
+
+ scatter_points(
+ read_subset,
+ store_subset,
+ transform,
+ values,
+ request.chunk_grid,
+ request.shape,
+ )
+
+
+class ZarristaEngine(_ZarristaEngineBase):
+ """Sync engine over `zarrista.Array`. No event loop involved."""
+
+ def __init__(
+ self, zarrista_array: Any, metadata: ArrayMetadata, config: ArrayConfig | None
+ ) -> None:
+ self._arr = zarrista_array
+ self._metadata = metadata
+ self._config = config
+
+ def with_metadata(self, metadata: ArrayMetadata) -> ZarristaEngine:
+ import zarrista
+
+ return ZarristaEngine(
+ zarrista.Array.from_metadata(_require_v3(metadata), self._arr.storage, self._arr.path),
+ metadata,
+ self._config,
+ )
+
+ def _source(self) -> ZarristaSource:
+ return ZarristaSource(self._arr, self._metadata.dtype.to_native_dtype())
+
+ def read_selection(
+ self,
+ request: SelectionRequest,
+ *,
+ prototype: BufferPrototype,
+ out: NDBuffer | None = None,
+ fields: Fields | None = None,
+ ) -> NDArrayLikeOrScalar:
+ return self._read(self._source(), request, out=out, fields=fields)
+
+ def write_selection(
+ self,
+ request: SelectionRequest,
+ value: npt.ArrayLike,
+ *,
+ prototype: BufferPrototype,
+ fields: Fields | None = None,
+ ) -> None:
+ arr = self._arr
+
+ def read_subset(key: Any) -> npt.NDArray[Any]:
+ # `np.array`, not `asarray`: zarrista hands back read-only
+ # Rust-owned memory, and the scatter patches this in place.
+ return np.array(tensor_to_numpy(arr.retrieve_array_subset(key)))
+
+ def store_subset(key: Any, data: npt.NDArray[Any]) -> None:
+ arr.store_array_subset(key, _as_contiguous(data))
+
+ self._write(
+ self._source(),
+ arr,
+ request,
+ value,
+ fields=fields,
+ store_subset=store_subset,
+ read_subset=read_subset,
+ )
+
+
+class ZarristaAsyncEngine(_ZarristaEngineBase):
+ """Async engine over `zarrista.AsyncArray`.
+
+ Store translation and construction of the underlying `zarrista.AsyncArray`
+ are deferred to the first read or write rather than done at construction.
+ `AsyncArray.__init__` always eagerly resolves *an* async engine — even for
+ a plain sync `Array`, which never touches it, since only the sync engine is
+ lazily resolved. Without this deferral, `engine="zarrista"` over a
+ sync-only store (e.g. `LocalStore`) would raise `UnsupportedEngineError`
+ just from opening the array, even when only ever accessed synchronously.
+ """
+
+ def __init__(
+ self, store: Store, path: str, metadata: ArrayMetadata, config: ArrayConfig | None
+ ) -> None:
+ self._store = store
+ self._path = path
+ self._metadata = metadata
+ self._config = config
+ self._arr: Any | None = None
+
+ def with_metadata(self, metadata: ArrayMetadata) -> ZarristaAsyncEngine:
+ _require_v3(metadata)
+ return ZarristaAsyncEngine(self._store, self._path, metadata, self._config)
+
+ def _ensure_arr(self) -> Any:
+ if self._arr is None:
+ import zarrista
+
+ self._arr = zarrista.AsyncArray.from_metadata(
+ _require_v3(self._metadata), translate_store_async(self._store), self._path
+ )
+ return self._arr
+
+ def _source(self, loop: asyncio.AbstractEventLoop) -> AsyncZarristaSource:
+ return AsyncZarristaSource(self._ensure_arr(), self._metadata.dtype.to_native_dtype(), loop)
+
+ async def read_selection(
+ self,
+ request: SelectionRequest,
+ *,
+ prototype: BufferPrototype,
+ out: NDBuffer | None = None,
+ fields: Fields | None = None,
+ ) -> NDArrayLikeOrScalar:
+ loop = asyncio.get_running_loop()
+ source = self._source(loop)
+ # `LazyArray` is synchronous and issues its reads back to this loop, so
+ # it must not resolve on the loop thread itself.
+ return await run_off_loop(lambda: self._read(source, request, out=out, fields=fields))
+
+ async def write_selection(
+ self,
+ request: SelectionRequest,
+ value: npt.ArrayLike,
+ *,
+ prototype: BufferPrototype,
+ fields: Fields | None = None,
+ ) -> None:
+ loop = asyncio.get_running_loop()
+ arr = self._ensure_arr()
+ source = self._source(loop)
+
+ def read_subset(key: Any) -> npt.NDArray[Any]:
+ return np.array(tensor_to_numpy(await_on(loop, lambda: arr.retrieve_array_subset(key))))
+
+ def store_subset(key: Any, data: npt.NDArray[Any]) -> None:
+ await_on(loop, lambda: arr.store_array_subset(key, _as_contiguous(data)))
+
+ await run_off_loop(
+ lambda: self._write(
+ source,
+ arr,
+ request,
+ value,
+ fields=fields,
+ store_subset=store_subset,
+ read_subset=read_subset,
+ )
+ )
+
+
+class ZarristaHierarchyEngine:
+ """Store-bound factory for sync zarrista engines (translates the store once)."""
+
+ def __init__(self, store: Store) -> None:
+ self._zarr_store = store
+ self._zstore = translate_store_sync(store)
+
+ def array_engine(
+ self, path: str, metadata: ArrayMetadata, config: ArrayConfig | None = None
+ ) -> ZarristaEngine:
+ """Mint a sync array engine bound to `path`/`metadata`."""
+ import zarrista
+
+ _reject_unenforceable_config(config)
+ return ZarristaEngine(
+ zarrista.Array.from_metadata(
+ _require_v3(metadata), self._zstore, "/" + path.strip("/")
+ ),
+ metadata,
+ config,
+ )
+
+
+class ZarristaAsyncHierarchyEngine:
+ """Store-bound factory for async zarrista engines.
+
+ Unlike `ZarristaHierarchyEngine`, this does *not* translate the store at
+ construction time — see `ZarristaAsyncEngine` for why.
+ """
+
+ def __init__(self, store: Store) -> None:
+ self._zarr_store = store
+
+ def array_engine(
+ self, path: str, metadata: ArrayMetadata, config: ArrayConfig | None = None
+ ) -> ZarristaAsyncEngine:
+ """Mint an async array engine bound to `path`/`metadata`.
+
+ `metadata` is validated as Zarr v3 eagerly (cheap, and lets an
+ unsupported-format error surface immediately); the store itself is only
+ translated lazily, on first I/O.
+ """
+ _reject_unenforceable_config(config)
+ _require_v3(metadata)
+ return ZarristaAsyncEngine(self._zarr_store, "/" + path.strip("/"), metadata, config)
diff --git a/src/zarr/zarrista/_resolve.py b/src/zarr/zarrista/_resolve.py
new file mode 100644
index 0000000000..a3b54e0376
--- /dev/null
+++ b/src/zarr/zarrista/_resolve.py
@@ -0,0 +1,198 @@
+"""Resolve a `SelectionRequest` against a zarrista array.
+
+Reads go through `zarr_indexing.LazyArray`, which grafts the whole NumPy
+indexing dialect onto a backend that natively offers only step-1 boxes.
+
+Writes have no `LazyArray` equivalent — it describes reads only — so they are
+resolved here, in two tiers:
+
+- A *dense forward box* (the shape of an ordinary `arr[a:b, c:d] = v`) is one
+ `store_array_subset` call. zarrista does any partial-chunk read-modify-write
+ internally, in Rust.
+- Anything else — strided, reversed, orthogonal, vectorized, masked — is
+ scattered pointwise: the transform is evaluated over its whole domain to get
+ one storage coordinate per value, the points are grouped by chunk, and each
+ touched chunk is read, patched, and written back once. Grouping by chunk is
+ what keeps this from degenerating into the hull rewrite that a bounding-box
+ strategy would perform (`oindex[[0, 999999]]` touches two chunks, not a
+ million rows).
+
+Both tiers take their transform from the `LazyArray` view rather than building
+one directly, because `LazyArray` re-bases every view's domain to origin 0.
+A transform built straight from `IndexTransform.from_shape(shape)[1:5]` has
+domain origin 1, and scattering a zero-origin value buffer through it would
+silently write to the wrong offsets.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+from zarr_indexing import LazyArray, unit_step_reader
+from zarr_indexing.output_map import ConstantMap, DimensionMap
+
+if TYPE_CHECKING:
+ import numpy.typing as npt
+ from zarr_indexing import IndexTransform
+
+ from zarr.abc.engine import SelectionRequest
+ from zarr.core.chunk_grids import ChunkGrid
+
+__all__ = ["dense_forward_box", "lazy_view", "scatter_points"]
+
+
+def lazy_view(source: Any, request: SelectionRequest, chunk_shape: tuple[int, ...] | None) -> Any:
+ """Build the `LazyArray` view `request` selects from `source`.
+
+ `chunk_shape`, when given, partitions the read so a fancy or strided
+ selection is covered by per-chunk boxes rather than one hull read.
+ """
+ lazy = LazyArray(source).with_reader(unit_step_reader)
+ if chunk_shape is not None:
+ lazy = lazy.with_parts(list(chunk_shape))
+ kind = request.kind
+ selection = request.selection
+ if kind == "basic":
+ return lazy.lazy[selection]
+ if kind == "orthogonal":
+ return lazy.lazy.oindex[selection]
+ if kind in ("coordinate", "mask"):
+ return lazy.lazy.vindex[selection]
+ if kind == "block":
+ # A block selection is per-axis slices once expanded against the grid;
+ # the indexer has already done that expansion.
+ return lazy.lazy[_block_slices(request)]
+ raise NotImplementedError(f"the zarrista engine cannot serve a {kind!r} selection")
+
+
+def _block_slices(request: SelectionRequest) -> tuple[slice, ...]:
+ """The element-space slices a block selection expands to."""
+ return tuple(
+ slice(d.start, d.stop, d.step) # BlockIndexer yields only SliceDimIndexers
+ for d in request.indexer.dim_indexers # type: ignore[attr-defined]
+ )
+
+
+def dense_forward_box(
+ transform: IndexTransform,
+) -> tuple[tuple[slice, ...], tuple[int, ...]] | None:
+ """The step-1 box `transform` selects, or `None` if it is not one.
+
+ Returns `(key, ndim_shape)` where `key` is a per-storage-axis step-1 slice
+ and `ndim_shape` is the ndim-preserving shape to reshape a value to (an
+ axis the selection dropped with an integer has extent 1). `None` means the
+ selection strides, reverses, transposes, broadcasts, or gathers, and so
+ must be scattered instead.
+ """
+ key: list[slice] = []
+ shape: list[int] = []
+ previous = -1
+ domain = transform.domain
+ for m in transform.output:
+ if isinstance(m, ConstantMap):
+ key.append(slice(m.offset, m.offset + 1))
+ shape.append(1)
+ elif isinstance(m, DimensionMap):
+ # stride 1 only: a reversal (-1) or a stride selects a sublattice
+ # of its hull, so the hull is not what the value fills. Requiring
+ # ascending input dimensions rejects a transposing selection,
+ # whose value layout does not match the box.
+ if m.stride != 1 or m.input_dimension <= previous:
+ return None
+ previous = m.input_dimension
+ d = m.input_dimension
+ lo = m.offset + domain.inclusive_min[d]
+ hi = m.offset + domain.exclusive_max[d]
+ key.append(slice(lo, hi))
+ shape.append(hi - lo)
+ else:
+ return None
+ # Every domain axis must be consumed exactly once; otherwise the value is
+ # broadcast across an axis rather than laid into a box.
+ dims = sorted(m.input_dimension for m in transform.output if isinstance(m, DimensionMap))
+ if dims != list(range(transform.input_rank)):
+ return None
+ return tuple(key), tuple(shape)
+
+
+def _domain_points(transform: IndexTransform) -> npt.NDArray[np.intp]:
+ """Every point of `transform`'s domain, in C order, as `(n, input_rank)`."""
+ domain = transform.domain
+ axes = [
+ np.arange(lo, hi) for lo, hi in zip(domain.inclusive_min, domain.exclusive_max, strict=True)
+ ]
+ if not axes:
+ return np.zeros((1, 0), dtype=np.intp)
+ grid = np.meshgrid(*axes, indexing="ij")
+ return np.stack(grid, axis=-1).reshape(-1, len(axes)).astype(np.intp, copy=False)
+
+
+def _chunk_indices(storage: npt.NDArray[np.intp], chunk_grid: ChunkGrid) -> npt.NDArray[np.intp]:
+ """Map storage coordinates to chunk-grid coordinates, per axis."""
+ dims = chunk_grid._dimensions
+ return np.stack(
+ [dim.indices_to_chunks(storage[:, axis]) for axis, dim in enumerate(dims)],
+ axis=-1,
+ )
+
+
+def scatter_points(
+ read_chunk: Any,
+ write_chunk: Any,
+ transform: IndexTransform,
+ value: npt.NDArray[Any],
+ chunk_grid: ChunkGrid,
+ shape: tuple[int, ...],
+) -> None:
+ """Write `value` through `transform`, one touched chunk at a time.
+
+ `read_chunk(key)` returns the decoded box at `key` as a writable, C-order
+ NumPy array; `write_chunk(key, data)` stores it back.
+
+ Duplicate coordinates (legal in a vectorized write) resolve last-wins,
+ matching `numpy.ndarray.__setitem__`, because the points keep their
+ original order within a chunk and NumPy's fancy assignment is last-wins.
+ """
+ points = _domain_points(transform)
+ if points.size == 0 and transform.input_rank > 0:
+ return
+ storage = np.asarray(transform.apply_many(points), dtype=np.intp)
+ values = np.asarray(value).reshape(-1)
+ if values.size != storage.shape[0]:
+ raise ValueError(
+ f"value has {values.size} elements but the selection covers {storage.shape[0]}"
+ )
+ if storage.shape[0] == 0:
+ return
+
+ chunk_ix = _chunk_indices(storage, chunk_grid)
+ # Group points by chunk. `lexsort` on the reversed key columns sorts by the
+ # leading axis first; a stable sort keeps each chunk's points in the user's
+ # original order, which is what makes last-wins match NumPy.
+ order = np.lexsort(tuple(chunk_ix[:, axis] for axis in range(chunk_ix.shape[1] - 1, -1, -1)))
+ storage, values, chunk_ix = storage[order], values[order], chunk_ix[order]
+
+ boundaries = np.flatnonzero(np.any(np.diff(chunk_ix, axis=0) != 0, axis=1)) + 1
+ starts = np.concatenate(([0], boundaries))
+ ends = np.concatenate((boundaries, [len(storage)]))
+
+ dims = chunk_grid._dimensions
+ for start, end in zip(starts, ends, strict=True):
+ coords = chunk_ix[start]
+ origin = np.array(
+ [dim.chunk_offset(int(c)) for dim, c in zip(dims, coords, strict=True)],
+ dtype=np.intp,
+ )
+ extent = np.array(
+ [dim.data_size(int(c)) for dim, c in zip(dims, coords, strict=True)],
+ dtype=np.intp,
+ )
+ # `data_size` is the chunk's *clipped* extent, so an edge chunk is read
+ # and written at its valid size rather than its nominal one.
+ stop = np.minimum(origin + extent, np.array(shape, dtype=np.intp))
+ key = tuple(slice(int(a), int(b)) for a, b in zip(origin, stop, strict=True))
+ buffer = read_chunk(key)
+ local = storage[start:end] - origin
+ buffer[tuple(local.T)] = values[start:end]
+ write_chunk(key, buffer)
diff --git a/src/zarr/zarrista/_source.py b/src/zarr/zarrista/_source.py
new file mode 100644
index 0000000000..8cfa2dbfc5
--- /dev/null
+++ b/src/zarr/zarrista/_source.py
@@ -0,0 +1,132 @@
+"""Adapt a `zarrista.Array` to what `zarr_indexing.LazyArray` reads from.
+
+`zarrista.Array.__getitem__` accepts exactly one selection dialect: integers,
+step-1 slices, and `Ellipsis`, ndim-preserving. That is precisely the contract
+`zarr_indexing.UnitStepReader` targets — it decomposes any transform into the
+smallest enclosing ascending unit-step slab plus a residual applied in memory —
+so the two compose with no translation beyond `shape`/`dtype` plumbing.
+
+Wrapping the array this way is what gives the engine the full NumPy dialect
+(orthogonal, vectorized, masks, negative steps, composition) over a backend
+that natively offers only boxes.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING, Any, cast
+
+if TYPE_CHECKING:
+ from collections.abc import Callable
+
+ import numpy as np
+ import numpy.typing as npt
+
+__all__ = [
+ "AsyncZarristaSource",
+ "ZarristaSource",
+ "await_on",
+ "run_off_loop",
+ "tensor_to_numpy",
+]
+
+
+def tensor_to_numpy(decoded: Any) -> npt.NDArray[Any]:
+ """Convert a zarrista decoded tensor to a NumPy array.
+
+ `zarrista.Tensor` is a union of four layouts. `FixedLengthTensor` and
+ `VariableLengthTensor` both export via `to_numpy()`. The two `Optional*`
+ layouts carry a validity mask and convert to `numpy.ma.MaskedArray`, which
+ has no zarr-python equivalent, so they are refused rather than silently
+ losing the mask.
+ """
+ type_name = type(decoded).__name__
+ if type_name in ("OptionalFixedLengthTensor", "OptionalVariableLengthTensor"):
+ raise NotImplementedError(
+ f"zarrista returned a {type_name}; masked layouts have no zarr-python "
+ "equivalent, so this array cannot be served by the zarrista engine"
+ )
+ return cast("npt.NDArray[Any]", decoded.to_numpy())
+
+
+class ZarristaSource:
+ """A `zarrista.Array` behind the narrow surface `LazyArray` reads from."""
+
+ __slots__ = ("_arr", "_dtype")
+
+ def __init__(self, zarrista_array: Any, dtype: np.dtype[Any]) -> None:
+ self._arr = zarrista_array
+ # Taken from zarr's own metadata rather than `zarrista.Array.dtype`,
+ # which is a zarrista `DataType`, not a NumPy one.
+ self._dtype = dtype
+
+ @property
+ def shape(self) -> tuple[int, ...]:
+ # zarrista reports shape as a list.
+ return tuple(self._arr.shape)
+
+ @property
+ def dtype(self) -> np.dtype[Any]:
+ return self._dtype
+
+ def __getitem__(self, key: Any) -> npt.NDArray[Any]:
+ return tensor_to_numpy(self._arr.retrieve_array_subset(key))
+
+
+class AsyncZarristaSource:
+ """`ZarristaSource` over an `AsyncArray`, driven from a worker thread.
+
+ `LazyArray` is synchronous, so the async engine resolves a selection on a
+ worker thread (see `run_off_loop`). Each read is handed back to the event
+ loop that owns the `zarrista.AsyncArray` with
+ `asyncio.run_coroutine_threadsafe`, and the worker blocks on the result.
+
+ Going through `zarr.core.sync.sync()` instead does not work: it would be
+ called from within a coroutine already running on that loop, and zarrista's
+ pyo3 futures are bound to the loop that awaits them, so the read would
+ either deadlock or fail with a cross-loop future error.
+ """
+
+ __slots__ = ("_arr", "_dtype", "_loop")
+
+ def __init__(
+ self, zarrista_array: Any, dtype: np.dtype[Any], loop: asyncio.AbstractEventLoop
+ ) -> None:
+ self._arr = zarrista_array
+ self._dtype = dtype
+ self._loop = loop
+
+ @property
+ def shape(self) -> tuple[int, ...]:
+ return tuple(self._arr.shape)
+
+ @property
+ def dtype(self) -> np.dtype[Any]:
+ return self._dtype
+
+ def __getitem__(self, key: Any) -> npt.NDArray[Any]:
+ return tensor_to_numpy(await_on(self._loop, lambda: self._arr.retrieve_array_subset(key)))
+
+
+def await_on(loop: asyncio.AbstractEventLoop, make_awaitable: Callable[[], Any]) -> Any:
+ """Call `make_awaitable` on `loop` and block the worker thread on its result.
+
+ A *factory*, not an awaitable: zarrista's pyo3 futures bind to the running
+ loop at the moment they are created, so calling
+ `arr.retrieve_array_subset(...)` on the worker thread raises "no running
+ event loop". Deferring the call into a coroutine that `loop` runs means the
+ future is both created and awaited there.
+ """
+
+ async def call() -> Any:
+ return await make_awaitable()
+
+ return asyncio.run_coroutine_threadsafe(call(), loop).result()
+
+
+async def run_off_loop[T](fn: Callable[[], T]) -> T:
+ """Run a blocking `LazyArray` resolution on a worker thread.
+
+ The caller's loop stays free to service the reads `fn` issues back to it.
+ """
+ return await asyncio.to_thread(fn)
diff --git a/src/zarr/zarrista/_translate.py b/src/zarr/zarrista/_translate.py
new file mode 100644
index 0000000000..f1bcfe1c4e
--- /dev/null
+++ b/src/zarr/zarrista/_translate.py
@@ -0,0 +1,46 @@
+"""Translate zarr-python stores into stores Zarrista can consume."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from zarr.errors import UnsupportedEngineError
+from zarr.storage import LocalStore
+
+if TYPE_CHECKING:
+ from zarr.abc.store import Store
+
+_SYNC_SUPPORTED = "LocalStore"
+_ASYNC_SUPPORTED = "zarr.storage.ObjectStore (obstore-backed) or an icechunk store"
+
+
+def translate_store_sync(store: Store) -> Any:
+ """zarr store -> zarrista sync store (`zarrista.store.FilesystemStore`)."""
+ import zarrista
+
+ if isinstance(store, LocalStore):
+ return zarrista.store.FilesystemStore(store.root)
+ raise UnsupportedEngineError(
+ f"the zarrista sync engine cannot serve a {type(store).__name__}; "
+ f"supported: {_SYNC_SUPPORTED}. Note: zarr's MemoryStore lives in the "
+ "Python process and cannot be shared with the Rust extension."
+ )
+
+
+def translate_store_async(store: Store) -> Any:
+ """zarr store -> zarrista async store (obstore `ObjectStore` or icechunk `Session`)."""
+ from zarr.storage import ObjectStore
+
+ if isinstance(store, ObjectStore):
+ return store.store # the underlying obstore instance
+ try:
+ from icechunk import IcechunkStore
+
+ if isinstance(store, IcechunkStore):
+ return store.session
+ except ImportError:
+ pass
+ raise UnsupportedEngineError(
+ f"the zarrista async engine cannot serve a {type(store).__name__}; "
+ f"supported: {_ASYNC_SUPPORTED}."
+ )
diff --git a/tests/engine/__init__.py b/tests/engine/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/engine/test_asyncarray_wiring.py b/tests/engine/test_asyncarray_wiring.py
new file mode 100644
index 0000000000..90dd04c7ba
--- /dev/null
+++ b/tests/engine/test_asyncarray_wiring.py
@@ -0,0 +1,147 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, cast
+
+import numpy as np
+
+import zarr
+from zarr.core.engine import DefaultAsyncArrayEngine
+from zarr.storage import MemoryStore
+
+if TYPE_CHECKING:
+ import numpy.typing as npt
+
+ from zarr.abc.engine import SelectionRequest
+ from zarr.core.buffer import BufferPrototype, NDArrayLikeOrScalar, NDBuffer
+ from zarr.core.indexing import BasicSelection, Fields
+ from zarr.core.metadata import ArrayMetadata
+
+
+class _SpyEngine:
+ """Wraps a real engine, recording the requests it is handed."""
+
+ def __init__(self, inner: DefaultAsyncArrayEngine) -> None:
+ self.inner = inner
+ self.reads: list[SelectionRequest] = []
+ self.writes: list[SelectionRequest] = []
+
+ async def read_selection(
+ self,
+ request: SelectionRequest,
+ *,
+ prototype: BufferPrototype,
+ out: NDBuffer | None = None,
+ fields: Fields | None = None,
+ ) -> NDArrayLikeOrScalar:
+ self.reads.append(request)
+ return await self.inner.read_selection(request, prototype=prototype, out=out, fields=fields)
+
+ async def write_selection(
+ self,
+ request: SelectionRequest,
+ value: npt.ArrayLike,
+ *,
+ prototype: BufferPrototype,
+ fields: Fields | None = None,
+ ) -> None:
+ self.writes.append(request)
+ return await self.inner.write_selection(request, value, prototype=prototype, fields=fields)
+
+ def with_metadata(self, metadata: ArrayMetadata) -> _SpyEngine:
+ return _SpyEngine(self.inner.with_metadata(metadata))
+
+
+async def test_asyncarray_routes_io_through_engine() -> None:
+ # NOTE: the async variant of the spy test. `Array` (the sync facade)
+ # resolves and calls its own sync engine (see tests/engine/test_sync_path.py);
+ # this test exercises the async `AsyncArray` methods directly, via its
+ # separate async engine.
+ z = zarr.create_array(MemoryStore(), shape=(10,), chunks=(3,), dtype="int16")
+ aa = z.async_array
+ spy = _SpyEngine(
+ DefaultAsyncArrayEngine(store_path=aa.store_path, metadata=aa.metadata, config=aa.config)
+ )
+ object.__setattr__(aa, "engine", spy)
+
+ await aa.setitem(slice(2, 8), np.arange(6, dtype="int16"))
+ data = await aa.getitem(slice(2, 8))
+
+ for request in (*spy.writes, *spy.reads):
+ # the selection reaches the engine as the user wrote it, tagged with
+ # its dialect and the context needed to resolve it
+ assert request.kind == "basic"
+ assert request.selection == slice(2, 8)
+ assert request.shape == (10,)
+ assert request.chunk_grid.chunk_shape == (3,)
+ assert len(spy.writes) == len(spy.reads) == 1
+ np.testing.assert_array_equal(np.asarray(data), np.arange(6, dtype="int16"))
+
+
+async def test_asyncarray_tags_each_selection_with_its_dialect() -> None:
+ # Every accessor has to name the dialect its selection is written in --
+ # the engine has no other way to tell an orthogonal tuple from a coordinate
+ # one.
+ z = zarr.create_array(MemoryStore(), shape=(8, 8), chunks=(4, 4), dtype="int16")
+ aa = z.async_array
+ spy = _SpyEngine(
+ DefaultAsyncArrayEngine(store_path=aa.store_path, metadata=aa.metadata, config=aa.config)
+ )
+ object.__setattr__(aa, "engine", spy)
+
+ await aa.getitem((slice(None), 1))
+ await aa.get_orthogonal_selection((np.array([0, 3]), slice(None)))
+ await aa.get_coordinate_selection((np.array([0, 3]), np.array([1, 2])))
+ await aa.get_mask_selection(np.zeros((8, 8), dtype=bool))
+
+ # no `block` here: `AsyncArray` exposes no block accessor, so that dialect
+ # only ever reaches an engine through the sync `Array`.
+ assert [r.kind for r in spy.reads] == ["basic", "orthogonal", "coordinate", "mask"]
+
+
+def test_asyncarray_default_engine_attribute() -> None:
+ z = zarr.create_array(MemoryStore(), shape=(4,), chunks=(2,), dtype="int8")
+ assert isinstance(z.async_array.engine, DefaultAsyncArrayEngine)
+
+
+def test_strided_read_preserves_fortran_order() -> None:
+ z = zarr.create_array(
+ MemoryStore(),
+ shape=(8, 8),
+ chunks=(4, 4),
+ dtype="float64",
+ config={"order": "F"},
+ )
+ z[:, :] = np.asfortranarray(np.arange(64, dtype="float64").reshape(8, 8))
+ full = np.asarray(z[:, :])
+ strided = np.asarray(z[::2, ::2])
+ assert full.flags.f_contiguous
+ assert strided.flags.f_contiguous
+ np.testing.assert_array_equal(strided, np.arange(64.0).reshape(8, 8)[::2, ::2])
+
+
+def test_basic_set_integer_axis_widens_value() -> None:
+ # Regression: a basic write whose dropped integer axis is *not* the leading
+ # axis (e.g. `arr[:, 0] = v`) has to widen the dimension-dropped value back
+ # to the selection's rank. A numpy integer scalar (not a Python int) keeps
+ # `__setitem__` routing on the basic facade rather than the orthogonal one,
+ # reproducing the original property-test failure.
+ expected = np.zeros((3, 3), dtype="int64")
+ z = zarr.create_array(MemoryStore(), shape=(3, 3), chunks=(3, 3), dtype="int64")
+ z[:, :] = expected
+ value = np.array([1, 2, 3], dtype="int64")
+ selection = cast("BasicSelection", (slice(None), np.int64(0)))
+ z.set_basic_selection(selection, value)
+ expected[:, 0] = value
+ np.testing.assert_array_equal(np.asarray(z[:, :]), expected)
+
+
+def test_empty_block_slice_reads_zero_length_box() -> None:
+ # Regression: an empty block slice (`blocks[1:0]`) yields a SliceDimIndexer
+ # with start > stop, which an engine that turns the block selection into a
+ # box must not read as a negative-length one.
+ data = np.arange(2, dtype="int64")
+ z = zarr.create_array(MemoryStore(), shape=(2,), chunks=(1,), dtype="int64")
+ z[:] = data
+ result: Any = z.get_block_selection((slice(1, 0),))
+ np.testing.assert_array_equal(np.asarray(result), data[2:2])
+ assert result.shape == (0,)
diff --git a/tests/engine/test_coordinate_out.py b/tests/engine/test_coordinate_out.py
new file mode 100644
index 0000000000..94a3761cf5
--- /dev/null
+++ b/tests/engine/test_coordinate_out.py
@@ -0,0 +1,68 @@
+"""`get_coordinate_selection(..., out=...)` must validate `out` against the
+shape the *engine* reads, and reshape only its return value.
+
+A coordinate selection is pointwise: the request the engine serves is flat, and
+the coordinate arrays' own (possibly multi-dimensional) shape is restored by the
+caller afterwards. So `out` is sized by the number of selected points, while the
+returned array carries the selection's shape. These tests pin both halves, and
+the error when `out` matches neither.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+import pytest
+
+import zarr
+from zarr.core.buffer import default_buffer_prototype
+from zarr.storage import MemoryStore
+
+
+def _filled_array() -> zarr.Array[Any]:
+ z = zarr.create_array(MemoryStore(), shape=(5, 5), chunks=(5, 5), dtype="int32")
+ z[:] = np.arange(25, dtype="int32").reshape(5, 5)
+ return z
+
+
+def test_coordinate_selection_out_multidim() -> None:
+ z = _filled_array()
+ coords = (np.array([[0, 1], [2, 3]]), np.array([[0, 1], [2, 3]]))
+ out = default_buffer_prototype().nd_buffer.from_numpy_array(np.zeros((4,), dtype="int32"))
+ result = z.get_coordinate_selection(coords, out=out)
+ expected = np.array([[0, 6], [12, 18]], dtype="int32")
+ # the result carries the coordinate arrays' shape; `out` holds the points
+ np.testing.assert_array_equal(np.asarray(result), expected)
+ np.testing.assert_array_equal(out.as_numpy_array(), expected.reshape(-1))
+
+
+async def test_coordinate_selection_out_multidim_async() -> None:
+ z = _filled_array()
+ coords = (np.array([[0, 1], [2, 3]]), np.array([[0, 1], [2, 3]]))
+ out = default_buffer_prototype().nd_buffer.from_numpy_array(np.zeros((4,), dtype="int32"))
+ result = await z.async_array.get_coordinate_selection(coords, out=out)
+ expected = np.array([[0, 6], [12, 18]], dtype="int32")
+ np.testing.assert_array_equal(np.asarray(result), expected)
+ np.testing.assert_array_equal(out.as_numpy_array(), expected.reshape(-1))
+
+
+def test_coordinate_selection_out_1d() -> None:
+ z = _filled_array()
+ coords = (np.array([0, 1, 4]), np.array([0, 1, 4]))
+ out = default_buffer_prototype().nd_buffer.from_numpy_array(np.zeros((3,), dtype="int32"))
+ result = z.get_coordinate_selection(coords, out=out)
+ expected = np.array([0, 6, 24], dtype="int32")
+ np.testing.assert_array_equal(np.asarray(result), expected)
+ np.testing.assert_array_equal(out.as_numpy_array(), expected)
+
+
+@pytest.mark.parametrize("out_shape", [(2, 2), (3,)])
+def test_coordinate_selection_out_shape_mismatch_raises(out_shape: tuple[int, ...]) -> None:
+ # neither the selection's own 2-d shape nor a wrong-length flat buffer is
+ # the shape the engine reads into
+ z = _filled_array()
+ coords = (np.array([[0, 1], [2, 3]]), np.array([[0, 1], [2, 3]]))
+ out = default_buffer_prototype().nd_buffer.from_numpy_array(np.zeros(out_shape, dtype="int32"))
+ with pytest.raises(ValueError, match="shape of out argument"):
+ z.get_coordinate_selection(coords, out=out)
diff --git a/tests/engine/test_default_engine.py b/tests/engine/test_default_engine.py
new file mode 100644
index 0000000000..8b154ffd78
--- /dev/null
+++ b/tests/engine/test_default_engine.py
@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import pytest
+
+import zarr
+from zarr.abc.engine import ArrayEngine, AsyncArrayEngine, SelectionRequest
+from zarr.core.buffer import default_buffer_prototype
+from zarr.core.chunk_grids import ChunkGrid
+from zarr.core.engine import DefaultArrayEngine, DefaultAsyncArrayEngine
+from zarr.core.sync import sync
+from zarr.errors import ChunkNotFoundError
+from zarr.storage import MemoryStore
+
+if TYPE_CHECKING:
+ from zarr.abc.engine import SelectionKind
+
+SHAPE = (10, 9)
+CHUNKS = (3, 4)
+
+
+def _make_array() -> zarr.Array[Any]:
+ z = zarr.create_array(MemoryStore(), shape=SHAPE, chunks=CHUNKS, dtype="int32", fill_value=0)
+ z[:, :] = np.arange(90, dtype="int32").reshape(SHAPE)
+ return z
+
+
+def _engine(z: zarr.Array[Any]) -> DefaultAsyncArrayEngine:
+ return DefaultAsyncArrayEngine(
+ store_path=z.async_array.store_path,
+ metadata=z.async_array.metadata,
+ config=z.async_array.config,
+ )
+
+
+def _request(z: zarr.Array[Any], kind: SelectionKind, selection: Any) -> SelectionRequest:
+ return SelectionRequest(
+ kind=kind,
+ selection=selection,
+ shape=z.shape,
+ chunk_grid=ChunkGrid.from_metadata(z.async_array.metadata),
+ )
+
+
+def test_default_async_engine_read_write_roundtrip() -> None:
+ z = _make_array()
+ eng = _engine(z)
+ assert isinstance(eng, AsyncArrayEngine)
+ proto = default_buffer_prototype()
+ request = _request(z, "basic", (slice(2, 7), slice(1, 5)))
+
+ out = sync(eng.read_selection(request, prototype=proto))
+ np.testing.assert_array_equal(np.asarray(out), np.asarray(z[2:7, 1:5]))
+
+ new: np.ndarray[Any, Any] = np.full((5, 4), -1, dtype="int32")
+ sync(eng.write_selection(request, new, prototype=proto))
+ np.testing.assert_array_equal(np.asarray(z[2:7, 1:5]), new)
+
+
+@pytest.mark.parametrize(
+ ("kind", "selection"),
+ [
+ ("basic", (slice(2, 7), slice(1, 5))),
+ ("orthogonal", (np.array([7, 1, 4]), np.array([0, 8]))),
+ ("coordinate", (np.array([9, 0, 3]), np.array([8, 0, 2]))),
+ ("block", (1, 2)),
+ ],
+)
+def test_default_engine_reads_every_dialect(kind: SelectionKind, selection: Any) -> None:
+ """Each dialect reaches the engine as a raw selection and reads what the
+ equivalent `Array` accessor reads."""
+ z = _make_array()
+ data = np.arange(90, dtype="int32").reshape(SHAPE)
+ expected = {
+ "basic": lambda: data[2:7, 1:5],
+ "orthogonal": lambda: data[np.ix_([7, 1, 4], [0, 8])],
+ "coordinate": lambda: data[[9, 0, 3], [8, 0, 2]],
+ "block": lambda: data[3:6, 8:9],
+ }[kind]()
+ result = sync(
+ _engine(z).read_selection(
+ _request(z, kind, selection), prototype=default_buffer_prototype()
+ )
+ )
+ np.testing.assert_array_equal(np.asarray(result), expected)
+
+
+def test_default_engine_rank_0_basic_read_is_a_scalar() -> None:
+ # The protocol mirrors `_get_selection`, including its scalar return, so an
+ # engine swap cannot change the type `get_basic_selection` hands back.
+ z = _make_array()
+ result = sync(
+ _engine(z).read_selection(
+ _request(z, "basic", (3, 4)), prototype=default_buffer_prototype()
+ )
+ )
+ assert isinstance(result, np.int32)
+ assert result == np.arange(90, dtype="int32").reshape(SHAPE)[3, 4]
+
+
+def test_default_sync_engine_matches_async() -> None:
+ z = _make_array()
+ eng = DefaultArrayEngine(_engine(z))
+ assert isinstance(eng, ArrayEngine)
+ request = _request(z, "basic", (slice(None), slice(None)))
+ np.testing.assert_array_equal(
+ np.asarray(eng.read_selection(request, prototype=default_buffer_prototype())),
+ np.asarray(z[:, :]),
+ )
+
+
+def test_with_metadata_rebinds() -> None:
+ z = _make_array()
+ eng = _engine(z)
+ assert eng.with_metadata(z.async_array.metadata) is not eng
+
+
+def test_read_missing_chunks_false_raises() -> None:
+ z = zarr.create_array(
+ MemoryStore(),
+ shape=(6,),
+ chunks=(2,),
+ dtype="int16",
+ config={"read_missing_chunks": False},
+ )
+ z[0:2] = np.arange(2, dtype="int16") # chunks 1 and 2 never written
+ with pytest.raises(ChunkNotFoundError):
+ sync(
+ _engine(z).read_selection(
+ _request(z, "basic", (slice(0, 6),)), prototype=default_buffer_prototype()
+ )
+ )
diff --git a/tests/engine/test_device_buffer_facade.py b/tests/engine/test_device_buffer_facade.py
new file mode 100644
index 0000000000..c75fb6e3fc
--- /dev/null
+++ b/tests/engine/test_device_buffer_facade.py
@@ -0,0 +1,261 @@
+"""Regression tests: neither `Array` nor `AsyncArray` may force a host (numpy)
+conversion on a device buffer (e.g. cupy/torch) whose implicit `np.asarray`
+coercion is refused.
+
+The engine owns the whole data path, so what these tests pin is the *facade's*
+end of the contract: the value handed to `write_selection` reaches the engine as
+the caller wrote it, and whatever `read_selection` returns reaches the caller
+unchanged. A `_DeviceArray` stand-in wraps a numpy array but raises `TypeError`
+from `__array__`, exactly like cupy refusing an implicit host copy. Everything
+else (indexing, `astype`, `copy`, and numpy's `__array_function__` protocol) is
+delegated to the wrapped array and re-wrapped, so operations that *stay in the
+array's own namespace* keep working while any `np.asarray`/`np.array` coercion
+raises.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import numpy.typing as npt
+import pytest
+
+import zarr
+from zarr.storage import MemoryStore
+
+if TYPE_CHECKING:
+ from zarr.abc.engine import SelectionRequest
+ from zarr.core.buffer import BufferPrototype, NDBuffer
+ from zarr.core.indexing import Fields
+ from zarr.core.metadata import ArrayMetadata
+
+
+class _DeviceArray:
+ """A minimal `NDArrayLike` whose implicit host conversion is refused.
+
+ Mimics a cupy array: indexing / `astype` / `copy` and numpy function
+ dispatch (`__array_function__`) all work in-namespace, but `__array__`
+ (what `np.asarray`/`np.array` call) raises, so any host coercion blows up.
+ """
+
+ def __init__(self, data: npt.NDArray[Any]) -> None:
+ self._a: npt.NDArray[Any] = data
+
+ # --- host coercion is forbidden -------------------------------------
+ def __array__(self, dtype: Any = None) -> npt.NDArray[Any]:
+ raise TypeError("implicit conversion to a host numpy array is not allowed")
+
+ # --- in-namespace numpy-function dispatch ---------------------------
+ def __array_function__(
+ self, func: Any, types: Any, args: tuple[Any, ...], kwargs: dict[str, Any]
+ ) -> Any:
+ unwrapped = tuple(a._a if isinstance(a, _DeviceArray) else a for a in args)
+ result = func(*unwrapped, **kwargs)
+ if isinstance(result, np.ndarray):
+ return _DeviceArray(result)
+ return result
+
+ # --- ndarray-like surface -------------------------------------------
+ @property
+ def dtype(self) -> np.dtype[Any]:
+ return self._a.dtype
+
+ @property
+ def shape(self) -> tuple[int, ...]:
+ return self._a.shape
+
+ @property
+ def ndim(self) -> int:
+ return self._a.ndim
+
+ @property
+ def size(self) -> int:
+ return self._a.size
+
+ def __len__(self) -> int:
+ return len(self._a)
+
+ def __getitem__(self, key: Any) -> _DeviceArray:
+ return _DeviceArray(self._a[key])
+
+ def __setitem__(self, key: Any, value: Any) -> None:
+ self._a[key] = value._a if isinstance(value, _DeviceArray) else value
+
+ def astype(self, dtype: Any, order: Any = "K", *, copy: bool = True) -> _DeviceArray:
+ return _DeviceArray(self._a.astype(dtype, order=order, copy=copy))
+
+ def copy(self) -> _DeviceArray:
+ return _DeviceArray(self._a.copy())
+
+ def reshape(self, *args: Any, **kwargs: Any) -> _DeviceArray:
+ return _DeviceArray(self._a.reshape(*args, **kwargs))
+
+
+def _numpy_key(request: SelectionRequest) -> Any:
+ """The numpy key equivalent to `request`'s raw selection.
+
+ These tests only use basic selections and orthogonal ones with a single
+ advanced index, for which numpy's own indexing already means what zarr's
+ does -- so the stub engine can apply the raw selection directly and stays
+ free of any resolution logic of its own.
+ """
+ assert request.kind in ("basic", "orthogonal")
+ return request.selection
+
+
+class _DeviceEngine:
+ """A synchronous `ArrayEngine` backed by numpy that hands the facade
+ `_DeviceArray` buffers (on read) and asserts it is handed them (on write)."""
+
+ def __init__(self, data: npt.NDArray[Any]) -> None:
+ self._data = data
+
+ def read_selection(
+ self,
+ request: SelectionRequest,
+ *,
+ prototype: BufferPrototype,
+ out: NDBuffer | None = None,
+ fields: Fields | None = None,
+ ) -> _DeviceArray:
+ # `np.asarray` keeps a rank-0 read a 0-d array rather than a numpy
+ # scalar, so even a scalar read stays in the device namespace.
+ return _DeviceArray(np.asarray(self._data[_numpy_key(request)]).copy())
+
+ def write_selection(
+ self,
+ request: SelectionRequest,
+ value: npt.ArrayLike,
+ *,
+ prototype: BufferPrototype,
+ fields: Fields | None = None,
+ ) -> None:
+ assert isinstance(value, _DeviceArray), (
+ "facade coerced the value off the device namespace before writing"
+ )
+ self._data[_numpy_key(request)] = value._a
+
+ def with_metadata(self, metadata: ArrayMetadata) -> _DeviceEngine:
+ return self
+
+
+class _AsyncDeviceEngine:
+ """Async mirror of `_DeviceEngine` for the `AsyncArray` data path."""
+
+ def __init__(self, data: npt.NDArray[Any]) -> None:
+ self._data = data
+ self._sync = _DeviceEngine(data)
+
+ async def read_selection(
+ self,
+ request: SelectionRequest,
+ *,
+ prototype: BufferPrototype,
+ out: NDBuffer | None = None,
+ fields: Fields | None = None,
+ ) -> _DeviceArray:
+ return self._sync.read_selection(request, prototype=prototype, out=out, fields=fields)
+
+ async def write_selection(
+ self,
+ request: SelectionRequest,
+ value: npt.ArrayLike,
+ *,
+ prototype: BufferPrototype,
+ fields: Fields | None = None,
+ ) -> None:
+ self._sync.write_selection(request, value, prototype=prototype, fields=fields)
+
+ def with_metadata(self, metadata: ArrayMetadata) -> _AsyncDeviceEngine:
+ return self
+
+
+def _array_on_device_engine() -> tuple[zarr.Array[Any], _DeviceEngine]:
+ z = zarr.create_array(MemoryStore(), shape=(8,), chunks=(4,), dtype="int64")
+ engine = _DeviceEngine(np.zeros(8, dtype="int64"))
+ # inject the device engine as this array's cached sync engine
+ z._engine = engine # type: ignore[assignment]
+ return z, engine
+
+
+def _array_2d_on_device_engine() -> tuple[zarr.Array[Any], _DeviceEngine]:
+ z = zarr.create_array(MemoryStore(), shape=(4, 4), chunks=(4, 4), dtype="int64")
+ engine = _DeviceEngine(np.arange(16, dtype="int64").reshape(4, 4))
+ z._engine = engine # type: ignore[assignment]
+ return z, engine
+
+
+def test_full_box_setitem_keeps_device_buffer() -> None:
+ z, engine = _array_on_device_engine()
+ z[:] = _DeviceArray(np.arange(8, dtype="int64"))
+ np.testing.assert_array_equal(engine._data, np.arange(8, dtype="int64"))
+
+
+def test_strided_setitem_keeps_device_buffer() -> None:
+ # a strided write is a read-modify-write inside the engine; the facade must
+ # still hand the value over untouched rather than normalizing it first.
+ z, engine = _array_on_device_engine()
+ z[::2] = _DeviceArray(np.array([10, 20, 30, 40], dtype="int64"))
+ expected = np.zeros(8, dtype="int64")
+ expected[::2] = [10, 20, 30, 40]
+ np.testing.assert_array_equal(engine._data, expected)
+
+
+def test_strided_getitem_keeps_device_buffer() -> None:
+ z, engine = _array_on_device_engine()
+ engine._data[:] = np.arange(8, dtype="int64")
+ result: object = z[::2]
+ assert isinstance(result, _DeviceArray), "strided read coerced off the device namespace"
+ np.testing.assert_array_equal(result._a, np.arange(8, dtype="int64")[::2])
+
+
+def test_basic_selection_scalar_read_keeps_device_buffer() -> None:
+ # an all-integer basic read is where zarr's own engine scalarizes; the
+ # facade must not repeat that on an engine that returned a device buffer.
+ z, engine = _array_2d_on_device_engine()
+ result: object = z.get_basic_selection((1, 2))
+ assert isinstance(result, _DeviceArray), "scalar read coerced off the device namespace"
+ np.testing.assert_array_equal(result._a, engine._data[1, 2])
+
+
+async def test_async_getitem_scalar_read_keeps_device_buffer() -> None:
+ z = zarr.create_array(MemoryStore(), shape=(4, 4), chunks=(4, 4), dtype="int64")
+ aa = z.async_array
+ engine = _AsyncDeviceEngine(np.arange(16, dtype="int64").reshape(4, 4))
+ object.__setattr__(aa, "engine", engine)
+ result: object = await aa.getitem((1, 2))
+ assert isinstance(result, _DeviceArray), "async scalar read coerced off the device namespace"
+ np.testing.assert_array_equal(result._a, engine._data[1, 2])
+
+
+def test_oindex_set_integer_and_array_axis_keeps_device_buffer() -> None:
+ # an orthogonal write with a dropped integer axis has to widen the value
+ # back to the selection's rank -- work the engine does, on a value the
+ # facade must not have converted on the way in.
+ z, engine = _array_2d_on_device_engine()
+ z.oindex[1, np.array([0, 2])] = _DeviceArray(np.array([10, 20], dtype="int64"))
+ expected = np.arange(16, dtype="int64").reshape(4, 4)
+ expected[1, [0, 2]] = [10, 20]
+ np.testing.assert_array_equal(engine._data, expected)
+
+
+@pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning")
+def test_scalar_bytes_read_returns_numpy_scalar() -> None:
+ # A fixed/vlen-bytes scalar read produces a numpy scalar (`np.bytes_`), not a
+ # 0-d ndarray; nothing in the read path may swallow it (`np.bytes_[()]`
+ # raises "byte indices must be integers"). Uses the real default engine, no
+ # device stand-in. (Distilled from a `test_basic_indexing` hypothesis
+ # failure.)
+ z = zarr.create_array(MemoryStore(), shape=(1,), chunks=(1,), dtype="S4")
+ z[:] = np.array([b"ab"], dtype="S4")
+ result = z.get_basic_selection(0)
+ assert isinstance(result, np.bytes_)
+ assert result == b"ab"
+
+
+def test_facade_never_coerces_device_buffer_to_host() -> None:
+ # Belt-and-braces: a bare `np.asarray` on the stand-in must raise, proving the
+ # tests above would trip any host coercion the facade performed.
+ with pytest.raises(TypeError):
+ np.asarray(_DeviceArray(np.arange(4)))
diff --git a/tests/engine/test_differential.py b/tests/engine/test_differential.py
new file mode 100644
index 0000000000..f2ce513847
--- /dev/null
+++ b/tests/engine/test_differential.py
@@ -0,0 +1,264 @@
+"""The same operations through both engines must agree with numpy and each other.
+
+The matrix covers every selection dialect (`basic`, `orthogonal`, `coordinate`,
+`mask`, `block`), reading and writing, over regular and sharded arrays. numpy is
+the oracle: each case is expressed as a zarr accessor call and the equivalent
+numpy key, so a divergence between the engines shows up as a divergence from
+numpy rather than as one engine's output being taken as the reference.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import pytest
+
+import zarr
+from zarr.errors import NegativeStepError
+from zarr.storage import LocalStore
+
+try:
+ import zarrista # noqa: F401
+
+ ENGINES = ["default", "zarrista"]
+except ImportError:
+ ENGINES = ["default"]
+
+if TYPE_CHECKING:
+ from pathlib import Path
+
+ import numpy.typing as npt
+
+ from zarr.abc.engine import SelectionKind
+ from zarr.core.engine import EngineName
+
+SHAPE = (10, 9)
+CHUNKS = (3, 4)
+
+MASK = np.zeros(SHAPE, dtype=bool)
+MASK[0, 0] = MASK[2, 2] = MASK[4, 5] = MASK[9, 8] = True
+
+
+@pytest.fixture
+def reference(tmp_path: Path) -> tuple[Path, npt.NDArray[np.float64]]:
+ z = zarr.create_array(LocalStore(tmp_path), shape=SHAPE, chunks=CHUNKS, dtype="float64")
+ data = np.arange(90, dtype="float64").reshape(SHAPE)
+ z[:, :] = data
+ return tmp_path, data
+
+
+@pytest.fixture
+def sharded(tmp_path: Path) -> tuple[Path, npt.NDArray[np.float64]]:
+ z = zarr.create_array(
+ LocalStore(tmp_path),
+ shape=SHAPE,
+ chunks=CHUNKS, # inner chunks
+ shards=(6, 8), # shard shape
+ dtype="float64",
+ )
+ data = np.arange(90, dtype="float64").reshape(SHAPE)
+ z[:, :] = data
+ return tmp_path, data
+
+
+def _read(z: zarr.Array[Any], kind: SelectionKind, selection: Any) -> npt.NDArray[Any]:
+ """Read `selection` through the accessor for its dialect."""
+ if kind == "basic":
+ return np.asarray(z[selection])
+ if kind == "orthogonal":
+ return np.asarray(z.oindex[selection])
+ if kind in ("coordinate", "mask"):
+ return np.asarray(z.vindex[selection])
+ return np.asarray(z.blocks[selection])
+
+
+def _write(z: zarr.Array[Any], kind: SelectionKind, selection: Any, value: Any) -> None:
+ if kind == "basic":
+ z[selection] = value
+ elif kind == "orthogonal":
+ z.oindex[selection] = value
+ elif kind in ("coordinate", "mask"):
+ z.vindex[selection] = value
+ else:
+ z.blocks[selection] = value
+
+
+def _numpy_key(kind: SelectionKind, selection: Any) -> Any:
+ """The numpy key that means what zarr's `kind` selection means."""
+ if kind == "orthogonal":
+ # With at most one advanced index numpy's own semantics *are* outer
+ # indexing; beyond that `np.ix_` is the equivalent, and a slice entry
+ # has to be spelled out as the indices it covers. Cases below keep to
+ # selections with no integer entry when they use more than one array,
+ # since `np.ix_` cannot express a dropped axis.
+ if sum(isinstance(s, np.ndarray) for s in selection) < 2:
+ return selection
+ return np.ix_(
+ *(
+ s if isinstance(s, np.ndarray) else np.arange(n)[s]
+ for s, n in zip(selection, SHAPE, strict=True)
+ )
+ )
+ if kind == "block":
+ return tuple(
+ slice(s * c, min((s + 1) * c, n))
+ if isinstance(s, int)
+ else slice((s.start or 0) * c, min(s.stop * c, n))
+ for s, c, n in zip(selection, CHUNKS, SHAPE, strict=True)
+ )
+ return selection
+
+
+READS: list[Any] = [
+ pytest.param("basic", (slice(None), slice(None)), id="basic-whole"),
+ pytest.param("basic", (slice(2, 7), slice(1, 5)), id="basic-box"),
+ pytest.param("basic", (slice(1, 9, 2), slice(None, None, 3)), id="basic-strided"),
+ pytest.param("basic", (3, slice(None)), id="basic-dropped-axis"),
+ pytest.param("basic", (-1, -2), id="basic-negative-ints"),
+ pytest.param("basic", (slice(4, 4), slice(None)), id="basic-empty"),
+ pytest.param("basic", (Ellipsis,), id="basic-ellipsis"),
+ pytest.param("orthogonal", (np.array([7, 1, 4]), np.array([0, 8])), id="oindex-arrays"),
+ pytest.param("orthogonal", (3, np.array([0, 8])), id="oindex-int-and-array"),
+ pytest.param("orthogonal", (np.array([7, 1, 4]), slice(1, 5)), id="oindex-array-and-slice"),
+ pytest.param(
+ "orthogonal",
+ (np.array([True] * 5 + [False] * 5), np.array([0, 8])),
+ id="oindex-bool-and-array",
+ ),
+ pytest.param("coordinate", (np.array([9, 0, 3]), np.array([8, 0, 2])), id="vindex-points"),
+ pytest.param(
+ "coordinate",
+ (np.array([[0, 1], [2, 3]]), np.array([[0, 1], [2, 3]])),
+ id="vindex-2d-points",
+ ),
+ pytest.param("mask", MASK, id="mask"),
+ pytest.param("block", (1, 2), id="block-ints"),
+ pytest.param("block", (slice(0, 2), 0), id="block-slice-and-int"),
+]
+
+
+@pytest.mark.parametrize("engine", ENGINES)
+@pytest.mark.parametrize(("kind", "selection"), READS)
+def test_reads_match_numpy(
+ reference: tuple[Path, npt.NDArray[np.float64]],
+ engine: EngineName,
+ kind: SelectionKind,
+ selection: Any,
+) -> None:
+ tmp_path, data = reference
+ z = zarr.open_array(LocalStore(tmp_path), engine=engine)
+ np.testing.assert_array_equal(_read(z, kind, selection), data[_numpy_key(kind, selection)])
+
+
+WRITES: list[Any] = [
+ pytest.param("basic", (slice(0, 3), slice(0, 4)), 7.0, id="basic-aligned-scalar"),
+ pytest.param(
+ "basic",
+ (slice(4, 6), slice(2, 9)),
+ np.arange(14, dtype="float64").reshape(2, 7),
+ id="basic-partial-chunks",
+ ),
+ pytest.param("basic", (slice(1, 9, 3), slice(None, None, 4)), -1.0, id="basic-strided"),
+ pytest.param("basic", (9, slice(0, 4)), -2.0, id="basic-edge-chunk-row"),
+ pytest.param("basic", (slice(4, 4), slice(None)), np.zeros((0, 9)), id="basic-empty"),
+ pytest.param(
+ "orthogonal",
+ (np.array([7, 1, 4]), np.array([0, 8])),
+ np.arange(6, dtype="float64").reshape(3, 2),
+ id="oindex-arrays",
+ ),
+ pytest.param(
+ "orthogonal", (3, np.array([0, 8])), np.array([11.0, 12.0]), id="oindex-int-and-array"
+ ),
+ pytest.param(
+ "orthogonal", (np.array([7, 1, 4]), slice(1, 5)), 5.0, id="oindex-array-and-slice"
+ ),
+ pytest.param(
+ "coordinate",
+ (np.array([9, 0, 3]), np.array([8, 0, 2])),
+ np.array([1.0, 2.0, 3.0]),
+ id="vindex-points",
+ ),
+ # duplicate coordinates are legal and resolve last-wins, as they do in numpy
+ pytest.param(
+ "coordinate",
+ (np.array([2, 2, 5]), np.array([3, 3, 6])),
+ np.array([1.0, 2.0, 3.0]),
+ id="vindex-duplicate-points",
+ ),
+ pytest.param("mask", MASK, np.array([1.0, 2.0, 3.0, 4.0]), id="mask"),
+ pytest.param(
+ "block", (1, 2), np.arange(3, dtype="float64").reshape(3, 1), id="block-edge-column"
+ ),
+ pytest.param("block", (slice(0, 2), 0), 8.0, id="block-slice-and-int"),
+]
+
+
+@pytest.mark.parametrize("engine", ENGINES)
+@pytest.mark.parametrize(("kind", "selection", "value"), WRITES)
+def test_writes_match_numpy(
+ reference: tuple[Path, npt.NDArray[np.float64]],
+ engine: EngineName,
+ kind: SelectionKind,
+ selection: Any,
+ value: Any,
+) -> None:
+ tmp_path, data = reference
+ z = zarr.open_array(LocalStore(tmp_path), engine=engine)
+ expected = data.copy()
+ expected[_numpy_key(kind, selection)] = value
+
+ _write(z, kind, selection, value)
+
+ # the whole array, so a write that spilled outside its selection is caught
+ np.testing.assert_array_equal(np.asarray(z[:, :]), expected)
+
+
+@pytest.mark.parametrize("engine", ENGINES)
+def test_sharded_read_write(
+ sharded: tuple[Path, npt.NDArray[np.float64]], engine: EngineName
+) -> None:
+ tmp_path, data = sharded
+ z = zarr.open_array(LocalStore(tmp_path), engine=engine)
+ np.testing.assert_array_equal(np.asarray(z[2:8, 3:9]), data[2:8, 3:9])
+
+ expected = data.copy()
+ z[1:9:2, ::3] = -5.0 # strided, so inner chunks are read-modify-written
+ expected[1:9:2, ::3] = -5.0
+ np.testing.assert_array_equal(np.asarray(z[:, :]), expected)
+
+
+@pytest.mark.parametrize("engine", ENGINES)
+def test_scalar_read_matches_numpy(
+ reference: tuple[Path, npt.NDArray[np.float64]], engine: EngineName
+) -> None:
+ tmp_path, data = reference
+ z = zarr.open_array(LocalStore(tmp_path), engine=engine)
+ assert float(np.asarray(z.get_basic_selection((3, 4)))) == data[3, 4]
+
+
+@pytest.mark.parametrize(
+ "selection", [(slice(None, None, -1), slice(None)), (slice(None), slice(None, None, -2))]
+)
+def test_negative_step_is_engine_specific(
+ reference: tuple[Path, npt.NDArray[np.float64]], selection: Any
+) -> None:
+ """The engines deliberately differ here, so this case stays out of the
+ shared matrix: zarr's own indexer refuses a negative step, while the
+ zarrista engine resolves the raw selection through `zarr_indexing` and
+ serves it.
+ """
+ tmp_path, data = reference
+ with pytest.raises(NegativeStepError):
+ zarr.open_array(LocalStore(tmp_path), engine="default")[selection]
+
+ if "zarrista" not in ENGINES:
+ pytest.skip("zarrista not installed")
+ z = zarr.open_array(LocalStore(tmp_path), engine="zarrista")
+ np.testing.assert_array_equal(np.asarray(z[selection]), data[selection])
+
+ expected = data.copy()
+ expected[selection] = -data[selection]
+ z[selection] = -data[selection]
+ np.testing.assert_array_equal(np.asarray(z[:, :]), expected)
diff --git a/tests/engine/test_engine_param.py b/tests/engine/test_engine_param.py
new file mode 100644
index 0000000000..5c130c6e8c
--- /dev/null
+++ b/tests/engine/test_engine_param.py
@@ -0,0 +1,80 @@
+"""`engine=` threading through `zarr.create_array` / `zarr.open_array` and their
+async counterparts (Task 7 of the array-engine-protocol plan)."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+import zarr
+from zarr.core.array import AsyncArray
+from zarr.core.engine import DefaultArrayEngine, DefaultAsyncArrayEngine
+from zarr.storage import MemoryStore
+
+
+def test_engine_param_combinations() -> None:
+ store = MemoryStore()
+ z = zarr.create_array(store, name="a", shape=(4,), chunks=(2,), dtype="int8", engine="default")
+ z[:] = np.arange(4, dtype="int8")
+ assert isinstance(z.engine, DefaultArrayEngine)
+
+ z2 = zarr.open_array(store, path="a", engine="default")
+ np.testing.assert_array_equal(np.asarray(z2[:]), np.arange(4, dtype="int8"))
+ assert isinstance(z2.async_array.engine, DefaultAsyncArrayEngine)
+
+ # user-provided sync instance
+ inst = z2.engine
+ z3 = zarr.open_array(store, path="a", engine=inst)
+ assert z3.engine is inst
+ # the wrapped AsyncArray keeps its own (default) engine -- a sync instance
+ # must never reach AsyncArray.
+ assert isinstance(z3.async_array.engine, DefaultAsyncArrayEngine)
+
+
+def test_engine_param_unknown_name() -> None:
+ with pytest.raises(ValueError, match="unknown engine"):
+ zarr.create_array(
+ MemoryStore(),
+ shape=(2,),
+ chunks=(2,),
+ dtype="int8",
+ engine="nope", # type: ignore[arg-type]
+ )
+
+
+def test_async_array_rejects_sync_engine_instance() -> None:
+ """A sync `ArrayEngine` instance passed to `AsyncArray` -- any construction
+ path, not just the public API -- must raise `TypeError` naming both
+ protocols."""
+ z = zarr.create_array(MemoryStore(), shape=(4,), chunks=(2,), dtype="int8")
+ sync_engine = z.engine
+ aa = z.async_array
+ with pytest.raises(TypeError, match="ArrayEngine"):
+ AsyncArray(
+ metadata=aa.metadata,
+ store_path=aa.store_path,
+ engine=sync_engine, # type: ignore[call-overload]
+ )
+
+
+def test_sync_api_rejects_async_engine_instance() -> None:
+ """An `AsyncArrayEngine` instance passed to the sync `create_array` entry
+ point must raise `TypeError` immediately (not lazily on first data
+ access)."""
+ z = zarr.create_array(MemoryStore(), shape=(4,), chunks=(2,), dtype="int8")
+ async_engine = z.async_array.engine
+ with pytest.raises(TypeError, match="ArrayEngine"):
+ zarr.create_array(MemoryStore(), shape=(4,), chunks=(2,), dtype="int8", engine=async_engine)
+
+
+def test_engine_missing_read_selection_raises_type_error() -> None:
+ """An object implementing neither engine protocol (no `read_selection`)
+ must raise `TypeError` naming the protocols, not some other error."""
+ with pytest.raises(TypeError, match="read_selection"):
+ zarr.create_array(
+ MemoryStore(),
+ shape=(4,),
+ chunks=(2,),
+ dtype="int8",
+ engine=object(), # type: ignore[arg-type]
+ )
diff --git a/tests/engine/test_engine_spec_preservation.py b/tests/engine/test_engine_spec_preservation.py
new file mode 100644
index 0000000000..0ee6fdd511
--- /dev/null
+++ b/tests/engine/test_engine_spec_preservation.py
@@ -0,0 +1,64 @@
+"""The original engine spec must survive `with_config` / `update_attributes`.
+
+`with_config` and `update_attributes` build a fresh array wrapper; each must
+re-thread the array's engine spec so a custom engine is not silently dropped in
+favour of the default one. Each path stores the ORIGINAL spec (name/instance),
+so a named engine is re-resolved against the new config and an instance is
+carried through unchanged -- these tests use instances and assert identity,
+which is exactly what regresses when the spec is dropped (the copy would fall
+back to a freshly-resolved default engine instead).
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import zarr
+from zarr.core.array import AsyncArray
+from zarr.core.engine import DefaultArrayEngine, DefaultAsyncArrayEngine
+from zarr.storage import MemoryStore
+
+
+def _async_array_with_custom_engine() -> tuple[AsyncArray[Any], DefaultAsyncArrayEngine]:
+ z = zarr.create_array(MemoryStore(), shape=(4,), chunks=(2,), dtype="int8")
+ aa = z.async_array
+ custom = DefaultAsyncArrayEngine(
+ store_path=aa.store_path, metadata=aa.metadata, config=aa.config
+ )
+ built = AsyncArray(
+ metadata=aa.metadata, store_path=aa.store_path, config=aa.config, engine=custom
+ )
+ return built, custom
+
+
+def test_asyncarray_with_config_preserves_engine() -> None:
+ built, custom = _async_array_with_custom_engine()
+ assert built.engine is custom
+ copy = built.with_config({"order": "F"})
+ # the custom engine spec is re-resolved (an instance is returned unchanged),
+ # so the copy keeps it instead of falling back to a default engine
+ assert copy.engine is custom
+
+
+def test_array_with_config_preserves_engine() -> None:
+ z = zarr.create_array(MemoryStore(), shape=(4,), chunks=(2,), dtype="int8")
+ aa = z.async_array
+ custom = DefaultArrayEngine(
+ DefaultAsyncArrayEngine(store_path=aa.store_path, metadata=aa.metadata, config=aa.config)
+ )
+ arr = zarr.Array(aa, engine_spec=custom)
+ assert arr.engine is custom
+ copy = arr.with_config({"order": "F"})
+ assert copy.engine is custom
+
+
+def test_array_update_attributes_preserves_engine() -> None:
+ z = zarr.create_array(MemoryStore(), shape=(4,), chunks=(2,), dtype="int8")
+ aa = z.async_array
+ custom = DefaultArrayEngine(
+ DefaultAsyncArrayEngine(store_path=aa.store_path, metadata=aa.metadata, config=aa.config)
+ )
+ arr = zarr.Array(aa, engine_spec=custom)
+ assert arr.engine is custom
+ updated = arr.update_attributes({"foo": "bar"})
+ assert updated.engine is custom
diff --git a/tests/engine/test_protocols.py b/tests/engine/test_protocols.py
new file mode 100644
index 0000000000..a5bc443a4f
--- /dev/null
+++ b/tests/engine/test_protocols.py
@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import pytest
+
+from zarr.abc.engine import ArrayEngine, AsyncArrayEngine, SelectionRequest
+from zarr.core.chunk_grids import ChunkGrid
+from zarr.core.indexing import (
+ BasicIndexer,
+ BlockIndexer,
+ CoordinateIndexer,
+ MaskIndexer,
+ OrthogonalIndexer,
+)
+from zarr.errors import UnsupportedEngineError
+
+if TYPE_CHECKING:
+ import numpy.typing as npt
+
+ from zarr.abc.engine import SelectionKind
+ from zarr.core.buffer import BufferPrototype, NDArrayLikeOrScalar, NDBuffer
+ from zarr.core.indexing import Fields, Indexer
+ from zarr.core.metadata import ArrayMetadata
+
+SHAPE = (10, 9)
+GRID = ChunkGrid.from_sizes(SHAPE, (3, 4))
+
+
+def _request(kind: SelectionKind, selection: Any) -> SelectionRequest:
+ return SelectionRequest(kind=kind, selection=selection, shape=SHAPE, chunk_grid=GRID)
+
+
+@pytest.mark.parametrize(
+ ("kind", "selection", "expected"),
+ [
+ ("basic", (slice(1, 5), slice(None)), BasicIndexer),
+ ("orthogonal", (np.array([0, 3]), slice(None)), OrthogonalIndexer),
+ ("coordinate", (np.array([0, 3]), np.array([1, 2])), CoordinateIndexer),
+ ("mask", np.zeros(SHAPE, dtype=bool), MaskIndexer),
+ ("block", (1, 2), BlockIndexer),
+ ],
+)
+def test_request_builds_the_indexer_for_its_kind(
+ kind: SelectionKind, selection: Any, expected: type[Indexer]
+) -> None:
+ """`kind` alone decides which indexing dialect the raw selection is read in."""
+ assert isinstance(_request(kind, selection).indexer, expected)
+
+
+def test_request_indexer_is_built_on_demand_and_cached() -> None:
+ # An engine that never asks for the zarr-native view must not pay to build
+ # it; one that asks twice must pay once.
+ request = _request("basic", (slice(1, 5), slice(None)))
+ assert "indexer" not in request.__dict__
+ assert request.indexer is request.indexer
+
+
+def test_request_keeps_the_selection_unresolved() -> None:
+ # The reason the request carries the raw selection: `CoordinateIndexer`
+ # stores its coordinates in chunk-sorted order, so an engine that
+ # reconstructed the selection from the indexer would silently reorder the
+ # points -- and with them the user's results.
+ coords = (np.array([9, 0, 3]), np.array([8, 0, 2]))
+ request = _request("coordinate", coords)
+ assert request.selection is coords
+ assert not np.array_equal(request.indexer.selection[0], coords[0]) # type: ignore[attr-defined]
+
+
+class _FakeSyncEngine:
+ def read_selection(
+ self,
+ request: SelectionRequest,
+ *,
+ prototype: BufferPrototype,
+ out: NDBuffer | None = None,
+ fields: Fields | None = None,
+ ) -> NDArrayLikeOrScalar:
+ return np.zeros(request.indexer.shape)
+
+ def write_selection(
+ self,
+ request: SelectionRequest,
+ value: npt.ArrayLike,
+ *,
+ prototype: BufferPrototype,
+ fields: Fields | None = None,
+ ) -> None:
+ return None
+
+ def with_metadata(self, metadata: ArrayMetadata) -> _FakeSyncEngine:
+ return self
+
+
+class _FakeAsyncEngine:
+ async def read_selection(
+ self,
+ request: SelectionRequest,
+ *,
+ prototype: BufferPrototype,
+ out: NDBuffer | None = None,
+ fields: Fields | None = None,
+ ) -> NDArrayLikeOrScalar:
+ return np.zeros(request.indexer.shape)
+
+ async def write_selection(
+ self,
+ request: SelectionRequest,
+ value: npt.ArrayLike,
+ *,
+ prototype: BufferPrototype,
+ fields: Fields | None = None,
+ ) -> None:
+ return None
+
+ def with_metadata(self, metadata: ArrayMetadata) -> _FakeAsyncEngine:
+ return self
+
+
+def test_runtime_checkable_protocols() -> None:
+ """isinstance checks verify method presence only; mypy is authoritative.
+
+ Both fakes therefore satisfy both protocols -- the sync/async distinction
+ is drawn at runtime by `classify_engine_arg`, not by isinstance.
+ """
+ assert isinstance(_FakeSyncEngine(), ArrayEngine)
+ assert isinstance(_FakeAsyncEngine(), AsyncArrayEngine)
+ assert not isinstance(object(), ArrayEngine)
+ assert not isinstance(object(), AsyncArrayEngine)
+
+
+def test_unsupported_engine_error_is_value_error() -> None:
+ with pytest.raises(ValueError):
+ raise UnsupportedEngineError("nope")
diff --git a/tests/engine/test_resolve.py b/tests/engine/test_resolve.py
new file mode 100644
index 0000000000..30bda6c694
--- /dev/null
+++ b/tests/engine/test_resolve.py
@@ -0,0 +1,168 @@
+from __future__ import annotations
+
+import gc
+from typing import Any
+
+import pytest
+
+import zarr
+from zarr.core.engine import (
+ DefaultArrayEngine,
+ DefaultAsyncArrayEngine,
+ resolve_async_engine,
+ resolve_sync_engine,
+)
+from zarr.core.engine._resolve import _hierarchy_cache
+from zarr.storage import MemoryStore
+
+
+def _array() -> zarr.Array[Any]:
+ return zarr.create_array(MemoryStore(), shape=(4,), chunks=(2,), dtype="int8")
+
+
+def test_list_engines() -> None:
+ # `list_engines` reports exactly the known engine names, sorted, and is
+ # re-exported at the top level for discoverability.
+ assert zarr.list_engines() == ["default", "zarrista"]
+ assert callable(zarr.list_engines)
+
+
+def test_resolution_combinations() -> None:
+ z = _array()
+ store = z.store
+ path = z.path
+ meta = z.async_array.metadata
+
+ # None and "default" produce default engines
+ for spec in (None, "default"):
+ assert isinstance(
+ resolve_async_engine(spec, store=store, path=path, metadata=meta),
+ DefaultAsyncArrayEngine,
+ )
+ assert isinstance(
+ resolve_sync_engine(spec, store=store, path=path, metadata=meta),
+ DefaultArrayEngine,
+ )
+
+ # instances pass through untouched
+ inst = resolve_sync_engine(None, store=store, path=path, metadata=meta)
+ assert resolve_sync_engine(inst, store=store, path=path, metadata=meta) is inst
+
+ # engines minted from the same store share a hierarchy engine
+ e1 = resolve_async_engine(None, store=store, path=path, metadata=meta)
+ e2 = resolve_async_engine(None, store=store, path="other", metadata=meta)
+ assert e1.store_path.store is e2.store_path.store # type: ignore[attr-defined]
+
+
+def test_hierarchy_cache_evicts_when_store_and_engines_are_unreferenced() -> None:
+ # A dedicated store (not shared with other tests) so the cache starts clean
+ # for this key. Measure the baseline *before* creating the array: creating it
+ # resolves the array's own engine, which already populates the (default,
+ # async, id(store)) cache entry that `e1`/`e2` then reuse.
+ store = MemoryStore()
+ # Collect any hierarchies left unreferenced by earlier tests first, so the
+ # baseline reflects only entries kept alive by still-referenced arrays.
+ gc.collect()
+ before = len(_hierarchy_cache)
+ z = zarr.create_array(store, shape=(4,), chunks=(2,), dtype="int8")
+ path = z.path
+ meta = z.async_array.metadata
+
+ e1 = resolve_async_engine(None, store=store, path=path, metadata=meta)
+ e2 = resolve_async_engine(None, store=store, path="other", metadata=meta)
+ assert len(_hierarchy_cache) == before + 1
+
+ # engines minted while the store is in concurrent use share one hierarchy
+ assert (
+ e1._resolve_hierarchy_keepalive # type: ignore[attr-defined]
+ is e2._resolve_hierarchy_keepalive # type: ignore[attr-defined]
+ )
+
+ del z, e1, e2, store
+ gc.collect()
+ assert len(_hierarchy_cache) == before
+
+
+def test_unknown_name_raises() -> None:
+ z = _array()
+ with pytest.raises(ValueError, match="unknown engine"):
+ resolve_async_engine(
+ "bogus", # type: ignore[arg-type]
+ store=z.store,
+ path=z.path,
+ metadata=z.async_array.metadata,
+ )
+
+
+def test_zarrista_missing_raises_import_error() -> None:
+ try:
+ import zarrista # noqa: F401
+
+ pytest.skip("zarrista installed; missing-module error not testable")
+ except ImportError:
+ pass
+ z = _array()
+ with pytest.raises(ImportError, match="zarrista"):
+ resolve_async_engine(
+ "zarrista", store=z.store, path=z.path, metadata=z.async_array.metadata
+ )
+
+
+def test_resolve_async_engine_rejects_sync_engine_instance() -> None:
+ """`resolve_async_engine` must reject a synchronous `ArrayEngine` instance --
+ `AsyncArray` can only be backed by an `AsyncArrayEngine`."""
+ z = _array()
+ sync_engine = resolve_sync_engine(
+ None, store=z.store, path=z.path, metadata=z.async_array.metadata
+ )
+ with pytest.raises(TypeError, match="AsyncArray"):
+ resolve_async_engine(
+ sync_engine, # type: ignore[arg-type]
+ store=z.store,
+ path=z.path,
+ metadata=z.async_array.metadata,
+ )
+
+
+def test_resolve_sync_engine_rejects_async_engine_instance() -> None:
+ """`resolve_sync_engine` must reject an asynchronous `AsyncArrayEngine`
+ instance -- `Array` can only be backed by a synchronous `ArrayEngine`."""
+ z = _array()
+ async_engine = resolve_async_engine(
+ None, store=z.store, path=z.path, metadata=z.async_array.metadata
+ )
+ with pytest.raises(TypeError, match="Array"):
+ resolve_sync_engine(
+ async_engine, # type: ignore[arg-type]
+ store=z.store,
+ path=z.path,
+ metadata=z.async_array.metadata,
+ )
+
+
+def test_resolve_async_engine_rejects_object_missing_read_selection() -> None:
+ """An object implementing neither engine protocol (no `read_selection`) must
+ raise `TypeError` naming the protocols, not some other error, when handed to
+ `resolve_async_engine`."""
+ z = _array()
+ with pytest.raises(TypeError, match="read_selection"):
+ resolve_async_engine(
+ object(), # type: ignore[arg-type]
+ store=z.store,
+ path=z.path,
+ metadata=z.async_array.metadata,
+ )
+
+
+def test_resolve_sync_engine_rejects_object_missing_read_selection() -> None:
+ """An object implementing neither engine protocol (no `read_selection`) must
+ raise `TypeError` naming the protocols, not some other error, when handed to
+ `resolve_sync_engine`."""
+ z = _array()
+ with pytest.raises(TypeError, match="read_selection"):
+ resolve_sync_engine(
+ object(), # type: ignore[arg-type]
+ store=z.store,
+ path=z.path,
+ metadata=z.async_array.metadata,
+ )
diff --git a/tests/engine/test_sync_path.py b/tests/engine/test_sync_path.py
new file mode 100644
index 0000000000..f3a3d4e5d8
--- /dev/null
+++ b/tests/engine/test_sync_path.py
@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING
+
+import numpy as np
+import pytest
+
+import zarr
+from zarr.abc.engine import ArrayEngine
+from zarr.core.engine import DefaultArrayEngine
+from zarr.storage import MemoryStore
+
+if TYPE_CHECKING:
+ import numpy.typing as npt
+
+ from zarr.abc.engine import SelectionRequest
+ from zarr.core.buffer import BufferPrototype, NDArrayLikeOrScalar, NDBuffer
+ from zarr.core.indexing import Fields
+ from zarr.core.metadata import ArrayMetadata
+
+
+def test_array_has_sync_engine() -> None:
+ z = zarr.create_array(MemoryStore(), shape=(6,), chunks=(2,), dtype="uint8")
+ assert isinstance(z.engine, ArrayEngine)
+ assert isinstance(z.engine, DefaultArrayEngine)
+
+
+def test_array_engine_is_cached() -> None:
+ z = zarr.create_array(MemoryStore(), shape=(6,), chunks=(2,), dtype="uint8")
+ assert z.engine is z.engine
+
+
+class _NoLoopEngine:
+ """A sync engine that asserts no event loop is running when called."""
+
+ def __init__(self, inner: ArrayEngine) -> None:
+ self._inner = inner
+ self.calls = 0
+
+ def read_selection(
+ self,
+ request: SelectionRequest,
+ *,
+ prototype: BufferPrototype,
+ out: NDBuffer | None = None,
+ fields: Fields | None = None,
+ ) -> NDArrayLikeOrScalar:
+ self.calls += 1
+ with pytest.raises(RuntimeError):
+ asyncio.get_running_loop()
+ return self._inner.read_selection(request, prototype=prototype, out=out, fields=fields)
+
+ def write_selection(
+ self,
+ request: SelectionRequest,
+ value: npt.ArrayLike,
+ *,
+ prototype: BufferPrototype,
+ fields: Fields | None = None,
+ ) -> None:
+ self.calls += 1
+ with pytest.raises(RuntimeError):
+ asyncio.get_running_loop()
+ return self._inner.write_selection(request, value, prototype=prototype, fields=fields)
+
+ def with_metadata(self, metadata: ArrayMetadata) -> _NoLoopEngine:
+ return _NoLoopEngine(self._inner.with_metadata(metadata))
+
+
+def test_sync_data_path_runs_without_event_loop_in_caller_thread() -> None:
+ z = zarr.create_array(MemoryStore(), shape=(6,), chunks=(2,), dtype="uint8")
+ probe = _NoLoopEngine(z.engine)
+ object.__setattr__(z, "_engine", probe) # match the attribute name used in impl
+
+ z[1:5] = np.arange(4, dtype="uint8")
+ out = z[1:5]
+
+ assert probe.calls == 2
+ np.testing.assert_array_equal(np.asarray(out), np.arange(4, dtype="uint8"))
+
+
+def test_resize_rebinds_cached_sync_engine() -> None:
+ """After `resize`, reads/writes beyond the old bounds must go through an
+ engine bound to the new metadata, not a stale cached one."""
+ z = zarr.create_array(MemoryStore(), shape=(4,), chunks=(2,), dtype="uint8")
+ z[:] = np.arange(4, dtype="uint8")
+ assert np.asarray(z[:]).tolist() == [0, 1, 2, 3]
+
+ # force engine resolution before the resize so the cache is populated
+ assert isinstance(z.engine, DefaultArrayEngine)
+
+ z.resize((8,))
+ z[4:8] = np.arange(4, 8, dtype="uint8")
+ out = z[:]
+
+ np.testing.assert_array_equal(np.asarray(out), np.arange(8, dtype="uint8"))
diff --git a/tests/test_array.py b/tests/test_array.py
index b1a7a3c0f2..ad4ae37cbe 100644
--- a/tests/test_array.py
+++ b/tests/test_array.py
@@ -70,7 +70,7 @@
from zarr.core.dtype.common import ENDIANNESS_STR, EndiannessStr
from zarr.core.dtype.npy.common import NUMPY_ENDIANNESS_STR, endianness_from_numpy_str
from zarr.core.group import AsyncGroup
-from zarr.core.indexing import BasicIndexer, _iter_grid, _iter_regions
+from zarr.core.indexing import _iter_grid, _iter_regions
from zarr.core.metadata.v2 import ArrayV2Metadata
from zarr.core.sync import sync
from zarr.errors import (
@@ -1598,10 +1598,7 @@ async def test_with_data(impl: Literal["sync", "async"], store: Store) -> None:
stored = arr[:]
elif impl == "async":
arr = await create_array(store, name=name, data=data, zarr_format=3)
- stored = await arr._get_selection(
- BasicIndexer(..., shape=arr.shape, chunk_grid=arr._chunk_grid),
- prototype=default_buffer_prototype(),
- )
+ stored = await arr.getitem(Ellipsis, prototype=default_buffer_prototype())
else:
raise ValueError(f"Invalid impl: {impl}")
diff --git a/tests/zarrista/__init__.py b/tests/zarrista/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/zarrista/test_engine.py b/tests/zarrista/test_engine.py
new file mode 100644
index 0000000000..2604ae6f24
--- /dev/null
+++ b/tests/zarrista/test_engine.py
@@ -0,0 +1,184 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import pytest
+
+pytest.importorskip("zarrista")
+
+import zarr
+from zarr.errors import UnsupportedEngineError
+from zarr.storage import LocalStore
+from zarr.zarrista import ZarristaAsyncEngine
+
+if TYPE_CHECKING:
+ from pathlib import Path
+
+ import numpy.typing as npt
+
+
+def _make(tmp_path: Path) -> zarr.Array[Any]:
+ z = zarr.create_array(LocalStore(tmp_path), shape=(10, 9), chunks=(3, 4), dtype="float32")
+ z[:, :] = np.arange(90, dtype="float32").reshape(10, 9)
+ return z
+
+
+def test_zarrista_engine_read_write_combinations(tmp_path: Path) -> None:
+ z = _make(tmp_path)
+ ze = zarr.open_array(LocalStore(tmp_path), engine="zarrista")
+
+ # contiguous read
+ np.testing.assert_array_equal(np.asarray(ze[2:7, 1:5]), np.asarray(z[2:7, 1:5]))
+ # strided read via bbox + post-index
+ np.testing.assert_array_equal(np.asarray(ze[1:9:2, ::3]), np.asarray(z[1:9:2, ::3]))
+ # full-chunk-aligned write
+ ze[0:3, 0:4] = np.zeros((3, 4), dtype="float32")
+ np.testing.assert_array_equal(np.asarray(z[0:3, 0:4]), np.zeros((3, 4), dtype="float32"))
+ # partial-chunk RMW write
+ ze[1:2, 1:2] = np.float32(99.0)
+ assert float(np.asarray(z[1, 1])) == 99.0
+ assert float(np.asarray(z[0, 0])) == 0.0 # neighbor in same chunk untouched
+
+
+def test_zarrista_engine_edge_chunk_full_write(tmp_path: Path) -> None:
+ # shape (10, 9) with chunks (3, 4): row-chunk index 3 (rows 9:10) is an
+ # edge chunk whose *clipped* extent (1 row) is smaller than the nominal
+ # chunk shape (3 rows). Writing the entirety of that chunk's valid
+ # (clipped) region is not the same as writing the full nominal chunk, so
+ # the engine must still take the RMW path (rather than treating it as a
+ # "full chunk" write and handing zarrista a wrongly-shaped buffer).
+ z = _make(tmp_path)
+ ze = zarr.open_array(LocalStore(tmp_path), engine="zarrista")
+
+ ze[9:10, 0:4] = -np.ones((1, 4), dtype="float32")
+ np.testing.assert_array_equal(np.asarray(z[9:10, 0:4]), -np.ones((1, 4), dtype="float32"))
+ # the preceding (non-edge) chunk must be untouched
+ np.testing.assert_array_equal(
+ np.asarray(z[6:9, 0:4]), np.arange(90, dtype="float32").reshape(10, 9)[6:9, 0:4]
+ )
+
+
+def test_zarrista_reads_are_writable(tmp_path: Path) -> None:
+ # zarrista's `Tensor` wraps Rust-owned memory that `np.asarray` exposes
+ # read-only. zarr-python reads have always returned writable arrays, and
+ # nothing downstream of the engine copies any more, so the engine itself
+ # must -- on a whole-array read as much as on a partial one.
+ _make(tmp_path)
+ ze = zarr.open_array(LocalStore(tmp_path), engine="zarrista")
+
+ full = np.asarray(ze[:, :])
+ assert full.flags.writeable
+ partial = np.asarray(ze[2:7, 1:5])
+ assert partial.flags.writeable
+
+
+def test_zarrista_rejects_v2(tmp_path: Path) -> None:
+ zarr.create_array(LocalStore(tmp_path), shape=(4,), chunks=(2,), dtype="int8", zarr_format=2)
+ with pytest.raises(UnsupportedEngineError, match="v3"):
+ zarr.open_array(LocalStore(tmp_path), engine="zarrista")
+
+
+def test_zarrista_vlen_read_matches_default_engine(tmp_path: Path) -> None:
+ # zarrista decodes a vlen dtype to a `VariableLengthTensor`, which the
+ # engine exports through `to_numpy()`. Reads must therefore agree with the
+ # default engine element for element, both whole and partial -- the shapes
+ # are what a raw `np.asarray` on the tensor would get wrong.
+ z = zarr.create_array(LocalStore(tmp_path), shape=(4,), chunks=(2,), dtype="str")
+ z[:] = np.array(["a", "bb", "ccc", "dddd"], dtype=object)
+ ze = zarr.open_array(LocalStore(tmp_path), engine="zarrista")
+
+ np.testing.assert_array_equal(np.asarray(ze[:]), np.asarray(z[:]))
+ np.testing.assert_array_equal(np.asarray(ze[1:3]), np.asarray(z[1:3]))
+
+
+def test_zarrista_sync_rejects_read_missing_chunks_false(tmp_path: Path) -> None:
+ # The zarrista engine cannot enforce read_missing_chunks=False (it fills
+ # missing chunks instead of raising), so minting a sync engine with that
+ # config must fail loudly rather than silently downgrade the semantics.
+ from zarr.core.array_spec import ArrayConfig
+ from zarr.zarrista._engine import ZarristaHierarchyEngine
+
+ z = _make(tmp_path)
+ config = ArrayConfig(order="C", write_empty_chunks=False, read_missing_chunks=False)
+ hierarchy = ZarristaHierarchyEngine(LocalStore(tmp_path))
+ with pytest.raises(UnsupportedEngineError, match="read_missing_chunks=False"):
+ hierarchy.array_engine("", z.metadata, config)
+
+
+async def test_zarrista_async_rejects_read_missing_chunks_false(tmp_path: Path) -> None:
+ # Same fail-loud contract as the sync engine, exercised on the async
+ # hierarchy engine's `array_engine` factory.
+ from zarr.core.array_spec import ArrayConfig
+ from zarr.zarrista._engine import ZarristaAsyncHierarchyEngine
+
+ z = _make(tmp_path)
+ config = ArrayConfig(order="C", write_empty_chunks=False, read_missing_chunks=False)
+ hierarchy = ZarristaAsyncHierarchyEngine(LocalStore(tmp_path))
+ with pytest.raises(UnsupportedEngineError, match="read_missing_chunks=False"):
+ hierarchy.array_engine("", z.metadata, config)
+
+
+async def _async_arrays(tmp_path: Path) -> tuple[Any, Any, npt.NDArray[np.float32]]:
+ """A filled array over an obstore-backed store, opened on both engines.
+
+ The async zarrista engine refuses a `LocalStore` (see `translate_store_async`),
+ so an `ObjectStore` is the store to exercise it with.
+ """
+ obstore = pytest.importorskip("obstore")
+ from zarr.api import asynchronous as async_api
+ from zarr.storage import ObjectStore
+
+ store = ObjectStore(obstore.store.LocalStore(str(tmp_path)))
+ z = await async_api.create_array(store=store, shape=(10, 9), chunks=(3, 4), dtype="float32")
+ data = np.arange(90, dtype="float32").reshape(10, 9)
+ await z.setitem((slice(None), slice(None)), data)
+ ze = await async_api.open_array(store=store, engine="zarrista")
+ # guard against a silent fallback making the assertions below vacuous
+ assert isinstance(ze.engine, ZarristaAsyncEngine)
+ return z, ze, data
+
+
+async def test_zarrista_async_engine_reads(tmp_path: Path) -> None:
+ # exercises `ZarristaAsyncEngine` directly, rather than through the sync
+ # `Array`/`ZarristaEngine` path the other tests in this module cover. The
+ # engine resolves its `LazyArray` off the event loop and hands each zarrista
+ # call back to it, so every dialect has to survive that round trip.
+ _, ze, data = await _async_arrays(tmp_path)
+
+ contiguous = await ze.getitem((slice(2, 7), slice(1, 5)))
+ np.testing.assert_array_equal(np.asarray(contiguous), data[2:7, 1:5])
+
+ strided = await ze.getitem((slice(1, 9, 2), slice(None, None, 3)))
+ np.testing.assert_array_equal(np.asarray(strided), data[1:9:2, ::3])
+
+ ortho = await ze.get_orthogonal_selection((np.array([7, 1, 4]), np.array([0, 8])))
+ np.testing.assert_array_equal(np.asarray(ortho), data[np.ix_([7, 1, 4], [0, 8])])
+
+ points = await ze.get_coordinate_selection((np.array([9, 0, 3]), np.array([8, 0, 2])))
+ np.testing.assert_array_equal(np.asarray(points), data[[9, 0, 3], [8, 0, 2]])
+
+ mask = np.zeros((10, 9), dtype=bool)
+ mask[0, 0] = mask[4, 5] = mask[9, 8] = True
+ np.testing.assert_array_equal(np.asarray(await ze.get_mask_selection(mask)), data[mask])
+
+
+async def test_zarrista_async_engine_writes(tmp_path: Path) -> None:
+ # `AsyncArray` exposes only basic writes, so this covers the box and
+ # read-modify-write tiers of the engine's write path; the fancy write tiers
+ # are covered through the sync engine in tests/engine/test_differential.py.
+ z, ze, data = await _async_arrays(tmp_path)
+ expected = data.copy()
+
+ await ze.setitem((slice(0, 3), slice(0, 4)), np.zeros((3, 4), dtype="float32"))
+ expected[0:3, 0:4] = 0.0
+ # a partial-chunk write, which the engine has to read-modify-write
+ await ze.setitem((slice(4, 6), slice(2, 9)), np.float32(99.0))
+ expected[4:6, 2:9] = 99.0
+ # strided, so the write scatters pointwise rather than laying down a box
+ await ze.setitem((slice(1, 9, 3), slice(None, None, 4)), np.float32(-1.0))
+ expected[1:9:3, ::4] = -1.0
+
+ # read back through the *default* engine: the bytes have to be right on
+ # disk, not merely round-trip through zarrista
+ np.testing.assert_array_equal(np.asarray(await z.getitem((slice(None), slice(None)))), expected)
diff --git a/tests/zarrista/test_translate.py b/tests/zarrista/test_translate.py
new file mode 100644
index 0000000000..0c47a8c924
--- /dev/null
+++ b/tests/zarrista/test_translate.py
@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+
+pytest.importorskip("zarrista")
+
+import zarrista
+
+from zarr.errors import UnsupportedEngineError
+from zarr.storage import LocalStore, MemoryStore
+from zarr.zarrista._translate import translate_store_async, translate_store_sync
+
+if TYPE_CHECKING:
+ from pathlib import Path
+
+
+def test_local_store_translates_to_filesystem_store(tmp_path: Path) -> None:
+ zs = translate_store_sync(LocalStore(tmp_path))
+ assert isinstance(zs, zarrista.store.FilesystemStore)
+
+
+def test_memory_store_rejected_sync(tmp_path: Path) -> None:
+ with pytest.raises(UnsupportedEngineError, match="MemoryStore"):
+ translate_store_sync(MemoryStore())
+
+
+def test_local_store_rejected_async(tmp_path: Path) -> None:
+ # async side wants obstore/icechunk; LocalStore is sync-only in v1
+ with pytest.raises(UnsupportedEngineError):
+ translate_store_async(LocalStore(tmp_path))
+
+
+def test_object_store_translates_to_obstore(tmp_path: Path) -> None:
+ obstore = pytest.importorskip("obstore")
+ from zarr.storage import ObjectStore
+
+ inner = obstore.store.LocalStore(prefix=str(tmp_path))
+ zstore = ObjectStore(inner)
+ assert translate_store_async(zstore) is inner
diff --git a/uv.lock b/uv.lock
index 5e3e33aefe..d60c9a6d52 100644
--- a/uv.lock
+++ b/uv.lock
@@ -1080,6 +1080,24 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/dd/08/681d4a272cd2812151581c3328e41a80a34e420d676e419a25b4b9dc2291/hypothesis-6.164.0-cp314-cp314t-win_amd64.whl", hash = "sha256:a845e59fae87bb47a6fb84e0d5adb5679b3b55042fc3f8791da91486103cfbf0", size = 660724, upload-time = "2026-07-30T12:38:40.341Z" },
]
+[[package]]
+name = "icechunk"
+version = "2.1.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "zarr" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/f9/ec/7d5188b499f43365af257ca2a254cd129b50fbfb60cb493a583a8dc39e9d/icechunk-2.1.2.tar.gz", hash = "sha256:762daf1ece116903b25c66555bc07bf9756de26f2e7d0cdbdb6d69530aa35f82", size = 3614380, upload-time = "2026-07-29T15:49:05.207Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/16/90/6b33d354836ef7bc3e4dbde5c3158e3994804481ed3739c6d64da7bae768/icechunk-2.1.2-cp312-abi3-macosx_10_12_x86_64.whl", hash = "sha256:1bc73fcea2411d1bab13851f46eed51cbe409b80c6b8754dbfc366f26a81a28f", size = 16422141, upload-time = "2026-07-29T15:49:11.531Z" },
+ { url = "https://files.pythonhosted.org/packages/d7/df/e5e473dff92d15488b85cf52152608ae236d5c60d1ee9ed8a0d1438dbb96/icechunk-2.1.2-cp312-abi3-macosx_11_0_arm64.whl", hash = "sha256:6e8039298886b73cb8d81fe36cab00a180cef31426f2ac3d5e7a42e47d4ff0fc", size = 15146547, upload-time = "2026-07-29T15:49:09.104Z" },
+ { url = "https://files.pythonhosted.org/packages/a7/2f/d6843aa47e282e55280dceb48196aa90e1f102a983b7086223cc4900c13f/icechunk-2.1.2-cp312-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:499d61e24487e318986b61371cb095e926d3d8f6cc1989cd5d394b1c7887614b", size = 17062708, upload-time = "2026-07-29T15:49:06.799Z" },
+ { url = "https://files.pythonhosted.org/packages/4c/0a/68e186bd8e80032f738e21b16f882cba2c65beb22e60187d21b3e5ad39f0/icechunk-2.1.2-cp312-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:f53f6d63da11846a1d32dd49b2cfb12f5bb6a5bca42b9b459748e8bdabb4df12", size = 16717011, upload-time = "2026-07-29T15:49:02.924Z" },
+ { url = "https://files.pythonhosted.org/packages/8b/1f/4e2b0058c2909734b10c8761179d35adfe6aaae776b5ca90609016deb999/icechunk-2.1.2-cp312-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:f409c6e17f89f85e26e2c227938c3dd23c9e873dc72997514356174779c19557", size = 16944698, upload-time = "2026-07-29T15:49:13.68Z" },
+ { url = "https://files.pythonhosted.org/packages/24/7f/626edb15fc21c9fd596f9212db1b93c3fe5d36fc8c2eeb0b97097dd48777/icechunk-2.1.2-cp312-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:41e3881cbe704ce605c071d9522a9246efacc358700bdaf75fff80fc366e8c76", size = 17519173, upload-time = "2026-07-29T15:49:15.94Z" },
+ { url = "https://files.pythonhosted.org/packages/11/e7/907e22a3dd9ecbbd2488e1901b55e6273925549a301e2203566348078002/icechunk-2.1.2-cp312-abi3-win_amd64.whl", hash = "sha256:bf796aaa99e3ad883d2125b9388237a253eb2d326475cc3e2833786ec62cb4a8", size = 16046437, upload-time = "2026-07-29T15:49:18.929Z" },
+]
+
[[package]]
name = "idna"
version = "3.15"
@@ -3544,6 +3562,24 @@ test = [
{ name = "tomlkit" },
{ name = "uv" },
]
+zarrista = [
+ { name = "coverage" },
+ { name = "hypothesis" },
+ { name = "icechunk" },
+ { name = "numpydoc" },
+ { name = "obstore" },
+ { name = "pytest" },
+ { name = "pytest-accept" },
+ { name = "pytest-asyncio" },
+ { name = "pytest-benchmark" },
+ { name = "pytest-codspeed" },
+ { name = "pytest-cov" },
+ { name = "pytest-xdist" },
+ { name = "tomlkit" },
+ { name = "uv" },
+ { name = "zarr-indexing" },
+ { name = "zarrista" },
+]
[package.metadata]
requires-dist = [
@@ -3648,3 +3684,53 @@ test = [
{ name = "tomlkit", specifier = "==0.15.1" },
{ name = "uv", specifier = "==0.12.0" },
]
+zarrista = [
+ { name = "coverage", specifier = "==7.15.2" },
+ { name = "hypothesis", specifier = "==6.164.0" },
+ { name = "icechunk", specifier = ">=1.1.21" },
+ { name = "numpydoc", specifier = "==1.10.0" },
+ { name = "obstore", specifier = ">=0.10.1" },
+ { name = "pytest", specifier = "==9.1.1" },
+ { name = "pytest-accept", specifier = "==0.3.0" },
+ { name = "pytest-asyncio", specifier = "==1.4.0" },
+ { name = "pytest-benchmark", specifier = "==5.2.3" },
+ { name = "pytest-codspeed", specifier = "==5.0.3" },
+ { name = "pytest-cov", specifier = "==7.1.0" },
+ { name = "pytest-xdist", specifier = "==3.8.0" },
+ { name = "tomlkit", specifier = "==0.15.1" },
+ { name = "uv", specifier = "==0.12.0" },
+ { name = "zarr-indexing", specifier = ">=0.2.1" },
+ { name = "zarrista", git = "https://github.com/developmentseed/zarrista" },
+]
+
+[[package]]
+name = "zarr-indexing"
+version = "0.2.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "numpy" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/45/b8/803a47caac58ea691a533c5ede821e6a772fc928963b5bae42bbfc82a8a8/zarr_indexing-0.2.1.tar.gz", hash = "sha256:ae8246d284d242d504698708f402db159a8ce75e8610b0a5754e03ccd7073613", size = 267699, upload-time = "2026-08-12T14:42:44.536Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/7d/62/d0e463ac800062a509779ffb2eec3a44dba74330a0c6630852d58226a336/zarr_indexing-0.2.1-py3-none-any.whl", hash = "sha256:e607e492876c82be6d5c22904e75060fadbf2340ab4364b5bf5a0821adee625e", size = 108080, upload-time = "2026-08-12T14:42:43.138Z" },
+]
+
+[[package]]
+name = "zarr-metadata"
+version = "0.4.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+ { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/04/76/ca9334554b9a875c2c68203e2ce8b0a3febe1c3e9efba00a57e9a12c4e40/zarr_metadata-0.4.0.tar.gz", hash = "sha256:56788af0b86ec653176410fd57e0f4815331ac61d61e562dfd020198e98ce3eb", size = 149567, upload-time = "2026-07-29T20:20:38.513Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/cb/8b/f8594d9160f0eaae3f988eb2969bafc9db96bf3e224b8199b1bb80fa1f7a/zarr_metadata-0.4.0-py3-none-any.whl", hash = "sha256:642106291b881284d08222973d0970dd950d254b2286f36fcc269b0b1a1f4be2", size = 69750, upload-time = "2026-07-29T20:20:37.059Z" },
+]
+
+[[package]]
+name = "zarrista"
+version = "0.1.0rc1"
+source = { git = "https://github.com/developmentseed/zarrista#e2bf68a58d55f9da1c248222b479610528c75e18" }
+dependencies = [
+ { name = "zarr-metadata" },
+]