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..ebaf18c11d --- /dev/null +++ b/changes/4181.feature.md @@ -0,0 +1,22 @@ +`Array` and `AsyncArray` now route data I/O through pluggable *array engines* +(`zarr.abc.engine.ArrayEngine` / `AsyncArrayEngine`). The default engine +preserves existing behavior on every store and format. Pass +`engine="zarrista"` (requires the `zarrista` package) to serve Zarr v3 arrays +on local-filesystem, obstore, or icechunk storage through the Rust `zarrs` +implementation. Call `zarr.list_engines()` to discover the available engine +names. + +As part of this rewiring, writing a single field of a structured dtype (e.g. +`z["x"] = ...`) now preserves the values of the array's other fields; it +previously clobbered them. + +Behavior note: strided and fancy (orthogonal/coordinate) selections are now +served internally via the contiguous bounding box that covers the selected +indices, with the requested elements picked out of that box afterwards. For +sparse selections this can transfer more data from storage than before, and +with the `array.read_missing_chunks` config set to `False`, a missing chunk +that lies within the bounding box will raise even if no selected element +falls inside that chunk. Relatedly, a partial (non-full-box) write now performs +an internal read of the bounding box before writing it back, so with +`array.read_missing_chunks` set to `False` such a write raises +`ChunkNotFoundError` when it covers an uninitialized chunk. 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/superpowers/plans/2026-06-12-zarrs-functional-api-phase1.md b/docs/superpowers/plans/2026-06-12-zarrs-functional-api-phase1.md new file mode 100644 index 0000000000..364de7cebc --- /dev/null +++ b/docs/superpowers/plans/2026-06-12-zarrs-functional-api-phase1.md @@ -0,0 +1,1796 @@ +# zarrs functional API (Phase 1) Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** A new in-repo PyO3 crate `zarrs-bindings` plus a `zarr.zarrs` subpackage exposing an async functional API (node lifecycle + whole-chunk I/O) that delegates to the Rust `zarrs` crate, working against any zarr-python `Store`. + +**Architecture:** Two layers. The Rust crate (`zarrs-bindings/`, maturin/PyO3 abi3-py312, native module `_zarrs_bindings`) is a thin binding over `zarrs` ≈0.23: functions take metadata as JSON strings, a store object, and a node path. The Python subpackage `src/zarr/zarrs/` owns the public API: dict metadata documents, `Store` adaptation (native `LocalStore` fast path + a generic sync-shim callback bridge), numpy conversion, and error translation. Spec: `docs/superpowers/specs/2026-06-11-zarrs-functional-api-design.md`. + +**Tech Stack:** Rust 1.91+ (1.96 installed), zarrs 0.23 (default features), pyo3 0.28 (abi3-py312), maturin build backend driven by uv (no maturin CLI needed), pytest with `asyncio_mode = "auto"`. + +--- + +## Environment notes (read first) + +- **Python/pytest/mypy always via `uv run`** (user preference). +- **Build/refresh the extension:** `uv sync --group zarrs --reinstall-package zarrs-bindings`. Plain `uv run --group zarrs ...` does NOT reliably rebuild after Rust edits — always re-sync with `--reinstall-package zarrs-bindings` after touching `zarrs-bindings/`. +- **Fast Rust feedback:** `cargo check --manifest-path zarrs-bindings/Cargo.toml` (compiles without packaging a wheel). +- Builds need network access (crates.io for cargo, PyPI for maturin). The Claude Code sandbox on this host fails at bwrap init, so run build commands with the sandbox disabled. +- Pre-commit hooks (ruff format/check, mypy, codespell) run on `git commit`. If a hook modifies files, `git add` the changes and commit again. +- The Rust snippets below were written against verified zarrs 0.23.13 / zarrs_storage 0.4.3 signatures. If `cargo check` reports a mismatch (most likely candidates: the exact signature of `zarrs::node::node_exists`, the re-export path of `store_set_partial_many`, or `TryInto for &NodePath`), check https://docs.rs/zarrs/latest — the primitives all exist; only spelling may need adjustment. +- Docstrings use **markdown** (mkdocs), single backticks — not RST. + +## File structure + +``` +zarrs-bindings/ # new Rust crate (own wheel: zarrs-bindings / _zarrs_bindings) + Cargo.toml + pyproject.toml # maturin backend + src/lib.rs # pymodule, exceptions, shared error helpers + src/store.rs # PyStore bridge + store resolution + src/node.rs # group/array creation, read_metadata, delete_node, list_children + src/chunk.rs # retrieve/store/erase chunk, retrieve_encoded_chunk +src/zarr/zarrs/ # new Python subpackage (public API) + __init__.py # import guard + re-exports + _bridge.py # StoreShim (sync adapter over async Store), resolve_store + _api.py # async functional API, numpy/JSON conversion, error translation +tests/zarrs/ # new test directory (skips when bindings missing) + __init__.py + conftest.py # store fixtures, array_metadata helper + test_bridge.py + test_node.py + test_chunk.py +.github/workflows/zarrs.yml # new CI job +pyproject.toml # modified: zarrs dependency group, uv source, sdist exclude +.gitignore # modified: zarrs-bindings/target/ +changes/+zarrs-bindings.feature.md +``` + +--- + +### Task 1: Rust crate scaffolding + uv wiring + +**Files:** +- Create: `zarrs-bindings/Cargo.toml` +- Create: `zarrs-bindings/pyproject.toml` +- Create: `zarrs-bindings/src/lib.rs` +- Modify: `pyproject.toml` (root) +- Modify: `.gitignore` + +- [ ] **Step 1: Create `zarrs-bindings/Cargo.toml`** + +```toml +[package] +name = "zarrs-bindings" +version = "0.1.0" +edition = "2024" +rust-version = "1.91" +publish = false + +[lib] +name = "_zarrs_bindings" +crate-type = ["cdylib"] + +[dependencies] +pyo3 = { version = "0.28", features = ["abi3-py312"] } +serde_json = "1" +zarrs = "0.23" + +[profile.release] +lto = "thin" +``` + +- [ ] **Step 2: Create `zarrs-bindings/pyproject.toml`** + +```toml +[build-system] +requires = ["maturin>=1.7,<2"] +build-backend = "maturin" + +[project] +name = "zarrs-bindings" +version = "0.1.0" +description = "PyO3 bindings to the zarrs Rust crate, consumed by zarr.zarrs" +requires-python = ">=3.12" +license = "MIT" + +[tool.maturin] +module-name = "_zarrs_bindings" +strip = true +``` + +- [ ] **Step 3: Create `zarrs-bindings/src/lib.rs`** (exceptions + version only for now) + +```rust +use pyo3::exceptions::{PyRuntimeError, PyValueError}; +use pyo3::prelude::*; + +pyo3::create_exception!( + _zarrs_bindings, + NodeExistsError, + PyValueError, + "A node already exists at the given path." +); +pyo3::create_exception!( + _zarrs_bindings, + NodeNotFoundError, + PyValueError, + "No node was found at the given path." +); + +pub(crate) fn runtime_err(err: impl std::fmt::Display) -> PyErr { + PyRuntimeError::new_err(err.to_string()) +} + +pub(crate) fn value_err(err: impl std::fmt::Display) -> PyErr { + PyValueError::new_err(err.to_string()) +} + +#[pyfunction] +fn version() -> &'static str { + env!("CARGO_PKG_VERSION") +} + +#[pymodule] +fn _zarrs_bindings(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add("NodeExistsError", m.py().get_type::())?; + m.add("NodeNotFoundError", m.py().get_type::())?; + m.add_function(wrap_pyfunction!(version, m)?)?; + Ok(()) +} +``` + +- [ ] **Step 4: Wire into the root `pyproject.toml`** + +Add to the `[dependency-groups]` table (after the `dev` group): + +```toml +zarrs = [ + {include-group = "test"}, + "zarrs-bindings", +] +``` + +Add a new section at the end of the file: + +```toml +[tool.uv.sources] +zarrs-bindings = { path = "zarrs-bindings" } +``` + +Add `"/zarrs-bindings",` to the `exclude` list under `[tool.hatch.build.targets.sdist]`. + +- [ ] **Step 5: Add `zarrs-bindings/target/` to `.gitignore`** + +- [ ] **Step 6: Lock, build, smoke-test** + +Run: `cargo check --manifest-path zarrs-bindings/Cargo.toml` +Expected: compiles clean (first run downloads ~zarrs dependency tree). + +Run: `uv lock && uv sync --group zarrs` +Expected: lockfile updated; `zarrs-bindings` builds via maturin and installs. + +Run: `uv run --group zarrs python -c "import _zarrs_bindings as z; print(z.version())"` +Expected: `0.1.0` + +- [ ] **Step 7: Commit** (include `zarrs-bindings/Cargo.lock`, which the build created) + +```bash +git add zarrs-bindings .gitignore pyproject.toml uv.lock +git commit -m "feat: scaffold zarrs-bindings PyO3 crate" +``` + +--- + +### Task 2: `zarr.zarrs` package skeleton + test scaffolding + +**Files:** +- Create: `src/zarr/zarrs/__init__.py` +- Create: `tests/zarrs/__init__.py` (empty) +- Create: `tests/zarrs/conftest.py` +- Test: `tests/zarrs/test_api.py` + +- [ ] **Step 1: Write the failing test** — `tests/zarrs/test_api.py` + +```python +from __future__ import annotations + + +def test_import() -> None: + import zarr.zarrs + + assert isinstance(zarr.zarrs.__version__, str) +``` + +- [ ] **Step 2: Create `tests/zarrs/__init__.py`** (empty file) **and `tests/zarrs/conftest.py`** + +```python +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import pytest + +pytest.importorskip("_zarrs_bindings", reason="zarrs-bindings is not installed") + +import zarr +from zarr.storage import LocalStore, MemoryStore + +if TYPE_CHECKING: + from pathlib import Path + + from zarr.abc.store import Store + + +@pytest.fixture(params=["memory", "local"]) +async def store(request: pytest.FixtureRequest, tmp_path: Path) -> Store: + """A writable store: MemoryStore exercises the generic Python-callback bridge, + LocalStore exercises the native zarrs filesystem store.""" + if request.param == "memory": + return await MemoryStore.open() + return await LocalStore.open(root=tmp_path / "store") + + +def array_metadata(**kwargs: Any) -> dict[str, Any]: + """Build an array metadata document using zarr-python itself, so the + documents fed to zarrs always match what zarr-python would write.""" + params: dict[str, Any] = { + "shape": (8, 8), + "chunks": (4, 4), + "dtype": "uint16", + "zarr_format": 3, + } | kwargs + arr = zarr.create_array(store=MemoryStore(), **params) + doc = dict(arr.metadata.to_dict()) + if params["zarr_format"] == 2: + # v2 attributes live in .zattrs, not in the .zarray document + doc.pop("attributes", None) + return doc +``` + +- [ ] **Step 3: Run the test to verify it fails** + +Run: `uv run --group zarrs pytest tests/zarrs -v` +Expected: FAIL with `ModuleNotFoundError: No module named 'zarr.zarrs'` + +- [ ] **Step 4: Create `src/zarr/zarrs/__init__.py`** + +```python +""" +Low-level functional API for zarr hierarchies, backed by the Rust +[`zarrs`](https://zarrs.dev) crate. + +This subpackage is experimental. It requires the `zarrs-bindings` package +(in-repo Rust crate; install for development with `uv sync --group zarrs`). + +All array routines take an explicit metadata document (a `dict` matching the +`zarr.json` / `.zarray` document) rather than reading metadata from the store, +which makes read-only and virtual views possible. +""" + +try: + import _zarrs_bindings +except ImportError as e: + raise ImportError( + "zarr.zarrs requires the `zarrs-bindings` package, which is not installed. " + "It is built from the zarr-python repository: run `uv sync --group zarrs`." + ) from e + +__version__: str = _zarrs_bindings.version() + +__all__ = ["__version__"] +``` + +- [ ] **Step 5: Run the test to verify it passes** + +Run: `uv run --group zarrs pytest tests/zarrs -v` +Expected: 1 passed. Also verify the skip path works in the default env: `uv run pytest tests/zarrs -v` → all skipped/deselected with "zarrs-bindings is not installed" (the default group lacks the bindings). + +- [ ] **Step 6: Commit** + +```bash +git add src/zarr/zarrs tests/zarrs +git commit -m "feat: add zarr.zarrs package skeleton and test scaffolding" +``` + +--- + +### Task 3: StoreShim — sync bridge over async stores (pure Python, TDD) + +**Files:** +- Create: `src/zarr/zarrs/_bridge.py` +- Test: `tests/zarrs/test_bridge.py` + +- [ ] **Step 1: Write the failing tests** — `tests/zarrs/test_bridge.py` + +```python +from __future__ import annotations + +from typing import TYPE_CHECKING + +from zarr.storage import LocalStore, MemoryStore +from zarr.zarrs._bridge import StoreShim, resolve_store + +if TYPE_CHECKING: + from pathlib import Path + + +def test_shim_get_set_delete() -> None: + shim = StoreShim(MemoryStore()) + assert shim.get("a/b") is None + shim.set("a/b", b"xyz") + assert shim.get("a/b") == b"xyz" + assert shim.get_range("a/b", 1, 1) == b"y" + assert shim.get_range("a/b", 1, None) == b"yz" + assert shim.get_suffix("a/b", 2) == b"yz" + assert shim.getsize("a/b") == 3 + assert shim.getsize("missing") is None + shim.delete("a/b") + assert shim.get("a/b") is None + + +def test_shim_listing() -> None: + shim = StoreShim(MemoryStore()) + shim.set("zarr.json", b"{}") + shim.set("a/zarr.json", b"{}") + shim.set("a/c/0/0", b"\x00") + assert shim.list() == ["a/c/0/0", "a/zarr.json", "zarr.json"] + assert shim.list_prefix("a/") == ["a/c/0/0", "a/zarr.json"] + assert shim.list_dir("a/") == (["a/zarr.json"], ["a/c/"]) + assert shim.list_dir("") == (["zarr.json"], ["a/"]) + assert shim.getsize_prefix("a/") == 3 + shim.delete_prefix("a/") + assert shim.list() == ["zarr.json"] + + +def test_resolve_store(tmp_path: Path) -> None: + local = LocalStore(tmp_path) + assert resolve_store(local) == {"filesystem": str(tmp_path)} + # read-only LocalStore must go through the shim so writes are rejected in Python + assert isinstance(resolve_store(LocalStore(tmp_path, read_only=True)), StoreShim) + assert isinstance(resolve_store(MemoryStore()), StoreShim) +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `uv run --group zarrs pytest tests/zarrs/test_bridge.py -v` +Expected: FAIL with `ModuleNotFoundError: No module named 'zarr.zarrs._bridge'` + +- [ ] **Step 3: Create `src/zarr/zarrs/_bridge.py`** + +```python +from __future__ import annotations + +from typing import TYPE_CHECKING + +from zarr.abc.store import OffsetByteRequest, RangeByteRequest, SuffixByteRequest +from zarr.core.buffer.core import default_buffer_prototype +from zarr.core.sync import _collect_aiterator, sync +from zarr.storage import LocalStore + +if TYPE_CHECKING: + from zarr.abc.store import Store + + +class StoreShim: + """ + Synchronous adapter over an async `Store`, called from Rust worker threads. + + Each method blocks the calling thread by submitting a coroutine to the zarr + event-loop thread (`zarr.core.sync`). Methods must never be called from the + zarr event-loop thread itself; the Rust bindings only call them from + `asyncio.to_thread` worker threads. + """ + + def __init__(self, store: Store) -> None: + self._store = store + self._prototype = default_buffer_prototype() + + def get(self, key: str) -> bytes | None: + buf = sync(self._store.get(key, prototype=self._prototype)) + return None if buf is None else buf.to_bytes() + + def get_range(self, key: str, offset: int, length: int | None) -> bytes | None: + byte_range = ( + RangeByteRequest(offset, offset + length) + if length is not None + else OffsetByteRequest(offset) + ) + buf = sync(self._store.get(key, prototype=self._prototype, byte_range=byte_range)) + return None if buf is None else buf.to_bytes() + + def get_suffix(self, key: str, suffix: int) -> bytes | None: + buf = sync( + self._store.get(key, prototype=self._prototype, byte_range=SuffixByteRequest(suffix)) + ) + return None if buf is None else buf.to_bytes() + + def set(self, key: str, value: bytes) -> None: + sync(self._store.set(key, self._prototype.buffer.from_bytes(value))) + + def delete(self, key: str) -> None: + sync(self._store.delete(key)) + + def delete_prefix(self, prefix: str) -> None: + sync(self._store.delete_dir(prefix.rstrip("/"))) + + def getsize(self, key: str) -> int | None: + try: + return sync(self._store.getsize(key)) + except FileNotFoundError: + return None + + def getsize_prefix(self, prefix: str) -> int: + return sync(self._store.getsize_prefix(prefix.rstrip("/"))) + + def list(self) -> list[str]: + return sorted(sync(_collect_aiterator(self._store.list()))) + + def list_prefix(self, prefix: str) -> list[str]: + return sorted(sync(_collect_aiterator(self._store.list_prefix(prefix)))) + + def list_dir(self, prefix: str) -> tuple[list[str], list[str]]: + """Return `(keys, prefixes)` directly under `prefix`, as zarrs expects: + full keys, and child prefixes ending in `/`.""" + stripped = prefix.rstrip("/") + children = sorted(sync(_collect_aiterator(self._store.list_dir(stripped)))) + keys: list[str] = [] + prefixes: list[str] = [] + for child in children: + full = f"{stripped}/{child}" if stripped else child + if sync(self._store.exists(full)): + keys.append(full) + else: + prefixes.append(full + "/") + return keys, prefixes + + +def resolve_store(store: Store) -> StoreShim | dict[str, str]: + """ + Convert a zarr `Store` into the representation `_zarrs_bindings` expects: + a config dict for stores with a native Rust implementation, otherwise a + `StoreShim` that Rust calls back into. + """ + if isinstance(store, LocalStore) and not store.read_only: + return {"filesystem": str(store.root)} + return StoreShim(store) +``` + +- [ ] **Step 4: Run tests to verify they pass** + +Run: `uv run --group zarrs pytest tests/zarrs/test_bridge.py -v` +Expected: 3 passed. (If `list_dir`/`delete_dir`/`getsize_prefix` choke on the stripped prefix, check the `Store` ABC docstrings in `src/zarr/abc/store.py:348-501` — these methods take prefixes without trailing slashes.) + +- [ ] **Step 5: Commit** + +```bash +git add src/zarr/zarrs/_bridge.py tests/zarrs/test_bridge.py +git commit -m "feat: sync store bridge for zarrs bindings" +``` + +--- + +### Task 4: Rust store bridge + group creation, end to end + +**Files:** +- Create: `zarrs-bindings/src/store.rs` +- Create: `zarrs-bindings/src/node.rs` +- Modify: `zarrs-bindings/src/lib.rs` +- Create: `src/zarr/zarrs/_api.py` +- Modify: `src/zarr/zarrs/__init__.py` +- Test: `tests/zarrs/test_node.py` + +- [ ] **Step 1: Write the failing tests** — `tests/zarrs/test_node.py` + +```python +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import pytest + +import zarr +from zarr.core.buffer.core import default_buffer_prototype +from zarr.zarrs import NodeExistsError, create_new_group, create_overwrite_group + +if TYPE_CHECKING: + from zarr.abc.store import Store + +GROUP_META: dict[str, Any] = { + "zarr_format": 3, + "node_type": "group", + "attributes": {"answer": 42}, +} + + +async def test_create_new_group(store: Store) -> None: + await create_new_group(GROUP_META, store, "foo") + group = zarr.open_group(store=store, path="foo", mode="r") + assert dict(group.attrs) == {"answer": 42} + + +async def test_create_new_group_at_root(store: Store) -> None: + await create_new_group(GROUP_META, store, "") + group = zarr.open_group(store=store, mode="r") + assert dict(group.attrs) == {"answer": 42} + + +async def test_create_new_group_existing_node(store: Store) -> None: + await create_new_group(GROUP_META, store, "foo") + with pytest.raises(NodeExistsError): + await create_new_group(GROUP_META, store, "foo") + + +async def test_create_overwrite_group(store: Store) -> None: + # an array and its chunks previously occupied the path; overwrite removes both + arr = zarr.create_array(store=store, name="foo", shape=(4,), chunks=(2,), dtype="uint8") + arr[:] = 1 + assert await store.exists("foo/c/0") + await create_overwrite_group(GROUP_META, store, "foo") + group = zarr.open_group(store=store, path="foo", mode="r") + assert dict(group.attrs) == {"answer": 42} + assert not await store.exists("foo/c/0") + assert await store.get("foo/zarr.json", prototype=default_buffer_prototype()) is not None +``` + +- [ ] **Step 2: Run tests to verify they fail** + +Run: `uv run --group zarrs pytest tests/zarrs/test_node.py -v` +Expected: FAIL with `ImportError: cannot import name 'NodeExistsError' from 'zarr.zarrs'` + +- [ ] **Step 3: Create `zarrs-bindings/src/store.rs`** + +```rust +use std::sync::Arc; + +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::types::{PyBytes, PyDict}; +use zarrs::filesystem::FilesystemStore; +use zarrs::storage::{ + Bytes, ByteRange, ByteRangeIterator, ListableStorageTraits, MaybeBytes, MaybeBytesIterator, + OffsetBytesIterator, ReadableStorageTraits, ReadableWritableListableStorage, StorageError, + StoreKey, StoreKeys, StoreKeysPrefixes, StorePrefix, WritableStorageTraits, +}; + +/// A zarrs store backed by a Python `zarr.zarrs._bridge.StoreShim`. +/// +/// Every method attaches to the Python interpreter and calls the shim, which +/// blocks on the zarr event loop. Blocking waits in Python release the GIL, so +/// the loop thread can make progress while a Rust worker waits here. +pub(crate) struct PyStore(Py); + +fn py_err(err: PyErr) -> StorageError { + StorageError::Other(err.to_string()) +} + +fn invalid(err: impl std::fmt::Display) -> StorageError { + StorageError::Other(err.to_string()) +} + +impl PyStore { + fn get_with_range( + &self, + key: &StoreKey, + range: Option<&ByteRange>, + ) -> Result { + Python::attach(|py| { + let shim = self.0.bind(py); + let result = match range { + None => shim.call_method1("get", (key.as_str(),)), + Some(ByteRange::FromStart(offset, length)) => { + shim.call_method1("get_range", (key.as_str(), *offset, *length)) + } + Some(ByteRange::Suffix(suffix)) => { + shim.call_method1("get_suffix", (key.as_str(), *suffix)) + } + } + .map_err(py_err)?; + if result.is_none() { + Ok(None) + } else { + let bytes: Vec = result.extract().map_err(py_err)?; + Ok(Some(Bytes::from(bytes))) + } + }) + } +} + +impl ReadableStorageTraits for PyStore { + fn get(&self, key: &StoreKey) -> Result { + self.get_with_range(key, None) + } + + fn get_partial_many<'a>( + &'a self, + key: &StoreKey, + byte_ranges: ByteRangeIterator<'a>, + ) -> Result, StorageError> { + let mut out = Vec::new(); + for byte_range in byte_ranges { + match self.get_with_range(key, Some(&byte_range))? { + Some(bytes) => out.push(Ok(bytes)), + None => return Ok(None), + } + } + Ok(Some(Box::new(out.into_iter()))) + } + + fn size_key(&self, key: &StoreKey) -> Result, StorageError> { + Python::attach(|py| { + self.0 + .bind(py) + .call_method1("getsize", (key.as_str(),)) + .map_err(py_err)? + .extract() + .map_err(py_err) + }) + } + + fn supports_get_partial(&self) -> bool { + true + } +} + +impl WritableStorageTraits for PyStore { + fn set(&self, key: &StoreKey, value: Bytes) -> Result<(), StorageError> { + Python::attach(|py| { + let data = PyBytes::new(py, &value); + self.0 + .bind(py) + .call_method1("set", (key.as_str(), data)) + .map_err(py_err)?; + Ok(()) + }) + } + + fn set_partial_many( + &self, + key: &StoreKey, + offset_values: OffsetBytesIterator, + ) -> Result<(), StorageError> { + // read-modify-write fallback provided by zarrs + zarrs::storage::store_set_partial_many(self, key, offset_values) + } + + fn supports_set_partial(&self) -> bool { + false + } + + fn erase(&self, key: &StoreKey) -> Result<(), StorageError> { + Python::attach(|py| { + self.0 + .bind(py) + .call_method1("delete", (key.as_str(),)) + .map_err(py_err)?; + Ok(()) + }) + } + + fn erase_prefix(&self, prefix: &StorePrefix) -> Result<(), StorageError> { + Python::attach(|py| { + self.0 + .bind(py) + .call_method1("delete_prefix", (prefix.as_str(),)) + .map_err(py_err)?; + Ok(()) + }) + } +} + +impl ListableStorageTraits for PyStore { + fn list(&self) -> Result { + Python::attach(|py| { + let keys: Vec = self + .0 + .bind(py) + .call_method0("list") + .map_err(py_err)? + .extract() + .map_err(py_err)?; + keys.into_iter() + .map(|k| StoreKey::new(k).map_err(invalid)) + .collect() + }) + } + + fn list_prefix(&self, prefix: &StorePrefix) -> Result { + Python::attach(|py| { + let keys: Vec = self + .0 + .bind(py) + .call_method1("list_prefix", (prefix.as_str(),)) + .map_err(py_err)? + .extract() + .map_err(py_err)?; + keys.into_iter() + .map(|k| StoreKey::new(k).map_err(invalid)) + .collect() + }) + } + + fn list_dir(&self, prefix: &StorePrefix) -> Result { + Python::attach(|py| { + let (keys, prefixes): (Vec, Vec) = self + .0 + .bind(py) + .call_method1("list_dir", (prefix.as_str(),)) + .map_err(py_err)? + .extract() + .map_err(py_err)?; + let keys = keys + .into_iter() + .map(|k| StoreKey::new(k).map_err(invalid)) + .collect::, StorageError>>()?; + let prefixes = prefixes + .into_iter() + .map(|p| StorePrefix::new(p).map_err(invalid)) + .collect::, StorageError>>()?; + Ok(StoreKeysPrefixes::new(keys, prefixes)) + }) + } + + fn size_prefix(&self, prefix: &StorePrefix) -> Result { + Python::attach(|py| { + self.0 + .bind(py) + .call_method1("getsize_prefix", (prefix.as_str(),)) + .map_err(py_err)? + .extract() + .map_err(py_err) + }) + } +} + +/// Convert the Python-side store representation (`zarr.zarrs._bridge.resolve_store` +/// output) into a zarrs storage handle. +pub(crate) fn resolve_store(obj: &Bound<'_, PyAny>) -> PyResult { + if let Ok(config) = obj.downcast::() { + if let Some(root) = config.get_item("filesystem")? { + let root: String = root.extract()?; + let store = + FilesystemStore::new(root).map_err(|e| PyValueError::new_err(e.to_string()))?; + return Ok(Arc::new(store)); + } + return Err(PyValueError::new_err("unrecognized store configuration")); + } + Ok(Arc::new(PyStore(obj.clone().unbind()))) +} +``` + +- [ ] **Step 4: Create `zarrs-bindings/src/node.rs`** (group functions only; later tasks extend this file) + +```rust +use pyo3::prelude::*; +use zarrs::group::Group; +use zarrs::metadata::GroupMetadata; +use zarrs::node::{node_exists, NodePath}; +use zarrs::storage::{ReadableWritableListableStorage, StorePrefix}; + +use crate::store::resolve_store; +use crate::{runtime_err, value_err, NodeExistsError}; + +pub(crate) fn parse_node_path(path: &str) -> PyResult { + NodePath::new(path).map_err(value_err) +} + +/// When a node exists at `node_path`: erase it (and everything under it) if +/// `overwrite`, otherwise raise `NodeExistsError`. +pub(crate) fn prepare_target( + storage: &ReadableWritableListableStorage, + node_path: &NodePath, + overwrite: bool, +) -> PyResult<()> { + if node_exists(storage, node_path).map_err(runtime_err)? { + if !overwrite { + return Err(NodeExistsError::new_err(format!( + "a node already exists at path {}", + node_path.as_str() + ))); + } + let prefix: StorePrefix = node_path.try_into().map_err(value_err)?; + storage.erase_prefix(&prefix).map_err(runtime_err)?; + } + Ok(()) +} + +#[pyfunction] +pub(crate) fn create_group( + py: Python<'_>, + store: &Bound<'_, PyAny>, + path: String, + metadata_json: String, + overwrite: bool, +) -> PyResult<()> { + let storage = resolve_store(store)?; + let metadata = GroupMetadata::try_from(metadata_json.as_str()).map_err(value_err)?; + py.detach(move || { + let node_path = parse_node_path(&path)?; + prepare_target(&storage, &node_path, overwrite)?; + let group = Group::new_with_metadata(storage, &path, metadata).map_err(value_err)?; + group.store_metadata().map_err(runtime_err) + }) +} +``` + +- [ ] **Step 5: Register in `zarrs-bindings/src/lib.rs`** + +Add after the `use` lines: + +```rust +mod node; +mod store; +``` + +Add to the `#[pymodule]` body before `Ok(())`: + +```rust + m.add_function(wrap_pyfunction!(node::create_group, m)?)?; +``` + +- [ ] **Step 6: Compile** + +Run: `cargo check --manifest-path zarrs-bindings/Cargo.toml` +Expected: success. If `node_exists` or `try_into::()` signatures mismatch, fix per https://docs.rs/zarrs/latest/zarrs/node/ (the helpers exist; argument form may differ, e.g. `node_exists(&storage, &node_path)` vs a `&Arc` receiver). + +- [ ] **Step 7: Create `src/zarr/zarrs/_api.py`** + +```python +from __future__ import annotations + +import asyncio +import json +from contextlib import contextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import _zarrs_bindings as _zb + +from zarr.errors import NodeNotFoundError +from zarr.zarrs._bridge import resolve_store + +if TYPE_CHECKING: + from collections.abc import Iterator, Mapping + + from zarr.abc.store import Store + from zarr.core.common import JSON + +NodeExistsError = _zb.NodeExistsError +"""Raised by `create_new_*` when a node already exists at the target path.""" + + +@dataclass(frozen=True, slots=True) +class ZarrsOptions: + """Options for zarrs-backed operations. + + Currently empty: fields (concurrency limits, checksum validation) arrive in + a later phase. Accepting it now keeps signatures stable. + """ + + +def _node_path(path: str) -> str: + """Convert a zarr-python node path (`""`, `"foo/bar"`) to a zarrs node path + (`"/"`, `"/foo/bar"`).""" + return f"/{path.strip('/')}" + + +@contextmanager +def _translate_errors() -> Iterator[None]: + try: + yield + except _zb.NodeNotFoundError as err: + raise NodeNotFoundError(str(err)) from err + + +async def create_new_group( + metadata: Mapping[str, JSON], + store: Store, + path: str, + *, + options: ZarrsOptions | None = None, +) -> None: + """Create a group at `path` from a group metadata document. + + Raises `NodeExistsError` if any node already exists at `path`. + """ + with _translate_errors(): + await asyncio.to_thread( + _zb.create_group, resolve_store(store), _node_path(path), json.dumps(metadata), False + ) + + +async def create_overwrite_group( + metadata: Mapping[str, JSON], + store: Store, + path: str, + *, + options: ZarrsOptions | None = None, +) -> None: + """Create a group at `path`, deleting any existing node (and its children) first.""" + with _translate_errors(): + await asyncio.to_thread( + _zb.create_group, resolve_store(store), _node_path(path), json.dumps(metadata), True + ) +``` + +- [ ] **Step 8: Re-export from `src/zarr/zarrs/__init__.py`** + +Replace the `__version__`/`__all__` lines at the end with: + +```python +__version__: str = _zarrs_bindings.version() + +from zarr.zarrs._api import ( + NodeExistsError, + ZarrsOptions, + create_new_group, + create_overwrite_group, +) + +__all__ = [ + "NodeExistsError", + "ZarrsOptions", + "__version__", + "create_new_group", + "create_overwrite_group", +] +``` + +- [ ] **Step 9: Rebuild and run the tests** + +Run: `uv sync --group zarrs --reinstall-package zarrs-bindings` +Run: `uv run --group zarrs pytest tests/zarrs/test_node.py -v` +Expected: 8 passed (4 tests × 2 store params). The MemoryStore param proves the full Rust→Python callback bridge; LocalStore proves the native path. + +- [ ] **Step 10: Commit** + +```bash +git add zarrs-bindings/src src/zarr/zarrs tests/zarrs/test_node.py +git commit -m "feat: zarrs store bridge and group creation" +``` + +--- + +### Task 5: Array creation + read_metadata + +**Files:** +- Modify: `zarrs-bindings/src/node.rs` +- Modify: `zarrs-bindings/src/lib.rs` +- Modify: `src/zarr/zarrs/_api.py`, `src/zarr/zarrs/__init__.py` +- Test: `tests/zarrs/test_node.py` + +- [ ] **Step 1: Add failing tests to `tests/zarrs/test_node.py`** + +Extend the imports: + +```python +import json + +import numpy as np + +from tests.zarrs.conftest import array_metadata +from zarr.errors import NodeNotFoundError +from zarr.zarrs import create_new_array, create_overwrite_array, read_metadata +``` + +(If `from tests.zarrs.conftest import ...` fails at collection, use a relative import `from .conftest import array_metadata` — `tests` is a package.) + +Add tests: + +```python +async def test_create_new_array(store: Store) -> None: + await create_new_array(array_metadata(), store, "arr") + arr = zarr.open_array(store=store, path="arr", mode="r") + assert arr.shape == (8, 8) + assert arr.chunks == (4, 4) + assert arr.dtype == np.dtype("uint16") + + +async def test_create_new_array_existing_node(store: Store) -> None: + await create_new_array(array_metadata(), store, "arr") + with pytest.raises(NodeExistsError): + await create_new_array(array_metadata(), store, "arr") + + +async def test_create_overwrite_array(store: Store) -> None: + zarr.create_group(store=store, path="arr") + await create_overwrite_array(array_metadata(), store, "arr") + arr = zarr.open_array(store=store, path="arr", mode="r") + assert arr.shape == (8, 8) + + +async def test_read_metadata_matches_stored_document(store: Store) -> None: + await create_new_array(array_metadata(), store, "arr") + observed = await read_metadata(store, "arr") + raw = await store.get("arr/zarr.json", prototype=default_buffer_prototype()) + assert raw is not None + assert observed == json.loads(raw.to_bytes()) + + +async def test_read_metadata_zarr_python_group(store: Store) -> None: + zarr.create_group(store=store, path="g", attributes={"a": 1}) + observed = await read_metadata(store, "g") + assert observed["node_type"] == "group" + assert observed["attributes"] == {"a": 1} + + +async def test_read_metadata_missing(store: Store) -> None: + with pytest.raises(NodeNotFoundError): + await read_metadata(store, "nope") +``` + +- [ ] **Step 2: Run to verify failure** + +Run: `uv run --group zarrs pytest tests/zarrs/test_node.py -v` +Expected: FAIL with `ImportError: cannot import name 'create_new_array'` + +- [ ] **Step 3: Add Rust functions to `zarrs-bindings/src/node.rs`** + +Extend the `use` block: + +```rust +use zarrs::array::Array; +use zarrs::metadata::ArrayMetadata; +use zarrs::node::Node; + +use crate::NodeNotFoundError; +``` + +Append: + +```rust +#[pyfunction] +pub(crate) fn create_array( + py: Python<'_>, + store: &Bound<'_, PyAny>, + path: String, + metadata_json: String, + overwrite: bool, +) -> PyResult<()> { + let storage = resolve_store(store)?; + let metadata = ArrayMetadata::try_from(metadata_json.as_str()).map_err(value_err)?; + py.detach(move || { + let node_path = parse_node_path(&path)?; + prepare_target(&storage, &node_path, overwrite)?; + let array = Array::new_with_metadata(storage, &path, metadata).map_err(value_err)?; + array.store_metadata().map_err(runtime_err) + }) +} + +#[pyfunction] +pub(crate) fn read_metadata( + py: Python<'_>, + store: &Bound<'_, PyAny>, + path: String, +) -> PyResult { + let storage = resolve_store(store)?; + py.detach(move || { + let node = Node::open(&storage, &path) + .map_err(|e| NodeNotFoundError::new_err(e.to_string()))?; + serde_json::to_string(node.metadata()).map_err(runtime_err) + }) +} +``` + +Register both in `lib.rs`: + +```rust + m.add_function(wrap_pyfunction!(node::create_array, m)?)?; + m.add_function(wrap_pyfunction!(node::read_metadata, m)?)?; +``` + +- [ ] **Step 4: Add Python wrappers to `src/zarr/zarrs/_api.py`** + +```python +async def create_new_array( + metadata: Mapping[str, JSON], + store: Store, + path: str, + *, + options: ZarrsOptions | None = None, +) -> None: + """Create an array at `path` from a v2 or v3 array metadata document. + + Raises `NodeExistsError` if any node already exists at `path`. + """ + with _translate_errors(): + await asyncio.to_thread( + _zb.create_array, resolve_store(store), _node_path(path), json.dumps(metadata), False + ) + + +async def create_overwrite_array( + metadata: Mapping[str, JSON], + store: Store, + path: str, + *, + options: ZarrsOptions | None = None, +) -> None: + """Create an array at `path`, deleting any existing node (and its children) first.""" + with _translate_errors(): + await asyncio.to_thread( + _zb.create_array, resolve_store(store), _node_path(path), json.dumps(metadata), True + ) + + +async def read_metadata( + store: Store, + path: str, + *, + options: ZarrsOptions | None = None, +) -> dict[str, JSON]: + """Read the metadata document of the array or group at `path`. + + Raises `zarr.errors.NodeNotFoundError` if no node exists there. + """ + with _translate_errors(): + raw = await asyncio.to_thread(_zb.read_metadata, resolve_store(store), _node_path(path)) + result: dict[str, JSON] = json.loads(raw) + return result +``` + +Add `create_new_array`, `create_overwrite_array`, `read_metadata` to the `__init__.py` import and `__all__`. + +- [ ] **Step 5: Rebuild and test** + +Run: `cargo check --manifest-path zarrs-bindings/Cargo.toml` → success +Run: `uv sync --group zarrs --reinstall-package zarrs-bindings` +Run: `uv run --group zarrs pytest tests/zarrs/test_node.py -v` +Expected: all pass (20 = 10 tests × 2 stores). Note: `test_read_metadata_matches_stored_document` asserts zarrs round-trips the document zarrs itself wrote; if zarrs normalizes a field zarr-python emits differently (e.g. drops a `null` `dimension_names`), adjust the *fixture* (`array_metadata`) to drop the field, not the assertion. + +- [ ] **Step 6: Commit** + +```bash +git add zarrs-bindings/src src/zarr/zarrs tests/zarrs +git commit -m "feat: zarrs-backed array creation and metadata reads" +``` + +--- + +### Task 6: delete_node + list_children + +**Files:** +- Modify: `zarrs-bindings/src/node.rs`, `zarrs-bindings/src/lib.rs` +- Modify: `src/zarr/zarrs/_api.py`, `src/zarr/zarrs/__init__.py` +- Test: `tests/zarrs/test_node.py` + +- [ ] **Step 1: Add failing tests to `tests/zarrs/test_node.py`** + +```python +from zarr.zarrs import delete_node, list_children + + +async def test_delete_node(store: Store) -> None: + arr = zarr.create_array(store=store, name="doomed", shape=(4,), chunks=(2,), dtype="uint8") + arr[:] = 1 + await delete_node(store, "doomed") + assert not await store.exists("doomed/zarr.json") + assert not await store.exists("doomed/c/0") + + +async def test_delete_node_missing(store: Store) -> None: + with pytest.raises(NodeNotFoundError): + await delete_node(store, "nope") + + +async def test_list_children(store: Store) -> None: + root = zarr.create_group(store=store) + root.create_group("sub_group", attributes={"kind": "group"}) + root.create_array("sub_array", shape=(4,), chunks=(2,), dtype="uint8") + children = await list_children(store, "") + by_path = dict(children) + assert set(by_path) == {"sub_group", "sub_array"} + assert by_path["sub_group"]["node_type"] == "group" + assert by_path["sub_array"]["node_type"] == "array" + + +async def test_list_children_missing(store: Store) -> None: + with pytest.raises(NodeNotFoundError): + await list_children(store, "nope") +``` + +- [ ] **Step 2: Run to verify failure** — `uv run --group zarrs pytest tests/zarrs/test_node.py -v` → ImportError. + +- [ ] **Step 3: Add Rust functions to `node.rs`** + +```rust +#[pyfunction] +pub(crate) fn delete_node( + py: Python<'_>, + store: &Bound<'_, PyAny>, + path: String, +) -> PyResult<()> { + let storage = resolve_store(store)?; + py.detach(move || { + let node_path = parse_node_path(&path)?; + if !node_exists(&storage, &node_path).map_err(runtime_err)? { + return Err(NodeNotFoundError::new_err(format!( + "no node found at path {}", + node_path.as_str() + ))); + } + let prefix: StorePrefix = (&node_path).try_into().map_err(value_err)?; + storage.erase_prefix(&prefix).map_err(runtime_err) + }) +} + +#[pyfunction] +pub(crate) fn list_children( + py: Python<'_>, + store: &Bound<'_, PyAny>, + path: String, +) -> PyResult> { + let storage = resolve_store(store)?; + py.detach(move || { + let group = Group::open(storage, &path) + .map_err(|e| NodeNotFoundError::new_err(e.to_string()))?; + let children = group.children(false).map_err(runtime_err)?; + children + .into_iter() + .map(|node| { + let metadata = serde_json::to_string(node.metadata()).map_err(runtime_err)?; + Ok((node.path().as_str().to_string(), metadata)) + }) + .collect() + }) +} +``` + +Register both in `lib.rs` as before. + +- [ ] **Step 4: Add Python wrappers to `_api.py`** + +```python +async def delete_node( + store: Store, + path: str, + *, + options: ZarrsOptions | None = None, +) -> None: + """Delete the node at `path`, including all keys and child nodes under it. + + Raises `zarr.errors.NodeNotFoundError` if no node exists there. Deleting the + root node (`path=""`) clears the entire store. + """ + with _translate_errors(): + await asyncio.to_thread(_zb.delete_node, resolve_store(store), _node_path(path)) + + +async def list_children( + store: Store, + path: str, + *, + options: ZarrsOptions | None = None, +) -> list[tuple[str, dict[str, JSON]]]: + """List the direct children of the group at `path` as + `(path, metadata_document)` pairs. Paths are store-relative (no leading `/`). + + Raises `zarr.errors.NodeNotFoundError` if no group exists at `path`. + """ + with _translate_errors(): + raw = await asyncio.to_thread(_zb.list_children, resolve_store(store), _node_path(path)) + return [(child_path.lstrip("/"), json.loads(doc)) for child_path, doc in raw] +``` + +Export both from `__init__.py`. + +- [ ] **Step 5: Rebuild and test** + +Run: `uv sync --group zarrs --reinstall-package zarrs-bindings && uv run --group zarrs pytest tests/zarrs/test_node.py -v` +Expected: all pass. + +- [ ] **Step 6: Commit** + +```bash +git add zarrs-bindings/src src/zarr/zarrs tests/zarrs +git commit -m "feat: zarrs-backed node deletion and child listing" +``` + +--- + +### Task 7: Whole-chunk I/O (decode/encode/raw/erase) + +**Files:** +- Create: `zarrs-bindings/src/chunk.rs` +- Modify: `zarrs-bindings/src/lib.rs` +- Modify: `src/zarr/zarrs/_api.py`, `src/zarr/zarrs/__init__.py` +- Test: `tests/zarrs/test_chunk.py` + +- [ ] **Step 1: Write the failing tests** — `tests/zarrs/test_chunk.py` + +```python +from __future__ import annotations + +import copy +from typing import TYPE_CHECKING, Any + +import numpy as np +import pytest + +import zarr +from tests.zarrs.conftest import array_metadata +from zarr.codecs import BloscCodec, GzipCodec, ZstdCodec +from zarr.core.buffer.core import default_buffer_prototype +from zarr.zarrs import ( + create_new_array, + decode_chunk, + encode_chunk, + erase_chunk, + read_encoded_chunk, +) + +if TYPE_CHECKING: + from zarr.abc.store import Store + + +def _filled( + store: Store, **kwargs: Any +) -> tuple[np.ndarray[Any, np.dtype[Any]], dict[str, Any]]: + """Create an 8x8 array named 'a' via zarr-python, fill it with a ramp, and + return (data, metadata_document).""" + params: dict[str, Any] = {"shape": (8, 8), "chunks": (4, 4), "dtype": "uint16"} | kwargs + arr = zarr.create_array(store=store, name="a", **params) + data = np.arange(64, dtype=params["dtype"]).reshape(8, 8) + arr[:, :] = data + doc = dict(arr.metadata.to_dict()) + if params.get("zarr_format") == 2: + # v2 attributes live in .zattrs, not in the .zarray document + doc.pop("attributes", None) + return data, doc + + +@pytest.mark.parametrize("dtype", ["uint8", "int32", "float64"]) +async def test_decode_chunk_differential(store: Store, dtype: str) -> None: + data, meta = _filled(store, dtype=dtype) + observed = await decode_chunk(meta, store, "a", (1, 0)) + np.testing.assert_array_equal(observed, data[4:8, 0:4]) + + +@pytest.mark.parametrize( + "compressors", [None, (GzipCodec(),), (ZstdCodec(),), (BloscCodec(cname="lz4"),)] +) +async def test_decode_chunk_codecs(store: Store, compressors: Any) -> None: + data, meta = _filled(store, compressors=compressors) + observed = await decode_chunk(meta, store, "a", (0, 1)) + np.testing.assert_array_equal(observed, data[0:4, 4:8]) + + +async def test_decode_chunk_v2(store: Store) -> None: + data, meta = _filled(store, zarr_format=2) + observed = await decode_chunk(meta, store, "a", (1, 1)) + np.testing.assert_array_equal(observed, data[4:8, 4:8]) + + +async def test_decode_chunk_sharding(store: Store) -> None: + # with sharding, the metadata chunk grid is the shard grid + data, meta = _filled(store, chunks=(2, 2), shards=(4, 4)) + observed = await decode_chunk(meta, store, "a", (1, 1)) + np.testing.assert_array_equal(observed, data[4:8, 4:8]) + + +async def test_decode_chunk_missing_returns_fill_value(store: Store) -> None: + arr = zarr.create_array( + store=store, name="a", shape=(8, 8), chunks=(4, 4), dtype="uint16", fill_value=7 + ) + meta = dict(arr.metadata.to_dict()) + observed = await decode_chunk(meta, store, "a", (0, 0)) + np.testing.assert_array_equal(observed, np.full((4, 4), 7, dtype="uint16")) + + +async def test_decode_chunk_selection_not_implemented(store: Store) -> None: + _, meta = _filled(store) + with pytest.raises(NotImplementedError): + await decode_chunk(meta, store, "a", (0, 0), selection=(slice(0, 2), slice(0, 2))) + + +async def test_decode_chunk_metadata_view(store: Store) -> None: + # the read-only-view case: decode with a metadata document the store never saw + data, meta = _filled(store, dtype="uint16", compressors=None) + view = copy.deepcopy(meta) + view["data_type"] = "uint8" + view["shape"] = [8, 16] + view["chunk_grid"]["configuration"]["chunk_shape"] = [4, 8] + observed = await decode_chunk(view, store, "a", (1, 0)) + np.testing.assert_array_equal(observed, data[4:8, 0:4].view("uint8")) + + +async def test_encode_chunk_differential(store: Store) -> None: + meta = array_metadata() + await create_new_array(meta, store, "a") + value = np.arange(16, dtype="uint16").reshape(4, 4) + await encode_chunk(meta, store, "a", (0, 1), value) + arr = zarr.open_array(store=store, path="a", mode="r") + np.testing.assert_array_equal(arr[0:4, 4:8], value) + + +async def test_encode_chunk_shape_mismatch(store: Store) -> None: + meta = array_metadata() + await create_new_array(meta, store, "a") + with pytest.raises(ValueError, match="chunk shape"): + await encode_chunk(meta, store, "a", (0, 0), np.zeros((2, 2), dtype="uint16")) + + +async def test_read_encoded_chunk_matches_store(store: Store) -> None: + _, meta = _filled(store) + raw = await read_encoded_chunk(meta, store, "a", (0, 0)) + expected = await store.get("a/c/0/0", prototype=default_buffer_prototype()) + assert expected is not None + assert raw == expected.to_bytes() + + +async def test_read_encoded_chunk_missing_returns_none(store: Store) -> None: + arr = zarr.create_array(store=store, name="empty", shape=(8, 8), chunks=(4, 4), dtype="uint16") + meta = dict(arr.metadata.to_dict()) + assert await read_encoded_chunk(meta, store, "empty", (0, 0)) is None + + +async def test_erase_chunk(store: Store) -> None: + data, meta = _filled(store) + assert await store.exists("a/c/0/0") + await erase_chunk(meta, store, "a", (0, 0)) + assert not await store.exists("a/c/0/0") + arr = zarr.open_array(store=store, path="a", mode="r") + np.testing.assert_array_equal(arr[0:4, 0:4], np.zeros((4, 4), dtype="uint16")) +``` + +- [ ] **Step 2: Run to verify failure** — `uv run --group zarrs pytest tests/zarrs/test_chunk.py -v` → ImportError. + +- [ ] **Step 3: Create `zarrs-bindings/src/chunk.rs`** + +```rust +use pyo3::exceptions::PyNotImplementedError; +use pyo3::prelude::*; +use pyo3::types::PyBytes; +use zarrs::array::{Array, ArrayBytes}; +use zarrs::metadata::ArrayMetadata; +use zarrs::storage::ReadableWritableListableStorage; + +use crate::store::resolve_store; +use crate::{runtime_err, value_err}; + +type DynArray = Array; + +/// Construct an Array view from an explicit metadata document, without +/// consulting the store for metadata. +fn array_view( + storage: ReadableWritableListableStorage, + path: &str, + metadata_json: &str, +) -> PyResult { + let metadata = ArrayMetadata::try_from(metadata_json).map_err(value_err)?; + Array::new_with_metadata(storage, path, metadata).map_err(value_err) +} + +#[pyfunction] +pub(crate) fn retrieve_chunk( + py: Python<'_>, + store: &Bound<'_, PyAny>, + path: String, + metadata_json: String, + chunk_coords: Vec, +) -> PyResult> { + let storage = resolve_store(store)?; + let data = py.detach(move || -> PyResult> { + let array = array_view(storage, &path, &metadata_json)?; + let bytes: ArrayBytes<'static> = + array.retrieve_chunk(&chunk_coords).map_err(runtime_err)?; + let fixed = bytes.into_fixed().map_err(|_| { + PyNotImplementedError::new_err("variable-length data types are not supported") + })?; + Ok(fixed.into_owned()) + })?; + Ok(PyBytes::new(py, &data).unbind()) +} + +#[pyfunction] +pub(crate) fn retrieve_encoded_chunk( + py: Python<'_>, + store: &Bound<'_, PyAny>, + path: String, + metadata_json: String, + chunk_coords: Vec, +) -> PyResult>> { + let storage = resolve_store(store)?; + let data = py.detach(move || -> PyResult>> { + let array = array_view(storage, &path, &metadata_json)?; + array + .retrieve_encoded_chunk(&chunk_coords) + .map_err(runtime_err) + })?; + Ok(data.map(|d| PyBytes::new(py, &d).unbind())) +} + +#[pyfunction] +pub(crate) fn store_chunk( + py: Python<'_>, + store: &Bound<'_, PyAny>, + path: String, + metadata_json: String, + chunk_coords: Vec, + data: Vec, +) -> PyResult<()> { + let storage = resolve_store(store)?; + py.detach(move || { + let array = array_view(storage, &path, &metadata_json)?; + array + .store_chunk(&chunk_coords, ArrayBytes::new_flen(data)) + .map_err(runtime_err) + }) +} + +#[pyfunction] +pub(crate) fn erase_chunk( + py: Python<'_>, + store: &Bound<'_, PyAny>, + path: String, + metadata_json: String, + chunk_coords: Vec, +) -> PyResult<()> { + let storage = resolve_store(store)?; + py.detach(move || { + let array = array_view(storage, &path, &metadata_json)?; + array.erase_chunk(&chunk_coords).map_err(runtime_err) + }) +} +``` + +Register in `lib.rs`: add `mod chunk;` and + +```rust + m.add_function(wrap_pyfunction!(chunk::retrieve_chunk, m)?)?; + m.add_function(wrap_pyfunction!(chunk::retrieve_encoded_chunk, m)?)?; + m.add_function(wrap_pyfunction!(chunk::store_chunk, m)?)?; + m.add_function(wrap_pyfunction!(chunk::erase_chunk, m)?)?; +``` + +- [ ] **Step 4: Add Python wrappers to `_api.py`** + +Extend imports: + +```python +from typing import Any + +import numpy as np +import numpy.typing as npt +``` + +Add: + +```python +def _chunk_dtype_and_shape( + metadata: Mapping[str, JSON], +) -> tuple[np.dtype[Any], tuple[int, ...]]: + """Resolve the numpy dtype and chunk shape from a metadata document, using + zarr-python's own metadata parsing.""" + from zarr.core.metadata.v2 import ArrayV2Metadata + from zarr.core.metadata.v3 import ArrayV3Metadata, RegularChunkGridMetadata + + if metadata.get("zarr_format") == 3: + meta3 = ArrayV3Metadata.from_dict(dict(metadata)) + grid = meta3.chunk_grid + if not isinstance(grid, RegularChunkGridMetadata): + raise NotImplementedError("only regular chunk grids are supported") + return meta3.data_type.to_native_dtype(), grid.chunk_shape + meta2 = ArrayV2Metadata.from_dict(dict(metadata)) + return meta2.dtype.to_native_dtype(), meta2.chunks + + +async def decode_chunk( + metadata: Mapping[str, JSON], + store: Store, + path: str, + chunk_coords: tuple[int, ...], + *, + selection: tuple[slice | int, ...] | None = None, + options: ZarrsOptions | None = None, +) -> np.ndarray[Any, np.dtype[Any]]: + """Read and decode the chunk at `chunk_coords` of the array described by + `metadata`, located at `path` in `store`. + + The metadata document is authoritative: it is not read from the store. + Missing chunks decode to the fill value. `selection` (a chunk-relative + subset) is not implemented yet. + """ + if selection is not None: + raise NotImplementedError("chunk subset selection is not implemented yet") + raw = await asyncio.to_thread( + _zb.retrieve_chunk, + resolve_store(store), + _node_path(path), + json.dumps(metadata), + list(chunk_coords), + ) + dtype, chunk_shape = _chunk_dtype_and_shape(metadata) + return np.frombuffer(raw, dtype=dtype).reshape(chunk_shape) + + +async def read_encoded_chunk( + metadata: Mapping[str, JSON], + store: Store, + path: str, + chunk_coords: tuple[int, ...], + *, + options: ZarrsOptions | None = None, +) -> bytes | None: + """Read the raw, still-encoded bytes of the chunk at `chunk_coords`, or + `None` if the chunk does not exist. No codecs are applied.""" + result: bytes | None = await asyncio.to_thread( + _zb.retrieve_encoded_chunk, + resolve_store(store), + _node_path(path), + json.dumps(metadata), + list(chunk_coords), + ) + return result + + +async def encode_chunk( + metadata: Mapping[str, JSON], + store: Store, + path: str, + chunk_coords: tuple[int, ...], + value: npt.ArrayLike, + *, + options: ZarrsOptions | None = None, +) -> None: + """Encode `value` with the codecs in `metadata` and store it as the chunk + at `chunk_coords`. `value` must match the chunk shape exactly.""" + dtype, chunk_shape = _chunk_dtype_and_shape(metadata) + arr = np.ascontiguousarray(np.asarray(value, dtype=dtype)) + if arr.shape != chunk_shape: + raise ValueError(f"value shape {arr.shape} does not match chunk shape {chunk_shape}") + await asyncio.to_thread( + _zb.store_chunk, + resolve_store(store), + _node_path(path), + json.dumps(metadata), + list(chunk_coords), + arr.tobytes(), + ) + + +async def erase_chunk( + metadata: Mapping[str, JSON], + store: Store, + path: str, + chunk_coords: tuple[int, ...], + *, + options: ZarrsOptions | None = None, +) -> None: + """Delete the chunk at `chunk_coords`. Deleting a missing chunk is a no-op.""" + await asyncio.to_thread( + _zb.erase_chunk, + resolve_store(store), + _node_path(path), + json.dumps(metadata), + list(chunk_coords), + ) +``` + +Export `decode_chunk`, `read_encoded_chunk`, `encode_chunk`, `erase_chunk` from `__init__.py`. + +- [ ] **Step 5: Rebuild and test** + +Run: `cargo check --manifest-path zarrs-bindings/Cargo.toml` → success +Run: `uv sync --group zarrs --reinstall-package zarrs-bindings` +Run: `uv run --group zarrs pytest tests/zarrs/test_chunk.py -v` +Expected: all pass. Likely first-run issues and their fixes: + - v2 differential test fails on dtype byte order → constrain the v2 test to `dtype=" None` annotations, which the code above has). + +- [ ] **Step 3: Re-run the full zarrs suite** — `uv run --group zarrs pytest tests/zarrs -v` → all pass. + +- [ ] **Step 4: Verify the rest of the test suite is unaffected** + +Run: `uv run pytest tests/test_array.py tests/test_group.py -x -q` +Expected: pass (no production code outside `src/zarr/zarrs/` changed). + +- [ ] **Step 5: Commit** + +```bash +git add -A +git commit -m "chore: lint fixes and changelog for zarr.zarrs" +``` + +--- + +### Task 9: CI workflow + +**Files:** +- Create: `.github/workflows/zarrs.yml` + +- [ ] **Step 1: Create `.github/workflows/zarrs.yml`** (action SHAs copied from `.github/workflows/test.yml` — keep them identical so dependabot groups them) + +```yaml +name: Zarrs bindings + +on: + push: + branches: [ main ] + pull_request: + branches: [ main ] + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + test: + runs-on: ubuntu-latest + 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 uv + uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # v8.2.0 + with: + python-version: '3.12' + - name: Run zarrs bindings tests + # the ubuntu runner image ships a Rust toolchain; the maturin build + # backend is fetched by uv on demand + run: uv run --group zarrs pytest tests/zarrs -v +``` + +- [ ] **Step 2: Validate the workflow** + +Run: `uvx zizmor .github/workflows/zarrs.yml` +Expected: no findings (matches the repo's zizmor policy). + +- [ ] **Step 3: Commit** + +```bash +git add .github/workflows/zarrs.yml +git commit -m "ci: test job for zarrs bindings" +``` + +--- + +## Out of scope for this plan (later phases, per spec) + +- `decode_region` / `encode_region` and chunk-subset `selection` (Phase 2: zarrs `retrieve_array_subset` / `partial_decoder`). +- `ZarrsOptions` fields (concurrency, checksum validation, direct IO), obstore native path, benchmarks (Phase 3). +- Variable-length data types, non-regular chunk grids, fancy indexing. +- Publishing the `zarrs-bindings` wheel / a `zarr[zarrs]` extra on PyPI. diff --git a/docs/superpowers/plans/2026-06-15-crud-backend-abstraction.md b/docs/superpowers/plans/2026-06-15-crud-backend-abstraction.md new file mode 100644 index 0000000000..985f325990 --- /dev/null +++ b/docs/superpowers/plans/2026-06-15-crud-backend-abstraction.md @@ -0,0 +1,1698 @@ +# Backend-agnostic CRUD layer Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Turn the low-level functional CRUD API into a backend-agnostic `zarr.crud` package with a pure-Python reference backend and the existing zarrs bindings as a second, interchangeable backend. + +**Architecture:** A narrow async `CrudBackend` protocol (byte/metadata level) plus a shared `zarr.crud` facade that holds all backend-neutral logic (selection normalization, numpy assembly, dtype handling, `read_encoded_chunk` via `store.get`). Two backends conform: `ReferenceBackend` (pure Python, wraps zarr-python's own codec pipeline / indexer / metadata machinery) and `ZarrsBackend` (wraps `_zarrs_bindings`). A registry + `zarr.config` key `crud.backend` (default `"reference"`) selects one; every facade function also takes `backend=`. + +**Tech Stack:** Python 3.12+, numpy, zarr-python internals (`BatchedCodecPipeline`, `AsyncArray`, `save_metadata`, `ArrayConfig`/`ArraySpec`, chunk-key encoding), the existing `_zarrs_bindings` Rust extension (unchanged — no Rust build needed). + +Spec: `docs/superpowers/specs/2026-06-15-crud-backend-abstraction-design.md`. + +--- + +## Environment notes (read first) + +- **Run python/pytest/mypy via `uv run`.** The zarrs backend needs the extension: `uv run --group zarrs pytest ...`. The reference backend works under plain `uv run pytest ...`. +- The Claude Code bash sandbox is broken on this host (`bwrap: loopback` error). Run commands with the sandbox **disabled**. +- **No Rust changes in this plan.** The `_zarrs_bindings` pyfunctions keep their existing names (`retrieve_chunk`, `store_chunk`, `erase_chunk`, `retrieve_array_subset`, `retrieve_encoded_chunk`, `create_array`, `create_group`, `read_metadata`, `delete_node`, `list_children`); `ZarrsBackend` adapts them to the contract's verb names. No `cargo` build or `uv sync --reinstall` is required, but the `zarrs` group must already be installed (`uv sync --group zarrs`) to run the zarrs-parametrized tests. +- Pre-commit hooks (ruff strict, mypy strict over `src`+`tests`, codespell) run on `git commit`. If a hook rewrites a file, `git add` and commit again. +- Docstrings use markdown (single backticks), not RST. +- pytest is configured with `asyncio_mode = "auto"` — async tests/fixtures need no decorator. + +## File structure + +``` +src/zarr/crud/ + __init__.py # public exports; registers the reference backend at import + _backend.py # CrudBackend Protocol + NodeExistsError + _registry.py # register_backend / get_backend + config default resolution + _reference.py # ReferenceBackend (pure Python) + _api.py # shared async facade (the public functions) + neutral helpers +src/zarr/zarrs/ + __init__.py # SHRINKS: version + register ZarrsBackend; no _api re-exports + _backend.py # ZarrsBackend (wraps _zarrs_bindings) — NEW + _bridge.py # unchanged + _api.py # DELETED +src/zarr/core/config.py # add "crud": {"backend": "reference"} +tests/crud/ + __init__.py + conftest.py # store fixture, backend fixture (reference+zarrs), metadata helpers + test_registry.py # registry + default + override + test_reference_backend.py # direct reference-backend smoke tests + test_crud.py # full differential suite, parametrized over backend x store +tests/zarrs/ + __init__.py # unchanged + conftest.py # unchanged (still used by test_bridge/test_cache) + test_bridge.py # unchanged + test_cache.py # imports updated to zarr.crud read_chunk/write_chunk, backend="zarrs" + test_node.py # DELETED (covered by tests/crud/test_crud.py) + test_chunk.py # DELETED (covered by tests/crud/test_crud.py) + test_api.py # DELETED (replaced by tests/crud import coverage) +changes/+zarrs-bindings.feature.md # reworded for zarr.crud +.github/workflows/zarrs.yml # run tests/crud tests/zarrs +``` + +--- + +### Task 1: `zarr.crud` skeleton — protocol, exceptions, registry, config + +**Files:** +- Create: `src/zarr/crud/__init__.py` +- Create: `src/zarr/crud/_backend.py` +- Create: `src/zarr/crud/_registry.py` +- Modify: `src/zarr/core/config.py` +- Create: `tests/crud/__init__.py` (empty) +- Test: `tests/crud/test_registry.py` + +- [ ] **Step 1: Write the failing test** — `tests/crud/test_registry.py` + +```python +from __future__ import annotations + +import pytest + +from zarr.crud import CrudBackend, NodeExistsError, get_backend, register_backend + + +def test_node_exists_error_is_value_error() -> None: + assert issubclass(NodeExistsError, ValueError) + + +def test_default_backend_is_reference() -> None: + # the reference backend is registered at import and is the configured default + be = get_backend() + assert be is get_backend("reference") + + +def test_get_unknown_backend_raises() -> None: + with pytest.raises(KeyError, match="no CRUD backend"): + get_backend("does-not-exist") + + +def test_register_and_resolve_instance() -> None: + class Dummy: + pass + + dummy = Dummy() + register_backend("dummy-test", dummy) # type: ignore[arg-type] + try: + assert get_backend("dummy-test") is dummy + finally: + from zarr.crud import _registry + + _registry._BACKENDS.pop("dummy-test", None) + + +def test_protocol_is_runtime_checkable() -> None: + # ReferenceBackend (registered as "reference") structurally satisfies the protocol + assert isinstance(get_backend("reference"), CrudBackend) +``` + +- [ ] **Step 2: Run it to verify failure** + +Run: `uv run pytest tests/crud/test_registry.py -v` +Expected: FAIL — `ModuleNotFoundError: No module named 'zarr.crud'` + +- [ ] **Step 3: Create `tests/crud/__init__.py`** (empty file) + +- [ ] **Step 4: Create `src/zarr/crud/_backend.py`** + +```python +from __future__ import annotations + +from typing import TYPE_CHECKING, Protocol, runtime_checkable + +if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + + from zarr.abc.store import Store + from zarr.core.common import JSON + + +class NodeExistsError(ValueError): + """Raised when a node already exists at a path and overwrite was not requested.""" + + +@runtime_checkable +class CrudBackend(Protocol): + """The byte/metadata-level contract a CRUD backend must implement. + + Methods take neutral types: the metadata document as a `dict`, a zarr + `Store`, and plain zarr paths (`""`, `"foo/bar"`). They return raw bytes, + parsed JSON documents, or `None`. The shared `zarr.crud` facade builds the + numpy- and selection-level API on top of these. + + `create_*` raise `zarr.crud.NodeExistsError` when a node exists and + `overwrite` is false. `read_metadata`/`delete_node`/`list_children` raise + `zarr.errors.NodeNotFoundError` when the target is missing. + """ + + async def create_array( + self, store: Store, path: str, metadata: Mapping[str, JSON], *, overwrite: bool + ) -> None: ... + + async def create_group( + self, store: Store, path: str, metadata: Mapping[str, JSON], *, overwrite: bool + ) -> None: ... + + async def read_metadata(self, store: Store, path: str) -> dict[str, JSON]: ... + + async def read_chunk( + self, store: Store, path: str, metadata: Mapping[str, JSON], coords: tuple[int, ...] + ) -> bytes: ... + + async def read_subset( + self, + store: Store, + path: str, + metadata: Mapping[str, JSON], + start: Sequence[int], + shape: Sequence[int], + ) -> bytes: ... + + async def write_chunk( + self, + store: Store, + path: str, + metadata: Mapping[str, JSON], + coords: tuple[int, ...], + data: bytes, + ) -> None: ... + + async def delete_chunk( + self, store: Store, path: str, metadata: Mapping[str, JSON], coords: tuple[int, ...] + ) -> None: ... + + async def delete_node(self, store: Store, path: str) -> None: ... + + async def list_children( + self, store: Store, path: str + ) -> list[tuple[str, dict[str, JSON]]]: ... +``` + +- [ ] **Step 5: Create `src/zarr/crud/_registry.py`** + +```python +from __future__ import annotations + +from typing import TYPE_CHECKING + +from zarr.core.config import config + +if TYPE_CHECKING: + from zarr.crud._backend import CrudBackend + +_BACKENDS: dict[str, CrudBackend] = {} + + +def register_backend(name: str, backend: CrudBackend) -> None: + """Register a CRUD backend instance under `name`.""" + _BACKENDS[name] = backend + + +def get_backend(name: str | None = None) -> CrudBackend: + """Resolve a backend by name, or the configured default when `name` is None. + + Selecting `"zarrs"` imports `zarr.zarrs` if needed so it can self-register. + """ + if name is None: + name = config.get("crud.backend") + if name not in _BACKENDS and name == "zarrs": + import zarr.zarrs # noqa: F401 (import registers the zarrs backend) + if name not in _BACKENDS: + raise KeyError( + f"no CRUD backend registered as {name!r}; registered: {sorted(_BACKENDS)}" + ) + return _BACKENDS[name] +``` + +- [ ] **Step 6: Create `src/zarr/crud/__init__.py`** (reference backend is added in Task 2; for now register nothing) + +```python +""" +Backend-agnostic low-level functional CRUD API for zarr hierarchies. + +The public functions delegate byte- and metadata-level work to a `CrudBackend`. +Two backends ship: a pure-Python reference backend (the default) and a +zarrs-accelerated backend (`zarr.zarrs`, requires the `zarrs-bindings` +extension). Select one with the `crud.backend` config key or a per-call +`backend=` argument. + +Array routines take an explicit metadata document (a `dict` matching the +`zarr.json` / `.zarray` document) rather than reading it from the store, which +makes read-only and virtual views possible. +""" + +from zarr.crud._backend import CrudBackend, NodeExistsError +from zarr.crud._registry import get_backend, register_backend + +__all__ = [ + "CrudBackend", + "NodeExistsError", + "get_backend", + "register_backend", +] +``` + +- [ ] **Step 7: Add the config default** — `src/zarr/core/config.py` + +Find the defaults mapping passed to the `Config(...)` constructor (it contains the `"codec_pipeline"` key). Add a sibling entry: + +```python + "crud": {"backend": "reference"}, +``` + +Run to confirm it loads: `uv run python -c "from zarr.core.config import config; print(config.get('crud.backend'))"` +Expected: `reference` + +- [ ] **Step 8: Run the test (note: `test_default_backend_is_reference` and the protocol test still fail — reference backend arrives in Task 2)** + +Run: `uv run pytest tests/crud/test_registry.py -v` +Expected: `test_node_exists_error_is_value_error`, `test_get_unknown_backend_raises`, `test_register_and_resolve_instance` PASS; `test_default_backend_is_reference` and `test_protocol_is_runtime_checkable` FAIL (KeyError: no backend `reference`). That is expected at this task boundary; they pass after Task 2. + +- [ ] **Step 9: Commit** + +```bash +git add src/zarr/crud/_backend.py src/zarr/crud/_registry.py src/zarr/crud/__init__.py src/zarr/core/config.py tests/crud/__init__.py tests/crud/test_registry.py +git commit -m "feat: zarr.crud skeleton — CrudBackend protocol, registry, config" +``` + +End every commit body in this plan with: +``` +Co-Authored-By: Claude Fable 5 +``` + +--- + +### Task 2: `ReferenceBackend` (pure Python) + +**Files:** +- Create: `src/zarr/crud/_reference.py` +- Modify: `src/zarr/crud/__init__.py` +- Test: `tests/crud/test_reference_backend.py` + +All snippets below are verified against the installed zarr-python. + +- [ ] **Step 1: Write the failing test** — `tests/crud/test_reference_backend.py` + +```python +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +import zarr +from zarr.crud import NodeExistsError, get_backend +from zarr.errors import NodeNotFoundError +from zarr.storage import MemoryStore + +if TYPE_CHECKING: + pass + +import pytest + + +def _array_meta() -> dict: + arr = zarr.create_array(store=MemoryStore(), shape=(8, 8), chunks=(4, 4), dtype="uint16") + return dict(arr.metadata.to_dict()) + + +async def test_reference_round_trip_chunk() -> None: + be = get_backend("reference") + store = MemoryStore() + meta = _array_meta() + await be.create_array(store, "a", meta, overwrite=False) + value = np.arange(16, dtype="uint16").reshape(4, 4) + await be.write_chunk(store, "a", meta, (0, 1), value.tobytes()) + raw = await be.read_chunk(store, "a", meta, (0, 1)) + np.testing.assert_array_equal(np.frombuffer(raw, dtype="uint16").reshape(4, 4), value) + + +async def test_reference_read_subset_spans_chunks() -> None: + be = get_backend("reference") + store = MemoryStore() + arr = zarr.create_array(store=store, name="a", shape=(8, 8), chunks=(4, 4), dtype="uint16") + data = np.arange(64, dtype="uint16").reshape(8, 8) + arr[:, :] = data + meta = dict(arr.metadata.to_dict()) + raw = await be.read_subset(store, "a", meta, (2, 1), (5, 4)) + np.testing.assert_array_equal( + np.frombuffer(raw, dtype="uint16").reshape(5, 4), data[2:7, 1:5] + ) + + +async def test_reference_create_exists_raises() -> None: + be = get_backend("reference") + store = MemoryStore() + meta = _array_meta() + await be.create_array(store, "a", meta, overwrite=False) + with pytest.raises(NodeExistsError): + await be.create_array(store, "a", meta, overwrite=False) + + +async def test_reference_read_metadata_missing_raises() -> None: + be = get_backend("reference") + with pytest.raises(NodeNotFoundError): + await be.read_metadata(MemoryStore(), "nope") +``` + +- [ ] **Step 2: Run it to verify failure** + +Run: `uv run pytest tests/crud/test_reference_backend.py -v` +Expected: FAIL — `KeyError: no CRUD backend registered as 'reference'` + +- [ ] **Step 3: Create `src/zarr/crud/_reference.py`** + +```python +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import numpy as np + +from zarr.core.array import AsyncArray, create_codec_pipeline +from zarr.core.array_spec import ArrayConfig, ArraySpec +from zarr.core.buffer.core import NDBuffer, default_buffer_prototype +from zarr.core.common import ZARR_JSON, ZARRAY_JSON, ZATTRS_JSON +from zarr.core.group import GroupMetadata +from zarr.core.metadata.io import save_metadata +from zarr.core.metadata.v2 import ArrayV2Metadata +from zarr.core.metadata.v3 import ArrayV3Metadata, RegularChunkGridMetadata +from zarr.crud._backend import NodeExistsError +from zarr.errors import NodeNotFoundError +from zarr.storage._common import StorePath + +if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + + from zarr.abc.store import Store + from zarr.core.common import JSON + + +def _parse_array_metadata( + metadata: Mapping[str, JSON], +) -> ArrayV3Metadata | ArrayV2Metadata: + """Parse a metadata document into a v2 or v3 array metadata object.""" + data = dict(metadata) + if data.get("zarr_format") == 3: + return ArrayV3Metadata.from_dict(data) + return ArrayV2Metadata.from_dict(data) + + +def _native_dtype(meta_obj: ArrayV3Metadata | ArrayV2Metadata) -> np.dtype[Any]: + """Numpy dtype in native byte order (zarrs and the facade assume native).""" + return meta_obj.dtype.to_native_dtype().newbyteorder("=") + + +def _chunk_shape(meta_obj: ArrayV3Metadata | ArrayV2Metadata) -> tuple[int, ...]: + if isinstance(meta_obj, ArrayV3Metadata): + grid = meta_obj.chunk_grid + if not isinstance(grid, RegularChunkGridMetadata): + raise NotImplementedError("only regular chunk grids are supported") + return tuple(grid.chunk_shape) + return tuple(meta_obj.chunks) + + +def _array_spec( + meta_obj: ArrayV3Metadata | ArrayV2Metadata, shape: tuple[int, ...] +) -> ArraySpec: + return ArraySpec( + shape=shape, + dtype=meta_obj.dtype, + fill_value=meta_obj.fill_value, + config=ArrayConfig.from_dict({}), + prototype=default_buffer_prototype(), + ) + + +def _meta_key(path: str, zarr_format: int) -> str: + fname = ZARR_JSON if zarr_format == 3 else ZARRAY_JSON + p = path.strip("/") + return f"{p}/{fname}" if p else fname + + +class ReferenceBackend: + """Pure-Python CRUD backend wrapping zarr-python's own machinery. + + Constructs no high-level `Array` for chunk operations (it drives the codec + pipeline directly); it does reuse `AsyncArray.getitem` for multi-chunk + subset reads, which is exactly the `BasicIndexer` + codec-pipeline read path. + """ + + async def create_array( + self, store: Store, path: str, metadata: Mapping[str, JSON], *, overwrite: bool + ) -> None: + meta_obj = _parse_array_metadata(metadata) + await self._create(store, path, meta_obj, meta_obj.zarr_format, overwrite=overwrite) + + async def create_group( + self, store: Store, path: str, metadata: Mapping[str, JSON], *, overwrite: bool + ) -> None: + meta_obj = GroupMetadata.from_dict(dict(metadata)) + await self._create(store, path, meta_obj, meta_obj.zarr_format, overwrite=overwrite) + + async def _create( + self, store: Store, path: str, meta_obj: Any, zarr_format: int, *, overwrite: bool + ) -> None: + sp = StorePath(store, path.strip("/")) + proto = default_buffer_prototype() + if overwrite: + await store.delete_dir(path.strip("/")) + else: + key = _meta_key(path, zarr_format) + if await store.get(key, prototype=proto) is not None: + raise NodeExistsError(f"a node already exists at path {path!r}") + await save_metadata(sp, meta_obj, ensure_parents=True) + + async def read_metadata(self, store: Store, path: str) -> dict[str, JSON]: + from zarr.core._json import buffer_to_json_object + + proto = default_buffer_prototype() + p = path.strip("/") + sp = StorePath(store, p) + buf = await (sp / ZARR_JSON).get(prototype=proto) + if buf is not None: + return buffer_to_json_object(buf) + buf2 = await (sp / ZARRAY_JSON).get(prototype=proto) + if buf2 is not None: + doc = buffer_to_json_object(buf2) + zattrs = await (sp / ZATTRS_JSON).get(prototype=proto) + if zattrs is not None: + doc["attributes"] = buffer_to_json_object(zattrs) + return doc + raise NodeNotFoundError(f"no node found at path {path!r}") + + async def read_chunk( + self, store: Store, path: str, metadata: Mapping[str, JSON], coords: tuple[int, ...] + ) -> bytes: + meta_obj = _parse_array_metadata(metadata) + shape = _chunk_shape(meta_obj) + np_dtype = _native_dtype(meta_obj) + sp = StorePath(store, path.strip("/")) + chunk_key = meta_obj.encode_chunk_key(coords) + buf = await (sp / chunk_key).get(prototype=default_buffer_prototype()) + if buf is None: + arr = np.full(shape, meta_obj.fill_value, dtype=np_dtype) + else: + pipeline = create_codec_pipeline(meta_obj) + spec = _array_spec(meta_obj, shape) + decoded = list(await pipeline.decode_batch([(buf, spec)])) + arr = np.asarray(decoded[0].as_numpy_array(), dtype=np_dtype) + return np.ascontiguousarray(arr).tobytes() + + async def read_subset( + self, + store: Store, + path: str, + metadata: Mapping[str, JSON], + start: Sequence[int], + shape: Sequence[int], + ) -> bytes: + meta_obj = _parse_array_metadata(metadata) + np_dtype = _native_dtype(meta_obj) + async_arr = AsyncArray(metadata=meta_obj, store_path=StorePath(store, path.strip("/"))) + selection = tuple(slice(s, s + length) for s, length in zip(start, shape, strict=True)) + result = await async_arr.getitem(selection) + return np.ascontiguousarray(np.asarray(result, dtype=np_dtype)).tobytes() + + async def write_chunk( + self, + store: Store, + path: str, + metadata: Mapping[str, JSON], + coords: tuple[int, ...], + data: bytes, + ) -> None: + meta_obj = _parse_array_metadata(metadata) + shape = _chunk_shape(meta_obj) + np_dtype = _native_dtype(meta_obj) + sp = StorePath(store, path.strip("/")) + chunk_key = meta_obj.encode_chunk_key(coords) + arr = np.frombuffer(data, dtype=np_dtype).reshape(shape) + pipeline = create_codec_pipeline(meta_obj) + spec = _array_spec(meta_obj, shape) + encoded = list(await pipeline.encode_batch([(NDBuffer.from_ndarray_like(arr), spec)])) + buf = encoded[0] + if buf is None: + await (sp / chunk_key).delete() + else: + await (sp / chunk_key).set(buf) + + async def delete_chunk( + self, store: Store, path: str, metadata: Mapping[str, JSON], coords: tuple[int, ...] + ) -> None: + meta_obj = _parse_array_metadata(metadata) + sp = StorePath(store, path.strip("/")) + await (sp / meta_obj.encode_chunk_key(coords)).delete() + + async def delete_node(self, store: Store, path: str) -> None: + proto = default_buffer_prototype() + p = path.strip("/") + sp = StorePath(store, p) + present = ( + await (sp / ZARR_JSON).get(prototype=proto) is not None + or await (sp / ZARRAY_JSON).get(prototype=proto) is not None + ) + if not present: + raise NodeNotFoundError(f"no node found at path {path!r}") + await store.delete_dir(p) + + async def list_children( + self, store: Store, path: str + ) -> list[tuple[str, dict[str, JSON]]]: + proto = default_buffer_prototype() + p = path.strip("/") + sp = StorePath(store, p) + if ( + await (sp / ZARR_JSON).get(prototype=proto) is None + and await (sp / ZARRAY_JSON).get(prototype=proto) is None + ): + raise NodeNotFoundError(f"no node found at path {path!r}") + prefix = f"{p}/" if p else "" + children: list[tuple[str, dict[str, JSON]]] = [] + async for name in store.list_dir(prefix): + child_path = f"{p}/{name}" if p else name + child_sp = StorePath(store, child_path) + if ( + await (child_sp / ZARR_JSON).get(prototype=proto) is not None + or await (child_sp / ZARRAY_JSON).get(prototype=proto) is not None + ): + children.append((name, await self.read_metadata(store, child_path))) + return children +``` + +Notes for the implementer: +- `decode_batch`/`encode_batch` are async and return iterables — wrap in `list(...)`. +- `ArraySpec.dtype` is the `ZDType` object (`meta_obj.dtype`), **not** a numpy dtype. +- `_native_dtype` byte-swaps to native order so both backends return identical + bytes through the facade (the facade reads them with a native dtype). +- `AsyncArray(metadata=meta_obj, store_path=...)` constructs from an explicit + document without reading the store. + +- [ ] **Step 4: Register the reference backend** — append to `src/zarr/crud/__init__.py` (after the imports, before `__all__`) + +```python +from zarr.crud._reference import ReferenceBackend + +register_backend("reference", ReferenceBackend()) +``` + +and add `"ReferenceBackend"` to `__all__`. + +- [ ] **Step 5: Run the tests** + +Run: `uv run pytest tests/crud/test_reference_backend.py tests/crud/test_registry.py -v` +Expected: all PASS (the two previously-failing registry tests now pass too). + +- [ ] **Step 6: Commit** + +```bash +git add src/zarr/crud/_reference.py src/zarr/crud/__init__.py tests/crud/test_reference_backend.py +git commit -m "feat: pure-Python ReferenceBackend for zarr.crud" +``` + +--- + +### Task 3: shared facade `zarr.crud._api` + differential suite (reference backend) + +**Files:** +- Create: `src/zarr/crud/_api.py` +- Modify: `src/zarr/crud/__init__.py` +- Create: `tests/crud/conftest.py` +- Test: `tests/crud/test_crud.py` + +- [ ] **Step 1: Create `tests/crud/conftest.py`** + +```python +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import numpy as np +import pytest + +import zarr +from zarr.storage import LocalStore, MemoryStore + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + from pathlib import Path + + from zarr.abc.store import Store + + +def _zarrs_available() -> bool: + try: + import _zarrs_bindings # noqa: F401 + except ImportError: + return False + return True + + +@pytest.fixture( + params=[ + "reference", + pytest.param( + "zarrs", + marks=pytest.mark.skipif( + not _zarrs_available(), reason="zarrs-bindings is not installed" + ), + ), + ] +) +def backend(request: pytest.FixtureRequest) -> str: + """A CRUD backend name. The zarrs param is skipped when the extension is absent.""" + import zarr.crud # noqa: F401 (ensures reference is registered) + + if request.param == "zarrs": + import zarr.zarrs # noqa: F401 (registers the zarrs backend) + return request.param + + +@pytest.fixture(params=["memory", "local"]) +async def store(request: pytest.FixtureRequest, tmp_path: Path) -> AsyncIterator[Store]: + if request.param == "memory": + s: Store = await MemoryStore.open() + else: + s = await LocalStore.open(root=tmp_path / "store") + try: + yield s + finally: + s.close() + + +def array_metadata(**kwargs: Any) -> dict[str, Any]: + """An array metadata document built via zarr-python itself.""" + params: dict[str, Any] = { + "shape": (8, 8), + "chunks": (4, 4), + "dtype": "uint16", + "zarr_format": 3, + } | kwargs + arr = zarr.create_array(store=MemoryStore(), **params) + doc = dict(arr.metadata.to_dict()) + if params["zarr_format"] == 2: + doc.pop("attributes", None) + return doc + + +def filled(store: Store, **kwargs: Any) -> tuple[np.ndarray[Any, np.dtype[Any]], dict[str, Any]]: + """Create an 8x8 array 'a', fill it with a ramp, return (data, metadata).""" + params: dict[str, Any] = {"shape": (8, 8), "chunks": (4, 4), "dtype": "uint16"} | kwargs + arr = zarr.create_array(store=store, name="a", **params) + data = np.arange(64, dtype=params["dtype"]).reshape(8, 8) + arr[:, :] = data + doc = dict(arr.metadata.to_dict()) + if params.get("zarr_format") == 2: + doc.pop("attributes", None) + return data, doc +``` + +- [ ] **Step 2: Write the failing test** — `tests/crud/test_crud.py` + +```python +from __future__ import annotations + +import copy +import json +from typing import TYPE_CHECKING, Any + +import numpy as np +import pytest + +import zarr +from tests.crud.conftest import array_metadata, filled +from zarr.codecs import BloscCodec, GzipCodec, ZstdCodec +from zarr.core.buffer.core import default_buffer_prototype +from zarr.crud import ( + NodeExistsError, + create_new_array, + create_new_group, + create_overwrite_array, + create_overwrite_group, + delete_chunk, + delete_node, + list_children, + read_chunk, + read_encoded_chunk, + read_metadata, + read_region, + write_chunk, +) +from zarr.errors import NodeNotFoundError + +if TYPE_CHECKING: + from zarr.abc.store import Store + +GROUP_META: dict[str, Any] = {"zarr_format": 3, "node_type": "group", "attributes": {"answer": 42}} + + +# --- node lifecycle --- + +async def test_create_new_group(backend: str, store: Store) -> None: + await create_new_group(GROUP_META, store, "foo", backend=backend) + assert dict(zarr.open_group(store=store, path="foo", mode="r").attrs) == {"answer": 42} + + +async def test_create_new_group_existing_raises(backend: str, store: Store) -> None: + await create_new_group(GROUP_META, store, "foo", backend=backend) + with pytest.raises(NodeExistsError): + await create_new_group(GROUP_META, store, "foo", backend=backend) + + +async def test_create_overwrite_group_replaces_array(backend: str, store: Store) -> None: + arr = zarr.create_array(store=store, name="foo", shape=(4,), chunks=(2,), dtype="uint8") + arr[:] = 1 + await create_overwrite_group(GROUP_META, store, "foo", backend=backend) + assert dict(zarr.open_group(store=store, path="foo", mode="r").attrs) == {"answer": 42} + assert not await store.exists("foo/c/0") + + +async def test_create_new_array(backend: str, store: Store) -> None: + await create_new_array(array_metadata(), store, "arr", backend=backend) + a = zarr.open_array(store=store, path="arr", mode="r") + assert a.shape == (8, 8) + assert a.dtype == np.dtype("uint16") + + +async def test_create_new_array_v2(backend: str, store: Store) -> None: + await create_new_array(array_metadata(zarr_format=2), store, "arr", backend=backend) + assert zarr.open_array(store=store, path="arr", mode="r").metadata.zarr_format == 2 + + +async def test_create_overwrite_array(backend: str, store: Store) -> None: + zarr.create_group(store=store, path="arr") + await create_overwrite_array(array_metadata(), store, "arr", backend=backend) + assert zarr.open_array(store=store, path="arr", mode="r").shape == (8, 8) + + +async def test_read_metadata(backend: str, store: Store) -> None: + await create_new_array(array_metadata(), store, "arr", backend=backend) + observed = await read_metadata(store, "arr", backend=backend) + raw = await store.get("arr/zarr.json", prototype=default_buffer_prototype()) + assert raw is not None + assert observed == json.loads(raw.to_bytes()) + + +async def test_read_metadata_missing(backend: str, store: Store) -> None: + with pytest.raises(NodeNotFoundError): + await read_metadata(store, "nope", backend=backend) + + +async def test_delete_node(backend: str, store: Store) -> None: + arr = zarr.create_array(store=store, name="doomed", shape=(4,), chunks=(2,), dtype="uint8") + arr[:] = 1 + await delete_node(store, "doomed", backend=backend) + assert not await store.exists("doomed/zarr.json") + assert not await store.exists("doomed/c/0") + + +async def test_delete_node_missing(backend: str, store: Store) -> None: + with pytest.raises(NodeNotFoundError): + await delete_node(store, "nope", backend=backend) + + +async def test_list_children(backend: str, store: Store) -> None: + root = zarr.create_group(store=store) + root.create_group("sub_group", attributes={"kind": "group"}) + root.create_array("sub_array", shape=(4,), chunks=(2,), dtype="uint8") + by_path = dict(await list_children(store, "", backend=backend)) + assert set(by_path) == {"sub_group", "sub_array"} + assert by_path["sub_group"]["node_type"] == "group" + assert by_path["sub_array"]["node_type"] == "array" + assert not any(p.startswith("/") for p in by_path) + + +# --- chunk I/O --- + +@pytest.mark.parametrize("dtype", ["uint8", "int32", "float64", "u2"]) +async def test_read_chunk_differential(backend: str, store: Store, dtype: str) -> None: + data, meta = filled(store, dtype=dtype) + observed = await read_chunk(meta, store, "a", (1, 0), backend=backend) + np.testing.assert_array_equal(observed, data[4:8, 0:4]) + + +@pytest.mark.parametrize( + "compressors", [None, (GzipCodec(),), (ZstdCodec(),), (BloscCodec(cname="lz4"),)] +) +async def test_read_chunk_codecs(backend: str, store: Store, compressors: Any) -> None: + data, meta = filled(store, compressors=compressors) + observed = await read_chunk(meta, store, "a", (0, 1), backend=backend) + np.testing.assert_array_equal(observed, data[0:4, 4:8]) + + +async def test_read_chunk_v2(backend: str, store: Store) -> None: + data, meta = filled(store, dtype=" None: + data, meta = filled(store, chunks=(2, 2), shards=(4, 4)) + observed = await read_chunk(meta, store, "a", (1, 1), backend=backend) + np.testing.assert_array_equal(observed, data[4:8, 4:8]) + + +async def test_read_chunk_missing_is_fill(backend: str, store: Store) -> None: + arr = zarr.create_array( + store=store, name="a", shape=(8, 8), chunks=(4, 4), dtype="uint16", fill_value=7 + ) + meta = dict(arr.metadata.to_dict()) + observed = await read_chunk(meta, store, "a", (0, 0), backend=backend) + np.testing.assert_array_equal(observed, np.full((4, 4), 7, dtype="uint16")) + + +async def test_read_chunk_metadata_view(backend: str, store: Store) -> None: + data, meta = filled(store, dtype="uint16", compressors=None) + view = copy.deepcopy(meta) + view["data_type"] = "uint8" + view["shape"] = [8, 16] + view["chunk_grid"]["configuration"]["chunk_shape"] = [4, 8] + observed = await read_chunk(view, store, "a", (1, 0), backend=backend) + np.testing.assert_array_equal(observed, data[4:8, 0:4].view("uint8")) + + +async def test_read_chunk_readonly(backend: str, store: Store) -> None: + _, meta = filled(store) + observed = await read_chunk(meta, store, "a", (0, 0), backend=backend) + assert not observed.flags.writeable + + +async def test_write_chunk_differential(backend: str, store: Store) -> None: + meta = array_metadata() + await create_new_array(meta, store, "a", backend=backend) + value = np.arange(16, dtype="uint16").reshape(4, 4) + await write_chunk(meta, store, "a", (0, 1), value, backend=backend) + np.testing.assert_array_equal(zarr.open_array(store=store, path="a", mode="r")[0:4, 4:8], value) + + +async def test_write_chunk_shape_mismatch(backend: str, store: Store) -> None: + meta = array_metadata() + await create_new_array(meta, store, "a", backend=backend) + with pytest.raises(ValueError, match="chunk shape"): + await write_chunk(meta, store, "a", (0, 0), np.zeros((2, 2), dtype="uint16"), backend=backend) + + +async def test_delete_chunk(backend: str, store: Store) -> None: + data, meta = filled(store) + assert await store.exists("a/c/0/0") + await delete_chunk(meta, store, "a", (0, 0), backend=backend) + assert not await store.exists("a/c/0/0") + + +async def test_read_encoded_chunk_matches_store(backend: str, store: Store) -> None: + _, meta = filled(store) + raw = await read_encoded_chunk(meta, store, "a", (0, 0), backend=backend) + expected = await store.get("a/c/0/0", prototype=default_buffer_prototype()) + assert expected is not None + assert raw == expected.to_bytes() + + +async def test_read_encoded_chunk_missing_is_none(backend: str, store: Store) -> None: + arr = zarr.create_array(store=store, name="e", shape=(8, 8), chunks=(4, 4), dtype="uint16") + meta = dict(arr.metadata.to_dict()) + assert await read_encoded_chunk(meta, store, "e", (0, 0), backend=backend) is None + + +# --- region I/O --- + +SELECTIONS: list[Any] = [ + (slice(None), slice(None)), + (slice(2, 7), slice(1, 5)), + (slice(None), 3), + (5, slice(None)), + (3, 4), + (slice(1, 8, 2), slice(None)), + (slice(None), slice(6, 1, -2)), + (slice(-3, None), slice(None, -1)), + ..., + (..., slice(2, 4)), + (slice(0, 0), slice(None)), + (slice(2, 6),), +] + + +@pytest.mark.parametrize("sel", SELECTIONS) +async def test_read_region_differential(backend: str, store: Store, sel: Any) -> None: + data, meta = filled(store) + observed = await read_region(meta, store, "a", sel, backend=backend) + np.testing.assert_array_equal(observed, data[sel]) + + +async def test_read_region_sharding(backend: str, store: Store) -> None: + data, meta = filled(store, chunks=(2, 2), shards=(4, 4)) + observed = await read_region(meta, store, "a", (slice(1, 7), slice(3, 8)), backend=backend) + np.testing.assert_array_equal(observed, data[1:7, 3:8]) + + +async def test_read_region_too_many_indices(backend: str, store: Store) -> None: + _, meta = filled(store) + with pytest.raises(IndexError, match="too many indices"): + await read_region(meta, store, "a", (0, 0, 0), backend=backend) + + +async def test_read_region_fancy_rejected(backend: str, store: Store) -> None: + _, meta = filled(store) + with pytest.raises(TypeError, match="only integers, slices"): + await read_region(meta, store, "a", ([0, 1], slice(None)), backend=backend) # type: ignore[arg-type] +``` + +- [ ] **Step 3: Run it to verify failure** + +Run: `uv run pytest tests/crud/test_crud.py -q` +Expected: collection error — `ImportError: cannot import name 'read_chunk' from 'zarr.crud'` + +- [ ] **Step 4: Create `src/zarr/crud/_api.py`** + +```python +from __future__ import annotations + +import operator +import types +from collections.abc import Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, cast + +import numpy as np + +from zarr.core.buffer.core import default_buffer_prototype +from zarr.crud._registry import get_backend + +if TYPE_CHECKING: + from collections.abc import Mapping + + import numpy.typing as npt + + from zarr.abc.store import Store + from zarr.core.common import JSON + from zarr.core.metadata.v2 import ArrayV2Metadata + from zarr.core.metadata.v3 import ArrayV3Metadata + from zarr.crud._backend import CrudBackend + + +@dataclass(frozen=True, slots=True) +class CrudOptions: + """Options for CRUD operations. + + Currently empty: fields (concurrency limits, checksum validation) arrive in + a later phase. Accepting it now keeps signatures stable. + """ + + +BasicIndex = int | slice | types.EllipsisType +BasicSelection = BasicIndex | tuple[BasicIndex, ...] + + +def _resolve_backend(backend: CrudBackend | str | None) -> CrudBackend: + if backend is None or isinstance(backend, str): + return get_backend(backend) + return backend + + +def _parse_array_metadata( + metadata: Mapping[str, JSON], +) -> ArrayV3Metadata | ArrayV2Metadata: + from zarr.core.metadata.v2 import ArrayV2Metadata + from zarr.core.metadata.v3 import ArrayV3Metadata + + data = dict(metadata) + if data.get("zarr_format") == 3: + return ArrayV3Metadata.from_dict(data) + return ArrayV2Metadata.from_dict(data) + + +def _chunk_dtype_and_shape( + metadata: Mapping[str, JSON], +) -> tuple[np.dtype[Any], tuple[int, ...]]: + """Resolve native-byte-order numpy dtype and regular chunk shape. + + Backends decode to (and encode from) the native in-memory representation, + applying any byte-order codec themselves, so the dtype is coerced to native. + """ + from zarr.core.metadata.v3 import ArrayV3Metadata, RegularChunkGridMetadata + + meta_obj = _parse_array_metadata(metadata) + if isinstance(meta_obj, ArrayV3Metadata): + grid = meta_obj.chunk_grid + if not isinstance(grid, RegularChunkGridMetadata): + raise NotImplementedError("only regular chunk grids are supported") + chunk_shape = tuple(grid.chunk_shape) + else: + chunk_shape = tuple(meta_obj.chunks) + return meta_obj.dtype.to_native_dtype().newbyteorder("="), chunk_shape + + +def _array_shape(metadata: Mapping[str, JSON]) -> tuple[int, ...]: + shape = metadata.get("shape") + if not isinstance(shape, Sequence) or isinstance(shape, str): + raise TypeError("metadata document has no valid 'shape'") + result: list[int] = [] + for s in shape: + if not isinstance(s, (int, float)): + raise TypeError(f"shape element {s!r} is not a number") + if isinstance(s, float) and not s.is_integer(): + raise TypeError(f"shape element {s!r} is not an integer") + result.append(int(s)) + return tuple(result) + + +def _chunk_key(metadata: Mapping[str, JSON], path: str, coords: tuple[int, ...]) -> str: + meta_obj = _parse_array_metadata(metadata) + rel = meta_obj.encode_chunk_key(coords) + p = path.strip("/") + return f"{p}/{rel}" if p else rel + + +def _normalize_selection( + selection: BasicSelection, shape: tuple[int, ...] +) -> tuple[list[int], list[int], tuple[slice | int, ...]]: + """Normalize a numpy basic-indexing selection to a step-1 bounding box. + + Returns `(start, bounding_shape, post_index)`: the box to fetch and the + numpy index to apply to it (strides, reversals, integer-axis removal). Only + integers, slices, and `Ellipsis` are supported; fancy indexing raises. + """ + sel_tuple = selection if isinstance(selection, tuple) else (selection,) + + n_ellipsis = sum(1 for s in sel_tuple if s is Ellipsis) + if n_ellipsis > 1: + raise IndexError("an index can only have a single ellipsis ('...')") + if n_ellipsis == 1: + i = sel_tuple.index(Ellipsis) + n_fill = len(shape) - (len(sel_tuple) - 1) + if n_fill < 0: + raise IndexError(f"too many indices for array: array is {len(shape)}-dimensional") + sel_tuple = sel_tuple[:i] + (slice(None),) * n_fill + sel_tuple[i + 1 :] + if len(sel_tuple) > len(shape): + raise IndexError(f"too many indices for array: array is {len(shape)}-dimensional") + sel_tuple = sel_tuple + (slice(None),) * (len(shape) - len(sel_tuple)) + + starts: list[int] = [] + lengths: list[int] = [] + post: list[slice | int] = [] + for dim, (sel, size) in enumerate(zip(sel_tuple, shape, strict=True)): + if isinstance(sel, slice): + start, stop, step = sel.indices(size) + n = len(range(start, stop, step)) + if n == 0: + starts.append(0) + lengths.append(0) + post.append(slice(None)) + elif step > 0: + last = start + (n - 1) * step + starts.append(start) + lengths.append(last - start + 1) + post.append(slice(None, None, step)) + else: + last = start + (n - 1) * step + starts.append(last) + lengths.append(start - last + 1) + post.append(slice(None, None, step)) + else: + assert not isinstance(sel, types.EllipsisType), "Ellipsis already expanded above" + try: + idx = operator.index(sel) + except TypeError: + raise TypeError( + "unsupported selection element " + f"{sel!r}: only integers, slices, and Ellipsis are supported" + ) from None + if idx < 0: + idx += size + if not 0 <= idx < size: + raise IndexError(f"index {sel} is out of bounds for axis {dim} with size {size}") + starts.append(idx) + lengths.append(1) + post.append(0) + return starts, lengths, tuple(post) + + +# --- node lifecycle --- + +async def create_new_group( + metadata: Mapping[str, JSON], + store: Store, + path: str, + *, + options: CrudOptions | None = None, + backend: CrudBackend | str | None = None, +) -> None: + """Create a group from a group metadata document. Raises `NodeExistsError` + if a node already exists at `path`. Not atomic against concurrent writers.""" + await _resolve_backend(backend).create_group(store, path, metadata, overwrite=False) + + +async def create_overwrite_group( + metadata: Mapping[str, JSON], + store: Store, + path: str, + *, + options: CrudOptions | None = None, + backend: CrudBackend | str | None = None, +) -> None: + """Create a group, deleting any existing node (and children) first. Not + atomic against concurrent writers.""" + await _resolve_backend(backend).create_group(store, path, metadata, overwrite=True) + + +async def create_new_array( + metadata: Mapping[str, JSON], + store: Store, + path: str, + *, + options: CrudOptions | None = None, + backend: CrudBackend | str | None = None, +) -> None: + """Create an array from a v2 or v3 metadata document. Raises + `NodeExistsError` if a node already exists. Not atomic against concurrent + writers.""" + await _resolve_backend(backend).create_array(store, path, metadata, overwrite=False) + + +async def create_overwrite_array( + metadata: Mapping[str, JSON], + store: Store, + path: str, + *, + options: CrudOptions | None = None, + backend: CrudBackend | str | None = None, +) -> None: + """Create an array, deleting any existing node (and children) first. Not + atomic against concurrent writers.""" + await _resolve_backend(backend).create_array(store, path, metadata, overwrite=True) + + +async def read_metadata( + store: Store, + path: str, + *, + options: CrudOptions | None = None, + backend: CrudBackend | str | None = None, +) -> dict[str, JSON]: + """Read the metadata document of the array or group at `path`. Raises + `zarr.errors.NodeNotFoundError` if no node exists there.""" + return await _resolve_backend(backend).read_metadata(store, path) + + +async def delete_node( + store: Store, + path: str, + *, + options: CrudOptions | None = None, + backend: CrudBackend | str | None = None, +) -> None: + """Delete the node at `path` and everything under it. Raises + `zarr.errors.NodeNotFoundError` if absent. `path=""` clears the store.""" + await _resolve_backend(backend).delete_node(store, path) + + +async def list_children( + store: Store, + path: str, + *, + options: CrudOptions | None = None, + backend: CrudBackend | str | None = None, +) -> list[tuple[str, dict[str, JSON]]]: + """List the direct children of the group at `path` as + `(path, metadata_document)` pairs (store-relative, no leading `/`). Raises + `zarr.errors.NodeNotFoundError` if no group exists there.""" + return await _resolve_backend(backend).list_children(store, path) + + +# --- chunk I/O --- + +async def read_chunk( + metadata: Mapping[str, JSON], + store: Store, + path: str, + chunk_coords: tuple[int, ...], + *, + options: CrudOptions | None = None, + backend: CrudBackend | str | None = None, +) -> np.ndarray[Any, np.dtype[Any]]: + """Read and decode the whole chunk at `chunk_coords`. The metadata document + is authoritative; missing chunks decode to the fill value. The result is a + read-only view (`.copy()` for a writable array).""" + be = _resolve_backend(backend) + raw = await be.read_chunk(store, path, metadata, tuple(chunk_coords)) + dtype, chunk_shape = _chunk_dtype_and_shape(metadata) + return np.frombuffer(raw, dtype=dtype).reshape(chunk_shape) + + +async def read_encoded_chunk( + metadata: Mapping[str, JSON], + store: Store, + path: str, + chunk_coords: tuple[int, ...], + *, + options: CrudOptions | None = None, + backend: CrudBackend | str | None = None, +) -> bytes | None: + """Read the raw, still-encoded bytes of the chunk at `chunk_coords`, or + `None` if absent. Pure store I/O (`store.get` on the chunk key): the + `backend` argument is accepted for signature uniformity but unused.""" + key = _chunk_key(metadata, path, tuple(chunk_coords)) + buf = await store.get(key, prototype=default_buffer_prototype()) + return None if buf is None else buf.to_bytes() + + +async def write_chunk( + metadata: Mapping[str, JSON], + store: Store, + path: str, + chunk_coords: tuple[int, ...], + value: npt.ArrayLike, + *, + options: CrudOptions | None = None, + backend: CrudBackend | str | None = None, +) -> None: + """Encode `value` with the codecs in `metadata` and store it as the chunk at + `chunk_coords`. `value` must match the chunk shape exactly.""" + be = _resolve_backend(backend) + dtype, chunk_shape = _chunk_dtype_and_shape(metadata) + arr = np.ascontiguousarray(np.asarray(value, dtype=dtype)) + if arr.shape != chunk_shape: + raise ValueError(f"value shape {arr.shape} does not match chunk shape {chunk_shape}") + await be.write_chunk(store, path, metadata, tuple(chunk_coords), arr.tobytes()) + + +async def delete_chunk( + metadata: Mapping[str, JSON], + store: Store, + path: str, + chunk_coords: tuple[int, ...], + *, + options: CrudOptions | None = None, + backend: CrudBackend | str | None = None, +) -> None: + """Delete the chunk at `chunk_coords`. Deleting a missing chunk is a no-op.""" + await _resolve_backend(backend).delete_chunk(store, path, metadata, tuple(chunk_coords)) + + +# --- region I/O --- + +async def read_region( + metadata: Mapping[str, JSON], + store: Store, + path: str, + selection: BasicSelection, + *, + options: CrudOptions | None = None, + backend: CrudBackend | str | None = None, +) -> np.ndarray[Any, np.dtype[Any]]: + """Read and decode a region given by a numpy basic-indexing `selection` + (integers, slices with steps, `Ellipsis`). One backend call fetches the + step-1 bounding box; strides/reversals/integer-axis removal are applied as + numpy views. Missing chunks decode to the fill value. Fancy indexing raises + `TypeError`. The result is a read-only view. + + Note: a `slice(0, N, step)` reads `O(N)` bytes even though `O(N / step)` are + returned; for sparse selections over large arrays prefer `read_chunk`.""" + be = _resolve_backend(backend) + dtype, _ = _chunk_dtype_and_shape(metadata) + shape = _array_shape(metadata) + starts, lengths, post_index = _normalize_selection(selection, shape) + if 0 in lengths: + block = np.empty(lengths, dtype=dtype) + block.flags.writeable = False + else: + raw = await be.read_subset(store, path, metadata, tuple(starts), tuple(lengths)) + block = np.frombuffer(raw, dtype=dtype).reshape(lengths) + return cast("np.ndarray[Any, np.dtype[Any]]", block[post_index]) +``` + +Note: `BackendArg` is a documentation alias only; use the literal +`CrudBackend | str | None` annotations as written above. + +- [ ] **Step 5: Export the facade from `src/zarr/crud/__init__.py`** + +Add to the imports and `__all__` (keep `__all__` sorted): + +```python +from zarr.crud._api import ( + CrudOptions, + create_new_array, + create_new_group, + create_overwrite_array, + create_overwrite_group, + delete_chunk, + delete_node, + list_children, + read_chunk, + read_encoded_chunk, + read_metadata, + read_region, + write_chunk, +) +``` + +Final `__all__`: + +```python +__all__ = [ + "CrudBackend", + "CrudOptions", + "NodeExistsError", + "ReferenceBackend", + "create_new_array", + "create_new_group", + "create_overwrite_array", + "create_overwrite_group", + "delete_chunk", + "delete_node", + "get_backend", + "list_children", + "read_chunk", + "read_encoded_chunk", + "read_metadata", + "read_region", + "register_backend", + "write_chunk", +] +``` + +- [ ] **Step 6: Run the suite against the reference backend** + +Run: `uv run pytest tests/crud/test_crud.py -q` +Expected: all PASS. The `backend` fixture's `zarrs` param is skipped (no `--group zarrs`), so every test runs once on `reference` × {memory, local}. If `test_read_chunk_differential[>u2-...]` fails, the byte-order coercion in `_reference._native_dtype` / `_chunk_dtype_and_shape` is wrong — both must end in `.newbyteorder("=")`; do not weaken the assertion. + +- [ ] **Step 7: Commit** + +```bash +git add src/zarr/crud/_api.py src/zarr/crud/__init__.py tests/crud/conftest.py tests/crud/test_crud.py +git commit -m "feat: zarr.crud shared facade + differential suite (reference backend)" +``` + +--- + +### Task 4: `ZarrsBackend` + shrink `zarr.zarrs` + migrate zarrs tests + +**Files:** +- Create: `src/zarr/zarrs/_backend.py` +- Modify: `src/zarr/zarrs/__init__.py` +- Delete: `src/zarr/zarrs/_api.py` +- Delete: `tests/zarrs/test_node.py`, `tests/zarrs/test_chunk.py`, `tests/zarrs/test_api.py` +- Modify: `tests/zarrs/test_cache.py` + +- [ ] **Step 1: Create `src/zarr/zarrs/_backend.py`** + +```python +from __future__ import annotations + +import asyncio +import json +from contextlib import contextmanager +from typing import TYPE_CHECKING, cast + +import _zarrs_bindings as _zb + +from zarr.crud import NodeExistsError +from zarr.errors import NodeNotFoundError +from zarr.zarrs._bridge import resolve_store + +if TYPE_CHECKING: + from collections.abc import Iterator, Mapping, Sequence + + from zarr.abc.store import Store + from zarr.core.common import JSON + + +def _node_path(path: str) -> str: + """Convert a zarr path (`""`, `"foo/bar"`) to a zarrs node path (`"/"`, + `"/foo/bar"`).""" + return f"/{path.strip('/')}" + + +@contextmanager +def _translate_errors() -> Iterator[None]: + try: + yield + except _zb.NodeNotFoundError as err: + raise NodeNotFoundError(str(err)) from err + except _zb.NodeExistsError as err: + raise NodeExistsError(str(err)) from err + + +class ZarrsBackend: + """CRUD backend backed by the Rust `zarrs` crate via `_zarrs_bindings`. + + Owns the zarrs-specific plumbing: JSON-serializing the metadata document, + the `/`-prefixed node-path form, store resolution, offloading the blocking + Rust calls to a worker thread, and translating binding exceptions to the + canonical `zarr.crud` / `zarr.errors` types. + """ + + async def create_array( + self, store: Store, path: str, metadata: Mapping[str, JSON], *, overwrite: bool + ) -> None: + with _translate_errors(): + await asyncio.to_thread( + _zb.create_array, + resolve_store(store), + _node_path(path), + json.dumps(metadata), + overwrite, + ) + + async def create_group( + self, store: Store, path: str, metadata: Mapping[str, JSON], *, overwrite: bool + ) -> None: + with _translate_errors(): + await asyncio.to_thread( + _zb.create_group, + resolve_store(store), + _node_path(path), + json.dumps(metadata), + overwrite, + ) + + async def read_metadata(self, store: Store, path: str) -> dict[str, JSON]: + with _translate_errors(): + raw = await asyncio.to_thread(_zb.read_metadata, resolve_store(store), _node_path(path)) + return cast("dict[str, JSON]", json.loads(raw)) + + async def read_chunk( + self, store: Store, path: str, metadata: Mapping[str, JSON], coords: tuple[int, ...] + ) -> bytes: + return await asyncio.to_thread( + _zb.retrieve_chunk, + resolve_store(store), + _node_path(path), + json.dumps(metadata), + list(coords), + ) + + async def read_subset( + self, + store: Store, + path: str, + metadata: Mapping[str, JSON], + start: Sequence[int], + shape: Sequence[int], + ) -> bytes: + return await asyncio.to_thread( + _zb.retrieve_array_subset, + resolve_store(store), + _node_path(path), + json.dumps(metadata), + list(start), + list(shape), + ) + + async def write_chunk( + self, + store: Store, + path: str, + metadata: Mapping[str, JSON], + coords: tuple[int, ...], + data: bytes, + ) -> None: + await asyncio.to_thread( + _zb.store_chunk, + resolve_store(store), + _node_path(path), + json.dumps(metadata), + list(coords), + data, + ) + + async def delete_chunk( + self, store: Store, path: str, metadata: Mapping[str, JSON], coords: tuple[int, ...] + ) -> None: + await asyncio.to_thread( + _zb.erase_chunk, + resolve_store(store), + _node_path(path), + json.dumps(metadata), + list(coords), + ) + + async def delete_node(self, store: Store, path: str) -> None: + with _translate_errors(): + await asyncio.to_thread(_zb.delete_node, resolve_store(store), _node_path(path)) + + async def list_children( + self, store: Store, path: str + ) -> list[tuple[str, dict[str, JSON]]]: + with _translate_errors(): + raw: list[tuple[str, str]] = await asyncio.to_thread( + _zb.list_children, resolve_store(store), _node_path(path) + ) + return [ + (child_path.lstrip("/"), cast("dict[str, JSON]", json.loads(doc))) + for child_path, doc in raw + ] +``` + +- [ ] **Step 2: Rewrite `src/zarr/zarrs/__init__.py`** + +```python +""" +The zarrs CRUD backend for `zarr.crud`, backed by the Rust +[`zarrs`](https://zarrs.dev) crate. + +Importing this module registers the `"zarrs"` backend. Requires the +`zarrs-bindings` extension (in-repo Rust crate; `uv sync --group zarrs`). Select +it with `zarr.config.set({"crud.backend": "zarrs"})` or per call via +`backend="zarrs"`. +""" + +try: + import _zarrs_bindings +except ImportError as e: + raise ImportError( + "zarr.zarrs requires the `zarrs-bindings` package, which is not installed. " + "It is built from the zarr-python repository: run `uv sync --group zarrs`." + ) from e + +from zarr.crud import register_backend +from zarr.zarrs._backend import ZarrsBackend + +__version__: str = _zarrs_bindings.version() + +register_backend("zarrs", ZarrsBackend()) + +__all__ = ["ZarrsBackend", "__version__"] +``` + +- [ ] **Step 3: Delete the moved module and obsolete tests** + +```bash +git rm src/zarr/zarrs/_api.py tests/zarrs/test_node.py tests/zarrs/test_chunk.py tests/zarrs/test_api.py +``` + +- [ ] **Step 4: Update `tests/zarrs/test_cache.py`** — change imports from the old `zarr.zarrs` functions to the `zarr.crud` facade with the zarrs backend. + +Replace the import block: + +```python +from zarr.zarrs import decode_chunk, encode_chunk +``` + +with: + +```python +from zarr.crud import read_chunk, write_chunk +``` + +Then in that file replace every `decode_chunk(` call with `read_chunk(` and every `encode_chunk(` call with `write_chunk(`, adding `backend="zarrs"` as the final keyword argument to each so they exercise the cached zarrs path. For example: + +```python + await read_chunk(meta, store, "a", (0, 0), backend="zarrs") +... + await write_chunk(meta, store, "a", (0, 0), new, backend="zarrs") +``` + +The cache assertions (`zb.array_cache_len()` / `zb.clear_array_cache()`) and the `import _zarrs_bindings as zb` line are unchanged. The module-level `pytest.importorskip("_zarrs_bindings", ...)` stays. + +- [ ] **Step 5: Add the zarrs param coverage — already wired** + +`tests/crud/conftest.py` already parametrizes `backend` over `["reference", "zarrs"]` with the zarrs case skipped when the extension is missing. No change needed; running with `--group zarrs` now exercises it. + +- [ ] **Step 6: Run everything with the zarrs extension** + +Run: `uv run --group zarrs pytest tests/crud tests/zarrs -q` +Expected: all PASS. `tests/crud/test_crud.py` now runs each test on both `reference` and `zarrs` × {memory, local}; `tests/zarrs/test_cache.py` and `test_bridge.py` pass. If a differential test passes on `reference` but fails on `zarrs` (or vice versa), the two backends disagree — investigate the backend, never weaken the assertion. + +- [ ] **Step 7: Run without the extension (reference-only path stays green)** + +Run: `uv run pytest tests/crud -q` +Expected: all PASS, zarrs params skipped. (`tests/zarrs` is not collectable without the extension; that's fine — its module-level `importorskip` skips it.) + +- [ ] **Step 8: Commit** + +```bash +git add src/zarr/zarrs tests/zarrs +git commit -m "feat: ZarrsBackend conforms to CrudBackend; zarr.zarrs is now a backend" +``` + +--- + +### Task 5: changelog, CI, and final verification + +**Files:** +- Modify: `changes/+zarrs-bindings.feature.md` +- Modify: `.github/workflows/zarrs.yml` + +- [ ] **Step 1: Reword the changelog fragment** — overwrite `changes/+zarrs-bindings.feature.md` + +```markdown +Added `zarr.crud`, an experimental backend-agnostic low-level functional API for +zarr hierarchy CRUD (`create_*`, `read_chunk`, `read_region`, `read_encoded_chunk`, +`write_chunk`, `delete_chunk`, `read_metadata`, `delete_node`, `list_children`). +Array routines take an explicit metadata document, enabling read-only views. +Operations delegate to a pluggable `CrudBackend`: a pure-Python reference backend +(the default) or the zarrs-accelerated backend in `zarr.zarrs`, backed by the Rust +[zarrs](https://zarrs.dev) crate via the in-repo `zarrs-bindings` PyO3 crate. +Select a backend with the `crud.backend` config key or a per-call `backend=` +argument. Build the zarrs backend for development with `uv sync --group zarrs`. +``` + +- [ ] **Step 2: Update the CI test command** — `.github/workflows/zarrs.yml` + +Change the test step's `run:` from: + +```yaml + run: uv run --group zarrs pytest tests/zarrs -v +``` + +to: + +```yaml + run: uv run --group zarrs pytest tests/crud tests/zarrs -v +``` + +Validate: `uvx zizmor .github/workflows/zarrs.yml` → no findings. + +- [ ] **Step 3: Lint and type-check the new code** + +Run: `uv run --group dev ruff format src/zarr/crud src/zarr/zarrs tests/crud tests/zarrs` +Run: `uv run --group dev ruff check --fix src/zarr/crud src/zarr/zarrs tests/crud tests/zarrs` +Run: `uv run --group dev --group zarrs mypy src/zarr/crud src/zarr/zarrs tests/crud tests/zarrs` +Expected: all clean. (mypy is strict; the facade and backends are fully annotated.) + +- [ ] **Step 4: Full suites, both with and without the extension** + +Run: `uv run --group zarrs pytest tests/crud tests/zarrs -q` → all pass +Run: `uv run pytest tests/crud -q` → all pass (zarrs skipped) +Run (regression — the rest of zarr-python is untouched): `uv run pytest tests/test_array.py tests/test_group.py -q` → pass + +- [ ] **Step 5: Commit** + +```bash +git add changes/+zarrs-bindings.feature.md .github/workflows/zarrs.yml +git commit -m "docs/ci: zarr.crud changelog and CI coverage" +``` + +--- + +## Out of scope (per spec) + +- Wiring `zarr.crud` under zarr-python's `Array`/`Group` classes. +- Entrypoint-based backend discovery (registration is explicit/import-time). +- A write-side region operation (`write_region`). +- Renaming the Rust `_zarrs_bindings` pyfunctions (private; adapted by `ZarrsBackend`). +- `CrudOptions` fields (concurrency, checksums) — still a placeholder. diff --git a/docs/superpowers/plans/2026-07-22-array-engine-protocol.md b/docs/superpowers/plans/2026-07-22-array-engine-protocol.md new file mode 100644 index 0000000000..2eae8dd902 --- /dev/null +++ b/docs/superpowers/plans/2026-07-22-array-engine-protocol.md @@ -0,0 +1,2205 @@ +# Array Engine Protocol Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** `Array`/`AsyncArray` route all data I/O through pluggable engine objects (`ArrayEngine`/`AsyncArrayEngine` protocols, `Region` interchange), with a default pure-Python engine and a Zarrista (zarrs-backed) engine, replacing `zarr.crud`/`zarr.zarrs`/`packages/zarrs-bindings`. + +**Architecture:** Engines are bound to `(store, path, metadata)` and speak only contiguous step-1 boxes (`Region(start, end_exclusive)`). The facade (`Array`/`AsyncArray`) normalizes every public selection kind to Region calls: contiguous basic selections map directly; strided/orthogonal/coordinate/mask selections go through their bounding box with numpy post-indexing (reads) or box-level read-patch-write (writes). Hierarchy engines are store-bound factories minting resource-sharing array engines. The sync `Array` calls a sync engine directly (no event loop on the data path). + +**Tech Stack:** Python protocols (`typing.Protocol`), existing codec-pipeline machinery (default engine), Zarrista main branch (git-pinned; PyO3/zarrs) for the accelerated engine. + +**Spec:** `docs/superpowers/specs/2026-07-22-array-engine-protocol-design.md` — read it before starting any task. + +## Global Constraints + +- Run all Python commands with `uv run` (e.g. `uv run pytest`, `uv run mypy`). +- Conventional commits with trailer `Assisted-by: ClaudeCode:`. +- Never `git add -A`; stage explicit paths only. Never commit `.claude/` or `.superpowers/`. +- Protocol names have no `Protocol` suffix: `ArrayEngine`, `AsyncArrayEngine`, `HierarchyEngine`, `AsyncHierarchyEngine`. +- The v1 engine contract has **no `out` parameter** and **no config key**; engine choice is explicit (`engine="default" | "zarrista" | `). +- Engine read results need only implement `__array__`. +- Fail loud: `UnsupportedEngineError` at engine construction; `NotImplementedError` for unsupported operations. No silent fallback. +- Docstrings are markdown (mkdocs), single backticks. +- Zarrista dependency: git-pinned to a `main` commit in an optional `zarrista` dependency group. + +## File Structure + +``` +src/zarr/abc/engine.py Region + the four protocols (new) +src/zarr/errors.py add UnsupportedEngineError +src/zarr/core/engine/__init__.py re-exports (new) +src/zarr/core/engine/_normalize.py selection-kind -> (Region, post_index) (new) +src/zarr/core/engine/_default.py DefaultArrayEngine / DefaultAsyncArrayEngine + hierarchy (new) +src/zarr/core/engine/_resolve.py engine-spec resolution + per-store hierarchy cache (new) +src/zarr/core/array.py facade rewiring (modify) +src/zarr/zarrista/__init__.py public zarrista engine surface (new) +src/zarr/zarrista/_translate.py zarr Store -> zarrista/obstore store translation (new) +src/zarr/zarrista/_engine.py ZarristaEngine / ZarristaAsyncEngine + hierarchy (new) +tests/engine/ protocol, normalization, default-engine, differential (new) +tests/zarrista/ translation + zarrista-only behavior (new) +DELETED: src/zarr/crud/, src/zarr/zarrs/, packages/zarrs-bindings/, tests/crud/, tests/zarrs/ +``` + +--- + +### Task 1: `Region`, the four protocols, `UnsupportedEngineError` + +**Files:** +- Create: `src/zarr/abc/engine.py` +- Modify: `src/zarr/errors.py` (append class) +- Test: `tests/engine/test_protocols.py`, `tests/engine/__init__.py` (empty) + +**Interfaces:** +- Produces: `Region(start: tuple[int, ...], end_exclusive: tuple[int, ...])` with property `shape -> tuple[int, ...]`; protocols `ArrayEngine`, `AsyncArrayEngine`, `HierarchyEngine`, `AsyncHierarchyEngine` exactly as below; `zarr.errors.UnsupportedEngineError(ValueError)`. Every later task imports these from `zarr.abc.engine` / `zarr.errors`. + +- [ ] **Step 1: Write the failing test** + +```python +# tests/engine/test_protocols.py +from __future__ import annotations + +import numpy as np +import pytest + +from zarr.abc.engine import ArrayEngine, AsyncArrayEngine, Region +from zarr.errors import UnsupportedEngineError + + +def test_region() -> None: + """Region carries start/end_exclusive and derives shape.""" + r = Region(start=(1, 2), end_exclusive=(4, 2)) + assert r.start == (1, 2) + assert r.end_exclusive == (4, 2) + assert r.shape == (3, 0) + + +class _FakeSyncEngine: + def read_selection(self, selection, *, prototype): + return np.zeros(selection.shape) + + def write_selection(self, selection, value, *, prototype): + return None + + def with_metadata(self, metadata): + return self + + +def test_runtime_checkable_protocols() -> None: + """isinstance checks verify method presence for the sync protocol.""" + assert isinstance(_FakeSyncEngine(), ArrayEngine) + assert not isinstance(object(), ArrayEngine) + assert not isinstance(_FakeSyncEngine(), AsyncArrayEngine) or True # names match; mypy is authoritative + + +def test_unsupported_engine_error_is_value_error() -> None: + with pytest.raises(ValueError): + raise UnsupportedEngineError("nope") +``` + +- [ ] **Step 2: Run it to make sure it fails** + +Run: `uv run pytest tests/engine/test_protocols.py -v` +Expected: FAIL (`ModuleNotFoundError: zarr.abc.engine`) + +- [ ] **Step 3: Implement** + +```python +# src/zarr/abc/engine.py +"""Array engine protocols. + +An *array engine* owns the data path of one open array: reading and writing +decoded data for contiguous regions. `Array` wraps an object satisfying +`ArrayEngine`; `AsyncArray` wraps an object satisfying `AsyncArrayEngine`. +A *hierarchy engine* is bound to a store and mints array engines that share +resources. See `docs/superpowers/specs/2026-07-22-array-engine-protocol-design.md`. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, NamedTuple, Protocol, runtime_checkable + +if TYPE_CHECKING: + from zarr.core.buffer import BufferPrototype, NDBuffer + from zarr.core.common import NDArrayLike + from zarr.core.metadata import ArrayMetadata + +__all__ = [ + "ArrayEngine", + "AsyncArrayEngine", + "AsyncHierarchyEngine", + "HierarchyEngine", + "Region", +] + + +class Region(NamedTuple): + """A contiguous, step-1 box in array-element coordinates. + + One entry per dimension; `start` is inclusive, `end_exclusive` exclusive. + Callers pass normalized values: non-negative and clipped to the array + shape. This is the only selection type that crosses the engine boundary. + """ + + start: tuple[int, ...] + end_exclusive: tuple[int, ...] + + @property + def shape(self) -> tuple[int, ...]: + """The ndim-preserving shape of the box.""" + return tuple(e - s for s, e in zip(self.start, self.end_exclusive, strict=True)) + + +@runtime_checkable +class ArrayEngine(Protocol): + """The synchronous data path of one open array. + + Bound to `(store, path, metadata)` at construction. Methods must not + require a running event loop. Read results are ndim-preserving with + shape `selection.shape` and need only implement `__array__`. + + Note: `runtime_checkable` isinstance checks only verify method names; + mypy is the authoritative conformance check. + """ + + def read_selection( + self, selection: Region, *, prototype: BufferPrototype + ) -> NDArrayLike: ... + + def write_selection( + self, selection: Region, value: NDBuffer, *, prototype: BufferPrototype + ) -> None: ... + + def with_metadata(self, metadata: ArrayMetadata) -> ArrayEngine: ... + + +@runtime_checkable +class AsyncArrayEngine(Protocol): + """The asynchronous data path of one open array. See `ArrayEngine`.""" + + async def read_selection( + self, selection: Region, *, prototype: BufferPrototype + ) -> NDArrayLike: ... + + async def write_selection( + self, selection: Region, value: NDBuffer, *, prototype: BufferPrototype + ) -> None: ... + + def with_metadata(self, metadata: ArrayMetadata) -> AsyncArrayEngine: ... + + +@runtime_checkable +class HierarchyEngine(Protocol): + """A store-bound factory for synchronous array engines.""" + + def array_engine(self, path: str, metadata: ArrayMetadata) -> ArrayEngine: ... + + +@runtime_checkable +class AsyncHierarchyEngine(Protocol): + """A store-bound factory for asynchronous array engines.""" + + def array_engine(self, path: str, metadata: ArrayMetadata) -> AsyncArrayEngine: ... +``` + +Append to `src/zarr/errors.py` (match the module's existing docstring style): + +```python +class UnsupportedEngineError(ValueError): + """Raised when an array engine cannot serve the requested store or array. + + Examples: a store the engine cannot translate, or metadata (e.g. Zarr v2) + the engine does not support. Raised at engine construction time. + """ +``` + +Add `"UnsupportedEngineError"` to `zarr.errors.__all__` if the module defines one. + +- [ ] **Step 4: Run tests, then mypy** + +Run: `uv run pytest tests/engine/test_protocols.py -v` — Expected: PASS +Run: `uv run mypy src/zarr/abc/engine.py` — Expected: no errors + +- [ ] **Step 5: Commit** + +```bash +git add src/zarr/abc/engine.py src/zarr/errors.py tests/engine +git commit -m "feat: Region interchange and array/hierarchy engine protocols" +``` + +--- + +### Task 2: Selection normalization (`_normalize.py`) + +Every public selection kind reduces to `(Region, post_index)` where +`post_index` is a numpy index applied to the ndim-preserving box result. +The basic-kind code is adapted from `src/zarr/crud/_api.py::_normalize_selection` +(lines 86-147) — copy it before Task 9 deletes that package. + +**Files:** +- Create: `src/zarr/core/engine/__init__.py`, `src/zarr/core/engine/_normalize.py` +- Test: `tests/engine/test_normalize.py` + +**Interfaces:** +- Produces (all in `zarr.core.engine._normalize`; re-exported from `zarr.core.engine`): + - `normalize_basic(selection, shape) -> tuple[Region, tuple[slice | int, ...]]` + - `normalize_orthogonal(selection, shape) -> tuple[Region, tuple[Any, ...]]` (post is per-axis arrays/slices for `np.ix_`-style outer indexing; the helper returns the ready-to-use index) + - `normalize_coordinate(selection, shape) -> tuple[Region, tuple[np.ndarray, ...]]` + - `normalize_block(selection, chunk_grid_shape, chunk_shape, shape) -> Region` (block selections are contiguous; no post index) + - mask selections are converted by callers via `np.nonzero(mask)` → `normalize_coordinate`. +- Invariant used by Tasks 5-6: for any array `a`, + `np.asarray(engine_read(region))[post] == a[original_selection]`. + +- [ ] **Step 1: Write the failing tests** — differential against numpy, per your test convention (one combo test per function + error cases): + +```python +# tests/engine/test_normalize.py +from __future__ import annotations + +import numpy as np +import pytest + +from zarr.core.engine import ( + normalize_basic, + normalize_block, + normalize_coordinate, + normalize_orthogonal, +) + +SHAPE = (10, 9) +ARR = np.arange(90).reshape(SHAPE) + + +def _read_box(region): + """Simulate an engine read: the ndim-preserving box.""" + return ARR[tuple(slice(s, e) for s, e in zip(region.start, region.end_exclusive, strict=True))] + + +@pytest.mark.parametrize( + "sel", + [ + (slice(2, 7), slice(0, 9)), + (slice(None), slice(3, 4)), + (slice(8, 2, -2), slice(None)), + (slice(1, 8, 3), slice(2, 9, 2)), + (3, slice(None)), + (-1, -2), + (Ellipsis, 4), + slice(5), + Ellipsis, + (slice(4, 4), slice(None)), # empty + ], +) +def test_normalize_basic_matches_numpy(sel) -> None: + region, post = normalize_basic(sel, SHAPE) + np.testing.assert_array_equal(_read_box(region)[post], ARR[sel]) + + +@pytest.mark.parametrize( + "sel", + [ + (np.array([1, 4, 7]), np.array([0, 8])), + (np.array([7, 1, 4]), slice(2, 6)), # unordered axis indices + (np.array([3, 3]), np.array([5, 5])), # repeats + (slice(1, 9, 2), np.array([2])), + (np.array([0]), 4), # int axis + ], +) +def test_normalize_orthogonal_matches_numpy_oindex(sel) -> None: + import numpy.lib.recfunctions # noqa: F401 (just ensures numpy fully loaded) + + region, post = normalize_orthogonal(sel, SHAPE) + # numpy oindex semantics == np.ix_ over per-axis index lists + axes = [] + for s, size in zip(sel if isinstance(sel, tuple) else (sel,), SHAPE, strict=False): + if isinstance(s, slice): + axes.append(np.arange(*s.indices(size))) + elif np.isscalar(s) or getattr(s, "ndim", 1) == 0: + axes.append(np.array([int(s)])) + else: + axes.append(np.asarray(s)) + expected = ARR[np.ix_(*axes)] + np.testing.assert_array_equal(_read_box(region)[post], expected) + + +def test_normalize_coordinate_matches_numpy_vindex() -> None: + coords = (np.array([9, 0, 3, 3]), np.array([8, 0, 2, 2])) + region, post = normalize_coordinate(coords, SHAPE) + np.testing.assert_array_equal(_read_box(region)[post], ARR[coords]) + + +def test_normalize_block_is_contiguous() -> None: + # chunk shape (3, 4) over SHAPE (10, 9): block (1, 2) spans rows 3:6, cols 8:9 + region = normalize_block((1, 2), chunk_grid_shape=(4, 3), chunk_shape=(3, 4), shape=SHAPE) + assert region.start == (3, 8) + assert region.end_exclusive == (6, 9) + + +def test_normalize_basic_rejects_fancy() -> None: + with pytest.raises(TypeError): + normalize_basic((np.array([1, 2]), slice(None)), SHAPE) + + +def test_normalize_basic_rejects_out_of_bounds_int() -> None: + with pytest.raises(IndexError): + normalize_basic((10, slice(None)), SHAPE) + + +def test_normalize_coordinate_rejects_out_of_bounds() -> None: + with pytest.raises(IndexError): + normalize_coordinate((np.array([10]), np.array([0])), SHAPE) + + +def test_normalize_orthogonal_rejects_boolean_ndim_mismatch() -> None: + with pytest.raises(IndexError): + normalize_orthogonal((np.zeros((2, 2), dtype=bool), slice(None)), SHAPE) +``` + +- [ ] **Step 2: Run to verify failure** + +Run: `uv run pytest tests/engine/test_normalize.py -v` +Expected: FAIL (`ImportError`) + +- [ ] **Step 3: Implement** + +```python +# src/zarr/core/engine/_normalize.py +"""Reduce public selection kinds to `(Region, post_index)` pairs. + +The engine boundary only speaks contiguous step-1 boxes (`Region`). Each +helper returns the box to transfer plus the numpy index that, applied to the +ndim-preserving box result, yields exactly `array[original_selection]`. +""" + +from __future__ import annotations + +import operator +import types +from typing import TYPE_CHECKING, Any + +import numpy as np + +from zarr.abc.engine import Region + +if TYPE_CHECKING: + from zarr.core.indexing import BasicSelection + + +def _expand(selection: Any, ndim: int) -> tuple[Any, ...]: + """Expand Ellipsis and pad missing trailing axes with full slices.""" + sel_tuple = selection if isinstance(selection, tuple) else (selection,) + n_ellipsis = sum(1 for s in sel_tuple if s is Ellipsis) + if n_ellipsis > 1: + raise IndexError("an index can only have a single ellipsis ('...')") + if n_ellipsis == 1: + i = sel_tuple.index(Ellipsis) + n_fill = ndim - (len(sel_tuple) - 1) + if n_fill < 0: + raise IndexError(f"too many indices for array: array is {ndim}-dimensional") + sel_tuple = sel_tuple[:i] + (slice(None),) * n_fill + sel_tuple[i + 1 :] + if len(sel_tuple) > ndim: + raise IndexError(f"too many indices for array: array is {ndim}-dimensional") + return sel_tuple + (slice(None),) * (ndim - len(sel_tuple)) + + +def _normalize_int(sel: Any, size: int, dim: int) -> int: + idx = operator.index(sel) + if idx < 0: + idx += size + if not 0 <= idx < size: + raise IndexError(f"index {sel} is out of bounds for axis {dim} with size {size}") + return idx + + +def normalize_basic( + selection: BasicSelection, shape: tuple[int, ...] +) -> tuple[Region, tuple[slice | int, ...]]: + """Normalize a numpy basic-indexing selection to a step-1 box. + + Supports integers, slices (any step), and `Ellipsis`. Integer axes get a + length-1 range in the box and `0` in the post index (dropping the axis, + matching numpy). Fancy elements raise `TypeError`. + """ + sel_tuple = _expand(selection, len(shape)) + starts: list[int] = [] + ends: list[int] = [] + post: list[slice | int] = [] + for dim, (sel, size) in enumerate(zip(sel_tuple, shape, strict=True)): + if isinstance(sel, slice): + start, stop, step = sel.indices(size) + n = len(range(start, stop, step)) + if n == 0: + starts.append(0) + ends.append(0) + post.append(slice(None)) + elif step > 0: + starts.append(start) + ends.append(start + (n - 1) * step + 1) + post.append(slice(None, None, step)) + else: + last = start + (n - 1) * step + starts.append(last) + ends.append(start + 1) + post.append(slice(None, None, step)) + elif isinstance(sel, types.EllipsisType): + raise AssertionError("Ellipsis expanded by _expand") + else: + try: + idx = _normalize_int(sel, size, dim) + except TypeError: + raise TypeError( + f"unsupported selection element {sel!r}: only integers, " + "slices, and Ellipsis are supported by basic indexing" + ) from None + starts.append(idx) + ends.append(idx + 1) + post.append(0) + return Region(start=tuple(starts), end_exclusive=tuple(ends)), tuple(post) + + +def normalize_orthogonal( + selection: Any, shape: tuple[int, ...] +) -> tuple[Region, tuple[Any, ...]]: + """Normalize an orthogonal (`oindex`) selection to a box + outer index. + + Each axis selector may be an integer, a slice, an integer array, or a 1-d + boolean mask for that axis. The post index is the `np.ix_`-broadcastable + tuple of per-axis integer arrays (relative to the box origin), with + integer axes dropped afterwards by the facade via the returned index + (integers appear as scalar entries produced by `np.ix_` inputs of shape + `(1,)` — the facade result keeps numpy `oindex` semantics). + """ + sel_tuple = _expand(selection, len(shape)) + starts: list[int] = [] + ends: list[int] = [] + axis_indices: list[np.ndarray] = [] + for dim, (sel, size) in enumerate(zip(sel_tuple, shape, strict=True)): + if isinstance(sel, slice): + idxs = np.arange(*sel.indices(size)) + elif isinstance(sel, np.ndarray) and sel.dtype == bool: + if sel.ndim != 1 or sel.shape[0] != size: + raise IndexError( + f"boolean index for axis {dim} must be 1-d with length {size}" + ) + idxs = np.nonzero(sel)[0] + elif isinstance(sel, (np.ndarray, list)): + idxs = np.asarray(sel, dtype=np.intp) + if idxs.ndim != 1: + raise IndexError(f"orthogonal index for axis {dim} must be 1-d") + idxs = np.where(idxs < 0, idxs + size, idxs) + if idxs.size and (idxs.min() < 0 or idxs.max() >= size): + raise IndexError(f"index out of bounds for axis {dim} with size {size}") + else: + idxs = np.array([_normalize_int(sel, size, dim)], dtype=np.intp) + if idxs.size == 0: + starts.append(0) + ends.append(0) + axis_indices.append(idxs) + else: + lo, hi = int(idxs.min()), int(idxs.max()) + starts.append(lo) + ends.append(hi + 1) + axis_indices.append(idxs - lo) + region = Region(start=tuple(starts), end_exclusive=tuple(ends)) + post = np.ix_(*axis_indices) if axis_indices else () + # integer axes were widened to length-1 arrays; the caller squeezes them + squeeze_axes = tuple( + i + for i, sel in enumerate(sel_tuple) + if not isinstance(sel, (slice, np.ndarray, list)) + ) + if squeeze_axes: + return region, (*post, _Squeeze(squeeze_axes)) # type: ignore[return-value] + return region, post + + +class _Squeeze: + """Marker appended to an orthogonal post index: squeeze these axes.""" + + def __init__(self, axes: tuple[int, ...]) -> None: + self.axes = axes + + +def apply_post_index(box: np.ndarray, post: tuple[Any, ...]) -> np.ndarray: + """Apply a post index produced by a `normalize_*` helper to a box read.""" + if post and isinstance(post[-1], _Squeeze): + result = box[post[:-1]] if len(post) > 1 else box + return np.squeeze(result, axis=post[-1].axes) + if post == (): + return box + return box[post] + + +def normalize_coordinate( + selection: tuple[Any, ...], shape: tuple[int, ...] +) -> tuple[Region, tuple[np.ndarray, ...]]: + """Normalize a coordinate (`vindex`) selection to a box + pointwise index.""" + coords = tuple(np.asarray(c, dtype=np.intp) for c in selection) + if len(coords) != len(shape): + raise IndexError( + f"coordinate selection needs {len(shape)} axis arrays, got {len(coords)}" + ) + coords = tuple( + np.where(c < 0, c + size, c) for c, size in zip(coords, shape, strict=True) + ) + for dim, (c, size) in enumerate(zip(coords, shape, strict=True)): + if c.size and (c.min() < 0 or c.max() >= size): + raise IndexError(f"index out of bounds for axis {dim} with size {size}") + starts = tuple(int(c.min()) if c.size else 0 for c in coords) + ends = tuple(int(c.max()) + 1 if c.size else 0 for c in coords) + post = tuple(c - s for c, s in zip(coords, starts, strict=True)) + return Region(start=starts, end_exclusive=ends), post + + +def normalize_block( + block_coords: tuple[int, ...], + *, + chunk_grid_shape: tuple[int, ...], + chunk_shape: tuple[int, ...], + shape: tuple[int, ...], +) -> Region: + """Normalize a block (chunk-grid) selection to its contiguous box.""" + starts = [] + ends = [] + for dim, (b, nblocks, csize, size) in enumerate( + zip(block_coords, chunk_grid_shape, chunk_shape, shape, strict=True) + ): + idx = _normalize_int(b, nblocks, dim) + starts.append(idx * csize) + ends.append(min((idx + 1) * csize, size)) + return Region(start=tuple(starts), end_exclusive=tuple(ends)) +``` + +```python +# src/zarr/core/engine/__init__.py +from zarr.core.engine._normalize import ( + apply_post_index, + normalize_basic, + normalize_block, + normalize_coordinate, + normalize_orthogonal, +) + +__all__ = [ + "apply_post_index", + "normalize_basic", + "normalize_block", + "normalize_coordinate", + "normalize_orthogonal", +] +``` + +Note for the implementer: the test uses `_read_box(region)[post]` for basic and +coordinate kinds and `apply_post_index` semantics for orthogonal. Update the +orthogonal test to call `apply_post_index(_read_box(region), post)` — plain +`box[post]` is wrong once the `_Squeeze` marker is present. + +- [ ] **Step 4: Run tests** + +Run: `uv run pytest tests/engine/test_normalize.py -v` — Expected: PASS +Run: `uv run mypy src/zarr/core/engine` — Expected: no errors + +- [ ] **Step 5: Commit** + +```bash +git add src/zarr/core/engine tests/engine/test_normalize.py +git commit -m "feat: normalize selection kinds to Region + post-index" +``` + +--- + +### Task 3: Default engines (`_default.py`) + +**Files:** +- Create: `src/zarr/core/engine/_default.py` +- Modify: `src/zarr/core/engine/__init__.py` (re-export) +- Test: `tests/engine/test_default_engine.py` + +**Interfaces:** +- Consumes: `Region` (Task 1). +- Produces: + - `DefaultAsyncArrayEngine(store_path: StorePath, metadata: ArrayMetadata, config: ArrayConfig)` satisfying `AsyncArrayEngine`. + - `DefaultArrayEngine(async_engine: DefaultAsyncArrayEngine)` satisfying `ArrayEngine` (wraps with `zarr.core.sync.sync`). + - `DefaultAsyncHierarchyEngine(store: Store)` / `DefaultHierarchyEngine(store: Store)` with `array_engine(path, metadata)`. +- Task 5 constructs these inside `AsyncArray.__init__`/`Array`. + +- [ ] **Step 1: Write the failing test** + +```python +# tests/engine/test_default_engine.py +from __future__ import annotations + +import numpy as np + +import zarr +from zarr.abc.engine import AsyncArrayEngine, Region +from zarr.core.buffer import default_buffer_prototype +from zarr.core.engine import DefaultArrayEngine, DefaultAsyncArrayEngine +from zarr.core.sync import sync +from zarr.storage import MemoryStore + + +def _make_array() -> zarr.Array: + z = zarr.create_array( + MemoryStore(), shape=(10, 9), chunks=(3, 4), dtype="int32", fill_value=0 + ) + z[:, :] = np.arange(90, dtype="int32").reshape(10, 9) + return z + + +def test_default_async_engine_read_write_roundtrip() -> None: + z = _make_array() + eng = DefaultAsyncArrayEngine( + store_path=z.async_array.store_path, + metadata=z.async_array.metadata, + config=z.async_array.config, + ) + assert isinstance(eng, AsyncArrayEngine) + proto = default_buffer_prototype() + region = Region(start=(2, 1), end_exclusive=(7, 5)) + + out = sync(eng.read_selection(region, prototype=proto)) + np.testing.assert_array_equal(np.asarray(out), np.asarray(z[2:7, 1:5])) + + new = np.full((5, 4), -1, dtype="int32") + value = proto.nd_buffer.from_ndarray_like(new) + sync(eng.write_selection(region, value, prototype=proto)) + np.testing.assert_array_equal(np.asarray(z[2:7, 1:5]), new) + + +def test_default_sync_engine_matches_async() -> None: + z = _make_array() + async_eng = DefaultAsyncArrayEngine( + store_path=z.async_array.store_path, + metadata=z.async_array.metadata, + config=z.async_array.config, + ) + eng = DefaultArrayEngine(async_eng) + proto = default_buffer_prototype() + region = Region(start=(0, 0), end_exclusive=(10, 9)) + np.testing.assert_array_equal( + np.asarray(eng.read_selection(region, prototype=proto)), np.asarray(z[:, :]) + ) + + +def test_with_metadata_rebinds() -> None: + z = _make_array() + eng = DefaultAsyncArrayEngine( + store_path=z.async_array.store_path, + metadata=z.async_array.metadata, + config=z.async_array.config, + ) + new_meta = z.async_array.metadata + assert eng.with_metadata(new_meta) is not eng +``` + +- [ ] **Step 2: Run to verify failure** + +Run: `uv run pytest tests/engine/test_default_engine.py -v` +Expected: FAIL (`ImportError: DefaultArrayEngine`) + +- [ ] **Step 3: Implement** + +```python +# src/zarr/core/engine/_default.py +"""The default engines: today's codec-pipeline machinery behind the protocol.""" + +from __future__ import annotations + +from dataclasses import replace +from typing import TYPE_CHECKING + +from zarr.abc.engine import Region +from zarr.core.array_spec import ArraySpec +from zarr.core.chunk_grids import ChunkGrid +from zarr.core.codec_pipeline import create_codec_pipeline +from zarr.core.indexing import BasicIndexer +from zarr.core.metadata import ArrayMetadata +from zarr.core.sync import sync + +if TYPE_CHECKING: + from zarr.abc.store import Store + from zarr.core.array_spec import ArrayConfig + from zarr.core.buffer import BufferPrototype, NDBuffer + from zarr.core.common import NDArrayLike + from zarr.core.storage import StorePath + + +def _region_to_slices(region: Region) -> tuple[slice, ...]: + return tuple( + slice(s, e) for s, e in zip(region.start, region.end_exclusive, strict=True) + ) + + +class DefaultAsyncArrayEngine: + """Codec-pipeline-backed engine. Any store, Zarr v2 and v3.""" + + def __init__( + self, store_path: StorePath, metadata: ArrayMetadata, config: ArrayConfig + ) -> None: + self.store_path = store_path + self.metadata = metadata + self.config = config + self._chunk_grid = ChunkGrid.from_metadata(metadata) + self._codec_pipeline = create_codec_pipeline( + metadata=metadata, store=store_path.store + ) + + def with_metadata(self, metadata: ArrayMetadata) -> DefaultAsyncArrayEngine: + return DefaultAsyncArrayEngine( + store_path=self.store_path, metadata=metadata, config=self.config + ) + + def _indexer(self, region: Region) -> BasicIndexer: + return BasicIndexer( + _region_to_slices(region), + shape=self.metadata.shape, + chunk_grid=self._chunk_grid, + ) + + def _chunk_spec(self, prototype: BufferPrototype) -> tuple[ArraySpec | None, ArrayConfig]: + # v2 honors metadata order; v3 honors config order (mirrors _get_selection) + config = self.config + if self.metadata.zarr_format == 2: + config = replace(config, order=self.metadata.order) + if self._chunk_grid.is_regular: + return ( + ArraySpec( + shape=self._chunk_grid.chunk_shape, + dtype=self.metadata.dtype, + fill_value=self.metadata.fill_value, + config=config, + prototype=prototype, + ), + config, + ) + return None, config + + async def read_selection( + self, selection: Region, *, prototype: BufferPrototype + ) -> NDArrayLike: + indexer = self._indexer(selection) + if self.metadata.zarr_format == 2: + dtype = self.metadata.dtype.to_native_dtype() + order = self.metadata.order + else: + dtype = self.metadata.data_type.to_native_dtype() + order = self.config.order + out_buffer = prototype.nd_buffer.empty( + shape=indexer.shape, dtype=dtype, order=order + ) + if all(e > s for s, e in zip(selection.start, selection.end_exclusive, strict=True)): + spec, config = self._chunk_spec(prototype) + await self._codec_pipeline.read( + [ + ( + self.store_path / self.metadata.encode_chunk_key(chunk_coords), + spec + if spec is not None + else _irregular_chunk_spec( + self.metadata, self._chunk_grid, chunk_coords, config, prototype + ), + chunk_selection, + out_selection, + is_complete_chunk, + ) + for chunk_coords, chunk_selection, out_selection, is_complete_chunk in indexer + ], + out_buffer, + drop_axes=indexer.drop_axes, + ) + return out_buffer.as_ndarray_like() + + async def write_selection( + self, selection: Region, value: NDBuffer, *, prototype: BufferPrototype + ) -> None: + indexer = self._indexer(selection) + spec, config = self._chunk_spec(prototype) + await self._codec_pipeline.write( + [ + ( + self.store_path / self.metadata.encode_chunk_key(chunk_coords), + spec + if spec is not None + else _irregular_chunk_spec( + self.metadata, self._chunk_grid, chunk_coords, config, prototype + ), + chunk_selection, + out_selection, + is_complete_chunk, + ) + for chunk_coords, chunk_selection, out_selection, is_complete_chunk in indexer + ], + value, + drop_axes=indexer.drop_axes, + ) + + +def _irregular_chunk_spec(metadata, chunk_grid, chunk_coords, config, prototype): + # same construction as array.py::_get_chunk_spec — import and delegate + from zarr.core.array import _get_chunk_spec + + return _get_chunk_spec(metadata, chunk_grid, chunk_coords, config, prototype) + + +class DefaultArrayEngine: + """Sync adapter over `DefaultAsyncArrayEngine` via `sync()`.""" + + def __init__(self, async_engine: DefaultAsyncArrayEngine) -> None: + self._async = async_engine + + def read_selection( + self, selection: Region, *, prototype: BufferPrototype + ) -> NDArrayLike: + return sync(self._async.read_selection(selection, prototype=prototype)) + + def write_selection( + self, selection: Region, value: NDBuffer, *, prototype: BufferPrototype + ) -> None: + sync(self._async.write_selection(selection, value, prototype=prototype)) + + def with_metadata(self, metadata: ArrayMetadata) -> DefaultArrayEngine: + return DefaultArrayEngine(self._async.with_metadata(metadata)) + + +class DefaultAsyncHierarchyEngine: + """Store-bound factory for default async engines.""" + + def __init__(self, store: Store) -> None: + self.store = store + + def array_engine(self, path: str, metadata: ArrayMetadata) -> DefaultAsyncArrayEngine: + from zarr.core.storage import StorePath + + from zarr.core.array_spec import parse_array_config + + return DefaultAsyncArrayEngine( + store_path=StorePath(self.store, path), + metadata=metadata, + config=parse_array_config(None), + ) + + +class DefaultHierarchyEngine: + """Store-bound factory for default sync engines.""" + + def __init__(self, store: Store) -> None: + self._async = DefaultAsyncHierarchyEngine(store) + + def array_engine(self, path: str, metadata: ArrayMetadata) -> DefaultArrayEngine: + return DefaultArrayEngine(self._async.array_engine(path, metadata)) +``` + +Implementation notes (verify against the live code, these matter): +- `StorePath` import path: check `src/zarr/core/array.py`'s own import and copy it. +- The empty-region guard in `read_selection` mirrors `_get_selection`'s + `product(indexer.shape) > 0` check; use `zarr.core.common.product` if simpler. +- v2 `metadata.dtype` / v3 `metadata.data_type` split is copied from + `_get_selection` (array.py:5392-5397). Do not invent a new path. + +Add to `src/zarr/core/engine/__init__.py`: + +```python +from zarr.core.engine._default import ( + DefaultArrayEngine, + DefaultAsyncArrayEngine, + DefaultAsyncHierarchyEngine, + DefaultHierarchyEngine, +) +``` +(and extend `__all__` accordingly). + +- [ ] **Step 4: Run tests** + +Run: `uv run pytest tests/engine/test_default_engine.py tests/engine/test_protocols.py -v` +Expected: PASS + +- [ ] **Step 5: Commit** + +```bash +git add src/zarr/core/engine tests/engine/test_default_engine.py +git commit -m "feat: default array engines over the codec pipeline" +``` + +--- + +### Task 4: Engine resolution (`_resolve.py`) + +**Files:** +- Create: `src/zarr/core/engine/_resolve.py` +- Modify: `src/zarr/core/engine/__init__.py` (re-export) +- Test: `tests/engine/test_resolve.py` + +**Interfaces:** +- Consumes: default engines (Task 3), protocols (Task 1). +- Produces (used by Task 5's `AsyncArray.__init__` / `Array` wiring): + - `EngineName = Literal["default", "zarrista"]` + - `resolve_async_engine(engine: AsyncArrayEngine | EngineName | None, *, store: Store, path: str, metadata: ArrayMetadata) -> AsyncArrayEngine` + - `resolve_sync_engine(engine: ArrayEngine | EngineName | None, *, store: Store, path: str, metadata: ArrayMetadata) -> ArrayEngine` + - `None` and `"default"` → default engine; `"zarrista"` → lazy import of `zarr.zarrista`, raising a clear `ImportError` if zarrista is missing; an instance → returned as-is. + - Hierarchy engines are cached per `(name, id(store))` in a module-level `WeakValueDictionary`-free plain dict keyed by a `weakref.ref` of the store (evicted via callback) so engines from one store share resources. + +- [ ] **Step 1: Write the failing test** + +```python +# tests/engine/test_resolve.py +from __future__ import annotations + +import numpy as np +import pytest + +import zarr +from zarr.core.engine import ( + DefaultArrayEngine, + DefaultAsyncArrayEngine, + resolve_async_engine, + resolve_sync_engine, +) +from zarr.storage import MemoryStore + + +def _array() -> zarr.Array: + return zarr.create_array(MemoryStore(), shape=(4,), chunks=(2,), dtype="int8") + + +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 + + +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: + pytest.importorskip("zarr.zarrista", reason="run only when zarrista absent") \ + if False else 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 + ) +``` + +- [ ] **Step 2: Run to verify failure** + +Run: `uv run pytest tests/engine/test_resolve.py -v` +Expected: FAIL (`ImportError: resolve_async_engine`) + +- [ ] **Step 3: Implement** + +```python +# src/zarr/core/engine/_resolve.py +"""Resolve an `engine=` argument to a bound array engine.""" + +from __future__ import annotations + +import weakref +from typing import TYPE_CHECKING, Literal + +from zarr.core.engine._default import ( + DefaultAsyncHierarchyEngine, + DefaultHierarchyEngine, +) + +if TYPE_CHECKING: + from zarr.abc.engine import ( + ArrayEngine, + AsyncArrayEngine, + AsyncHierarchyEngine, + HierarchyEngine, + ) + from zarr.abc.store import Store + from zarr.core.metadata import ArrayMetadata + +EngineName = Literal["default", "zarrista"] + +# (name, kind, id(store)) -> hierarchy engine; entries evicted when the store dies +_hierarchy_cache: dict[tuple[str, str, int], object] = {} + + +def _cached_hierarchy(name: str, kind: str, store: Store, factory) -> object: + key = (name, kind, id(store)) + if key not in _hierarchy_cache: + _hierarchy_cache[key] = factory(store) + weakref.finalize(store, _hierarchy_cache.pop, key, None) + return _hierarchy_cache[key] + + +def _hierarchy_factory(name: str, *, sync: bool): + 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, +) -> AsyncArrayEngine: + 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 hierarchy.array_engine(path, metadata) + return engine + + +def resolve_sync_engine( + engine: ArrayEngine | EngineName | None, + *, + store: Store, + path: str, + metadata: ArrayMetadata, +) -> ArrayEngine: + 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 hierarchy.array_engine(path, metadata) + return engine +``` + +Note: if `weakref.finalize` on a `Store` fails (stores with `__slots__` and no +`__weakref__`), fall back to a `weakref.WeakKeyDictionary[Store, dict[str, object]]` +keyed on the store; adjust the test only if the sharing assertion still holds. + +Re-export `EngineName`, `resolve_async_engine`, `resolve_sync_engine` from +`zarr.core.engine.__init__` and extend `__all__`. + +- [ ] **Step 4: Run tests** + +Run: `uv run pytest tests/engine/test_resolve.py -v` — Expected: PASS + +- [ ] **Step 5: Commit** + +```bash +git add src/zarr/core/engine tests/engine/test_resolve.py +git commit -m "feat: engine-spec resolution with per-store hierarchy cache" +``` + +--- + +### Task 5: Wire `AsyncArray` through its engine + +The public behavior contract for this task and Task 6: **the whole existing +array API test suite keeps passing** (`tests/test_array.py`, `tests/test_indexing.py`). +That suite is the real acceptance test; the new unit tests below only pin the +wiring itself. + +**Files:** +- Modify: `src/zarr/core/array.py` +- Test: `tests/engine/test_asyncarray_wiring.py` + +**Interfaces:** +- Consumes: `resolve_async_engine` (Task 4), `normalize_*`/`apply_post_index` (Task 2). +- Produces: `AsyncArray.__init__(..., engine: AsyncArrayEngine | EngineName | None = None)`; attribute `AsyncArray.engine`; rewritten module-level `_get_selection(engine, indexer_or_selection…)`/`_set_selection` (final signatures below). Task 6 mirrors this for `Array`; Task 7 threads `engine=` through creation APIs. + +- [ ] **Step 1: Write the failing test** + +```python +# tests/engine/test_asyncarray_wiring.py +from __future__ import annotations + +import numpy as np + +import zarr +from zarr.abc.engine import Region +from zarr.core.engine import DefaultAsyncArrayEngine +from zarr.storage import MemoryStore + + +class _SpyEngine: + """Wraps a real engine, recording regions.""" + + def __init__(self, inner: DefaultAsyncArrayEngine) -> None: + self.inner = inner + self.read_regions: list[Region] = [] + self.write_regions: list[Region] = [] + + async def read_selection(self, selection, *, prototype): + self.read_regions.append(selection) + return await self.inner.read_selection(selection, prototype=prototype) + + async def write_selection(self, selection, value, *, prototype): + self.write_regions.append(selection) + return await self.inner.write_selection(selection, value, prototype=prototype) + + def with_metadata(self, metadata): + return _SpyEngine(self.inner.with_metadata(metadata)) + + +def test_asyncarray_routes_io_through_engine() -> None: + 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) + + z[2:8] = np.arange(6, dtype="int16") + data = z[2:8] + + assert spy.write_regions == [Region(start=(2,), end_exclusive=(8,))] + assert spy.read_regions == [Region(start=(2,), end_exclusive=(8,))] + np.testing.assert_array_equal(np.asarray(data), np.arange(6, dtype="int16")) + + +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) +``` + +Note: `test_asyncarray_routes_io_through_engine` depends on Task 6 (the sync +`Array` path shares the async engine attribute until Task 6 gives `Array` its +own sync engine; when implementing Task 5 alone, exercise `await aa.getitem(...)` +instead of `z[...]` — final test content above is what must pass after Task 6). + +- [ ] **Step 2: Implement the `AsyncArray` changes** + +In `src/zarr/core/array.py`: + +1. `AsyncArray.__init__` (array.py:349): add keyword + `engine: AsyncArrayEngine | EngineName | None = None` to all three + signatures (both overloads + implementation), and set + +```python +object.__setattr__( + self, + "engine", + resolve_async_engine( + engine, + store=store_path.store, + path=store_path.path, + metadata=metadata_parsed, + ), +) +``` + + with attribute declaration `engine: AsyncArrayEngine = field(init=False)` + next to `codec_pipeline` (array.py:329). Import `resolve_async_engine` + and `EngineName` from `zarr.core.engine`, `AsyncArrayEngine` from + `zarr.abc.engine` (TYPE_CHECKING block for the types). + +2. Metadata-mutating paths: wherever `AsyncArray` constructs a successor + `AsyncArray` with new metadata (`resize`, `append`, `update_attributes`, + `with_config` — find them with + `grep -n "type(self)(" src/zarr/core/array.py` and + `grep -n "replace(self" src/zarr/core/array.py`), pass + `engine=self.engine.with_metadata(new_metadata)` if the constructor + accepts it, or call `object.__setattr__(new_array, "engine", + self.engine.with_metadata(new_metadata))` right after construction. + +3. Rewrite the two internal I/O helpers. Replace the bodies of + `AsyncArray._get_selection` (array.py:1412) and + `AsyncArray._set_selection` (array.py:1550) and the module-level + `_get_selection` (array.py:5353) / `_set_selection` (array.py:5702) with + engine-routed versions. The module-level functions get new signatures + (all callers are inside this file and the sync `Array` methods — + update every call site found via + `grep -n "_get_selection\|_set_selection" src/zarr/core/array.py`): + +```python +async def _get_selection( + engine: AsyncArrayEngine, + metadata: ArrayMetadata, + config: ArrayConfig, + region: Region, + post_index: tuple[Any, ...], + *, + prototype: BufferPrototype, + out: NDBuffer | None = None, + fields: Fields | None = None, +) -> NDArrayLikeOrScalar: + dtype = ( + metadata.dtype if metadata.zarr_format == 2 else metadata.data_type + ).to_native_dtype() + out_dtype = check_fields(fields, dtype) + if product(region.shape) == 0: + empty = np.empty( + apply_post_index(np.empty(region.shape, dtype=out_dtype), post_index).shape, + dtype=out_dtype, + ) + return _finalize_result(empty, out) + box = np.asarray(await engine.read_selection(region, prototype=prototype)) + if fields is not None: + box = box[fields] + result = apply_post_index(box, post_index) + return _finalize_result(result, out) + + +def _finalize_result(result: np.ndarray, out: NDBuffer | None) -> NDArrayLikeOrScalar: + if out is not None: + out.as_ndarray_like()[...] = result + return out.as_ndarray_like() + if result.shape == (): + return result[()] # scalar extraction, matching current behavior + return result +``` + +```python +async def _set_selection( + engine: AsyncArrayEngine, + metadata: ArrayMetadata, + config: ArrayConfig, + region: Region, + post_index: tuple[Any, ...], + value: npt.ArrayLike, + *, + prototype: BufferPrototype, + fields: Fields | None = None, +) -> None: + dtype = ( + metadata.dtype if metadata.zarr_format == 2 else metadata.data_type + ).to_native_dtype() + check_fields(fields, dtype) + fields = check_no_multi_fields(fields) + if product(region.shape) == 0: + return + value_np = np.asanyarray(value, dtype=None if fields else dtype) + + identity_post = all( + isinstance(p, slice) and p == slice(None) for p in post_index + ) and fields is None + if identity_post: + # broadcast the value to the full ndim-preserving box + box = np.broadcast_to(value_np, region.shape).astype(dtype, copy=False) + else: + # facade-level read-modify-write for strided/fancy/fields writes + box = np.array( + np.asarray(await engine.read_selection(region, prototype=prototype)) + ) + target = box[fields] if fields is not None else box + if post_index == (): + target[...] = value_np + else: + # _Squeeze markers only occur on reads (orthogonal); writes use + # the raw np.ix_/pointwise index without the marker + target[post_index] = value_np + value_buffer = prototype.nd_buffer.from_ndarray_like( + np.ascontiguousarray(box) + ) + await engine.write_selection(region, value_buffer, prototype=prototype) +``` + + Implementation notes: + - Keep the existing scalar/array-like `value` coercion logic from the old + `_set_selection` (array.py:5749-5769) ahead of the code above where it is + richer than `np.asanyarray` (non-numpy array types); preserve it rather + than the simplified line if any existing test fails. + - Writes for orthogonal selections must pass the *raw* per-axis outer index + (`np.ix_(...)` result without the `_Squeeze` marker); add a + `strip_squeeze(post) -> post` helper in `_normalize.py` and use it here. + - The old `regular_chunk_spec` optimization lives in the default engine + now (Task 3); nothing here touches chunks. + +4. Update `AsyncArray` methods that build indexers + (`getitem`, `setitem`, `get_basic_selection`, `set_basic_selection`, + `get_orthogonal_selection`, `set_orthogonal_selection`, + `get_mask_selection`, `set_mask_selection`, + `get_coordinate_selection`, `set_coordinate_selection`, + `get_block_selection`, `set_block_selection` — find each with + `grep -n "Indexer(" src/zarr/core/array.py`). Replace + `indexer = BasicIndexer(...)` + `self._get_selection(indexer, ...)` with: + +```python +region, post = normalize_basic(selection, self.shape) +return await _get_selection( + self.engine, self.metadata, self.config, region, post, + prototype=prototype, out=out, fields=fields, +) +``` + + and analogously `normalize_orthogonal` for orthogonal, `normalize_coordinate` + for coordinate, `np.nonzero(mask)` → `normalize_coordinate` for mask, and + `normalize_block(...)` (post index `()`) for block selections, passing the + chunk-grid shape from `self._chunk_grid`. Mask/coordinate *reads* return the + pointwise result (already numpy `vindex` semantics via the pointwise post + index). Delete the now-unused `Indexer` imports only if nothing else in the + file uses them (resize/`nchunks_initialized` may). + +- [ ] **Step 3: Run the wiring test and the array suites** + +Run: `uv run pytest tests/engine/test_asyncarray_wiring.py -v` — Expected: PASS (the async variant of the spy test if Task 6 not yet done) +Run: `uv run pytest tests/test_array.py tests/test_indexing.py -x -q` — Expected: PASS. Any failure here is a behavior regression: fix it before proceeding, favoring the old code's semantics. + +- [ ] **Step 4: Commit** + +```bash +git add src/zarr/core/array.py src/zarr/core/engine tests/engine/test_asyncarray_wiring.py +git commit -m "feat: AsyncArray routes selection I/O through its engine" +``` + +--- + +### Task 6: Sync `Array` data path (no event loop) + +**Files:** +- Modify: `src/zarr/core/array.py` +- Test: `tests/engine/test_sync_path.py` + +**Interfaces:** +- Consumes: `resolve_sync_engine` (Task 4), normalization helpers (Task 2), the `_finalize_result`/RMW patterns from Task 5. +- Produces: `Array.engine: ArrayEngine` property; sync module-level `_get_selection_sync`/`_set_selection_sync` mirroring Task 5's helpers with `engine.read_selection(...)` (no await). `Array.__init__`/constructors accept `engine: ArrayEngine | EngineName | None = None`. + +- [ ] **Step 1: Write the failing test** + +```python +# tests/engine/test_sync_path.py +from __future__ import annotations + +import asyncio + +import numpy as np +import pytest + +import zarr +from zarr.abc.engine import ArrayEngine, Region +from zarr.core.engine import DefaultArrayEngine +from zarr.storage import MemoryStore + + +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) + + +class _NoLoopEngine: + """A sync engine that asserts no event loop is running when called.""" + + def __init__(self, inner) -> None: + self._inner = inner + self.calls = 0 + + def read_selection(self, selection, *, prototype): + self.calls += 1 + with pytest.raises(RuntimeError): + asyncio.get_running_loop() + return self._inner.read_selection(selection, prototype=prototype) + + def write_selection(self, selection, value, *, prototype): + self.calls += 1 + with pytest.raises(RuntimeError): + asyncio.get_running_loop() + return self._inner.write_selection(selection, value, prototype=prototype) + + def with_metadata(self, metadata): + 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")) +``` + +(The engine methods themselves run in the caller thread; the *default* engine +internally uses `sync()` which runs coroutines on zarr's loop thread — that is +allowed. The probe asserts the caller thread has no running loop, i.e. `Array` +did not wrap the engine call in a coroutine.) + +- [ ] **Step 2: Implement** + +1. `Array` gains a lazily-resolved sync engine: store the creation-time spec + and resolve on first use. + +```python +# on class Array (near _async_array declaration, array.py:1803) +_engine: ArrayEngine | None = None +_engine_spec: ArrayEngine | EngineName | None = None + +@property +def engine(self) -> ArrayEngine: + """The synchronous array engine serving this array's data path.""" + if self._engine is None: + aa = self._async_array + object.__setattr__( + self, + "_engine", + resolve_sync_engine( + self._engine_spec, + store=aa.store_path.store, + path=aa.store_path.path, + metadata=aa.metadata, + ), + ) + object.__setattr__(self, "_engine_spec", None) + return self._engine +``` + + Whatever constructor `Array` uses (dataclass or `__init__` — check + `grep -n "def __init__\|@dataclass" src/zarr/core/array.py` around line + 1798) gains the optional `engine` keyword storing `_engine_spec`. When + `Array` wraps an `AsyncArray` that was itself given a *name* spec, reuse + that name so sync and async engines come from the same family. + +2. Every sync data method that currently does + `sync(self.async_array._get_selection(indexer, ...))` + (sites: array.py:2830, 2939, 3068, 3186, 3274, 3362, 3451, 3563, 3664, 3763 + — re-grep, they will have shifted after Task 5) is rewritten to the sync + mirror: + +```python +region, post = normalize_basic(selection, self.shape) # kind-appropriate helper +return _get_selection_sync( + self.engine, self._async_array.metadata, self._async_array.config, + region, post, prototype=prototype, out=out, fields=fields, +) +``` + +3. Add module-level `_get_selection_sync` / `_set_selection_sync` beside the + async versions from Task 5 — same bodies with `engine.read_selection(...)` / + `engine.write_selection(...)` un-awaited. Factor the shared pure logic + (dtype/fields resolution, empty short-circuit, box patching, result + finalization) into small helper functions used by both so the async and + sync variants differ only in the two engine calls. + +4. Metadata mutation on `Array` (`resize`, `append`) invalidates the cached + sync engine: `object.__setattr__(self, "_engine", self._engine.with_metadata(new_meta))` + when `_engine` is not None (sites via `grep -n "def resize\|def append" src/zarr/core/array.py`). + +- [ ] **Step 3: Run tests** + +Run: `uv run pytest tests/engine/ -v` — Expected: PASS (including the full Task 5 spy test now) +Run: `uv run pytest tests/test_array.py tests/test_indexing.py -q` — Expected: PASS +Run: `uv run mypy src/zarr` — Expected: no new errors + +- [ ] **Step 4: Commit** + +```bash +git add src/zarr/core/array.py src/zarr/core/engine tests/engine/test_sync_path.py +git commit -m "feat: Array sync data path calls its engine without the event loop" +``` + +--- + +### Task 7: `engine=` on creation/open APIs + +**Files:** +- Modify: `src/zarr/core/array.py` (`AsyncArray.open`, `Array.open`, `_create` paths), `src/zarr/api/asynchronous.py` and `src/zarr/api/synchronous.py` (`create_array`, `open_array`), `src/zarr/core/array_creation.py` if `create_array` lives there (find with `grep -rn "def create_array" src/zarr`). +- Test: `tests/engine/test_engine_param.py` + +**Interfaces:** +- Consumes: `EngineName`, engines from earlier tasks. +- Produces: public keyword `engine: ArrayEngine | AsyncArrayEngine | EngineName | None = None` on `zarr.create_array`, `zarr.open_array`, `zarr.api.asynchronous.create_array`, `zarr.api.asynchronous.open_array`, threaded to `AsyncArray.__init__` (async instances/names) and `Array._engine_spec` (sync instances/names). Passing a *sync instance* through an async API raises `TypeError`, and vice versa; a *name* is valid everywhere. + +- [ ] **Step 1: Write the failing test** + +```python +# tests/engine/test_engine_param.py +from __future__ import annotations + +import numpy as np +import pytest + +import zarr +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 + + +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" + ) +``` + +- [ ] **Step 2: Implement** + +Thread the keyword: each `create_array`/`open_array` entry point passes +`engine=engine` down to where `AsyncArray(...)`/`Array(...)` is constructed. +Rules: +- Async entry points pass names and `AsyncArrayEngine` instances to + `AsyncArray.__init__`; a sync `ArrayEngine` instance there raises + `TypeError("a sync ArrayEngine cannot serve an AsyncArray; pass a name or an AsyncArrayEngine")`. +- Sync entry points pass names to *both* layers (the `Array` stores the name + for its sync engine; the inner `AsyncArray` also gets the name so async + access stays consistent), and sync instances only to `Array`. +- Detection: `isinstance(engine, str) or engine is None` → both layers; + otherwise check for a coroutine `read_selection` via + `inspect.iscoroutinefunction(engine.read_selection)`. + +- [ ] **Step 3: Run tests** + +Run: `uv run pytest tests/engine/test_engine_param.py -v` — Expected: PASS +Run: `uv run pytest tests/test_api.py -q` — Expected: PASS + +- [ ] **Step 4: Commit** + +```bash +git add src/zarr/core/array.py src/zarr/api tests/engine/test_engine_param.py +git commit -m "feat: engine= parameter on array creation and open APIs" +``` + +--- + +### Task 8: Zarrista engines (`zarr.zarrista`) + +**Files:** +- Create: `src/zarr/zarrista/__init__.py`, `src/zarr/zarrista/_translate.py`, `src/zarr/zarrista/_engine.py` +- Modify: `pyproject.toml` (add `zarrista` dependency group, git-pinned) +- Test: `tests/zarrista/__init__.py` (empty), `tests/zarrista/test_translate.py`, `tests/zarrista/test_engine.py` + +**Interfaces:** +- Consumes: protocols/`Region` (Task 1), `UnsupportedEngineError` (Task 1). +- Produces: `ZarristaHierarchyEngine(store: Store)`, `ZarristaAsyncHierarchyEngine(store: Store)`, `ZarristaEngine`, `ZarristaAsyncEngine` (imported lazily by Task 4's resolver). +- Zarrista API used (from its `.pyi` stubs, main branch): sync + `zarrista.Array.from_metadata(metadata: dict, store, path)`, + `.retrieve_array_subset(selection) -> DecodedArray`, + `.store_chunk(chunk_indices, ArrayBytes)`, `.retrieve_chunk(chunk_indices)`, + `.chunk_subset(chunk_indices) -> tuple[slice, ...]`, `.chunk_grid_shape`, + `.chunk_shape(chunk_indices)`; async twins on `zarrista.AsyncArray` + (`from_metadata`, awaitable retrieve/store). Stores: + `zarrista.FilesystemStore(path)` (sync), obstore `ObjectStore` / icechunk + `Session` (async). `zarrista.ArrayBytes(bytes=...)` wraps decoded chunk bytes. + +- [ ] **Step 1: Add the dependency group** + +In `pyproject.toml`, next to the existing dependency groups (see +`[dependency-groups]`), add (pin the current zarrista main commit hash — get it +with `git ls-remote https://github.com/developmentseed/zarrista HEAD`): + +```toml +zarrista = [ + "zarrista @ git+https://github.com/developmentseed/zarrista@", +] +``` + +Run: `uv sync --group zarrista` — Expected: builds/installs zarrista (Rust +toolchain required; if the build fails locally, note it and let CI cover it — +zarrista publishes no sdist-independent wheels from git). + +- [ ] **Step 2: Write the failing tests** + +```python +# tests/zarrista/test_translate.py +from __future__ import annotations + +import pytest + +zarrista = pytest.importorskip("zarrista") + +from zarr.errors import UnsupportedEngineError +from zarr.storage import LocalStore, MemoryStore +from zarr.zarrista._translate import translate_store_async, translate_store_sync + + +def test_local_store_translates_to_filesystem_store(tmp_path) -> None: + zs = translate_store_sync(LocalStore(tmp_path)) + assert isinstance(zs, zarrista.FilesystemStore) + + +def test_memory_store_rejected_sync(tmp_path) -> None: + with pytest.raises(UnsupportedEngineError, match="MemoryStore"): + translate_store_sync(MemoryStore()) + + +def test_local_store_rejected_async(tmp_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) -> 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 +``` + +```python +# tests/zarrista/test_engine.py +from __future__ import annotations + +import numpy as np +import pytest + +zarrista = pytest.importorskip("zarrista") + +import zarr +from zarr.errors import UnsupportedEngineError +from zarr.storage import LocalStore + + +def _make(tmp_path, **kwargs): + z = zarr.create_array( + LocalStore(tmp_path), shape=(10, 9), chunks=(3, 4), dtype="float32", **kwargs + ) + z[:, :] = np.arange(90, dtype="float32").reshape(10, 9) + return z + + +def test_zarrista_engine_read_write_combinations(tmp_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(z[1, 1]) == 99.0 + assert float(z[0, 0]) == 0.0 # neighbor in same chunk untouched + + +def test_zarrista_rejects_v2(tmp_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") +``` + +Run: `uv run --group zarrista pytest tests/zarrista -v` +Expected: FAIL (`ModuleNotFoundError: zarr.zarrista`) + +- [ ] **Step 3: Implement translation** + +```python +# src/zarr/zarrista/_translate.py +"""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 (`FilesystemStore`).""" + import zarrista + + if isinstance(store, LocalStore): + return zarrista.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 # type: ignore[import-not-found] + + 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}." + ) +``` + +(Verify the icechunk attribute: `IcechunkStore.session` — check with +`uv run python -c "import icechunk, inspect; print([m for m in dir(icechunk.IcechunkStore) if 'session' in m])"` +if icechunk is installed; otherwise leave the guarded branch as written.) + +- [ ] **Step 4: Implement the engines** + +```python +# src/zarr/zarrista/_engine.py +"""Zarrista-backed array engines.""" + +from __future__ import annotations + +import itertools +from typing import TYPE_CHECKING, Any + +import numpy as np + +from zarr.abc.engine import Region +from zarr.errors import UnsupportedEngineError +from zarr.zarrista._translate import translate_store_async, translate_store_sync + +if TYPE_CHECKING: + from zarr.abc.store import Store + from zarr.core.buffer import BufferPrototype, NDBuffer + from zarr.core.metadata import ArrayMetadata + + +def _require_v3(metadata: ArrayMetadata) -> dict[str, Any]: + if metadata.zarr_format != 3: + raise UnsupportedEngineError( + "the zarrista engine supports Zarr v3 only; this array is " + f"format v{metadata.zarr_format}" + ) + return metadata.to_dict() + + +def _region_to_selection(region: Region) -> tuple[slice, ...]: + return tuple( + slice(s, e) for s, e in zip(region.start, region.end_exclusive, strict=True) + ) + + +def _decoded_to_numpy(decoded: Any) -> np.ndarray: + # Tensor / (future) VariableArray implement __array__; masked layouts do not + # map to a zarr-python dtype. + type_name = type(decoded).__name__ + if type_name in ("MaskedTensor", "MaskedVariableArray"): + raise NotImplementedError( + f"zarrista returned {type_name}; masked layouts have no " + "zarr-python equivalent" + ) + return np.asarray(decoded) + + +def _chunks_overlapping(region: Region, chunk_shape: tuple[int, ...]) -> Any: + ranges = [ + range(s // c, (e + c - 1) // c) if e > s else range(0) + for s, e, c in zip(region.start, region.end_exclusive, chunk_shape, strict=True) + ] + return itertools.product(*ranges) + + +class ZarristaEngine: + """Sync engine over `zarrista.Array`. No event loop involved.""" + + def __init__(self, zarrista_array: Any) -> None: + self._arr = zarrista_array + + @classmethod + def from_zarr_metadata(cls, store: Store, path: str, metadata: ArrayMetadata) -> ZarristaEngine: + import zarrista + + meta_dict = _require_v3(metadata) + zstore = translate_store_sync(store) + return cls(zarrista.Array.from_metadata(meta_dict, zstore, "/" + path.strip("/"))) + + def with_metadata(self, metadata: ArrayMetadata) -> ZarristaEngine: + import zarrista + + return ZarristaEngine( + zarrista.Array.from_metadata( + _require_v3(metadata), self._arr.store, self._arr.path + ) + ) + + def read_selection(self, selection: Region, *, prototype: BufferPrototype) -> Any: + return _decoded_to_numpy( + self._arr.retrieve_array_subset(_region_to_selection(selection)) + ) + + def write_selection( + self, selection: Region, value: NDBuffer, *, prototype: BufferPrototype + ) -> None: + import zarrista + + value_np = np.ascontiguousarray(value.as_ndarray_like()) + # regular grids: the origin chunk has the full (unclipped) chunk shape + chunk_shape = tuple(self._arr.chunk_shape([0] * len(selection.start))) + for chunk_idx in _chunks_overlapping(selection, chunk_shape): + chunk_idx = list(chunk_idx) + chunk_slices = self._arr.chunk_subset(chunk_idx) + # overlap of the write region with this chunk, in array coords + lo = tuple( + max(cs.start, s) + for cs, s in zip(chunk_slices, selection.start, strict=True) + ) + hi = tuple( + min(cs.stop, e) + for cs, e in zip(chunk_slices, selection.end_exclusive, strict=True) + ) + in_value = tuple( + slice(a - s, b - s) + for a, b, s in zip(lo, hi, selection.start, strict=True) + ) + in_chunk = tuple( + slice(a - cs.start, b - cs.start) + for a, b, cs in zip(lo, hi, chunk_slices, strict=True) + ) + full_chunk = all( + sl.start == 0 and sl.stop == (cs.stop - cs.start) + for sl, cs in zip(in_chunk, chunk_slices, strict=True) + ) + if full_chunk: + chunk_np = np.ascontiguousarray(value_np[in_value]) + else: + chunk_np = np.array(_decoded_to_numpy(self._arr.retrieve_chunk(chunk_idx))) + chunk_np[in_chunk] = value_np[in_value] + chunk_np = np.ascontiguousarray(chunk_np) + self._arr.store_chunk(chunk_idx, zarrista.ArrayBytes(chunk_np)) + + +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) -> ZarristaEngine: + import zarrista + + return ZarristaEngine( + zarrista.Array.from_metadata( + _require_v3(metadata), self._zstore, "/" + path.strip("/") + ) + ) +``` + +`ZarristaAsyncEngine` / `ZarristaAsyncHierarchyEngine`: same shape with +`zarrista.AsyncArray.from_metadata`, `translate_store_async`, and `await` on +`retrieve_array_subset` / `retrieve_chunk` / `store_chunk`. Write the full +async twin classes — do not factor a shared base for v1 (the sync class must +never touch a loop; keeping them textually parallel is clearer). + +Implementation notes: +- `zarrista.ArrayBytes(chunk_np)` — the constructor takes `bytes: Buffer`; + pass the numpy array (it implements the buffer protocol). If zarrista + rejects multi-dimensional buffers, pass + `chunk_np.reshape(-1).view(np.uint8)` — decide by running the test. +- `store_chunk` expects the *full* chunk shape including edge-chunk overhang; + `chunk_subset` returns clipped slices at array edges. For edge chunks the + decoded chunk from `retrieve_chunk` already has the full chunk shape — + patch using `in_chunk` offsets (clipped coords are correct because + `chunk_subset` slices start at the chunk origin). For a *full-chunk* write + on an edge chunk (`full_chunk` computed against the clipped extent), the + written array is smaller than the chunk shape — in that case fall through + to the RMW branch. Adjust `full_chunk` to also require + `(cs.stop - cs.start) == chunk_shape[dim]`. +- `zarrista.Array.from_metadata` path convention is `/`-prefixed (`"/"` is + the root), matching the stub default `path="/"`. +- vlen dtypes: `retrieve_array_subset` returns `VariableArray`; if + `np.asarray` on it fails under the pinned commit, raise + `NotImplementedError("vlen read support pending zarrista numpy export")` + and mark the corresponding differential test `xfail(strict=False)`. + +```python +# src/zarr/zarrista/__init__.py +"""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", +] +``` + +- [ ] **Step 5: Run tests** + +Run: `uv run --group zarrista pytest tests/zarrista -v` — Expected: PASS +Run: `uv run pytest tests/zarrista -v` (no group) — Expected: all SKIPPED (importorskip) + +- [ ] **Step 6: Commit** + +```bash +git add src/zarr/zarrista tests/zarrista pyproject.toml uv.lock +git commit -m "feat: zarrista-backed array engines with store translation" +``` + +--- + +### Task 9: Delete `zarr.crud`, `zarr.zarrs`, `packages/zarrs-bindings` + +**Files:** +- Delete: `src/zarr/crud/`, `src/zarr/zarrs/`, `packages/zarrs-bindings/`, `tests/crud/`, `tests/zarrs/`, `.github/workflows/zarrs.yml`, the `changes/` fragments describing zarr.crud/zarr.zarrs (find: `grep -rl "crud\|zarrs" changes/`) +- Modify: `pyproject.toml` — remove the `zarrs` dependency group (line ~138), the `zarrs-bindings` uv source (line ~502), the `/packages/zarrs-bindings` exclude (line ~10), and the `--ignore=src/zarr/zarrs` pytest flag (line ~444). + +- [ ] **Step 1: Delete and unwire** + +```bash +git rm -r src/zarr/crud src/zarr/zarrs packages/zarrs-bindings tests/crud tests/zarrs .github/workflows/zarrs.yml +``` + +Edit `pyproject.toml` as listed above. Search for stragglers: + +```bash +grep -rn "zarr.crud\|zarr\.zarrs\|zarrs_bindings\|zarrs-bindings" src tests docs pyproject.toml .github --include="*" | grep -v superpowers +``` + +Expected: no hits outside `docs/superpowers/` (specs/plans are historical records — leave them). + +- [ ] **Step 2: Full test run** + +Run: `uv sync && uv run pytest tests -x -q --ignore=tests/zarrista` +Expected: PASS (no imports of deleted modules anywhere) +Run: `uv run mypy src/zarr` — Expected: no errors + +- [ ] **Step 3: Commit** + +```bash +git add -u src tests pyproject.toml uv.lock .github +git commit -m "refactor!: remove zarr.crud, zarr.zarrs, and the zarrs-bindings crate" +``` + +--- + +### Task 10: Differential suite, CI, changelog + +**Files:** +- Create: `tests/engine/test_differential.py`, `.github/workflows/engine.yml`, `changes/.feature.md` (get the PR number from the branch's open PR via `gh pr view --json number -q .number`; if no PR exists yet, use the next number after `gh api repos/{owner}/{repo}/issues?state=all&per_page=1 -q '.[0].number'` and note it must match the eventual PR) + +**Interfaces:** +- Consumes: everything. + +- [ ] **Step 1: Differential test** + +```python +# tests/engine/test_differential.py +"""The same operations through both engines must agree with numpy and each other.""" + +from __future__ import annotations + +import numpy as np +import pytest + +import zarr +from zarr.storage import LocalStore + +try: + import zarrista # noqa: F401 + + ENGINES = ["default", "zarrista"] +except ImportError: + ENGINES = ["default"] + +SHAPE = (10, 9) +CHUNKS = (3, 4) + + +@pytest.fixture +def reference(tmp_path): + 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 + + +READS = [ + (slice(None), slice(None)), + (slice(2, 7), slice(1, 5)), + (slice(8, 2, -2), slice(None, None, 3)), + (3, slice(None)), + (-1, -2), + (slice(4, 4), slice(None)), +] + + +@pytest.mark.parametrize("engine", ENGINES) +@pytest.mark.parametrize("sel", READS) +def test_reads_match_numpy(reference, engine, sel) -> None: + tmp_path, data = reference + z = zarr.open_array(LocalStore(tmp_path), engine=engine) + np.testing.assert_array_equal(np.asarray(z[sel]), data[sel]) + + +@pytest.mark.parametrize("engine", ENGINES) +def test_fancy_reads_match_numpy(reference, engine) -> None: + tmp_path, data = reference + z = zarr.open_array(LocalStore(tmp_path), engine=engine) + np.testing.assert_array_equal( + np.asarray(z.oindex[np.array([7, 1, 4]), np.array([0, 8])]), + data[np.ix_([7, 1, 4], [0, 8])], + ) + np.testing.assert_array_equal( + np.asarray(z.vindex[np.array([9, 0, 3]), np.array([8, 0, 2])]), + data[np.array([9, 0, 3]), np.array([8, 0, 2])], + ) + np.testing.assert_array_equal(np.asarray(z.blocks[1, 2]), data[3:6, 8:9]) + + +@pytest.mark.parametrize("engine", ENGINES) +def test_writes_match_numpy(reference, engine) -> None: + tmp_path, data = reference + z = zarr.open_array(LocalStore(tmp_path), engine=engine) + expected = data.copy() + + z[0:3, 0:4] = 7.0 # aligned full chunk + expected[0:3, 0:4] = 7.0 + z[4:6, 2:9] = np.arange(14, dtype="float64").reshape(2, 7) # partial chunks + expected[4:6, 2:9] = np.arange(14, dtype="float64").reshape(2, 7) + z[1:9:3, ::4] = -1.0 # strided write (facade RMW) + expected[1:9:3, ::4] = -1.0 + + np.testing.assert_array_equal(np.asarray(z[:, :]), expected) + + +@pytest.mark.parametrize("engine", ENGINES) +def test_sharded_reads(reference, engine, tmp_path_factory) -> None: + path = tmp_path_factory.mktemp("sharded") + z = zarr.create_array( + LocalStore(path), + shape=SHAPE, + chunks=(3, 4), # inner chunks + shards=(6, 8), # shard shape + dtype="int32", + ) + data = np.arange(90, dtype="int32").reshape(SHAPE) + z[:, :] = data + zr = zarr.open_array(LocalStore(path), engine=engine) + np.testing.assert_array_equal(np.asarray(zr[2:8, 3:9]), data[2:8, 3:9]) +``` + +Run: `uv run --group zarrista pytest tests/engine/test_differential.py -v` — Expected: PASS (both engines) +Run: `uv run pytest tests/engine/test_differential.py -v` — Expected: PASS (default only) + +- [ ] **Step 2: CI workflow** + +Create `.github/workflows/engine.yml` modeled on the deleted `zarrs.yml` +(check its triggers/paths in git history: `git show HEAD~1:.github/workflows/zarrs.yml`): +a single job on `ubuntu-latest` that installs Rust (`dtolnay/rust-toolchain@stable`), +sets up uv (`astral-sh/setup-uv`), runs +`uv sync --group zarrista` and +`uv run --group zarrista pytest tests/engine tests/zarrista -v`. +Trigger on pull requests touching `src/zarr/**`, `tests/engine/**`, +`tests/zarrista/**`, `pyproject.toml`, and the workflow file itself. Keep +`permissions: {}` at top level and pin action SHAs, matching the repo's other +workflows (zizmor enforces this). + +- [ ] **Step 3: Changelog fragment** + +```markdown +# changes/.feature.md +`Array` and `AsyncArray` now route data I/O through pluggable *array engines* +(`zarr.abc.engine.ArrayEngine` / `AsyncArrayEngine`). The default engine +preserves existing behavior on every store and format. Pass +`engine="zarrista"` (requires the `zarrista` package) to serve Zarr v3 arrays +on local-filesystem, obstore, or icechunk storage through the Rust `zarrs` +implementation. The experimental `zarr.crud` and `zarr.zarrs` modules and the +bundled `zarrs-bindings` crate are removed in favor of this interface. +``` + +- [ ] **Step 4: Full local gate** + +Run: `uv run --group zarrista pytest tests -q` — Expected: PASS +Run: `uv run mypy src/zarr` — Expected: no errors +Run: `uv run pre-commit run --all-files` — Expected: PASS + +- [ ] **Step 5: Commit** + +```bash +git add tests/engine/test_differential.py .github/workflows/engine.yml changes/ +git commit -m "test/ci: engine differential suite and zarrista CI job" +``` + +--- + +## Self-Review Notes (already applied) + +- Spec coverage: protocols+Region (T1), normalization/facade rules (T2, T5, T6), + default engines incl. sync adapter (T3), explicit engine selection without a + config key (T4, T7), truly-sync path (T6), zarrista store translation / + v3-only / RMW writes / buffer handling (T8), fail-loud errors (T1, T4, T8), + deletions (T9), differential tests + CI + changelog (T10). `with_metadata` + rebinding covered in T3/T5/T6. +- Deliberately out of scope, per spec non-goals: group operations, engine-owned + metadata mutation, `out=` in the engine contract (facade copies instead). +- Known judgment calls the implementer may hit: exact `Array` constructor shape + (T6 step 1), zarrista `ArrayBytes` buffer dimensionality (T8 note), icechunk + session attribute (T8 note), vlen export under the pinned commit (T8 note). + Each has a decision rule written at the point of use. diff --git a/docs/superpowers/specs/2026-06-11-zarrs-functional-api-design.md b/docs/superpowers/specs/2026-06-11-zarrs-functional-api-design.md new file mode 100644 index 0000000000..22cb0b5785 --- /dev/null +++ b/docs/superpowers/specs/2026-06-11-zarrs-functional-api-design.md @@ -0,0 +1,210 @@ +# zarrs-backed low-level functional API for zarr-python + +Date: 2026-06-11 +Status: approved +Branch: `zarrs-bindings` + +## Goal + +Give zarr-python a low-level, functional API for zarr hierarchy CRUD whose +implementation delegates to the Rust [`zarrs`](https://docs.rs/zarrs) crate via +new PyO3 bindings. Every array routine takes a metadata document as an explicit +parameter, so callers can operate on read-only or virtual views of arrays +(e.g. decode a chunk with metadata the store never saw, or read a chunk as raw +bytes without decoding). + +Non-goals for this work: rewiring zarr-python's `Array`/`Group` classes or the +codec-pipeline registry through this API (possible later), fancy +(non-slice) indexing, and use of zarrs's experimental async feature. + +## Background + +- zarr-python is pure Python (hatchling). Its `Store` ABC + (`src/zarr/abc/store.py`) is async; metadata classes live under + `src/zarr/core/metadata/`. +- The Rust `zarrs` crate (~0.23) supports exactly the metadata-driven shape we + need: `Array::new_with_metadata(storage, path, metadata)` and + `Group::new_with_metadata(...)` construct nodes from a metadata document + without touching the store; `store_metadata()` persists separately. Chunk and + region I/O: `retrieve_chunk`, `retrieve_encoded_chunk` (raw bytes), + `retrieve_array_subset`, `partial_decoder` (sharding-aware), and the + corresponding `store_*` methods. `ArrayMetadata`/`GroupMetadata` parse + directly from JSON strings (v2 or v3; v2 converts internally). +- The existing `zarrs` PyPI package (github.com/zarrs/zarrs-python) exposes only + a codec pipeline (`CodecPipelineImpl`) and supports only a fixed set of + native stores. It cannot provide the API designed here, but its build setup + (maturin, PyO3 abi3, tokio/rayon) is the reference for ours. + +## Architecture + +Two distributions in this repo, hard boundary between them: + +1. **Rust crate `zarrs-bindings`** under `packages/` (`packages/zarrs-bindings/`, + alongside the existing `zarr-metadata` subpackage), + built with maturin (PyO3, `abi3-py312`), publishing wheel `zarrs-bindings` + with native module `_zarrs_bindings`. It is a thin, mechanical binding over + `zarrs`: functions/pyclasses take metadata as a **JSON string**, a + store-config object, a node path, and return bytes / numpy arrays. It knows + nothing about zarr-python except the store sniffing described below. +2. **Python subpackage `zarr.zarrs`** in zarr-python: the public functional + API. Owns conversion between zarr-python types (`dict` metadata documents, + `zarr.abc.store.Store`, numpy arrays) and the binding layer, plus + validation, ergonomics, and error translation. Imports `_zarrs_bindings` + lazily and raises a helpful `ImportError` naming the `zarr[zarrs]` extra if + it is missing. + +zarr-python's own wheel remains pure Python; `zarrs-bindings` becomes an +optional dependency (`zarr[zarrs]`). + +## Public API (`zarr.zarrs`) + +All functions are `async def`. Parameters: + +- `metadata`: `dict[str, JSON]` — the literal metadata document (`zarr.json`, + or v2 `.zarray`/`.zgroup` equivalents). Never read from the store by the + array routines. +- `store`: `zarr.abc.store.Store`. +- `path`: node path within the store (str, `""` = root). +- `chunk_coords`: `tuple[int, ...]` grid coordinates. +- `selection`: numpy-style basic indexing — integers, slices (including steps; strided/reversed selections fetch the step-1 bounding box in one call and apply numpy views), and `Ellipsis`. Fancy indexing (integer/boolean arrays) and `np.newaxis` are not supported. +- `options`: every function also accepts keyword-only + `options: ZarrsOptions | None = None` (omitted from the signatures below for + brevity) — a dataclass holding concurrency limits and checksum validation + flags. Defaults are applied when omitted; in Phase 1 the dataclass exists + but carries only defaults (fields become meaningful in Phase 3). + +```python +# node lifecycle +async def create_new_group(metadata, store, path) -> None # error if node exists +async def create_overwrite_group(metadata, store, path) -> None +async def create_new_array(metadata, store, path) -> None +async def create_overwrite_array(metadata, store, path) -> None +async def read_metadata(store, path) -> dict[str, JSON] # array or group doc +async def delete_node(store, path) -> None +async def list_children(store, path) -> list[tuple[str, dict]] # (path, metadata) + +# chunk-level I/O +async def decode_chunk(metadata, store, path, chunk_coords, *, selection=None) -> np.ndarray +async def read_encoded_chunk(metadata, store, path, chunk_coords) -> bytes | None +async def encode_chunk(metadata, store, path, chunk_coords, value) -> None +async def erase_chunk(metadata, store, path, chunk_coords) -> None + +# region-level I/O (selection in array coordinates, may span chunks) +async def decode_region(metadata, store, path, selection) -> np.ndarray +async def encode_region(metadata, store, path, selection, value) -> None +``` + +Mapping to zarrs primitives: + +| API function | zarrs primitive | +|---|---| +| `create_new_group` / `create_overwrite_group` | `Group::new_with_metadata` + `store_metadata` (existence check first for `new`) | +| `create_new_array` / `create_overwrite_array` | `Array::new_with_metadata` + `store_metadata` | +| `read_metadata` | `Array::open` / `Group::open` metadata retrieval | +| `delete_node` | `erase_metadata` + chunk erasure / prefix delete | +| `list_children` | `Group::children` / `traverse` | +| `decode_chunk` (no selection) | `retrieve_chunk` | +| `decode_chunk` (selection) | `partial_decoder(chunk).partial_decode` (sharding-aware) | +| `read_encoded_chunk` | `retrieve_encoded_chunk` | +| `encode_chunk` | `store_chunk` | +| `erase_chunk` | `erase_chunk` | +| `decode_region` | `retrieve_array_subset` | +| `encode_region` | `store_array_subset` | + +## Store bridge + +A Rust-side `StoreConfig` resolver, tried in priority order: + +1. `zarr.storage.LocalStore` → native `zarrs_filesystem` store. +2. obstore-backed `ObjectStore` → `zarrs_object_store` (Phase 3). +3. **Anything else** → generic `PyStore`: a Rust struct implementing + `ReadableStorageTraits` / `WritableStorageTraits` / + `ListableStorageTraits` over a Python callback object. + +The callback path: the async API function wraps the user's `Store` in a small +sync Python shim whose methods submit coroutines to zarr-python's existing +sync event-loop thread (`zarr.core.sync`, +`asyncio.run_coroutine_threadsafe(...)` + blocking result). Rust calls the +shim while holding no locks of its own. This makes any conformant `Store` +(Memory, Zip, Logging, Wrapper, user-defined) work without Rust knowing its +type. Deadlock safety relies on the existing invariant that code running on +the zarr sync loop never blocks on these Rust entry points. + +## Sync/async seam + +The public API is async to match zarr-python conventions. Internally each +function calls a blocking Rust entry point via `asyncio.to_thread`; the Rust +side releases the GIL during I/O and compute (reacquiring it only inside +`PyStore` callbacks). zarrs's experimental async feature is not used. + +## Array construction cache + +`Array::new_with_metadata` (serde-parsing the metadata document and building the +codec chain) is the dominant per-call cost on the native path — measured at +~20µs for a bytes-only array up to ~80µs for sharded+blosc, against single-digit +µs of actual chunk I/O on a warm filesystem. To amortize it across the common +"open one array, then do many chunk operations" pattern, the chunk/region +routines memoize the constructed `Array` in a process-global LRU cache +(capacity 128) keyed on `(filesystem root, node path, metadata JSON)`. + +This is safe because a zarrs `Array` caches no chunk data — it is metadata plus +codec chain plus a storage handle — so every read/write still goes through to +the store, and a correctly-keyed hit is behaviorally identical to a fresh build. +The key must include all three components: the same document at a different path +or store is a different array. Only native filesystem stores are cached; the +generic `PyStore` callback path has no stable cross-call identity to key on and +is left uncached (a future change may cache it if a store can supply a stable +value-based token). No invalidation hook is needed: delete/overwrite with +different metadata yields a different key, and an entry for a deleted-and-rebuilt +array with identical metadata stays valid because reads go through to the store. +A poisoned cache mutex is recovered rather than propagated, so the cache can +never wedge array I/O. Measured win: 14–20% faster per repeated call on a local +store, free on every hit. + +## Error handling + +The binding layer raises a small set of typed exceptions defined in one place: +`NodeExistsError`, `NodeNotFoundError`, and `ValueError` subclasses for +metadata-parse failures. In Phase 1 the translation surface is deliberately +small: `zarr.zarrs` re-raises the bindings' `NodeNotFoundError` as +`zarr.errors.NodeNotFoundError`; `NodeExistsError` is exposed as +`zarr.zarrs.NodeExistsError`. Exceptions raised by Python store callbacks are +flattened to a `RuntimeError` carrying the original message — the original +exception type and traceback are lost crossing the Rust boundary. Faithful +propagation of store-callback exceptions (and richer mapping onto +`zarr.errors` types) is deferred to a later phase. + +## Testing + +`tests/zarrs/`, module-level skip when `_zarrs_bindings` is not importable. + +- **Differential tests** are the core: every operation checked against + zarr-python's own implementation on the same store — write with zarr-python, + read with zarrs, and vice versa; metadata documents produced by both must + round-trip. +- Parametrized over: `MemoryStore` (exercises generic bridge) and `LocalStore` + (native path); zarr formats v2 and v3; a codec matrix including + `sharding_indexed`. +- Read-only-view tests: decode a chunk using a metadata dict not present in + the store; `read_encoded_chunk` returns bytes identical to `store.get`. +- A CI job builds the crate with `maturin develop` and runs `tests/zarrs/`. + Existing CI jobs are untouched (the suite skips without the extension). + +## Phasing + +1. **Phase 1**: crate scaffolding (maturin, CI build), store bridge (native + LocalStore + generic PyStore), node lifecycle functions, whole-chunk + `decode_chunk` / `read_encoded_chunk` / `encode_chunk` / `erase_chunk`. +2. **Phase 2**: `decode_region` (read side of region I/O) is implemented on + this branch. `encode_region` and chunk-subset `selection` for `decode_chunk` + via partial decoders remain Phase 2. +3. **Phase 3**: `ZarrsOptions` surface (concurrency, checksum validation, + direct IO), obstore native path, benchmarks vs. the pure-Python pipeline. + +## Naming decisions + +- Python API: `zarr.zarrs`. +- Rust crate / PyPI distribution: `zarrs-bindings` (PyPI name `zarrs` is taken + by the existing project); native module `_zarrs_bindings`. +- Function names follow the requested `create_new_*` / `create_overwrite_*` + pattern; reads are `decode_*` / `read_*`, writes `encode_*`. diff --git a/docs/superpowers/specs/2026-06-15-crud-backend-abstraction-design.md b/docs/superpowers/specs/2026-06-15-crud-backend-abstraction-design.md new file mode 100644 index 0000000000..bf8043513b --- /dev/null +++ b/docs/superpowers/specs/2026-06-15-crud-backend-abstraction-design.md @@ -0,0 +1,259 @@ +# Backend-agnostic CRUD layer for zarr-python + +Date: 2026-06-15 +Status: approved +Branch: `zarrs-bindings` + +## Goal + +Turn the low-level functional CRUD API (introduced as `zarr.zarrs` earlier on +this branch) into a backend-agnostic layer, with the Rust zarrs bindings as one +of several interchangeable implementations. Define the CRUD contract abstractly, +provide a pure-Python reference backend (no Rust required), and make the zarrs +bindings conform to the same contract. + +This validates the abstraction by having two real backends agree with each other +and with zarr-python, and it gives users a no-Rust fallback. + +Non-goals for this change (deliberately deferred): + +- Wiring the CRUD layer under zarr-python's own `Array`/`Group` classes. +- Entrypoint-based backend discovery (this change uses explicit import-time + registration). +- A write-side region operation (`write_region`) remains future work. + +## Background + +The current `zarr.zarrs._api` is a flat module of 13 async functions that +delegate to the `_zarrs_bindings` Rust extension. It already separates two +concerns that this design formalizes into a hard boundary: + +- **Backend-neutral glue:** `_normalize_selection`, `_array_shape`, + `_chunk_dtype_and_shape`, numpy assembly (`np.frombuffer`/reshape/strided + views), native-dtype coercion, options handling, error translation. +- **Genuinely zarrs-specific work:** producing/consuming raw chunk bytes, + reading array subsets as bytes, writing metadata documents — all via + `_zarrs_bindings` and the `_bridge.StoreShim`/`resolve_store` plumbing. + +The public surface (`zarr.zarrs.decode_region`, etc.) is unreleased on this +branch, so it can move without backward-compatibility constraints. + +zarr-python already contains everything a pure-Python backend needs: +`BatchedCodecPipeline` (`src/zarr/core/codec_pipeline.py`), `BasicIndexer` +(`src/zarr/core/indexing.py`), `save_metadata` (`src/zarr/core/metadata/io.py`), +metadata parsing (`ArrayV3Metadata.from_dict` / `ArrayV2Metadata.from_dict`), +and chunk-key encoding (`src/zarr/core/chunk_key_encodings.py`). + +## Architecture + +Two packages with a hard boundary. + +### `zarr.crud` (new, backend-neutral) + +- `_backend.py` — the `CrudBackend` `Protocol` (the narrow byte/metadata + contract below) plus the canonical exceptions. +- `_api.py` — the shared async facade: the 13 public functions moved out of + `zarr.zarrs`, holding all backend-neutral logic. Each function resolves a + backend (from the `backend` argument or the registry default) and calls its + byte/metadata methods, then does selection normalization, dtype handling, and + numpy assembly. +- `_reference.py` — `ReferenceBackend`, pure Python, wrapping zarr-python's own + codec/indexing/metadata machinery. Always importable; the default backend. +- `_registry.py` — `register_backend(name, backend)`, `get_backend(name)`, and + the config-driven default resolution. +- `__init__.py` — re-exports the facade functions, `CrudBackend`, + `ZarrsOptions`, the exceptions, and `register_backend`. + +### `zarr.zarrs` (shrinks to the zarrs provider) + +- `_backend.py` — `ZarrsBackend`, implementing `CrudBackend` by wrapping + `_zarrs_bindings`. Owns the zarrs-isms that move out of the facade: + `json.dumps` of the metadata dict, the `/`-prefixed zarrs node-path form + (formerly `_node_path`), `_bridge.resolve_store`, and translation of + `_zarrs_bindings` exceptions into the `zarr.crud` canonical exceptions. +- `_bridge.py` — unchanged (`StoreShim`, `resolve_store`). +- the Rust crate `zarrs-bindings/` and the construction cache — unchanged. +- registers itself as backend `"zarrs"` at import time. + +## The `CrudBackend` contract + +Narrow, byte/metadata level. Methods pass neutral types — the metadata document +as a `dict`, the zarr `Store`, and plain zarr paths (`""`, `"foo/bar"`) — and +return raw bytes / JSON-as-dict / `None`. Each backend serializes and bridges as +it needs. + +```python +class CrudBackend(Protocol): + async def create_array(self, store, path, metadata, *, overwrite: bool) -> None: ... + async def create_group(self, store, path, metadata, *, overwrite: bool) -> None: ... + async def read_metadata(self, store, path) -> dict[str, JSON]: ... + async def read_chunk(self, store, path, metadata, coords) -> bytes: ... + async def read_subset(self, store, path, metadata, start, shape) -> bytes: ... + async def write_chunk(self, store, path, metadata, coords, data: bytes) -> None: ... + async def delete_chunk(self, store, path, metadata, coords) -> None: ... + async def delete_node(self, store, path) -> None: ... + async def list_children(self, store, path) -> list[tuple[str, dict[str, JSON]]]: ... +``` + +Nine methods. `read_encoded_chunk` is deliberately **not** a backend method — +see below; it is a backend-independent facade helper over `store.get`. + +Byte conventions: `read_chunk`/`read_subset` return C-contiguous raw bytes in the +array's native byte order for the requested chunk / step-1 bounding box; +`write_chunk` takes the same. `read_metadata`/`list_children` return parsed JSON +documents as dicts. + +### Two read-addressing axes (no overlap) + +Reads are addressed in one of two coordinate spaces, and the two never overlap: + +- **Chunk-grid coordinates** — `read_chunk(coords)` / `read_encoded_chunk(coords)` + return a whole chunk addressed by its grid position. `read_chunk` (a backend + method) decodes and returns the *full* chunk shape, including the fill-padded + overhang of edge chunks; `read_encoded_chunk` (a facade helper, not a backend + method) returns the raw stored bytes or `None`. These pair with `write_chunk` / + `delete_chunk`, which are chunk-grid-addressed backend methods. +- **Array-element coordinates** — `read_subset(start, shape)` returns an + arbitrary box in array space, which generally spans multiple chunks and is + clipped to the array bounds. The facade's `read_region(selection)` normalizes + a numpy selection to a step-1 bounding box and calls it. + +`read_chunk` takes no `selection` parameter. A sub-region *within* a single +chunk is simply a `read_region` whose bounding box lies inside one chunk; the +backend already decodes only the overlapping chunk(s) (sharding-aware in the +zarrs backend), so a chunk-relative partial-read needs no separate API. The +`Store.get(key, byte_range=)` analogue is therefore `read_region` over a +single-chunk box, not a parameter on `read_chunk`; `read_subset` itself has no +single-`get` analogue — it is closer to "`get_partial_values` across many keys, +stitched into one array." + +### `read_encoded_chunk` is facade-level, not a backend method + +Reading a chunk's raw stored bytes is just `store.get(chunk_key)`, and the chunk +key is computable from the metadata document alone via zarr-python's +`chunk_key_encoding` — no decoding, no codec pipeline, nothing backend-specific. +Both backends would implement it identically. So the facade implements +`read_encoded_chunk(store, path, metadata, coords)` directly as: encode the chunk +key from the metadata, `store.get` it, return the bytes or `None`. It works the +same regardless of which backend (or none) is selected, which is correct since it +is pure store I/O. Under sharding the chunk key holds the whole shard blob, and +this returns exactly that raw object. + +This raw read can also be *expressed* through `read_chunk` by supplying a view +metadata document (`data_type: uint8` + a single `bytes` codec, identity decode) +— a nice demonstration that the read-only-view mechanism is general — but that +route requires knowing the encoded byte length up front to set the chunk shape (a +`store.getsize` round-trip) and would synthesize a fill-valued array for a missing +chunk instead of returning `None`. So `store.get` is the correct implementation +for fetching stored bytes; the view trick is the general tool for *reinterpreting* +decoded data under a different dtype/shape, which `read_chunk` already supports. + +## Method naming + +Both the public facade and the backend contract use a single, consistent verb +set: **create / read / write / delete / list**. No `decode`/`encode`/`retrieve`/ +`store`/`erase` synonyms. + +Public facade (`zarr.crud`): + +| Function | Verb | Notes | +|---|---|---| +| `create_new_group` / `create_overwrite_group` | create | node lifecycle | +| `create_new_array` / `create_overwrite_array` | create | node lifecycle | +| `read_metadata` | read | array or group document | +| `read_chunk` | read | decoded chunk → `ndarray` | +| `read_encoded_chunk` | read | raw stored bytes, no decode (facade-only, `store.get`) | +| `read_region` | read | numpy basic-indexing selection → `ndarray` | +| `write_chunk` | write | encode + store a chunk | +| `delete_chunk` | delete | remove one chunk | +| `delete_node` | delete | remove a node + descendants | +| `list_children` | list | direct children of a group | + +Facade → backend mapping for the byte-level methods: `read_chunk` → +`backend.read_chunk`, `read_region` → `backend.read_subset` (the facade +normalizes the selection to a step-1 bounding box `(start, shape)`), +`write_chunk` → `backend.write_chunk`, `delete_chunk` → `backend.delete_chunk`. +`read_encoded_chunk` maps to no backend method — the facade serves it from +`store.get`. The two distinct names `read_region` (selection-based, public) and +`read_subset` (bounding-box bytes, backend) are intentional: they have different +signatures and the facade is the adapter between them. + +## Facade / backend split + +What stays in the `zarr.crud` facade (written once, backend-neutral): + +- selection normalization (`_normalize_selection`), shape/dtype resolution + (`_array_shape`, `_chunk_dtype_and_shape`), native-dtype coercion; +- numpy assembly: `np.frombuffer(...).reshape(...)` and the strided/reversed/ + integer-axis post-index views; read-only result guarantee; +- the empty-selection short circuit (no backend call); +- `read_encoded_chunk`: encode the chunk key from the metadata and `store.get` + it (no backend involved); +- `ZarrsOptions` acceptance (still a placeholder) and backend resolution. + +What moves into each backend: + +- `ZarrsBackend`: `json.dumps`, the `/`-prefixed node-path form, + `resolve_store`, calling `_zarrs_bindings`, exception translation. +- `ReferenceBackend`: `ArrayV3Metadata.from_dict`/`ArrayV2Metadata.from_dict`, + building a `BatchedCodecPipeline` and `ChunkGrid`/`BasicIndexer`, assembling + `batch_info` and calling `codec_pipeline.read`/`write`, `save_metadata`, + `store.delete_dir`, and `list_dir` + per-child metadata reads. + +## Backend selection + +- A registry in `zarr.crud._registry`: `register_backend(name, backend)`, + `get_backend(name) -> CrudBackend`. +- A `zarr.config` key `crud.backend`, default `"reference"`. The pure-Python + backend always works and is predictable; `"zarrs"` opts into the accelerator + and is registered when `zarr.zarrs` is imported. +- Every facade function accepts `backend: CrudBackend | str | None = None`. + `None` → registry default; a string → registry lookup; an instance → used + directly. This enables side-by-side testing of backends. + +## Error handling + +`zarr.crud` defines the canonical exceptions: reuse +`zarr.errors.NodeNotFoundError`, and keep a `NodeExistsError` (exposed as +`zarr.crud.NodeExistsError`). Each backend raises these directly: + +- `ReferenceBackend` raises them at the point of detection. +- `ZarrsBackend` translates `_zarrs_bindings.NodeExistsError` / + `_zarrs_bindings.NodeNotFoundError` into the canonical types. + +The facade therefore no longer needs the `_translate_errors` shim. Phase-1 +fidelity limits (store-callback exceptions flattened to `RuntimeError` across the +Rust boundary) are unchanged for the zarrs backend; the reference backend +surfaces native exceptions directly. + +## Testing + +- Shared differential suite moves to `tests/crud/`, parametrized over + `backend ∈ {reference, zarrs}` × `store ∈ {memory, local}`. Each test writes + with zarr-python and reads through the facade (and vice versa), so the two + backends are checked against zarr-python *and*, transitively, against each + other. The zarrs-parametrized cases skip when `_zarrs_bindings` is not + installed (xdist-safe module-level `importorskip` in a zarrs-only conftest + helper, or a skip marker on the zarrs param). +- Zarrs-only tests stay in `tests/zarrs/`: the construction cache + (`test_cache.py`) and the store bridge (`test_bridge.py`). +- A focused `tests/crud/test_registry.py`: default resolution, `register_backend`, + string vs instance `backend=` override. +- `uv run --group zarrs pytest tests/crud tests/zarrs` is the full local check; + `uv run pytest tests/crud` (no zarrs group) must pass with the reference + backend alone and skip the zarrs params. + +## Migration notes + +- Move the 13 functions and the neutral helpers from `zarr.zarrs._api` into + `zarr.crud._api`; delete `zarr.zarrs._api`. No aliases in `zarr.zarrs`. +- Rename to the consistent verb set in the move (no compatibility aliases, since + the surface is unreleased): `decode_chunk` → `read_chunk`, `decode_region` → + `read_region`, `encode_chunk` → `write_chunk`, `erase_chunk` → `delete_chunk`. + `read_metadata`, `read_encoded_chunk`, `delete_node`, `list_children`, and the + `create_*` functions keep their names. +- `zarr.zarrs.__init__` exports only what is needed to register and identify the + zarrs backend (`ZarrsBackend`, and re-registers `"zarrs"` on import). +- The changelog fragment is updated to describe `zarr.crud` as the public CRUD + surface with pluggable backends, and `zarr.zarrs` as the zarrs backend. +- The CI job continues to build the crate and now runs `tests/crud tests/zarrs`. diff --git a/docs/superpowers/specs/2026-07-22-array-engine-protocol-design.md b/docs/superpowers/specs/2026-07-22-array-engine-protocol-design.md new file mode 100644 index 0000000000..f56c7e0859 --- /dev/null +++ b/docs/superpowers/specs/2026-07-22-array-engine-protocol-design.md @@ -0,0 +1,331 @@ +# Array engine protocol for zarr-python + +Date: 2026-07-22 +Status: approved +Branch: `claude/zarr-array-engine-protocol-239cfa` +Reviewed-by: @kylebarron (Zarrista developer), 2026-07-22 — feedback integrated + +## Goal + +Move the pluggable-backend boundary inside zarr-python's array classes: +`Array` wraps an object satisfying the `ArrayEngine` protocol, and +`AsyncArray` wraps an object satisfying the `AsyncArrayEngine` protocol (the +same contract with async methods). The engine owns the I/O data path — reading and writing +decoded data for selections — while `Array`/`AsyncArray` keep metadata +management, indexing normalization, and resize/append logic. + +Two engine families ship: + +- **Default engines** — the existing codec-pipeline machinery, moved behind + the protocol. Always available, every store, Zarr v2 and v3. Behavior and + performance are unchanged. +- **Zarrista engines** — powered by [Zarrista](https://pypi.org/project/zarrista/) + (Development Seed's zarrs-backed low-level Zarr API). We build against + Zarrista's `main` branch (unreleased; approved for use), pinned to a git + commit; the Zarrista team can publish a new beta release on request, at + which point the pin becomes `>=` that beta. + +This strategy **replaces** the `zarr.crud` layer and the homemade Rust +bindings from earlier on this branch. `zarr.crud`, `zarr.zarrs`, and the +`packages/zarrs-bindings` crate are deleted (see Deletions). + +Non-goals for this change: + +- Group-level engine operations (listing, group creation). The hierarchy + engine exists only to mint array engines in v1. +- Engine-owned metadata mutation (resize/append/update_attributes stay in + `Array`/`AsyncArray`). +- Silent fallback between engines (see Error handling). + +## Background: what Zarrista provides (main branch) + +Sync `Array` and async `AsyncArray` are separate native classes: + +- Construction: `open(store, path)` and `from_metadata(metadata, store, path)` + (constructs without writing; `store_metadata()` persists). Metadata is typed + as `zarr_metadata.ArrayMetadataV3` — our in-repo `packages/zarr-metadata` + package, so `ArrayV3Metadata.to_dict()` output is already the right currency. +- Reads: `retrieve_array_subset(selection)` (numpy basic indexing, step-1 + slices, ndim-preserving), `retrieve_chunk`, `retrieve_encoded_chunk`, and + sharding-aware `retrieve_subchunk` variants. +- Writes: `store_chunk` (decoded input as `ArrayBytes`; drops fill-value + chunks), `store_encoded_chunk`, `erase_chunk`, `compact_chunk`. **No + multi-chunk write** (`store_array_subset` is not exposed yet). +- Results: `DecodedArray = Tensor | VariableArray | MaskedTensor | + MaskedVariableArray`. `Tensor` is zero-copy to numpy (buffer protocol, + `__array__`, DLPack); `VariableArray` exports Arrow capsules. +- Stores: sync side takes Rust-native `FilesystemStore` / `MemoryStore`; + async side takes an obstore `ObjectStore` or an icechunk `Session`. There is + no bridge for arbitrary Python store objects. +- Zarr v3 only. + +## Architecture + +### Protocols (`zarr.abc.engine`) + +Four runtime-checkable protocols, named without a `Protocol` suffix +(`ArrayEngine(Protocol)` is enough; implementations get names like +`ZarristaEngine`). The async pair: + +```python +class Region(NamedTuple): + """A contiguous, step-1 box in array-element coordinates. + + One entry per dimension; `start` inclusive, `end_exclusive` exclusive, + already normalized (non-negative, clipped to the array shape). + """ + + start: tuple[int, ...] + end_exclusive: tuple[int, ...] + + +class AsyncArrayEngine(Protocol): + async def read_selection( + self, + selection: Region, + *, + prototype: BufferPrototype, + ) -> NDArrayLike: ... + + async def write_selection( + self, + selection: Region, + value: NDBuffer, + *, + prototype: BufferPrototype, + ) -> None: ... + + def with_metadata(self, metadata: ArrayMetadata) -> AsyncArrayEngine: ... + + +class AsyncHierarchyEngine(Protocol): + def array_engine( + self, path: str, metadata: ArrayMetadata, config: ArrayConfig | None = None + ) -> AsyncArrayEngine: ... +``` + +`ArrayEngine` and `HierarchyEngine` are identical with sync methods and sync +return types. + +Contract details: + +- An **array engine is bound** to `(store, path, metadata)` at construction. + Methods take only selection-level arguments, so engines can hold expensive + per-array state (a constructed `zarrista.Array`, codec pipelines, caches). +- A **hierarchy engine is bound to a store** and mints array engines that + share resources (the translated store handle, runtime, caches). This is the + factory the registry uses; users may also construct an array engine directly + and pass it to array creation/open. +- `read_selection` returns a buffer of the engine's own making: any object + implementing at least `__array__` (numpy coercion), zero-copy where the + backend allows. Additional export protocols (DLPack, Arrow) are welcome on + concrete engines but are not part of the contract. There is **no `out` + parameter** in v1; one may be added later as a performance improvement, and + the facade meanwhile serves user-supplied `out=` arguments by copying the + engine's result into them. +- The **engine speaks `Region`**: indexing across the engine boundary is + restricted to contiguous, step-1 boxes, interchanged as the `Region` + namedtuple above. Results and `value` inputs are ndim-preserving with shape + `end_exclusive - start` per dimension (matching Zarrista's ndim-preserving + `Selection` semantics). Engines never see raw user selections, steps, + fancy indices, or zarr-python `Indexer` objects. `write_selection` must + handle boxes that partially overlap chunks (read-modify-write is the + engine's concern; the default engine's codec pipeline already does this). +- `with_metadata` returns a rebound engine for the same store/path with new + metadata. `resize`/`append`/`update_attributes` use it, so rebinding works + uniformly for factory-made and user-provided engines. +- Facade-side responsibilities (not the engine's): normalizing every public + selection kind down to `Region` calls, `fields` handling, output + dtype/order resolution, scalar extraction and integer-axis squeezing, + copying into a user-supplied `out=` buffer, and the empty-selection short + circuit. Normalization rules: + - contiguous basic selections (slices with step 1, integers as length-1 + ranges, Ellipsis expansion) map to one `Region` directly; + - everything else — strided slices, orthogonal/coordinate/mask/block + selections — is served through its step-1 **bounding box**: reads fetch + the box and post-index the result; writes read the box, patch it with + numpy indexing, and write the box back (facade-level RMW). + - Consequence, accepted for now: sparse fancy selections over a large + extent transfer the whole bounding box through the engine. Revisit via a + richer interchange if it bites. + +### Wiring into `Array`/`AsyncArray` + +- `AsyncArray` gains an `engine: AsyncArrayEngine` attribute. The + module-level `_get_selection`/`_set_selection` shrink to: normalize the + selection to `Region` calls (see normalization rules above), call the + engine, post-process. The default async engine implements + `read_selection`/`write_selection` by building a `BasicIndexer` for the + region and running today's codec-pipeline machinery — for contiguous + selections the work done is identical to today. For strided and fancy + selections the facade's bounding-box normalization replaces today's + per-chunk gather/scatter for **all** engines; this is a deliberate + simplification (see the accepted consequence above). +- `Array` gains `engine: ArrayEngine`. Its data methods + (`__getitem__`/`__setitem__`, `get/set_basic_selection`, + `get/set_orthogonal_selection`, `get/set_mask_selection`, + `get/set_coordinate_selection`, `get/set_block_selection`, and the + `oindex`/`vindex`/`blocks` accessors) call the sync engine directly — **the + sync data path no longer touches the event loop**. Non-I/O operations + (metadata, resize, create) keep the existing sync-over-async route. +- `DefaultAsyncArrayEngine` wraps store_path + metadata + codec pipeline + + config. `DefaultArrayEngine` adapts it with `sync()` — so the default sync + path is exactly as fast as today, while engines with native sync + implementations (Zarrista on local files) skip asyncio entirely. + +### Engine selection + +- `engine=` parameter on array creation/open APIs, with three accepted forms: + the string `"default"`, the string `"zarrista"`, or an + `ArrayEngine`/`AsyncArrayEngine` instance the user constructed themselves + (used as-is, sync vs async matching the entry point): + + ```python + engine = ZarristaEngine(...) # user-constructed, full control + zarr.open_array(store, engine=engine) + zarr.open_array(store, engine="zarrista") # named, we construct it + ``` + +- **No config key in v1.** Omitting `engine=` means the default (pure-Python) + engine — today's behavior. A global config default can be added later. +- Name resolution is a small internal mapping from name to hierarchy-engine + factory taking the zarr store; the resulting hierarchy engine is cached per + `(name, store)` so array engines opened from one store share resources. + This replaces the `zarr.crud` registry. + +### Zarrista engines (`zarr.zarrista`) + +- Array engines `ZarristaEngine` (sync) / `ZarristaAsyncEngine`, minted by + `ZarristaHierarchyEngine` / `ZarristaAsyncHierarchyEngine`: the hierarchy + engine is constructed with a zarr-python store, translating it at + construction time: + - sync: `LocalStore → zarrista.FilesystemStore`. + - async: `zarr.storage.ObjectStore →` its underlying obstore instance; + an icechunk store/session → `icechunk.Session`. + - anything else (including zarr's `MemoryStore`, whose contents live in the + Python process): raise immediately with a message naming the store type + and the supported set. +- Array engines are built with + `zarrista.Array.from_metadata(metadata.to_dict(), store, path)` (async: + `AsyncArray.from_metadata`). Zarr v2 metadata raises immediately (Zarrista + is v3-only). +- **Reads**: a `Region` maps directly to a zarrista `Selection` (a tuple of + step-1 slices) and one `retrieve_array_subset` call — all-Rust, + sharding-aware, parallel. No selection-kind dispatch in the engine at all. +- **Writes**: `write_selection` decomposes the region over the chunk grid + internally: a chunk fully covered by the region encodes the value directly + via `store_chunk`; a partially covered chunk does read-modify-write — + `retrieve_chunk` → patch with numpy → `store_chunk`. Kyle confirmed + exposing `store_array_subset` (multi-chunk write) is easy on the Zarrista + side; when it lands, `write_selection` collapses to one Rust call with no + protocol change. +- **Buffers**: `Tensor` → zero-copy numpy (`to_numpy`); `VariableArray` → + numpy via Zarrista's upcoming numpy export (a copy for now — confirmed + planned); `MaskedTensor`/`MaskedVariableArray` unsupported (zarr-python has + no masked dtype) — raise. +- Selected by the engine resolver under the name `"zarrista"`, which lazily + imports `zarr.zarrista` only when that name is resolved; the import errors + cleanly when `zarrista` is not installed. +- Dependency: optional dependency group `zarrista`, git-pinned to a Zarrista + `main` commit until its next release, then `>=` that release. + +## Error handling + +Fail loud, no silent fallback: + +- Engine/hierarchy-engine construction raises `UnsupportedEngineError` (new, + in `zarr.errors`) when the engine cannot serve the store or metadata + (untranslatable store, v2 metadata, a dtype that decodes to a masked + layout). The message names the incompatibility and the supported set. +- Operations an engine cannot perform raise `NotImplementedError`. +- Engines raise zarr-python's canonical exceptions where they exist + (`zarr.errors` types); Zarrista exceptions are translated at the engine + boundary. + +## Testing + +- `tests/engine/` differential suite: the same operations executed through + the default and zarrista engines with results compared — reads/writes over + basic (with and without steps), orthogonal, coordinate, mask, and block + selections; partial-chunk RMW; fill-value handling; sharded arrays; + vlen dtypes. Sync engines on `LocalStore`; async engines on an + obstore local store. Zarrista params skip when `zarrista` is not installed + (xdist-safe importorskip). +- Per the testing convention: one test function per operation family covering + reasonable input combinations, plus one test function per error case + (untranslatable store, v2 metadata, masked dtype, unsupported operation). +- Protocol conformance is checked statically: mypy verifies both default and + zarrista engines satisfy the protocols (`runtime_checkable` isinstance + checks only verify method presence). +- The sync zarrista path gets a no-event-loop regression test (engine methods + run with no running loop and never create one). +- CI: drop the Rust-crate build job; add a job installing the pinned Zarrista + and running `tests/engine`. + +## Deletions + +- `src/zarr/crud/` (protocol, facade, reference backend, registry). +- `src/zarr/zarrs/` and `packages/zarrs-bindings/` (the homemade PyO3 crate, + its construction cache, store shim, and bridge). +- `tests/crud/`, `tests/zarrs/`, and the crate-build CI wiring. +- Changelog fragments describing `zarr.crud`/`zarr.zarrs` are replaced by one + describing the engine protocol and the zarrista engine. + +The differential-testing idea from `zarr.crud` survives as `tests/engine/`; +the store-translation and metadata-handoff learnings carry into +`zarr.zarrista`. + +## Feature requests for Zarrista (tracked, not blocking) + +Kyle has reviewed this list and is open to all of them; (1) is confirmed easy +and (a copying) numpy export for `VariableArray` is already planned. + +1. **`store_array_subset`** — zarrs has it natively; exposing it moves + partial-chunk read-modify-write into Rust and collapses the per-chunk + write path for basic selections. +2. **Python store bridge** — would lift the store-translation restriction and + let any zarr-python `Store` back a zarrista engine. +3. **Step > 1 slices in `Selection`** — widens the read fast path to strided + basic selections. +4. **Zarr v2 read support** — zarrs supports a v2 subset; exposing it would + let the zarrista engine serve v2 arrays. + +## Decision log + +- Engine protocols replace `zarr.crud`/`CrudBackend` entirely; homemade Rust + bindings dropped in favor of Zarrista (main branch approved). +- Engine scope is I/O only; metadata management stays in the array classes. +- Both acquisition paths: hierarchy engine as factory (resource sharing) and + user-provided array engines at creation time. +- Sync path is truly sync: `Array` calls its sync engine directly. +- Store bridging by translating known stores; fail loud otherwise. +- No fallback between engines; errors are immediate and specific. +- Region writes: protocol is selection-level; zarrista engine does + Python-side per-chunk RMW now, upgradeable when `store_array_subset` lands. +- Protocol granularity: engines speak `Indexer` (approach A), keeping the + default path unchanged and giving engines full selection information. + **Superseded** — see the post-review revision below: the interchange is now + the contiguous-box `Region`. + +Post-review (kylebarron, on the shared gist): + +- Protocol names drop the `Protocol` suffix: `ArrayEngine`, + `AsyncArrayEngine`, `HierarchyEngine`, `AsyncHierarchyEngine`. +- No `out` parameter in the v1 protocol; revisit later for performance. +- Engine read results only need `__array__`; other export protocols + (DLPack, Arrow) are optional extras on implementations. +- No config key in v1: engine choice is explicit via + `engine="default" | "zarrista" | `; a global config default may + come later. +- Zarrista can cut a beta release on request, so the git pin is temporary. +- Kyle's open question — chunk-level indexing as a simpler v1 engine + contract — initially answered with indexer-level (approach A). + +Post-review revision (supersedes approach A's interchange): + +- The engine interchange is restricted to contiguous, step-1 boxes: the + `Region` namedtuple (`start`, `end_exclusive`). Engines never see raw + selections or `Indexer` objects. All selection normalization (and + bounding-box + post-index / facade-RMW for strided and fancy selections) + moves to the facade. The bounding-box amplification for sparse fancy + selections is an accepted trade-off for now. 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..0209babcb8 --- /dev/null +++ b/examples/open_with_engine/README.md @@ -0,0 +1,52 @@ +# 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 + +This example demonstrates an **unreleased** feature and the `"zarrista"` engine +depends on the optional `zarrista` package (pulled in through the `zarrista` +dependency group). It therefore does **not** carry a PEP 723 inline-dependency +header like some other examples: such a header would install a zarr without this +feature and fail at runtime. + +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..f7bdfbb378 --- /dev/null +++ b/examples/open_with_engine/open_with_engine.py @@ -0,0 +1,129 @@ +# NOTE on dependencies / how to run +# ---------------------------------- +# Unlike some other examples in this directory, this one does NOT carry a PEP 723 +# `# /// script` inline-dependency header. +# +# The feature it demonstrates -- selecting a per-array execution "engine" -- is +# unreleased: it is not present in any published release. The `"zarrista"` engine +# additionally requires the `zarrista` package, which this repo pulls in through +# the `zarrista` dependency *group* (something a single inline `dependencies` +# pin cannot express cleanly). A PEP 723 header pinning a published zarr would +# therefore silently install a zarr WITHOUT this feature and fail at runtime, so +# we deliberately omit it. +# +# To run this example from a checkout of this branch: +# +# uv run python examples/open_with_engine/open_with_engine.py +# +# and, with the Rust-backed engine available: +# +# uv run --group zarrista python examples/open_with_engine/open_with_engine.py +# +# The zarrista portion is guarded: if the `zarrista` engine cannot be imported, +# the example still runs and demonstrates the default engine. + +""" +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 7a4bfa35ef..6387c23b59 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.ipynb - API Reference: - api/zarr/index.md @@ -48,10 +49,12 @@ nav: - api/zarr/registry.md - api/zarr/storage.md - api/zarr/experimental.md + - api/zarr/zarrista.md - ABC: - api/zarr/abc/index.md - api/zarr/abc/buffer.md - api/zarr/abc/codec.md + - api/zarr/abc/engine.md - api/zarr/abc/numcodec.md - api/zarr/abc/metadata.md - api/zarr/abc/store.md diff --git a/pyproject.toml b/pyproject.toml index 9f6005f981..4d930efe30 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -134,6 +134,15 @@ dev = [ "universal-pathlib", "mypy", ] +zarrista = [ + {include-group = "test"}, + "zarrista @ git+https://github.com/developmentseed/zarrista@95e47ad4c414c5920f0cf15550f923039641da8e", + # 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..f7ef6260c5 --- /dev/null +++ b/src/zarr/abc/engine.py @@ -0,0 +1,96 @@ +"""Array engine protocols. + +An *array engine* owns the data path of one open array: reading and writing +decoded data for contiguous regions. `Array` wraps an object satisfying +`ArrayEngine`; `AsyncArray` wraps an object satisfying `AsyncArrayEngine`. +A *hierarchy engine* is bound to a store and mints array engines that share +resources. See `docs/superpowers/specs/2026-07-22-array-engine-protocol-design.md`. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, NamedTuple, Protocol, runtime_checkable + +if TYPE_CHECKING: + from zarr.core.array_spec import ArrayConfig + from zarr.core.buffer import BufferPrototype, NDArrayLike, NDBuffer + from zarr.core.metadata import ArrayMetadata + +__all__ = [ + "ArrayEngine", + "AsyncArrayEngine", + "AsyncHierarchyEngine", + "HierarchyEngine", + "Region", +] + + +class Region(NamedTuple): + """A contiguous, step-1 box in array-element coordinates. + + One entry per dimension; `start` is inclusive, `end_exclusive` exclusive. + Callers pass normalized values: non-negative and clipped to the array + shape. This is the only selection type that crosses the engine boundary. + """ + + start: tuple[int, ...] + end_exclusive: tuple[int, ...] + + @property + def shape(self) -> tuple[int, ...]: + """The ndim-preserving shape of the box.""" + return tuple(e - s for s, e in zip(self.start, self.end_exclusive, strict=True)) + + +@runtime_checkable +class ArrayEngine(Protocol): + """The synchronous data path of one open array. + + Bound to `(store, path, metadata)` at construction. Methods must not + require a running event loop. Read results are ndim-preserving with + shape `selection.shape` and need only implement `__array__`. + + Note: `runtime_checkable` isinstance checks only verify method names; + mypy is the authoritative conformance check. + """ + + def read_selection(self, selection: Region, *, prototype: BufferPrototype) -> NDArrayLike: ... + + def write_selection( + self, selection: Region, value: NDBuffer, *, prototype: BufferPrototype + ) -> None: ... + + def with_metadata(self, metadata: ArrayMetadata) -> ArrayEngine: ... + + +@runtime_checkable +class AsyncArrayEngine(Protocol): + """The asynchronous data path of one open array. See `ArrayEngine`.""" + + async def read_selection( + self, selection: Region, *, prototype: BufferPrototype + ) -> NDArrayLike: ... + + async def write_selection( + self, selection: Region, value: NDBuffer, *, prototype: BufferPrototype + ) -> None: ... + + def with_metadata(self, metadata: ArrayMetadata) -> 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: ... + + +@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: ... diff --git a/src/zarr/api/asynchronous.py b/src/zarr/api/asynchronous.py index 7f185535df..d563879350 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 @@ -881,6 +883,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. @@ -1001,6 +1004,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 ------- @@ -1064,6 +1073,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, ) @@ -1204,6 +1217,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. @@ -1221,6 +1235,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]. @@ -1237,7 +1258,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) @@ -1246,6 +1274,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 8386427b3f..d39db0ed4e 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 @@ -634,6 +637,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. @@ -754,12 +758,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( @@ -790,9 +801,11 @@ def create( dimension_names=dimension_names, storage_options=storage_options, config=config, + engine=engine_for_async, **kwargs, ) - ) + ), + engine_spec=engine_for_array, ) @@ -818,6 +831,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. @@ -924,6 +938,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 ------- @@ -944,6 +964,7 @@ def create_array( # ``` """ + engine_for_async, engine_for_array = route_sync_engine_arg(engine) return Array( sync( zarr.core.array.create_array( @@ -967,8 +988,10 @@ def create_array( overwrite=overwrite, config=config, write_data=write_data, + engine=engine_for_async, ) - ) + ), + engine_spec=engine_for_array, ) @@ -1330,6 +1353,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. @@ -1347,6 +1371,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]. @@ -1356,6 +1387,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( @@ -1363,9 +1395,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 d15c70064b..7dfd9d65ad 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -3,13 +3,14 @@ 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 ( TYPE_CHECKING, Any, Literal, + NamedTuple, TypedDict, cast, overload, @@ -21,6 +22,7 @@ import zarr from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec, Codec +from zarr.abc.engine import Region from zarr.abc.numcodec import Numcodec, _is_numcodec from zarr.codecs._v2 import V2Codec from zarr.codecs.bytes import BytesCodec @@ -81,32 +83,49 @@ parse_dtype, ) from zarr.core.dtype.common import HasEndianness, HasItemSize, HasObjectCodec +from zarr.core.engine import ( + apply_post_index, + normalize_basic, + normalize_coordinate, + normalize_orthogonal, + resolve_async_engine, + resolve_sync_engine, + route_sync_engine_arg, + squeeze_axes, + strip_squeeze, +) from zarr.core.indexing import ( AsyncOIndex, AsyncVIndex, - BasicIndexer, BasicSelection, BlockIndex, BlockIndexer, - CoordinateIndexer, CoordinateSelection, Fields, - Indexer, - MaskIndexer, MaskSelection, OIndex, - OrthogonalIndexer, OrthogonalSelection, Selection, VIndex, _iter_grid, _iter_regions, + boundscheck_indices, check_fields, check_no_multi_fields, + ensure_tuple, + is_bool_array, + is_coordinate_selection, + is_integer, + is_integer_array, + is_mask_selection, is_pure_fancy_indexing, is_pure_orthogonal_indexing, is_scalar, + normalize_integer_selection, pop_fields, + replace_ellipsis, + replace_lists, + wraparound_indices, ) from zarr.core.metadata import ( ArrayMetadata, @@ -131,8 +150,8 @@ from zarr.core.sync import sync from zarr.errors import ( ArrayNotFoundError, - ChunkNotFoundError, MetadataValidationError, + NegativeStepError, ZarrDeprecationWarning, ZarrUserWarning, ) @@ -152,9 +171,11 @@ 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 ShardingCodecIndexLocation from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar + from zarr.core.engine import EngineName from zarr.storage import StoreLike from zarr.types import AnyArray, AnyAsyncArray, ArrayV2, ArrayV3, AsyncArrayV2, AsyncArrayV3 @@ -327,6 +348,16 @@ class AsyncArray[T_ArrayMetadata: (ArrayV2Metadata, ArrayV3Metadata)]: metadata: T_ArrayMetadata store_path: StorePath codec_pipeline: CodecPipeline = field(init=False) + # derived, per-array data path; excluded from equality/repr like other + # engine instances lack value semantics (two default engines are never `==`) + engine: AsyncArrayEngine = field(init=False, compare=False, repr=False) + # the ORIGINAL engine spec (name/instance/None) this array was built with, + # kept so `with_config` can re-resolve the same engine family against the new + # config instead of freezing the already-resolved engine (which would carry + # the old config); excluded from equality/repr for the same reason as `engine` + _engine_spec: AsyncArrayEngine | EngineName | None = field( + init=False, compare=False, repr=False + ) _chunk_grid: ChunkGrid = field(init=False) config: ArrayConfig @@ -336,6 +367,7 @@ def __init__( metadata: ArrayV2Metadata | ArrayV2MetadataDict, store_path: StorePath, config: ArrayConfigLike | None = None, + engine: AsyncArrayEngine | EngineName | None = None, ) -> None: ... @overload @@ -344,6 +376,7 @@ def __init__( metadata: ArrayV3Metadata | ArrayMetadataJSON_V3, store_path: StorePath, config: ArrayConfigLike | None = None, + engine: AsyncArrayEngine | EngineName | None = None, ) -> None: ... def __init__( @@ -351,6 +384,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) @@ -364,6 +398,18 @@ def __init__( "codec_pipeline", create_codec_pipeline(metadata=metadata_parsed, store=store_path.store), ) + object.__setattr__(self, "_engine_spec", engine) + object.__setattr__( + self, + "engine", + resolve_async_engine( + engine, + store=store_path.store, + path=store_path.path, + metadata=metadata_parsed, + config=config_parsed, + ), + ) @classmethod async def _create( @@ -396,6 +442,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`][]. @@ -452,6 +499,7 @@ async def _create( overwrite=overwrite, config=config_parsed, chunk_grid=chunk_grid, + engine=engine, ) elif zarr_format == 2: if codecs is not None: @@ -496,6 +544,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 @@ -575,6 +624,7 @@ async def _create_v3( dimension_names: DimensionNamesLike = None, attributes: dict[str, JSON] | None = None, overwrite: bool = False, + engine: AsyncArrayEngine | EngineName | None = None, ) -> AsyncArrayV3: if overwrite: if store_path.store.supports_deletes: @@ -602,7 +652,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 @@ -656,6 +706,7 @@ async def _create_v2( compressor: CompressorLike = "auto", attributes: dict[str, JSON] | None = None, overwrite: bool = False, + engine: AsyncArrayEngine | EngineName | None = None, ) -> AsyncArrayV2: if overwrite: if store_path.store.supports_deletes: @@ -691,7 +742,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 @@ -732,6 +783,7 @@ 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. @@ -744,6 +796,10 @@ async def open( for a description of all valid StoreLike values. zarr_format : ZarrFormat | None, optional The Zarr format version (default is 3). + engine : AsyncArrayEngine | Literal["default", "zarrista"] | None, optional + The data-path engine backing the opened array: a name (`"default"`, + `"zarrista"`) or a pre-built engine instance. When omitted, the + `"default"` engine is used. Returns ------- @@ -775,7 +831,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: @@ -1170,7 +1226,15 @@ 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) + # re-resolve the ORIGINAL engine spec against the new config (a named + # engine gets rebuilt with `new_config`; an instance is returned as-is). + # Passing `self.engine` here 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: """ @@ -1411,22 +1475,25 @@ def nbytes(self) -> int: async def _get_selection( self, - indexer: Indexer, + region: Region, + post_index: tuple[Any, ...], *, prototype: BufferPrototype, out: NDBuffer | None = None, fields: Fields | None = None, + scalarize: bool = False, ) -> NDArrayLikeOrScalar: + """Route a normalized `(region, post_index)` read through this array's engine.""" return await _get_selection( - self.store_path, + self.engine, self.metadata, - self.codec_pipeline, self.config, - self._chunk_grid, - indexer, + region, + post_index, prototype=prototype, out=out, fields=fields, + scalarize=scalarize, ) async def getitem( @@ -1471,15 +1538,10 @@ 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() + region, post = _basic_region_post(selection, self.metadata.shape) + return await self._get_selection(region, post, prototype=prototype, scalarize=True) async def get_orthogonal_selection( self, @@ -1489,17 +1551,10 @@ async def get_orthogonal_selection( fields: Fields | None = None, prototype: BufferPrototype | None = None, ) -> NDArrayLikeOrScalar: - return await _get_orthogonal_selection( - self.store_path, - self.metadata, - self.codec_pipeline, - self.config, - self._chunk_grid, - selection, - out=out, - fields=fields, - prototype=prototype, - ) + if prototype is None: + prototype = default_buffer_prototype() + region, post = _orthogonal_region_post(selection, self.metadata.shape) + return await self._get_selection(region, post, prototype=prototype, out=out, fields=fields) async def get_mask_selection( self, @@ -1509,17 +1564,10 @@ async def get_mask_selection( fields: Fields | None = None, prototype: BufferPrototype | None = None, ) -> NDArrayLikeOrScalar: - return await _get_mask_selection( - self.store_path, - self.metadata, - self.codec_pipeline, - self.config, - self._chunk_grid, - mask, - out=out, - fields=fields, - prototype=prototype, - ) + if prototype is None: + prototype = default_buffer_prototype() + region, post = _mask_region_post(mask, self.metadata.shape) + return await self._get_selection(region, post, prototype=prototype, out=out, fields=fields) async def get_coordinate_selection( self, @@ -1529,17 +1577,14 @@ async def get_coordinate_selection( fields: Fields | None = None, prototype: BufferPrototype | None = None, ) -> NDArrayLikeOrScalar: - return await _get_coordinate_selection( - self.store_path, - self.metadata, - self.codec_pipeline, - self.config, - self._chunk_grid, - selection, - out=out, - fields=fields, - prototype=prototype, - ) + if prototype is None: + prototype = default_buffer_prototype() + region, post, sel_shape = _coordinate_region_post(selection, self.metadata.shape) + # `out` is validated/filled against the (possibly multi-dimensional) + # selection shape, not the flattened pointwise read, so it is handled by + # `_finalize_coordinate_result` rather than the flat-shaped inner read. + out_array = await self._get_selection(region, post, prototype=prototype, fields=fields) + return _finalize_coordinate_result(out_array, sel_shape, out) async def _save_metadata(self, metadata: ArrayMetadata, ensure_parents: bool = False) -> None: """ @@ -1549,19 +1594,20 @@ async def _save_metadata(self, metadata: ArrayMetadata, ensure_parents: bool = F async def _set_selection( self, - indexer: Indexer, + region: Region, + post_index: tuple[Any, ...], value: npt.ArrayLike, *, prototype: BufferPrototype, fields: Fields | None = None, ) -> None: + """Route a normalized `(region, post_index)` write through this array's engine.""" return await _set_selection( - self.store_path, + self.engine, self.metadata, - self.codec_pipeline, self.config, - self._chunk_grid, - indexer, + region, + post_index, value, prototype=prototype, fields=fields, @@ -1606,16 +1652,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() + region, post = _basic_region_post(selection, self.metadata.shape) + return await self._set_selection(region, post, value, prototype=prototype) @property def oindex(self) -> AsyncOIndex[T_ArrayMetadata]: @@ -1801,6 +1841,18 @@ class Array[T_ArrayMetadata: (ArrayV2Metadata, ArrayV3Metadata)]: """ _async_array: AsyncArray[T_ArrayMetadata] + # `engine_spec` seeds `_engine`'s lazy resolution (see the `engine` property + # below); it is not itself a stored attribute (`InitVar`), so it cannot + # share the `engine` property's name without the class-body assignment of + # one clobbering the other in `Array.__dict__`. + engine_spec: InitVar[ArrayEngine | EngineName | None] = None + # derived, per-array sync data path; lazily resolved on first `.engine` + # access (see the `engine` property) so wrapping an `AsyncArray` never + # eagerly binds a second engine that might go unused. + _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 async_array(self) -> AsyncArray[T_ArrayMetadata]: @@ -1812,6 +1864,39 @@ def async_array(self) -> AsyncArray[T_ArrayMetadata]: """ return self._async_array + @property + def engine(self) -> ArrayEngine: + """The synchronous array engine serving this array's data path. + + Resolved lazily on first access and cached; every `get_*_selection`/ + `set_*_selection` method (and `__getitem__`/`__setitem__`) calls the + engine directly, without wrapping the call in a coroutine, so this + array's data path never touches the event loop in the caller's thread + -- the default engine still uses `sync()` internally, but that runs on + zarr's dedicated loop thread, not the caller's. + """ + 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, + ) + return self._engine + + def _rebind_engine(self) -> None: + """Rebind the cached sync engine to this array's current metadata. + + Called after `resize`/`append` mutate `self._async_array.metadata` in + place, so a subsequent read/write through `self.engine` sees the new + shape instead of a stale cached engine. A no-op if `.engine` was never + accessed (nothing cached to rebind). + """ + if self._engine is not None: + self._engine = self._engine.with_metadata(self._async_array.metadata) + @property def config(self) -> ArrayConfig: """ @@ -1860,10 +1945,12 @@ def _create( # runtime overwrite: bool = False, config: ArrayConfigLike | None = None, + engine: ArrayEngine | 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, @@ -1883,9 +1970,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( @@ -1922,6 +2010,7 @@ def from_dict( def open( cls, store: StoreLike, + engine: ArrayEngine | EngineName | None = None, ) -> Self: """Opens an existing Array from a store. @@ -1931,14 +2020,19 @@ def open( Store containing the Array. See the [storage documentation in the user guide][user-guide-store-like] for a description of all valid StoreLike values. + engine : ArrayEngine | Literal["default", "zarrista"] | None, optional + The data-path engine backing the opened array: a name (`"default"`, + `"zarrista"`) or a pre-built engine instance. When omitted, the + `"default"` behavior is unchanged. Returns ------- 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: @@ -2226,7 +2320,9 @@ def with_config(self, config: ArrayConfigLike) -> Self: ------- A new Array """ - return type(self)(self._async_array.with_config(config)) + # preserve this array's original engine spec so the new copy re-resolves + # the same engine family (the async array carries its own spec too) + return type(self)(self._async_array.with_config(config), engine_spec=self._engine_spec) @property def nbytes(self) -> int: @@ -2826,13 +2922,17 @@ 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, - ) + region, post = _basic_region_post(selection, self.shape) + return _get_selection_sync( + self.engine, + self.metadata, + self.config, + region, + post, + out=out, + fields=fields, + prototype=prototype, + scalarize=True, ) def set_basic_selection( @@ -2935,8 +3035,17 @@ 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)) + region, post = _basic_region_post(selection, self.shape) + _set_selection_sync( + self.engine, + self.metadata, + self.config, + region, + post, + value, + fields=fields, + prototype=prototype, + ) def get_orthogonal_selection( self, @@ -3063,11 +3172,16 @@ 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 - ) + region, post = _orthogonal_region_post(selection, self.shape) + return _get_selection_sync( + self.engine, + self.metadata, + self.config, + region, + post, + out=out, + fields=fields, + prototype=prototype, ) def set_orthogonal_selection( @@ -3181,9 +3295,16 @@ 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) + region, post = _orthogonal_region_post(selection, self.shape) + _set_selection_sync( + self.engine, + self.metadata, + self.config, + region, + post, + value, + fields=fields, + prototype=prototype, ) def get_mask_selection( @@ -3269,11 +3390,16 @@ 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 - ) + region, post = _mask_region_post(mask, self.shape) + return _get_selection_sync( + self.engine, + self.metadata, + self.config, + region, + post, + out=out, + fields=fields, + prototype=prototype, ) def set_mask_selection( @@ -3358,8 +3484,17 @@ 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)) + region, post = _mask_region_post(mask, self.shape) + _set_selection_sync( + self.engine, + self.metadata, + self.config, + region, + post, + value, + fields=fields, + prototype=prototype, + ) def get_coordinate_selection( self, @@ -3446,17 +3581,20 @@ 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 - ) + region, post, sel_shape = _coordinate_region_post(selection, self.shape) + # `out` is validated/filled against the (possibly multi-dimensional) + # selection shape, not the flattened pointwise read, so it is handled by + # `_finalize_coordinate_result` rather than the flat-shaped inner read. + out_array = _get_selection_sync( + self.engine, + self.metadata, + self.config, + region, + post, + fields=fields, + prototype=prototype, ) - - if hasattr(out_array, "shape"): - # restore shape - out_array = np.array(out_array).reshape(indexer.sel_shape) - return out_array + return _finalize_coordinate_result(out_array, sel_shape, out) def set_coordinate_selection( self, @@ -3537,8 +3675,9 @@ def set_coordinate_selection( """ if prototype is None: prototype = default_buffer_prototype() - # setup indexer - indexer = CoordinateIndexer(selection, self.shape, self._chunk_grid) + # normalize the coordinate selection to a contiguous box + pointwise index + region, post, sel_shape = _coordinate_region_post(selection, self.shape) + flat_shape = (int(product(sel_shape)),) # handle value - need ndarray-like flatten value if not is_scalar(value, self.dtype): @@ -3553,14 +3692,23 @@ def set_coordinate_selection( value = np.array(value).reshape(-1) if not is_scalar(value, self.dtype) and ( - isinstance(value, NDArrayLike) and indexer.shape != value.shape + isinstance(value, NDArrayLike) and flat_shape != value.shape ): raise ValueError( - f"Attempting to set a selection of {indexer.sel_shape[0]} " + f"Attempting to set a selection of {sel_shape[0]} " f"elements with an array of {value.shape[0]} elements." ) - sync(self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype)) + _set_selection_sync( + self.engine, + self.metadata, + self.config, + region, + post, + value, + fields=fields, + prototype=prototype, + ) def get_block_selection( self, @@ -3659,11 +3807,16 @@ 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 - ) + region = _block_region(selection, self.shape, self._chunk_grid) + return _get_selection_sync( + self.engine, + self.metadata, + self.config, + region, + (), + out=out, + fields=fields, + prototype=prototype, ) def set_block_selection( @@ -3759,8 +3912,17 @@ 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)) + region = _block_region(selection, self.shape, self._chunk_grid) + _set_selection_sync( + self.engine, + self.metadata, + self.config, + region, + (), + value, + fields=fields, + prototype=prototype, + ) @property def vindex(self) -> VIndex: @@ -3825,6 +3987,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`. @@ -3860,7 +4023,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)) + new_shape = sync(self.async_array.append(data, axis=axis)) + self._rebind_engine() + return new_shape def update_attributes(self, new_attributes: dict[str, JSON]) -> Self: """ @@ -3888,7 +4053,9 @@ 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) + # carry the original engine spec so the new wrapper re-resolves the same + # engine family rather than silently falling back to the default engine + return type(self)(new_array, engine_spec=self._engine_spec) def __repr__(self) -> str: return f"" @@ -4032,6 +4199,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. @@ -4256,6 +4424,7 @@ async def from_array( dimension_names=dimension_names, overwrite=overwrite, config=config_parsed, + engine=engine, ) if write_data: @@ -4305,6 +4474,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. @@ -4510,7 +4680,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 @@ -4537,6 +4707,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. @@ -4641,6 +4812,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 ------- @@ -4685,6 +4862,10 @@ async def create_array( storage_options=storage_options, overwrite=overwrite, config=config, + # `from_array` 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), ) else: mode: Literal["a"] = "a" @@ -4709,6 +4890,8 @@ async def create_array( dimension_names=dimension_names, overwrite=overwrite, config=config, + # see the comment above the analogous `from_array` call. + engine=cast("AsyncArrayEngine | EngineName | None", engine), ) @@ -5350,520 +5533,641 @@ def _get_chunk_spec( ) -async def _get_selection( - store_path: StorePath, - metadata: ArrayMetadata, - codec_pipeline: CodecPipeline, - config: ArrayConfig, - chunk_grid: ChunkGrid, - indexer: Indexer, - *, - prototype: BufferPrototype, - out: NDBuffer | None = None, - fields: Fields | None = None, -) -> NDArrayLikeOrScalar: - """ - Get a selection from an array. +def _native_dtype(metadata: ArrayMetadata) -> np.dtype[Any]: + """The array's native numpy dtype, for both Zarr v2 and v3 metadata.""" + zdtype = metadata.dtype if metadata.zarr_format == 2 else metadata.data_type + return zdtype.to_native_dtype() - Parameters - ---------- - store_path : StorePath - The store path of the array. - metadata : ArrayMetadata - The array metadata. - codec_pipeline : CodecPipeline - The codec pipeline for encoding/decoding. - config : ArrayConfig - The array configuration. - indexer : Indexer - The indexer specifying the selection. - prototype : BufferPrototype - A buffer prototype to use for the retrieved data. - out : NDBuffer | None, optional - An output buffer to write the data to. - fields : Fields | None, optional - Fields to select from structured arrays. - Returns - ------- - NDArrayLikeOrScalar - The selected data. +def _selection_result_shape( + region: Region, post_index: tuple[Any, ...], out_dtype: np.dtype[Any] +) -> tuple[int, ...]: + """Shape of the post-indexed result for an ndim-preserving box read.""" + return apply_post_index(np.empty(region.shape, dtype=out_dtype), post_index).shape + + +def _axis_covers_full_ordered(p: Any, n: int, axis: int, ndim: int) -> bool: + """Whether post entry `p` selects axis `axis` (length `n`) fully and in order. + + True for a full-range step-1 slice, and for an `np.ix_` outer-product array + for this axis -- shape `n` at `axis` and 1 elsewhere, values `arange(n)`. A + pointwise (coordinate) array has a different shape (`(N,)` on every axis), so + it is correctly rejected: `box[([0,1,2],[0,1,2])]` is the diagonal, not the + whole box. """ - # Get dtype from metadata - if metadata.zarr_format == 2: - zdtype = metadata.dtype - else: - zdtype = metadata.data_type - dtype = zdtype.to_native_dtype() + if isinstance(p, slice): + return p.start is None and p.stop is None and p.step in (None, 1) + if isinstance(p, np.ndarray): + expected_shape = tuple(n if j == axis else 1 for j in range(ndim)) + return p.shape == expected_shape and np.array_equal(p.reshape(-1), np.arange(n)) + return False + + +def _is_identity_read(post_index: tuple[Any, ...], box_shape: tuple[int, ...]) -> bool: + """Whether the post index leaves the ndim-preserving box read unchanged. + + True for an empty post (block selection), a per-axis tuple of full step-1 + slices (contiguous basic selections), and full-axis orthogonal outer slices + (`np.ix_(arange(n))`). In those cases the box the engine returned already *is* + the result, so it is returned without numpy coercion, preserving a custom + `NDArrayLike` type. Integer axes and the `_Squeeze` marker drop dimensions, so + they make the length differ or fail the per-axis check and are excluded. + """ + if post_index == (): + return True + if len(post_index) != len(box_shape): + return False + ndim = len(box_shape) + return all( + _axis_covers_full_ordered(p, n, i, ndim) + for i, (p, n) in enumerate(zip(post_index, box_shape, strict=True)) + ) - # Determine memory order - if metadata.zarr_format == 2: - order = metadata.order - else: - order = config.order - # check fields are sensible - out_dtype = check_fields(fields, dtype) +def _finalize_result( + result: NDArrayLike, out: NDBuffer | None, *, scalarize: bool +) -> NDArrayLikeOrScalar: + """Apply the shared post-read finalization: `out` copy, or scalar extraction.""" + if out is not None: + out.as_ndarray_like()[...] = result # type: ignore[index] + return out.as_ndarray_like() + if scalarize and result.shape == (): + # basic indexing collapses to a scalar (matches the historical + # `out_buffer.as_scalar()` return for all-integer basic selections). + # Any numpy result -- a 0-d `ndarray` or an already-collapsed numpy + # scalar (`np.generic`, e.g. `np.bytes_` from a vlen/fixed-bytes read) -- + # goes through `np.asarray(...)[()]`, bit-identical to the old behaviour; + # only a genuine device buffer (cupy/torch), which is neither, is indexed + # in its own namespace to avoid a host coercion. + if isinstance(result, (np.ndarray, np.generic)): + return cast("NDArrayLikeOrScalar", np.asarray(result)[()]) + device_result: Any = result + return cast("NDArrayLikeOrScalar", device_result[()]) + return result + - # setup output buffer +def _finalize_coordinate_result( + out_array: NDArrayLikeOrScalar, sel_shape: tuple[int, ...], out: NDBuffer | None +) -> NDArrayLikeOrScalar: + """Reshape a flat pointwise (coordinate/mask) read back to the selection + shape and, if an `out` buffer was supplied, validate and fill it. + + `_coordinate_region_post` flattens the (possibly multi-dimensional) + coordinate arrays before handing the engine a pointwise index, so the raw + read is 1-d. `out` is validated against the SELECTION shape -- which may be + multi-dimensional -- not that flattened shape, matching the pre-engine + behaviour where a 2-d coordinate array accepted a matching 2-d `out`. + """ + if hasattr(out_array, "shape"): + # restore the (possibly multi-dimensional) selection shape + out_array = cast("NDArrayLikeOrScalar", np.asarray(out_array).reshape(sel_shape)) if out is not None: - if isinstance(out, NDBuffer): - out_buffer = out - else: + if not isinstance(out, NDBuffer): raise TypeError(f"out argument needs to be an NDBuffer. Got {type(out)!r}") - if out_buffer.shape != indexer.shape: + if out.shape != sel_shape: raise ValueError( - f"shape of out argument doesn't match. Expected {indexer.shape}, got {out.shape}" - ) - else: - out_buffer = prototype.nd_buffer.empty( - shape=indexer.shape, - dtype=out_dtype, - order=order, - ) - if product(indexer.shape) > 0: - # need to use the order from the metadata for v2 - _config = config - if metadata.zarr_format == 2: - _config = replace(_config, order=order) - - # reading chunks and decoding them - indexed_chunks = list(indexer) - # For regular grids, all chunks share the same ArraySpec, so build it once - # and reuse it to avoid per-chunk ChunkGrid lookups and ArraySpec construction. - regular_grid = chunk_grid.is_regular - if regular_grid: - regular_chunk_spec = ArraySpec( - shape=chunk_grid.chunk_shape, - dtype=metadata.dtype, - fill_value=metadata.fill_value, - config=_config, - prototype=prototype, + f"shape of out argument doesn't match. Expected {sel_shape}, got {out.shape}" ) - results = await codec_pipeline.read( - [ - ( - store_path / metadata.encode_chunk_key(chunk_coords), - regular_chunk_spec - if regular_grid - else _get_chunk_spec(metadata, chunk_grid, chunk_coords, _config, prototype), - chunk_selection, - out_selection, - is_complete_chunk, - ) - for chunk_coords, chunk_selection, out_selection, is_complete_chunk in indexed_chunks - ], - out_buffer, - drop_axes=indexer.drop_axes, - ) - if _config.read_missing_chunks is False: - missing_info = [] - for i, result in enumerate(results): - if result["status"] == "missing": - coords = indexed_chunks[i][0] - key = metadata.encode_chunk_key(coords) - missing_info.append(f" chunk '{key}' (grid position {coords})") - if missing_info: - chunks_str = "\n".join(missing_info) - raise ChunkNotFoundError( - f"{len(missing_info)} chunk(s) not found in store '{store_path}'.\n" - f"Set the 'array.read_missing_chunks' config to True to fill " - f"missing chunks with the fill value.\n" - f"Missing chunks:\n{chunks_str}" - ) - if isinstance(indexer, BasicIndexer) and indexer.shape == (): - return out_buffer.as_scalar() - return out_buffer.as_ndarray_like() + out.as_ndarray_like()[...] = out_array # type: ignore[index] + return out.as_ndarray_like() + return out_array -async def _getitem( - store_path: StorePath, +def _get_selection_prepare( metadata: ArrayMetadata, - codec_pipeline: CodecPipeline, config: ArrayConfig, - chunk_grid: ChunkGrid, - selection: BasicSelection, + region: Region, + post_index: tuple[Any, ...], *, - prototype: BufferPrototype | None = None, -) -> NDArrayLikeOrScalar: + out: NDBuffer | None, + fields: Fields | None, +) -> tuple[np.dtype[Any], tuple[int, ...], MemoryOrder]: + """Shared preamble for `_get_selection`/`_get_selection_sync`: compute the + read's output dtype/shape/memory order and validate `out` against them. + + Split out so the async and sync selection helpers share this pure logic + and differ only in whether `engine.read_selection` is awaited. """ - Retrieve a subset of the array's data based on the provided selection. + dtype = _native_dtype(metadata) + out_dtype = check_fields(fields, dtype) + result_shape = _selection_result_shape(region, post_index, out_dtype) + order = metadata.order if metadata.zarr_format == 2 else config.order - Parameters - ---------- - store_path : StorePath - The store path of the array. - metadata : ArrayMetadata - The array metadata. - codec_pipeline : CodecPipeline - The codec pipeline for encoding/decoding. - config : ArrayConfig - The array configuration. - chunk_grid : ChunkGrid - The chunk grid. - selection : BasicSelection - A selection object specifying the subset of data to retrieve. - prototype : BufferPrototype, optional - A buffer prototype to use for the retrieved data (default is None). + if out is not None: + if not isinstance(out, NDBuffer): + raise TypeError(f"out argument needs to be an NDBuffer. Got {type(out)!r}") + if out.shape != result_shape: + raise ValueError( + f"shape of out argument doesn't match. Expected {result_shape}, got {out.shape}" + ) + return out_dtype, result_shape, order - Returns - ------- - NDArrayLikeOrScalar - The retrieved subset of the array's data. + +def _finish_get_selection( + raw: NDArrayLike, + region: Region, + post_index: tuple[Any, ...], + *, + fields: Fields | None, + order: MemoryOrder, + out: NDBuffer | None, + scalarize: bool, +) -> NDArrayLikeOrScalar: + """Shared postamble for `_get_selection`/`_get_selection_sync`: apply + `fields`/`post_index` to the engine's raw box read, then finalize into + `out`/scalar -- everything that happens after the (a)waited + `engine.read_selection` call. """ - if prototype is None: - prototype = default_buffer_prototype() - indexer = BasicIndexer( - selection, - shape=metadata.shape, - chunk_grid=chunk_grid, - ) - return await _get_selection( - store_path, metadata, codec_pipeline, config, chunk_grid, indexer, prototype=prototype - ) + if not fields and _is_identity_read(post_index, region.shape): + # a full-box, step-1 read is exactly the engine result; return it + # unchanged so a custom `NDArrayLike` type (e.g. GPU/torch buffers) + # survives instead of being coerced to numpy by `np.asarray`. + # zarr-python reads have always returned writable arrays; the non-identity + # paths below copy (via `astype(copy=True)`), but this fast path hands the + # engine buffer straight back. An engine may expose read-only memory (e.g. + # `zarrista` wraps Rust-owned bytes that `np.asarray` sees as + # non-writable), so copy a numpy-visible, non-writable result to restore + # the writability guarantee. Device buffers (cupy/torch) have no `.flags` + # and are left untouched, so this never forces them onto the host. + flags = getattr(raw, "flags", None) + if flags is not None and not flags.writeable: + raw = raw.copy() + return _finalize_result(raw, out, scalarize=scalarize) + # Stay in the engine result's own array namespace: field selection and + # `apply_post_index` are plain indexing (which cupy/torch implement), so a + # device buffer is never forced onto the host by an `np.asarray` here. + box = raw + if fields: + # non-empty `fields` selects structured sub-fields; an empty list/tuple + # (from `pop_fields` on a field-free selection) means "all fields" + box = box[fields] # type: ignore[index] + result = apply_post_index(box, post_index) + # Materialize the post-indexed (strided/fancy) view in the array's effective + # memory order without leaving the result's namespace -- `astype` is a method + # on the array itself, so a device buffer stays on-device (unlike + # `np.asarray(..., order=...)`, which would coerce it to numpy). + result = result.astype(result.dtype, order=order, copy=True) + return _finalize_result(result, out, scalarize=scalarize) -async def _get_orthogonal_selection( - store_path: StorePath, +async def _get_selection( + engine: AsyncArrayEngine, metadata: ArrayMetadata, - codec_pipeline: CodecPipeline, config: ArrayConfig, - chunk_grid: ChunkGrid, - selection: OrthogonalSelection, + region: Region, + post_index: tuple[Any, ...], *, + prototype: BufferPrototype, out: NDBuffer | None = None, fields: Fields | None = None, - prototype: BufferPrototype | None = None, + scalarize: bool = False, ) -> NDArrayLikeOrScalar: - """ - Get an orthogonal selection from the array. - - Parameters - ---------- - store_path : StorePath - The store path of the array. - metadata : ArrayMetadata - The array metadata. - codec_pipeline : CodecPipeline - The codec pipeline for encoding/decoding. - config : ArrayConfig - The array configuration. - chunk_grid : ChunkGrid - The chunk grid. - selection : OrthogonalSelection - The orthogonal selection specification. - out : NDBuffer | None, optional - An output buffer to write the data to. - fields : Fields | None, optional - Fields to select from structured arrays. - prototype : BufferPrototype | None, optional - A buffer prototype to use for the retrieved data. + """Read a contiguous `region` via the engine and re-index it to a selection. - Returns - ------- - NDArrayLikeOrScalar - The selected data. + The engine only speaks contiguous boxes; `normalize_*` produced `region` + (the box to transfer) and `post_index` (the numpy index that turns the + ndim-preserving box read into the requested selection). This applies + `fields`, then `post_index`, and finally optional scalar extraction / `out` + handling. """ - if prototype is None: - prototype = default_buffer_prototype() - indexer = OrthogonalIndexer(selection, metadata.shape, chunk_grid) - return await _get_selection( - store_path, - metadata, - codec_pipeline, - config, - chunk_grid, - indexer=indexer, - out=out, - fields=fields, - prototype=prototype, + out_dtype, result_shape, order = _get_selection_prepare( + metadata, config, region, post_index, out=out, fields=fields + ) + if product(region.shape) == 0: + # allocate through the prototype so an empty read honours the requested + # buffer type (e.g. a GPU prototype) instead of a host numpy array + empty = prototype.nd_buffer.empty(result_shape, out_dtype).as_ndarray_like() + return _finalize_result(empty, out, scalarize=scalarize) + + raw = await engine.read_selection(region, prototype=prototype) + return _finish_get_selection( + raw, region, post_index, fields=fields, order=order, out=out, scalarize=scalarize ) -async def _get_mask_selection( - store_path: StorePath, +def _get_selection_sync( + engine: ArrayEngine, metadata: ArrayMetadata, - codec_pipeline: CodecPipeline, config: ArrayConfig, - chunk_grid: ChunkGrid, - mask: MaskSelection, + region: Region, + post_index: tuple[Any, ...], *, + prototype: BufferPrototype, out: NDBuffer | None = None, fields: Fields | None = None, - prototype: BufferPrototype | None = None, + scalarize: bool = False, ) -> NDArrayLikeOrScalar: + """Synchronous mirror of `_get_selection`: same logic, calling + `engine.read_selection` directly instead of awaiting it, so `Array`'s data + path never runs a coroutine on the caller's thread. """ - Get a mask selection from the array. + out_dtype, result_shape, order = _get_selection_prepare( + metadata, config, region, post_index, out=out, fields=fields + ) + if product(region.shape) == 0: + # allocate through the prototype so an empty read honours the requested + # buffer type (e.g. a GPU prototype) instead of a host numpy array + empty = prototype.nd_buffer.empty(result_shape, out_dtype).as_ndarray_like() + return _finalize_result(empty, out, scalarize=scalarize) + + raw = engine.read_selection(region, prototype=prototype) + return _finish_get_selection( + raw, region, post_index, fields=fields, order=order, out=out, scalarize=scalarize + ) - Parameters - ---------- - store_path : StorePath - The store path of the array. - metadata : ArrayMetadata - The array metadata. - codec_pipeline : CodecPipeline - The codec pipeline for encoding/decoding. - config : ArrayConfig - The array configuration. - chunk_grid : ChunkGrid - The chunk grid. - mask : MaskSelection - The boolean mask specifying the selection. - out : NDBuffer | None, optional - An output buffer to write the data to. - fields : Fields | None, optional - Fields to select from structured arrays. - prototype : BufferPrototype | None, optional - A buffer prototype to use for the retrieved data. - Returns - ------- - NDArrayLikeOrScalar - The selected data. +def _coerce_write_value( + value: npt.ArrayLike, dtype: np.dtype[Any], prototype: BufferPrototype, fields: Fields | None +) -> NDArrayLike: + """Coerce a user-supplied write `value` to an ndarray-like of the array dtype. + + Ported from the historical `_set_selection`: scalars are materialized with + the prototype's buffer type (so GPU prototypes stay on-device), and + array-likes are cast to the array dtype. When `fields` is set the value + keeps its own (structured sub-field) dtype. + """ + # empty `fields` (list/tuple) means "no field selection"; treat it like None + target_dtype = None if fields else dtype + if np.isscalar(value): + array_like = prototype.buffer.create_zero_length().as_array_like() + # only NDArrayLike implements __array_function__ (e.g. numpy, cupy) + if isinstance(array_like, np._typing._SupportsArrayFunc): + array_like_ = cast("np._typing._SupportsArrayFunc", array_like) + value = np.asanyarray(value, dtype=target_dtype, like=array_like_) + else: + if not hasattr(value, "shape"): + value = np.asarray(value, dtype=target_dtype) + if target_dtype is not None and ( + not hasattr(value, "dtype") or value.dtype.name != dtype.name + ): + if hasattr(value, "astype"): + # Handle things that are already NDArrayLike more efficiently + value = value.astype(dtype=dtype, order="A") + else: + value = np.array(value, dtype=dtype, order="A") + return cast("NDArrayLike", value) + + +class _SetSelectionPrep(NamedTuple): + """Engine-agnostic state prepared for a selection write, shared by + `_set_selection`/`_set_selection_sync`. + + `identity_box`, when not `None`, is the fully-assembled box for a + full-box overwrite -- no read is needed, it is ready to write as-is. When + `None`, the write is a read-modify-write: `fields`/`stripped`/ + `assign_value` are the arguments `_patch_selection_box` needs to patch the + box the engine reads back. """ - if prototype is None: - prototype = default_buffer_prototype() - indexer = MaskIndexer(mask, metadata.shape, chunk_grid) - return await _get_selection( - store_path, - metadata, - codec_pipeline, - config, - chunk_grid, - indexer=indexer, - out=out, + + identity_box: NDArrayLike | None + fields: Fields | None + stripped: tuple[Any, ...] + assign_value: NDArrayLike + + +def _prepare_set_selection( + metadata: ArrayMetadata, + region: Region, + post_index: tuple[Any, ...], + value: npt.ArrayLike, + *, + prototype: BufferPrototype, + fields: Fields | None, +) -> _SetSelectionPrep | None: + """Shared preamble for `_set_selection`/`_set_selection_sync`: validate + `fields`, short-circuit an empty selection (returning `None`), coerce + `value`, and decide whether the write is a full-box overwrite or a + read-modify-write. + """ + dtype = _native_dtype(metadata) + check_fields(fields, dtype) + fields = check_no_multi_fields(fields) + if product(region.shape) == 0: + return None + + value_np = _coerce_write_value(value, dtype, prototype, fields) + + stripped = strip_squeeze(post_index) + sq_axes = squeeze_axes(post_index) + # a "full box" write addresses every element of the box exactly once, in + # order, so the value can be broadcast straight in without a read. This + # covers contiguous basic slices (post `slice(None, None, 1)`), full-axis + # orthogonal slices (post `np.ix_(arange(n))`), integer axes (length-1 + # boxes), and block selections (post `()`). + identity_post = not fields and _is_full_box_write(stripped, region.shape) + + # orthogonal selections drop integer axes; widen the value with np.newaxis at + # those axes so it aligns with the (un-squeezed) box (matching `oindex_set`) + assign_value = _widen_value_for_squeeze(value_np, sq_axes, len(region.shape)) + + identity_box = None + if identity_post: + # full-box write: broadcast the value into the box without a read. Basic + # integer axes drop a dimension from the value (their post entry is a + # plain int, not an orthogonal `_Squeeze` marker), so re-insert a size-1 + # axis at each dropped position before broadcasting -- otherwise a + # non-trailing dropped axis (e.g. `arr[:, 0] = v`) leaves the value's + # shape misaligned with the ndim-preserving box. + int_axes = tuple(i for i, p in enumerate(stripped) if isinstance(p, (int, np.integer))) + broadcast_value = _widen_value_for_squeeze(assign_value, int_axes, len(region.shape)) + # `np.broadcast_to` dispatches through `__array_function__`, so a device + # buffer stays in its own namespace; `astype` (a method on the array) + # then materializes a writable, dtype-correct box without an + # `np.asarray` host coercion. + identity_box = np.broadcast_to(broadcast_value, region.shape).astype(dtype, copy=True) + return _SetSelectionPrep( + identity_box=identity_box, fields=fields, - prototype=prototype, + stripped=stripped, + assign_value=assign_value, ) -async def _get_coordinate_selection( - store_path: StorePath, +def _patch_selection_box(raw: NDArrayLike, prep: _SetSelectionPrep) -> NDArrayLike: + """Patch a box read back from the engine for the read-modify-write path + (`prep.identity_box is None`), using the same `post_index` the read path + uses (dropping the orthogonal `_Squeeze` marker and widening the value at + dropped integer axes, matching `oindex_set`).""" + # `.copy()` yields a writable box in the engine result's own namespace (a + # device buffer is patched on-device); `np.asarray` would coerce it to host. + box = raw.copy() + # `target` indexing/assignment uses tuple/ellipsis keys the `NDArrayLike` + # protocol does not type (it only declares slice-key access), so it is + # `Any`-typed here; at runtime numpy/cupy/torch all support these keys. + target: Any = box[prep.fields] if prep.fields else box # type: ignore[index] + if prep.stripped == (): + target[...] = prep.assign_value + else: + target[prep.stripped] = prep.assign_value + return box + + +def _finish_set_selection(box: NDArrayLike, prototype: BufferPrototype) -> NDBuffer: + """Shared postamble for `_set_selection`/`_set_selection_sync`: make the + box C-contiguous and wrap it for `engine.write_selection`. + + `astype` (a method on the array) forces C-order in the box's own namespace + -- copying only when needed, exactly like `np.ascontiguousarray` -- so a + device buffer (cupy/torch) stays on-device instead of being coerced to host. + 0-d arrays pass through unchanged (`np.ascontiguousarray` would have promoted + them to shape `(1,)`).""" + contiguous_box = box if box.ndim == 0 else box.astype(box.dtype, order="C", copy=False) + return prototype.nd_buffer.from_ndarray_like(contiguous_box) + + +async def _set_selection( + engine: AsyncArrayEngine, metadata: ArrayMetadata, - codec_pipeline: CodecPipeline, config: ArrayConfig, - chunk_grid: ChunkGrid, - selection: CoordinateSelection, + region: Region, + post_index: tuple[Any, ...], + value: npt.ArrayLike, *, - out: NDBuffer | None = None, + prototype: BufferPrototype, fields: Fields | None = None, - prototype: BufferPrototype | None = None, -) -> NDArrayLikeOrScalar: - """ - Get a coordinate selection from the array. - - Parameters - ---------- - store_path : StorePath - The store path of the array. - metadata : ArrayMetadata - The array metadata. - codec_pipeline : CodecPipeline - The codec pipeline for encoding/decoding. - config : ArrayConfig - The array configuration. - chunk_grid : ChunkGrid - The chunk grid. - selection : CoordinateSelection - The coordinate selection specification. - out : NDBuffer | None, optional - An output buffer to write the data to. - fields : Fields | None, optional - Fields to select from structured arrays. - prototype : BufferPrototype | None, optional - A buffer prototype to use for the retrieved data. +) -> None: + """Write a selection through the engine as a contiguous-box read-modify-write. - Returns - ------- - NDArrayLikeOrScalar - The selected data. + An identity `post_index` (a full-box write with no `fields`) broadcasts the + value straight into the box and writes it. Any strided/fancy/orthogonal/ + fields write instead reads the box, patches it with numpy using the same + `post_index` the read path uses, then writes the whole box back. """ - if prototype is None: - prototype = default_buffer_prototype() - indexer = CoordinateIndexer(selection, metadata.shape, chunk_grid) - out_array = await _get_selection( - store_path, - metadata, - codec_pipeline, - config, - chunk_grid, - indexer=indexer, - out=out, - fields=fields, - prototype=prototype, + prep = _prepare_set_selection( + metadata, region, post_index, value, prototype=prototype, fields=fields ) + if prep is None: + return - if hasattr(out_array, "shape"): - # restore shape - out_array = cast("NDArrayLikeOrScalar", np.array(out_array).reshape(indexer.sel_shape)) - return out_array + if prep.identity_box is not None: + box = prep.identity_box + else: + # facade-level read-modify-write for strided/fancy/orthogonal/fields writes + raw = await engine.read_selection(region, prototype=prototype) + box = _patch_selection_box(raw, prep) + value_buffer = _finish_set_selection(box, prototype) + await engine.write_selection(region, value_buffer, prototype=prototype) -async def _set_selection( - store_path: StorePath, + +def _set_selection_sync( + engine: ArrayEngine, metadata: ArrayMetadata, - codec_pipeline: CodecPipeline, config: ArrayConfig, - chunk_grid: ChunkGrid, - indexer: Indexer, + region: Region, + post_index: tuple[Any, ...], value: npt.ArrayLike, *, prototype: BufferPrototype, fields: Fields | None = None, ) -> None: + """Synchronous mirror of `_set_selection`: same logic, calling + `engine.read_selection`/`engine.write_selection` directly instead of + awaiting them, so `Array`'s data path never runs a coroutine on the + caller's thread. """ - Set a selection in an array. + prep = _prepare_set_selection( + metadata, region, post_index, value, prototype=prototype, fields=fields + ) + if prep is None: + return - Parameters - ---------- - store_path : StorePath - The store path of the array. - metadata : ArrayMetadata - The array metadata. - codec_pipeline : CodecPipeline - The codec pipeline for encoding/decoding. - config : ArrayConfig - The array configuration. - chunk_grid : ChunkGrid - The chunk grid. - indexer : Indexer - The indexer specifying the selection. - value : npt.ArrayLike - The values to write. - prototype : BufferPrototype - A buffer prototype to use. - fields : Fields | None, optional - Fields to select from structured arrays. - """ - # Get dtype from metadata - if metadata.zarr_format == 2: - zdtype = metadata.dtype + if prep.identity_box is not None: + box = prep.identity_box else: - zdtype = metadata.data_type - dtype = zdtype.to_native_dtype() + raw = engine.read_selection(region, prototype=prototype) + box = _patch_selection_box(raw, prep) - # check fields are sensible - check_fields(fields, dtype) - fields = check_no_multi_fields(fields) + value_buffer = _finish_set_selection(box, prototype) + engine.write_selection(region, value_buffer, prototype=prototype) - # check value shape - if np.isscalar(value): - array_like = prototype.buffer.create_zero_length().as_array_like() - if isinstance(array_like, np._typing._SupportsArrayFunc): - # TODO: need to handle array types that don't support __array_function__ - # like PyTorch and JAX - array_like_ = cast("np._typing._SupportsArrayFunc", array_like) - value = np.asanyarray(value, dtype=dtype, like=array_like_) - else: - if not hasattr(value, "shape"): - value = np.asarray(value, dtype) - # assert ( - # value.shape == indexer.shape - # ), f"shape of value doesn't match indexer shape. Expected {indexer.shape}, got {value.shape}" - if not hasattr(value, "dtype") or value.dtype.name != dtype.name: - if hasattr(value, "astype"): - # Handle things that are already NDArrayLike more efficiently - value = value.astype(dtype=dtype, order="A") - else: - value = np.array(value, dtype=dtype, order="A") - value = cast("NDArrayLike", value) - # We accept any ndarray like object from the user and convert it - # to an NDBuffer (or subclass). From this point onwards, we only pass - # Buffer and NDBuffer between components. - value_buffer = prototype.nd_buffer.from_ndarray_like(value) +def _basic_region_post( + selection: BasicSelection, shape: tuple[int, ...] +) -> tuple[Region, tuple[slice | int, ...]]: + """Validate a basic selection with the historical rules, then normalize it. - # Determine memory order - if metadata.zarr_format == 2: - order = metadata.order - else: - order = config.order - - # need to use the order from the metadata for v2 - _config = config - if metadata.zarr_format == 2: - _config = replace(_config, order=order) - - # merging with existing data and encoding chunks - # For regular grids, all chunks share the same ArraySpec, so build it once - # and reuse it to avoid per-chunk ChunkGrid lookups and ArraySpec construction. - regular_grid = chunk_grid.is_regular - if regular_grid: - regular_chunk_spec = ArraySpec( - shape=chunk_grid.chunk_shape, - dtype=metadata.dtype, - fill_value=metadata.fill_value, - config=_config, - prototype=prototype, - ) - await codec_pipeline.write( - [ - ( - store_path / metadata.encode_chunk_key(chunk_coords), - regular_chunk_spec - if regular_grid - else _get_chunk_spec(metadata, chunk_grid, chunk_coords, _config, prototype), - chunk_selection, - out_selection, - is_complete_chunk, + `normalize_basic` (the engine helper) accepts negative-step slices and raises + `TypeError` for non-basic elements, but the public zarr basic-indexing API + has always rejected both with `IndexError` (via `BasicIndexer`). This + re-imposes those rules -- too-many-indices, negative steps, and unsupported + element types -- before delegating the box computation to `normalize_basic`. + """ + sel = replace_ellipsis(selection, shape) + for dim_sel, dim_len in zip(sel, shape, strict=True): + if is_integer(dim_sel): + normalize_integer_selection(dim_sel, dim_len) + elif isinstance(dim_sel, slice): + if dim_sel.step is not None and dim_sel.step < 1: + raise NegativeStepError("only slices with step >= 1 are supported.") + else: + raise IndexError( + "unsupported selection item for basic indexing; " + f"expected integer or slice, got {type(dim_sel)!r}" ) - for chunk_coords, chunk_selection, out_selection, is_complete_chunk in indexer - ], - value_buffer, - drop_axes=indexer.drop_axes, - ) + return normalize_basic(cast("BasicSelection", sel), shape) -async def _setitem( - store_path: StorePath, - metadata: ArrayMetadata, - codec_pipeline: CodecPipeline, - config: ArrayConfig, - chunk_grid: ChunkGrid, - selection: BasicSelection, - value: npt.ArrayLike, - prototype: BufferPrototype | None = None, -) -> None: +def _orthogonal_region_post( + selection: OrthogonalSelection, shape: tuple[int, ...] +) -> tuple[Region, tuple[Any, ...]]: + """Validate an orthogonal selection with the historical rules, then normalize it. + + Mirrors `OrthogonalIndexer`'s per-axis validation (and its exact error + messages) -- integer bounds, negative-step slices, 1-d integer/boolean + arrays of the right length, and unsupported element types -- before handing + the box computation to `normalize_orthogonal`. """ - Set values in the array using basic indexing. + sel = replace_lists(replace_ellipsis(selection, shape)) + for dim_sel, dim_len in zip(sel, shape, strict=True): + if is_integer(dim_sel): + normalize_integer_selection(dim_sel, dim_len) + elif isinstance(dim_sel, slice): + if dim_sel.step is not None and dim_sel.step < 1: + raise NegativeStepError("only slices with step >= 1 are supported.") + elif is_integer_array(dim_sel): + if not is_integer_array(dim_sel, 1): + raise IndexError( + "integer arrays in an orthogonal selection must be 1-dimensional only" + ) + checked = np.asanyarray(dim_sel).copy() + wraparound_indices(checked, dim_len) + boundscheck_indices(checked, dim_len) + elif is_bool_array(dim_sel): + if not is_bool_array(dim_sel, 1): + raise IndexError( + "Boolean arrays in an orthogonal selection must be 1-dimensional only" + ) + if dim_sel.shape[0] != dim_len: + raise IndexError( + f"Boolean array has the wrong length for dimension; expected {dim_len}, " + f"got {dim_sel.shape[0]}" + ) + else: + raise IndexError( + "unsupported selection item for orthogonal indexing; " + "expected integer, slice, integer array or Boolean " + f"array, got {type(dim_sel)!r}" + ) + return normalize_orthogonal(sel, shape) - Parameters - ---------- - store_path : StorePath - The store path of the array. - metadata : ArrayMetadata - The array metadata. - codec_pipeline : CodecPipeline - The codec pipeline for encoding/decoding. - config : ArrayConfig - The array configuration. - chunk_grid : ChunkGrid - The chunk grid. - selection : BasicSelection - The selection defining the region of the array to set. - value : npt.ArrayLike - The values to be written into the selected region of the array. - prototype : BufferPrototype or None, optional - A prototype buffer that defines the structure and properties of the array chunks being modified. - If None, the default buffer prototype is used. + +def _widen_value_for_squeeze( + value: NDArrayLike, squeeze_axes_: tuple[int, ...], ndim: int +) -> NDArrayLike: + """Insert `np.newaxis` at dropped integer axes so a value aligns with the box. + + Orthogonal selections drop integer axes from the result, so a non-scalar + value has fewer dimensions than the ndim-preserving box. Re-inserting a + size-1 axis at each dropped position lets the value broadcast/assign against + the box, matching `zarr.core.indexing.oindex_set`. Scalars and 0-d values are + returned unchanged (they broadcast on their own). """ - if prototype is None: - prototype = default_buffer_prototype() - indexer = BasicIndexer( - selection, - shape=metadata.shape, - chunk_grid=chunk_grid, - ) - return await _set_selection( - store_path, - metadata, - codec_pipeline, - config, - chunk_grid, - indexer, - value, - prototype=prototype, - ) + if squeeze_axes_ and not np.isscalar(value) and getattr(value, "ndim", 0) > 0: + value_selection: list[Any] = [slice(None)] * ndim + for ax in squeeze_axes_: + value_selection[ax] = np.newaxis + # index the value in its own array namespace so a device buffer + # (cupy/torch) is widened on-device instead of coerced to host + indexed: Any = value + return cast("NDArrayLike", indexed[tuple(value_selection)]) + return value + + +def _is_full_box_write(stripped: tuple[Any, ...], box_shape: tuple[int, ...]) -> bool: + """Whether a write's post index addresses every box element exactly once, in order. + + When true, the value fills the whole box and the write needs no + read-modify-write. Handles the post-index forms the `normalize_*` helpers + emit: an empty tuple (block selection), step-1 basic slices, single integer + axes (length-1 boxes), and `np.ix_` arrays equal to `arange(n)` (full-axis + orthogonal slices). The trailing `_Squeeze` marker must already be removed. + """ + if stripped == (): + return True + if len(stripped) != len(box_shape): + return False + ndim = len(box_shape) + for i, (p, n) in enumerate(zip(stripped, box_shape, strict=True)): + if isinstance(p, (int, np.integer)): + # an integer axis writes a length-1 box axis fully + if n != 1: + return False + elif not _axis_covers_full_ordered(p, n, i, ndim): + return False + return True + + +def _block_region( + selection: BasicSelection, shape: tuple[int, ...], chunk_grid: ChunkGrid +) -> Region: + """Map a block (chunk-grid) selection to its contiguous element-space box. + + Block selections are always contiguous (a slice of whole chunks spans a + contiguous element range), so the post index is empty. `BlockIndexer` is + reused purely for its coordinate math, preserving every existing behavior: + integer and step-1 slice block indices, irregular (rectilinear) grids, and + bounds errors. + """ + indexer = BlockIndexer(selection, shape, chunk_grid) + starts = tuple(di.start for di in indexer.dim_indexers) + # an empty block slice (e.g. `blocks[1:0]`) yields `start > stop` on that + # axis; clamp so the box is zero-length rather than negative-length. + ends = tuple(max(di.start, di.stop) for di in indexer.dim_indexers) + return Region(start=starts, end_exclusive=ends) + + +def _coordinate_region_post( + selection: CoordinateSelection, shape: tuple[int, ...] +) -> tuple[Region, tuple[np.ndarray[Any, Any], ...], tuple[int, ...]]: + """Normalize a coordinate (vindex) selection to `(region, post, sel_shape)`. + + Replicates `CoordinateIndexer`'s normalization and validation: integer axes + are widened to length-1 arrays, lists become arrays, the selection is checked + to be one integer array per dimension (raising the same `IndexError`), and + per-axis wraparound / bounds checks are applied (raising the same "dimension + with length" `IndexError`). The per-axis arrays are then broadcast to a common + shape (`sel_shape`) and flattened to the pointwise index the engine box read + is re-indexed with; the caller reshapes the flat result back to `sel_shape`. + """ + sel = ensure_tuple(selection) + sel = tuple(np.asarray([i]) if is_integer(i) else i for i in sel) + sel = replace_lists(sel) + if not is_coordinate_selection(sel, shape): + raise IndexError( + "invalid coordinate selection; expected one integer " + "(coordinate) array per dimension of the target array, " + f"got {selection!r}" + ) + for dim_sel, dim_len in zip(sel, shape, strict=True): + checked = np.asanyarray(dim_sel).copy() + wraparound_indices(checked, dim_len) + boundscheck_indices(checked, dim_len) + broadcast = np.broadcast_arrays(*(np.asarray(s) for s in sel)) + sel_shape = broadcast[0].shape or (1,) + flat = tuple(np.reshape(b, -1) for b in broadcast) + region, post = normalize_coordinate(flat, shape) + return region, post, sel_shape + + +def _mask_region_post( + mask: MaskSelection, shape: tuple[int, ...] +) -> tuple[Region, tuple[np.ndarray[Any, Any], ...]]: + """Normalize a boolean-mask (vindex) selection to `(region, post)`. + + Equivalent to coordinate indexing: the mask is converted to coordinate + arrays with `np.nonzero`, matching `MaskIndexer`. The result is 1-d in + row-major order, so no reshape is needed afterwards. + """ + sel = ensure_tuple(mask) + sel = replace_lists(sel) + if not is_mask_selection(sel, shape): + raise IndexError( + "invalid mask selection; expected one Boolean (mask)" + f"array with the same shape as the target array, got {sel!r}" + ) + coords = np.nonzero(sel[0]) + return normalize_coordinate(coords, shape) async def _resize( @@ -5913,9 +6217,10 @@ async def _delete_key(key: str) -> None: # Write new metadata await save_metadata(array.store_path, new_metadata) - # Update metadata and chunk_grid (in place) + # Update metadata, chunk_grid, and engine (in place) object.__setattr__(array, "metadata", new_metadata) object.__setattr__(array, "_chunk_grid", new_chunk_grid) + object.__setattr__(array, "engine", array.engine.with_metadata(new_metadata)) async def _append( @@ -5976,14 +6281,15 @@ async def _append( slice(None) if i != axis else slice(old_shape[i], new_shape[i]) for i in range(len(array.shape)) ) - await _setitem( - array.store_path, + region, post = _basic_region_post(append_selection, array.metadata.shape) + await _set_selection( + array.engine, array.metadata, - array.codec_pipeline, array.config, - array._chunk_grid, - append_selection, + region, + post, data, + prototype=default_buffer_prototype(), ) return new_shape diff --git a/src/zarr/core/engine/__init__.py b/src/zarr/core/engine/__init__.py new file mode 100644 index 0000000000..ade232b6c3 --- /dev/null +++ b/src/zarr/core/engine/__init__.py @@ -0,0 +1,41 @@ +from zarr.core.engine._default import ( + DefaultArrayEngine, + DefaultAsyncArrayEngine, + DefaultAsyncHierarchyEngine, + DefaultHierarchyEngine, +) +from zarr.core.engine._normalize import ( + apply_post_index, + normalize_basic, + normalize_coordinate, + normalize_orthogonal, + squeeze_axes, + strip_squeeze, +) +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", + "apply_post_index", + "classify_engine_arg", + "list_engines", + "normalize_basic", + "normalize_coordinate", + "normalize_orthogonal", + "resolve_async_engine", + "resolve_sync_engine", + "route_sync_engine_arg", + "squeeze_axes", + "strip_squeeze", +] diff --git a/src/zarr/core/engine/_default.py b/src/zarr/core/engine/_default.py new file mode 100644 index 0000000000..a3078c3a5c --- /dev/null +++ b/src/zarr/core/engine/_default.py @@ -0,0 +1,205 @@ +"""The default engines: today's codec-pipeline machinery behind the protocol.""" + +from __future__ import annotations + +from dataclasses import replace +from typing import TYPE_CHECKING + +from zarr.core.array_spec import ArraySpec, parse_array_config +from zarr.core.chunk_grids import ChunkGrid +from zarr.core.common import product +from zarr.core.indexing import BasicIndexer +from zarr.core.sync import sync +from zarr.errors import ChunkNotFoundError +from zarr.storage._common import StorePath + +if TYPE_CHECKING: + from zarr.abc.engine import Region + from zarr.abc.store import Store + from zarr.core.array_spec import ArrayConfig + from zarr.core.buffer import BufferPrototype, NDArrayLike, NDBuffer + from zarr.core.metadata import ArrayMetadata + +__all__ = [ + "DefaultArrayEngine", + "DefaultAsyncArrayEngine", + "DefaultAsyncHierarchyEngine", + "DefaultHierarchyEngine", +] + + +def _region_to_slices(region: Region) -> tuple[slice, ...]: + return tuple(slice(s, e) for s, e in zip(region.start, region.end_exclusive, strict=True)) + + +class DefaultAsyncArrayEngine: + """Codec-pipeline-backed engine. Any store, Zarr v2 and v3.""" + + def __init__(self, store_path: StorePath, metadata: ArrayMetadata, config: ArrayConfig) -> None: + # Imported lazily to avoid an import cycle: `zarr.core.array` imports the + # engine package (for engine resolution) while the default engine reuses + # `zarr.core.array`'s codec-pipeline machinery. + from zarr.core.array import create_codec_pipeline + + self.store_path = store_path + self.metadata = metadata + self.config = config + self._chunk_grid = ChunkGrid.from_metadata(metadata) + self._codec_pipeline = create_codec_pipeline(metadata=metadata, store=store_path.store) + + def with_metadata(self, metadata: ArrayMetadata) -> DefaultAsyncArrayEngine: + return DefaultAsyncArrayEngine( + store_path=self.store_path, metadata=metadata, config=self.config + ) + + def _indexer(self, region: Region) -> BasicIndexer: + return BasicIndexer( + _region_to_slices(region), + shape=self.metadata.shape, + chunk_grid=self._chunk_grid, + ) + + def _regular_chunk_spec( + self, config: ArrayConfig, prototype: BufferPrototype + ) -> ArraySpec | None: + """Same optimization as `_get_selection`/`_set_selection`: build one shared + `ArraySpec` for regular chunk grids instead of a per-chunk lookup. + """ + if not self._chunk_grid.is_regular: + return None + return ArraySpec( + shape=self._chunk_grid.chunk_shape, + dtype=self.metadata.dtype, + fill_value=self.metadata.fill_value, + config=config, + prototype=prototype, + ) + + def _v2_order_config(self) -> ArrayConfig: + # need to use the order from the metadata for v2 (mirrors _get_selection/_set_selection) + if self.metadata.zarr_format == 2: + return replace(self.config, order=self.metadata.order) + return self.config + + async def read_selection(self, selection: Region, *, prototype: BufferPrototype) -> NDArrayLike: + from zarr.core.array import _get_chunk_spec + + indexer = self._indexer(selection) + if self.metadata.zarr_format == 2: + dtype = self.metadata.dtype.to_native_dtype() + order = self.metadata.order + else: + dtype = self.metadata.data_type.to_native_dtype() + order = self.config.order + out_buffer = prototype.nd_buffer.empty(shape=indexer.shape, dtype=dtype, order=order) + if product(indexer.shape) > 0: + _config = self._v2_order_config() + regular_chunk_spec = self._regular_chunk_spec(_config, prototype) + indexed_chunks = list(indexer) + results = await self._codec_pipeline.read( + [ + ( + self.store_path / self.metadata.encode_chunk_key(chunk_coords), + regular_chunk_spec + if regular_chunk_spec is not None + else _get_chunk_spec( + self.metadata, self._chunk_grid, chunk_coords, _config, prototype + ), + chunk_selection, + out_selection, + is_complete_chunk, + ) + for chunk_coords, chunk_selection, out_selection, is_complete_chunk in indexed_chunks + ], + out_buffer, + drop_axes=indexer.drop_axes, + ) + if _config.read_missing_chunks is False: + missing_info = [] + for i, result in enumerate(results): + if result["status"] == "missing": + coords = indexed_chunks[i][0] + key = self.metadata.encode_chunk_key(coords) + missing_info.append(f" chunk '{key}' (grid position {coords})") + if missing_info: + chunks_str = "\n".join(missing_info) + raise ChunkNotFoundError( + f"{len(missing_info)} chunk(s) not found in store '{self.store_path}'.\n" + f"Set the 'array.read_missing_chunks' config to True to fill " + f"missing chunks with the fill value.\n" + f"Missing chunks:\n{chunks_str}" + ) + return out_buffer.as_ndarray_like() + + async def write_selection( + self, selection: Region, value: NDBuffer, *, prototype: BufferPrototype + ) -> None: + from zarr.core.array import _get_chunk_spec + + indexer = self._indexer(selection) + _config = self._v2_order_config() + regular_chunk_spec = self._regular_chunk_spec(_config, prototype) + await self._codec_pipeline.write( + [ + ( + self.store_path / self.metadata.encode_chunk_key(chunk_coords), + regular_chunk_spec + if regular_chunk_spec is not None + else _get_chunk_spec( + self.metadata, self._chunk_grid, chunk_coords, _config, prototype + ), + chunk_selection, + out_selection, + is_complete_chunk, + ) + for chunk_coords, chunk_selection, out_selection, is_complete_chunk in indexer + ], + value, + drop_axes=indexer.drop_axes, + ) + + +class DefaultArrayEngine: + """Sync adapter over `DefaultAsyncArrayEngine` via `sync()`.""" + + def __init__(self, async_engine: DefaultAsyncArrayEngine) -> None: + self._async = async_engine + + def read_selection(self, selection: Region, *, prototype: BufferPrototype) -> NDArrayLike: + return sync(self._async.read_selection(selection, prototype=prototype)) + + def write_selection( + self, selection: Region, value: NDBuffer, *, prototype: BufferPrototype + ) -> None: + sync(self._async.write_selection(selection, value, prototype=prototype)) + + def with_metadata(self, metadata: ArrayMetadata) -> DefaultArrayEngine: + return DefaultArrayEngine(self._async.with_metadata(metadata)) + + +class DefaultAsyncHierarchyEngine: + """Store-bound factory for default 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 default 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/_normalize.py b/src/zarr/core/engine/_normalize.py new file mode 100644 index 0000000000..a6c11d5ea0 --- /dev/null +++ b/src/zarr/core/engine/_normalize.py @@ -0,0 +1,213 @@ +"""Reduce public selection kinds to `(Region, post_index)` pairs. + +The engine boundary only speaks contiguous step-1 boxes (`Region`). Each +helper returns the box to transfer plus the numpy index that, applied to the +ndim-preserving box result, yields exactly `array[original_selection]`. +""" + +from __future__ import annotations + +import operator +import types +from typing import TYPE_CHECKING, Any, cast + +import numpy as np + +from zarr.abc.engine import Region + +if TYPE_CHECKING: + from zarr.core.buffer import NDArrayLike + from zarr.core.indexing import BasicSelection + + +def _expand(selection: Any, ndim: int) -> tuple[Any, ...]: + """Expand Ellipsis and pad missing trailing axes with full slices.""" + sel_tuple = selection if isinstance(selection, tuple) else (selection,) + n_ellipsis = sum(1 for s in sel_tuple if s is Ellipsis) + if n_ellipsis > 1: + raise IndexError("an index can only have a single ellipsis ('...')") + if n_ellipsis == 1: + i = sel_tuple.index(Ellipsis) + n_fill = ndim - (len(sel_tuple) - 1) + if n_fill < 0: + raise IndexError(f"too many indices for array: array is {ndim}-dimensional") + sel_tuple = sel_tuple[:i] + (slice(None),) * n_fill + sel_tuple[i + 1 :] + if len(sel_tuple) > ndim: + raise IndexError(f"too many indices for array: array is {ndim}-dimensional") + return sel_tuple + (slice(None),) * (ndim - len(sel_tuple)) + + +def _normalize_int(sel: Any, size: int, dim: int) -> int: + idx = operator.index(sel) + if idx < 0: + idx += size + if not 0 <= idx < size: + raise IndexError(f"index {sel} is out of bounds for axis {dim} with size {size}") + return idx + + +def normalize_basic( + selection: BasicSelection, shape: tuple[int, ...] +) -> tuple[Region, tuple[slice | int, ...]]: + """Normalize a numpy basic-indexing selection to a step-1 box. + + Supports integers, slices (any step), and `Ellipsis`. Integer axes get a + length-1 range in the box and `0` in the post index (dropping the axis, + matching numpy). Fancy elements raise `TypeError`. + """ + sel_tuple = _expand(selection, len(shape)) + starts: list[int] = [] + ends: list[int] = [] + post: list[slice | int] = [] + for dim, (sel, size) in enumerate(zip(sel_tuple, shape, strict=True)): + if isinstance(sel, slice): + start, stop, step = sel.indices(size) + n = len(range(start, stop, step)) + if n == 0: + starts.append(0) + ends.append(0) + post.append(slice(None)) + elif step > 0: + starts.append(start) + ends.append(start + (n - 1) * step + 1) + post.append(slice(None, None, step)) + else: + last = start + (n - 1) * step + starts.append(last) + ends.append(start + 1) + post.append(slice(None, None, step)) + elif isinstance(sel, types.EllipsisType): + raise AssertionError("Ellipsis expanded by _expand") # noqa: TRY004 + else: + try: + idx = _normalize_int(sel, size, dim) + except TypeError: + raise TypeError( + f"unsupported selection element {sel!r}: only integers, " + "slices, and Ellipsis are supported by basic indexing" + ) from None + starts.append(idx) + ends.append(idx + 1) + post.append(0) + return Region(start=tuple(starts), end_exclusive=tuple(ends)), tuple(post) + + +class _Squeeze: + """Marker appended to an orthogonal post index: squeeze these axes.""" + + def __init__(self, axes: tuple[int, ...]) -> None: + self.axes = axes + + +def normalize_orthogonal(selection: Any, shape: tuple[int, ...]) -> tuple[Region, tuple[Any, ...]]: + """Normalize an orthogonal (`oindex`) selection to a box + outer index. + + Each axis selector may be an integer, a slice, an integer array, or a 1-d + boolean mask for that axis. The post index is the `np.ix_`-broadcastable + tuple of per-axis integer arrays (relative to the box origin); if any axis + was an integer, a trailing `_Squeeze` marker records which axes to drop + afterwards. Use `apply_post_index` to apply the result correctly. + """ + sel_tuple = _expand(selection, len(shape)) + starts: list[int] = [] + ends: list[int] = [] + axis_indices: list[np.ndarray] = [] + for dim, (sel, size) in enumerate(zip(sel_tuple, shape, strict=True)): + if isinstance(sel, slice): + idxs = np.arange(*sel.indices(size)) + elif isinstance(sel, (np.ndarray, list)): + # Convert to array first to check dtype (handles list of bools) + sel_array = np.asarray(sel) + if sel_array.dtype == bool: + # Boolean mask path + if sel_array.ndim != 1 or sel_array.shape[0] != size: + raise IndexError(f"boolean index for axis {dim} must be 1-d with length {size}") + idxs = np.nonzero(sel_array)[0] + else: + # Integer index path + idxs = np.asarray(sel, dtype=np.intp) + if idxs.ndim != 1: + raise IndexError(f"orthogonal index for axis {dim} must be 1-d") + idxs = np.where(idxs < 0, idxs + size, idxs) + if idxs.size and (idxs.min() < 0 or idxs.max() >= size): + raise IndexError(f"index out of bounds for axis {dim} with size {size}") + else: + idxs = np.array([_normalize_int(sel, size, dim)], dtype=np.intp) + if idxs.size == 0: + starts.append(0) + ends.append(0) + axis_indices.append(idxs) + else: + lo, hi = int(idxs.min()), int(idxs.max()) + starts.append(lo) + ends.append(hi + 1) + axis_indices.append(idxs - lo) + region = Region(start=tuple(starts), end_exclusive=tuple(ends)) + post = np.ix_(*axis_indices) if axis_indices else () + # integer axes were widened to length-1 arrays; the caller squeezes them + squeeze_axes = tuple( + i for i, sel in enumerate(sel_tuple) if not isinstance(sel, (slice, np.ndarray, list)) + ) + if squeeze_axes: + return region, (*post, _Squeeze(squeeze_axes)) + return region, post + + +def apply_post_index(box: NDArrayLike, post: tuple[Any, ...]) -> NDArrayLike: + """Apply a post index produced by a `normalize_*` helper to a box read. + + Indexing stays in `box`'s own array namespace (numpy, cupy, torch, ...) so a + device buffer is never forced onto the host; the local is `Any`-typed because + the `NDArrayLike` protocol only types slice-key indexing, not the tuple keys + the `normalize_*` helpers emit. + """ + indexed: Any = box + if post and isinstance(post[-1], _Squeeze): + result = indexed[post[:-1]] if len(post) > 1 else indexed + return cast("NDArrayLike", np.squeeze(result, axis=post[-1].axes)) + if post == (): + return box + return cast("NDArrayLike", indexed[post]) + + +def strip_squeeze(post: tuple[Any, ...]) -> tuple[Any, ...]: + """Return `post` without a trailing `_Squeeze` marker, if present. + + Writes apply the post index directly to a target array (not through + `apply_post_index`), so callers that only need the "real" numpy index + strip the `_Squeeze` marker first. Returns `post` unchanged when there is + no marker. + """ + if post and isinstance(post[-1], _Squeeze): + return post[:-1] + return post + + +def squeeze_axes(post: tuple[Any, ...]) -> tuple[int, ...]: + """Return the axes recorded by a trailing `_Squeeze` marker, else `()`. + + Orthogonal *writes* assign into the unsqueezed box using + `strip_squeeze(post)`; when the selection dropped integer axes, the value + must be widened with `np.newaxis` at exactly these axes (matching + `oindex_set`). Returns `()` when `post` carries no marker. + """ + if post and isinstance(post[-1], _Squeeze): + return post[-1].axes + return () + + +def normalize_coordinate( + selection: tuple[Any, ...], shape: tuple[int, ...] +) -> tuple[Region, tuple[np.ndarray, ...]]: + """Normalize a coordinate (`vindex`) selection to a box + pointwise index.""" + coords = tuple(np.asarray(c, dtype=np.intp) for c in selection) + if len(coords) != len(shape): + raise IndexError(f"coordinate selection needs {len(shape)} axis arrays, got {len(coords)}") + coords = tuple(np.where(c < 0, c + size, c) for c, size in zip(coords, shape, strict=True)) + for dim, (c, size) in enumerate(zip(coords, shape, strict=True)): + if c.size and (c.min() < 0 or c.max() >= size): + raise IndexError(f"index out of bounds for axis {dim} with size {size}") + starts = tuple(int(c.min()) if c.size else 0 for c in coords) + ends = tuple(int(c.max()) + 1 if c.size else 0 for c in coords) + post = tuple(c - s for c, s in zip(coords, starts, strict=True)) + return Region(start=starts, end_exclusive=ends), post 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 52eaa3e144..c75a86f2a4 100644 --- a/src/zarr/core/group.py +++ b/src/zarr/core/group.py @@ -47,6 +47,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") @@ -1029,6 +1032,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. @@ -1122,6 +1126,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 ------- @@ -1152,6 +1162,7 @@ async def create_array( overwrite=overwrite, config=config, write_data=write_data, + engine=engine, ) async def require_array( @@ -2543,6 +2554,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. @@ -2638,6 +2650,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 ------- @@ -2646,6 +2664,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( @@ -2667,8 +2686,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..76f7807b88 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,11 @@ 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, or metadata (e.g. Zarr v2) + the engine does not support. Raised at engine construction time. + """ 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..4e25549f31 --- /dev/null +++ b/src/zarr/zarrista/_engine.py @@ -0,0 +1,282 @@ +"""Zarrista-backed array engines.""" + +from __future__ import annotations + +import itertools +from typing import TYPE_CHECKING, Any, cast + +import numpy as np + +from zarr.errors import UnsupportedEngineError +from zarr.zarrista._translate import translate_store_async, translate_store_sync + +if TYPE_CHECKING: + from zarr_metadata import ArrayMetadataV3 + + from zarr.abc.engine import Region + from zarr.abc.store import Store + from zarr.core.array_spec import ArrayConfig + from zarr.core.buffer import BufferPrototype, NDBuffer + from zarr.core.metadata import ArrayMetadata + +__all__ = [ + "ZarristaAsyncEngine", + "ZarristaAsyncHierarchyEngine", + "ZarristaEngine", + "ZarristaHierarchyEngine", +] + + +def _require_v3(metadata: ArrayMetadata) -> ArrayMetadataV3: + 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 `ArrayMetadataV3` TypedDict, which the + # dict's runtime shape matches by construction. + return cast("ArrayMetadataV3", metadata.to_dict()) + + +def _reject_unenforceable_config(config: ArrayConfig | None) -> None: + """Reject an `ArrayConfig` the zarrista engine cannot honour. + + The zarrista engine owns its own codec options and does not consult + zarr-python's `ArrayConfig`. Most fields (e.g. `order`) only affect the + in-memory layout of the returned array, which the facade normalizes, so + ignoring them is safe. `read_missing_chunks=False`, however, changes + *semantics*: the default engine raises `ChunkNotFoundError` for a missing + chunk, whereas zarrista silently fills it with the fill value. Per the + project's fail-loud rule, refuse rather than silently downgrade to + fill-value reads. + """ + if config is not None and 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" + ) + + +def _region_to_selection(region: Region) -> tuple[slice, ...]: + return tuple(slice(s, e) for s, e in zip(region.start, region.end_exclusive, strict=True)) + + +def _decoded_to_numpy(decoded: Any) -> np.ndarray[Any, Any]: + # Only `Tensor` (fixed-width, unmasked) implements `__array__`/the buffer + # protocol in a way that produces a correct zarr-python array. `VariableArray` + # (vlen data) only exposes the Arrow C Data interface, so `np.asarray` on it + # does *not* raise -- it silently wraps the opaque object in a 0-d `object` + # array -- which is why this is a type-name check rather than a try/except + # around `np.asarray`. `MaskedTensor`/`MaskedVariableArray` do convert (the + # former to a `numpy.ma.MaskedArray`), but masked layouts have no + # zarr-python equivalent to hold that result. + type_name = type(decoded).__name__ + if type_name != "Tensor": + raise NotImplementedError( + f"zarrista returned a {type_name}; only fixed-width, unmasked reads " + "(zarrista `Tensor`) are supported by this engine. vlen reads " + "(`VariableArray`) are pending zarrista numpy export; masked layouts " + "(`MaskedTensor`/`MaskedVariableArray`) have no zarr-python equivalent." + ) + return np.asarray(decoded) + + +def _chunks_overlapping( + region: Region, chunk_shape: tuple[int, ...] +) -> itertools.product[tuple[int, ...]]: + ranges = [ + range(s // c, (e + c - 1) // c) if e > s else range(0) + for s, e, c in zip(region.start, region.end_exclusive, chunk_shape, strict=True) + ] + return itertools.product(*ranges) + + +class ZarristaEngine: + """Sync engine over `zarrista.Array`. No event loop involved.""" + + def __init__(self, zarrista_array: Any) -> None: + self._arr = zarrista_array + + def with_metadata(self, metadata: ArrayMetadata) -> ZarristaEngine: + import zarrista + + return ZarristaEngine( + zarrista.Array.from_metadata(_require_v3(metadata), self._arr.store, self._arr.path) + ) + + def read_selection(self, selection: Region, *, prototype: BufferPrototype) -> Any: + return _decoded_to_numpy(self._arr.retrieve_array_subset(_region_to_selection(selection))) + + def write_selection( + self, selection: Region, value: NDBuffer, *, prototype: BufferPrototype + ) -> None: + import zarrista + + value_np = np.ascontiguousarray(value.as_ndarray_like()) + # regular grids: the origin chunk has the full (unclipped) chunk shape + chunk_shape = tuple(self._arr.chunk_shape([0] * len(selection.start))) + for chunk_idx_tuple in _chunks_overlapping(selection, chunk_shape): + chunk_idx = list(chunk_idx_tuple) + chunk_slices = self._arr.chunk_subset(chunk_idx) + # overlap of the write region with this chunk, in array coords + lo = tuple( + max(cs.start, s) for cs, s in zip(chunk_slices, selection.start, strict=True) + ) + hi = tuple( + min(cs.stop, e) for cs, e in zip(chunk_slices, selection.end_exclusive, strict=True) + ) + in_value = tuple( + slice(a - s, b - s) for a, b, s in zip(lo, hi, selection.start, strict=True) + ) + in_chunk = tuple( + slice(a - cs.start, b - cs.start) + for a, b, cs in zip(lo, hi, chunk_slices, strict=True) + ) + full_chunk = all( + sl.start == 0 and sl.stop == (cs.stop - cs.start) and (cs.stop - cs.start) == c + for sl, cs, c in zip(in_chunk, chunk_slices, chunk_shape, strict=True) + ) + if full_chunk: + chunk_np = np.ascontiguousarray(value_np[in_value]) + else: + chunk_np = np.array(_decoded_to_numpy(self._arr.retrieve_chunk(chunk_idx))) + chunk_np[in_chunk] = value_np[in_value] + chunk_np = np.ascontiguousarray(chunk_np) + self._arr.store_chunk( + chunk_idx, zarrista.ArrayBytes(chunk_np.reshape(-1).view(np.uint8)) + ) + + +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`. + + `config` is accepted for protocol conformance with `HierarchyEngine` + and is otherwise unused -- zarrista owns its own codec options and does + not read zarr-python's `ArrayConfig` (e.g. `order`) -- with one + exception: `config.read_missing_chunks=False` is rejected with an + `UnsupportedEngineError`, since zarrista cannot enforce it and would + silently fill missing chunks instead of raising. + """ + import zarrista + + _reject_unenforceable_config(config) + return ZarristaEngine( + zarrista.Array.from_metadata(_require_v3(metadata), self._zstore, "/" + path.strip("/")) + ) + + +class ZarristaAsyncEngine: + """Async engine over `zarrista.AsyncArray`. + + Store translation, and construction of the underlying `zarrista.AsyncArray`, + are deferred to the first `read_selection`/`write_selection` call rather + than done at construction time. `AsyncArray.__init__` always eagerly + resolves *an* async engine -- even for a plain sync `Array` that never + touches it, since only the sync engine is lazily resolved there (see + `Array.engine`). 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) -> None: + self._store = store + self._path = path + self._metadata = metadata + self._arr: Any | None = None + + def with_metadata(self, metadata: ArrayMetadata) -> ZarristaAsyncEngine: + _require_v3(metadata) + return ZarristaAsyncEngine(self._store, self._path, metadata) + + async def _ensure_arr(self) -> Any: + if self._arr is None: + import zarrista + + meta_dict = _require_v3(self._metadata) + zstore = translate_store_async(self._store) + self._arr = zarrista.AsyncArray.from_metadata(meta_dict, zstore, self._path) + return self._arr + + async def read_selection(self, selection: Region, *, prototype: BufferPrototype) -> Any: + arr = await self._ensure_arr() + return _decoded_to_numpy(await arr.retrieve_array_subset(_region_to_selection(selection))) + + async def write_selection( + self, selection: Region, value: NDBuffer, *, prototype: BufferPrototype + ) -> None: + import zarrista + + arr = await self._ensure_arr() + value_np = np.ascontiguousarray(value.as_ndarray_like()) + # regular grids: the origin chunk has the full (unclipped) chunk shape + chunk_shape = tuple(arr.chunk_shape([0] * len(selection.start))) + for chunk_idx_tuple in _chunks_overlapping(selection, chunk_shape): + chunk_idx = list(chunk_idx_tuple) + chunk_slices = arr.chunk_subset(chunk_idx) + lo = tuple( + max(cs.start, s) for cs, s in zip(chunk_slices, selection.start, strict=True) + ) + hi = tuple( + min(cs.stop, e) for cs, e in zip(chunk_slices, selection.end_exclusive, strict=True) + ) + in_value = tuple( + slice(a - s, b - s) for a, b, s in zip(lo, hi, selection.start, strict=True) + ) + in_chunk = tuple( + slice(a - cs.start, b - cs.start) + for a, b, cs in zip(lo, hi, chunk_slices, strict=True) + ) + full_chunk = all( + sl.start == 0 and sl.stop == (cs.stop - cs.start) and (cs.stop - cs.start) == c + for sl, cs, c in zip(in_chunk, chunk_slices, chunk_shape, strict=True) + ) + if full_chunk: + chunk_np = np.ascontiguousarray(value_np[in_value]) + else: + chunk_np = np.array(_decoded_to_numpy(await arr.retrieve_chunk(chunk_idx))) + chunk_np[in_chunk] = value_np[in_value] + chunk_np = np.ascontiguousarray(chunk_np) + await arr.store_chunk( + chunk_idx, zarrista.ArrayBytes(chunk_np.reshape(-1).view(np.uint8)) + ) + + +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`. + + `config` is accepted for protocol conformance with + `AsyncHierarchyEngine` and is otherwise unused -- zarrista owns its own + codec options and does not read zarr-python's `ArrayConfig` (e.g. + `order`) -- with one exception: `config.read_missing_chunks=False` is + rejected with an `UnsupportedEngineError`, since zarrista cannot enforce + it and would silently fill missing chunks instead of raising. `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) 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..dd234acbc7 --- /dev/null +++ b/tests/engine/test_asyncarray_wiring.py @@ -0,0 +1,152 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, cast + +import numpy as np + +import zarr +from zarr.abc.engine import Region +from zarr.core.engine import DefaultAsyncArrayEngine +from zarr.storage import MemoryStore + +if TYPE_CHECKING: + from zarr.core.buffer import BufferPrototype, NDArrayLike, NDBuffer + from zarr.core.indexing import BasicSelection + from zarr.core.metadata import ArrayMetadata + + +class _SpyEngine: + """Wraps a real engine, recording regions.""" + + def __init__(self, inner: DefaultAsyncArrayEngine) -> None: + self.inner = inner + self.read_regions: list[Region] = [] + self.write_regions: list[Region] = [] + + async def read_selection(self, selection: Region, *, prototype: BufferPrototype) -> NDArrayLike: + self.read_regions.append(selection) + return await self.inner.read_selection(selection, prototype=prototype) + + async def write_selection( + self, selection: Region, value: NDBuffer, *, prototype: BufferPrototype + ) -> None: + self.write_regions.append(selection) + return await self.inner.write_selection(selection, value, prototype=prototype) + + def with_metadata(self, metadata: ArrayMetadata) -> _SpyEngine: + return _SpyEngine(self.inner.with_metadata(metadata)) + + +class _ReadOnlyEngine: + """Wraps a real engine but returns a read-only numpy view from reads. + + Mimics an engine (e.g. `zarrista`) that exposes foreign-owned memory as a + non-writable numpy array. The facade must copy such a result so a read + still yields a writable array. + """ + + def __init__(self, inner: DefaultAsyncArrayEngine) -> None: + self.inner = inner + + async def read_selection(self, selection: Region, *, prototype: BufferPrototype) -> NDArrayLike: + result = await self.inner.read_selection(selection, prototype=prototype) + view = np.asarray(result) + view.flags.writeable = False + return cast("NDArrayLike", view) + + async def write_selection( + self, selection: Region, value: NDBuffer, *, prototype: BufferPrototype + ) -> None: + return await self.inner.write_selection(selection, value, prototype=prototype) + + def with_metadata(self, metadata: ArrayMetadata) -> _ReadOnlyEngine: + return _ReadOnlyEngine(self.inner.with_metadata(metadata)) + + +async def test_read_only_engine_result_is_copied_to_writable() -> None: + # Regression: the identity-read fast path returned the engine buffer + # unchanged. When an engine returns non-writable numpy memory, the facade + # must copy it so reads keep zarr-python's writable-array guarantee. + z = zarr.create_array(MemoryStore(), shape=(6,), chunks=(3,), dtype="int16") + aa = z.async_array + inner = DefaultAsyncArrayEngine( + store_path=aa.store_path, metadata=aa.metadata, config=aa.config + ) + object.__setattr__(aa, "engine", _ReadOnlyEngine(inner)) + await aa.setitem(slice(None), np.arange(6, dtype="int16")) + + result = np.asarray(await aa.getitem(slice(None))) + assert result.flags.writeable + np.testing.assert_array_equal(result, np.arange(6, dtype="int16")) + + +async def test_asyncarray_routes_io_through_engine() -> None: + # NOTE: the async variant of the spy test. `Array` (the sync facade) now + # 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)) + + assert spy.write_regions == [Region(start=(2,), end_exclusive=(8,))] + assert spy.read_regions == [Region(start=(2,), end_exclusive=(8,))] + np.testing.assert_array_equal(np.asarray(data), np.arange(6, dtype="int16")) + + +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`) took the full-box fast path, which broadcast the + # dimension-dropped value straight into the ndim-preserving box -- (3,) could + # not broadcast to box shape (3, 1). A numpy integer scalar keeps the routing + # in the basic (not orthogonal) facade, matching the property-test example. + 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") + # a numpy integer scalar (not a Python int) keeps `__setitem__` routing on the + # basic facade rather than the orthogonal one, reproducing the failing example. + 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]`) produced a + # SliceDimIndexer with start > stop, which `_block_region` mapped to a + # negative-length box (start=1, end_exclusive=0) and crashed with + # "negative dimensions are not allowed". + data = np.arange(2, dtype="int64") + z = zarr.create_array(MemoryStore(), shape=(2,), chunks=(1,), dtype="int64") + z[:] = data + result = np.asarray(z.get_block_selection((slice(1, 0),))) + np.testing.assert_array_equal(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..20b6a66605 --- /dev/null +++ b/tests/engine/test_coordinate_out.py @@ -0,0 +1,54 @@ +"""`get_coordinate_selection(..., out=...)` must validate `out` against the +selection shape, which may be multi-dimensional. + +The engine facade flattens the coordinate arrays into a pointwise index, so the +raw read is 1-d; a naive `out`-shape check against that flattened shape rejected +a multi-dimensional `out` that the pre-engine implementation accepted. These +tests pin both the multi-dimensional and the flat cases. +""" + +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((2, 2), dtype="int32")) + result = z.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) + + +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) + + +def test_coordinate_selection_out_shape_mismatch_raises() -> None: + z = _filled_array() + coords = (np.array([[0, 1], [2, 3]]), np.array([[0, 1], [2, 3]])) + # a flat out no longer matches the 2-d selection shape + out = default_buffer_prototype().nd_buffer.from_numpy_array(np.zeros((4,), 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..80431c4b7d --- /dev/null +++ b/tests/engine/test_default_engine.py @@ -0,0 +1,88 @@ +from __future__ import annotations + +from typing import Any + +import numpy as np +import pytest + +import zarr +from zarr.abc.engine import AsyncArrayEngine, Region +from zarr.core.buffer import default_buffer_prototype +from zarr.core.engine import DefaultArrayEngine, DefaultAsyncArrayEngine +from zarr.core.sync import sync +from zarr.errors import ChunkNotFoundError +from zarr.storage import MemoryStore + + +def _make_array() -> zarr.Array[Any]: + z = zarr.create_array(MemoryStore(), shape=(10, 9), chunks=(3, 4), dtype="int32", fill_value=0) + z[:, :] = np.arange(90, dtype="int32").reshape(10, 9) + return z + + +def test_default_async_engine_read_write_roundtrip() -> None: + z = _make_array() + eng = DefaultAsyncArrayEngine( + store_path=z.async_array.store_path, + metadata=z.async_array.metadata, + config=z.async_array.config, + ) + assert isinstance(eng, AsyncArrayEngine) + proto = default_buffer_prototype() + region = Region(start=(2, 1), end_exclusive=(7, 5)) + + out = sync(eng.read_selection(region, 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") + value = proto.nd_buffer.from_ndarray_like(new) + sync(eng.write_selection(region, value, prototype=proto)) + np.testing.assert_array_equal(np.asarray(z[2:7, 1:5]), new) + + +def test_default_sync_engine_matches_async() -> None: + z = _make_array() + async_eng = DefaultAsyncArrayEngine( + store_path=z.async_array.store_path, + metadata=z.async_array.metadata, + config=z.async_array.config, + ) + eng = DefaultArrayEngine(async_eng) + proto = default_buffer_prototype() + region = Region(start=(0, 0), end_exclusive=(10, 9)) + np.testing.assert_array_equal( + np.asarray(eng.read_selection(region, prototype=proto)), np.asarray(z[:, :]) + ) + + +def test_with_metadata_rebinds() -> None: + z = _make_array() + eng = DefaultAsyncArrayEngine( + store_path=z.async_array.store_path, + metadata=z.async_array.metadata, + config=z.async_array.config, + ) + new_meta = z.async_array.metadata + assert eng.with_metadata(new_meta) 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 + eng = DefaultAsyncArrayEngine( + store_path=z.async_array.store_path, + metadata=z.async_array.metadata, + config=z.async_array.config, + ) + with pytest.raises(ChunkNotFoundError): + sync( + eng.read_selection( + Region(start=(0,), end_exclusive=(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..a86e841419 --- /dev/null +++ b/tests/engine/test_device_buffer_facade.py @@ -0,0 +1,240 @@ +"""Regression tests: the engine facade must not force a host (numpy) conversion +on a device buffer (e.g. cupy/torch) whose implicit `np.asarray` coercion is +refused. + +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. Driving `Array.__setitem__`/`__getitem__` through a stub engine that +returns such buffers therefore fails before the facade's namespace fixes and +passes after them. +""" + +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 Region + from zarr.core.buffer import BufferPrototype, NDBuffer + 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 _region_to_index(region: Region) -> tuple[slice, ...]: + return tuple(slice(s, e) for s, e in zip(region.start, region.end_exclusive, strict=True)) + + +class _DeviceEngine: + """A synchronous `ArrayEngine` backed by numpy that hands the facade + `_DeviceArray` buffers (on read) and unwraps them (on write).""" + + def __init__(self, data: np.ndarray[Any, Any]) -> None: + self._data = data + + def read_selection(self, selection: Region, *, prototype: BufferPrototype) -> _DeviceArray: + return _DeviceArray(self._data[_region_to_index(selection)].copy()) + + def write_selection( + self, selection: Region, value: NDBuffer, *, prototype: BufferPrototype + ) -> None: + arr: object = value.as_ndarray_like() + assert isinstance(arr, _DeviceArray), ( + "facade coerced the box off the device namespace before writing" + ) + self._data[_region_to_index(selection)] = arr._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 + + async def read_selection( + self, selection: Region, *, prototype: BufferPrototype + ) -> _DeviceArray: + return _DeviceArray(self._data[_region_to_index(selection)].copy()) + + async def write_selection( + self, selection: Region, value: NDBuffer, *, prototype: BufferPrototype + ) -> None: + arr: object = value.as_ndarray_like() + assert isinstance(arr, _DeviceArray), ( + "facade coerced the box off the device namespace before writing" + ) + self._data[_region_to_index(selection)] = arr._a + + 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_identity_setitem_keeps_device_buffer() -> None: + # full-box write: `_prepare_set_selection` broadcasts the value into the box + # without a read; it must not `np.asarray` the device value. + z, engine = _array_on_device_engine() + value = _DeviceArray(np.arange(8, dtype="int64")) + z[:] = value + np.testing.assert_array_equal(engine._data, np.arange(8, dtype="int64")) + + +def test_strided_setitem_reads_and_patches_on_device() -> None: + # strided write is a read-modify-write: the engine box read comes back as a + # `_DeviceArray`, and `_patch_selection_box` must patch it in-namespace + # (`raw.copy()`), never `np.asarray(raw)`. + z, engine = _array_on_device_engine() + z[::2] = 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: + # strided read: `_finish_get_selection` applies the post index and reorders + # the result in the device namespace, returning a `_DeviceArray` untouched + # by `np.asarray`. + 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: + # all-integer basic read scalarizes a 0-d result; `_finalize_result` must + # extract the element in the device namespace, not via `np.asarray(result)[()]`. + 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: + # orthogonal write with a dropped integer axis widens the value with a + # `np.newaxis`; `_widen_value_for_squeeze` must index the value in its own + # namespace instead of `np.asarray(value)[...]`. + z, engine = _array_2d_on_device_engine() + value = _DeviceArray(np.array([10, 20], dtype="int64")) + z.oindex[1, np.array([0, 2])] = value + 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; the device-namespace `result[()]` branch must not swallow it + # (`np.bytes_[()]` raises "byte indices must be integers"). numpy scalars go + # through the `np.asarray` path, staying bit-identical to the historical + # `as_scalar()` return. (Distilled from a `test_basic_indexing` hypothesis + # failure; uses the real default engine, no device stand-in.) + 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..5b08110b54 --- /dev/null +++ b/tests/engine/test_differential.py @@ -0,0 +1,116 @@ +"""The same operations through both engines must agree with numpy and each other.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import pytest + +import zarr +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.core.engine import EngineName + +SHAPE = (10, 9) +CHUNKS = (3, 4) + + +@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 + + +READS: list[tuple[int | slice, int | slice]] = [ + (slice(None), slice(None)), + (slice(2, 7), slice(1, 5)), + # negative-step slices are rejected by `Array.__getitem__` (see + # `NegativeStepError`); only positive step > 1 is supported, so this + # covers strided reads on both axes without hitting that restriction. + (slice(1, 9, 2), slice(None, None, 3)), + (3, slice(None)), + (-1, -2), + (slice(4, 4), slice(None)), +] + + +@pytest.mark.parametrize("engine", ENGINES) +@pytest.mark.parametrize("sel", READS) +def test_reads_match_numpy( + reference: tuple[Path, npt.NDArray[np.float64]], + engine: EngineName, + sel: tuple[int | slice, int | slice], +) -> None: + tmp_path, data = reference + z = zarr.open_array(LocalStore(tmp_path), engine=engine) + np.testing.assert_array_equal(np.asarray(z[sel]), data[sel]) + + +@pytest.mark.parametrize("engine", ENGINES) +def test_fancy_reads_match_numpy( + reference: tuple[Path, npt.NDArray[np.float64]], engine: EngineName +) -> None: + tmp_path, data = reference + z = zarr.open_array(LocalStore(tmp_path), engine=engine) + np.testing.assert_array_equal( + np.asarray(z.oindex[np.array([7, 1, 4]), np.array([0, 8])]), + data[np.ix_([7, 1, 4], [0, 8])], + ) + np.testing.assert_array_equal( + np.asarray(z.vindex[np.array([9, 0, 3]), np.array([8, 0, 2])]), + data[np.array([9, 0, 3]), np.array([8, 0, 2])], + ) + np.testing.assert_array_equal(np.asarray(z.blocks[1, 2]), data[3:6, 8:9]) + + +@pytest.mark.parametrize("engine", ENGINES) +def test_writes_match_numpy( + reference: tuple[Path, npt.NDArray[np.float64]], engine: EngineName +) -> None: + tmp_path, data = reference + z = zarr.open_array(LocalStore(tmp_path), engine=engine) + expected = data.copy() + + z[0:3, 0:4] = 7.0 # aligned full chunk + expected[0:3, 0:4] = 7.0 + z[4:6, 2:9] = np.arange(14, dtype="float64").reshape(2, 7) # partial chunks + expected[4:6, 2:9] = np.arange(14, dtype="float64").reshape(2, 7) + z[1:9:3, ::4] = -1.0 # strided write (facade RMW) + expected[1:9:3, ::4] = -1.0 + + np.testing.assert_array_equal(np.asarray(z[:, :]), expected) + + +@pytest.mark.parametrize("engine", ENGINES) +def test_sharded_reads( + reference: tuple[Path, npt.NDArray[np.float64]], + engine: EngineName, + tmp_path_factory: pytest.TempPathFactory, +) -> None: + path = tmp_path_factory.mktemp("sharded") + z = zarr.create_array( + LocalStore(path), + shape=SHAPE, + chunks=(3, 4), # inner chunks + shards=(6, 8), # shard shape + dtype="int32", + ) + data = np.arange(90, dtype="int32").reshape(SHAPE) + z[:, :] = data + zr = zarr.open_array(LocalStore(path), engine=engine) + np.testing.assert_array_equal(np.asarray(zr[2:8, 3:9]), data[2:8, 3:9]) 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_normalize.py b/tests/engine/test_normalize.py new file mode 100644 index 0000000000..587a4c2a05 --- /dev/null +++ b/tests/engine/test_normalize.py @@ -0,0 +1,140 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import numpy as np +import numpy.typing as npt +import pytest + +from zarr.core.engine import ( + apply_post_index, + normalize_basic, + normalize_coordinate, + normalize_orthogonal, + strip_squeeze, +) +from zarr.core.engine._normalize import _Squeeze + +if TYPE_CHECKING: + from zarr.abc.engine import Region + +SHAPE = (10, 9) +ARR = np.arange(90).reshape(SHAPE) + + +def _read_box(region: Region) -> npt.NDArray[np.int_]: + """Simulate an engine read: the ndim-preserving box.""" + return ARR[tuple(slice(s, e) for s, e in zip(region.start, region.end_exclusive, strict=True))] + + +@pytest.mark.parametrize( + "sel", + [ + (slice(2, 7), slice(0, 9)), + (slice(None), slice(3, 4)), + (slice(8, 2, -2), slice(None)), + (slice(1, 8, 3), slice(2, 9, 2)), + (3, slice(None)), + (-1, -2), + (Ellipsis, 4), + slice(5), + Ellipsis, + (slice(4, 4), slice(None)), # empty + ], +) +def test_normalize_basic_matches_numpy(sel: Any) -> None: + region, post = normalize_basic(sel, SHAPE) + np.testing.assert_array_equal(apply_post_index(_read_box(region), post), ARR[sel]) + + +@pytest.mark.parametrize( + "sel", + [ + (np.array([1, 4, 7]), np.array([0, 8])), + (np.array([7, 1, 4]), slice(2, 6)), # unordered axis indices + (np.array([3, 3]), np.array([5, 5])), # repeats + (slice(1, 9, 2), np.array([2])), + (np.array([0]), 4), # int axis + ], +) +def test_normalize_orthogonal_matches_numpy_oindex(sel: Any) -> None: + from zarr.core.indexing import oindex as _reference_oindex + + region, post = normalize_orthogonal(sel, SHAPE) + # `zarr.core.indexing.oindex` is the existing, well-tested reference + # implementation of orthogonal indexing (np.ix_ outer indexing with + # integer axes dropped afterwards) -- use it as the ground truth rather + # than reimplementing it by hand. + expected = _reference_oindex(ARR, sel if isinstance(sel, tuple) else (sel,)) + np.testing.assert_array_equal(apply_post_index(_read_box(region), post), expected) + + +def test_normalize_coordinate_matches_numpy_vindex() -> None: + coords = (np.array([9, 0, 3, 3]), np.array([8, 0, 2, 2])) + region, post = normalize_coordinate(coords, SHAPE) + np.testing.assert_array_equal(apply_post_index(_read_box(region), post), ARR[coords]) + + +def test_normalize_basic_rejects_fancy() -> None: + with pytest.raises(TypeError): + normalize_basic((np.array([1, 2]), slice(None)), SHAPE) # type: ignore[arg-type] + + +def test_normalize_basic_rejects_out_of_bounds_int() -> None: + with pytest.raises(IndexError): + normalize_basic((10, slice(None)), SHAPE) + + +def test_normalize_coordinate_rejects_out_of_bounds() -> None: + with pytest.raises(IndexError): + normalize_coordinate((np.array([10]), np.array([0])), SHAPE) + + +def test_normalize_orthogonal_rejects_boolean_ndim_mismatch() -> None: + with pytest.raises(IndexError): + normalize_orthogonal((np.zeros((2, 2), dtype=bool), slice(None)), SHAPE) + + +def test_strip_squeeze_removes_trailing_marker() -> None: + post_with_marker = (slice(None), np.array([0]), _Squeeze((1,))) + post_without = (slice(None), np.array([0])) + stripped = strip_squeeze(post_with_marker) + assert len(stripped) == len(post_without) + assert stripped[0] == post_without[0] + np.testing.assert_array_equal(stripped[1], post_without[1]) + + +def test_strip_squeeze_identity_when_no_marker() -> None: + post = (slice(None), slice(2, 4)) + assert strip_squeeze(post) == post + + +def test_normalize_orthogonal_list_of_bools_matches_array_bool_mask() -> None: + """List of booleans should produce the same result as np.ndarray bool mask.""" + from zarr.core.indexing import oindex as _reference_oindex + + # Test case from the defect report: list of 10 booleans for axis 0 + bool_list = [True, False, True] + [False] * 7 + bool_array = np.array(bool_list, dtype=bool) + + # Both should produce the same result when used as selectors + region_list, post_list = normalize_orthogonal((bool_list, slice(None)), SHAPE) + region_array, post_array = normalize_orthogonal((bool_array, slice(None)), SHAPE) + + # Reference implementation using zarr.core.indexing.oindex + expected = _reference_oindex(ARR, (bool_array, slice(None))) + + # List of bools should match the array bool mask result + result_list = apply_post_index(_read_box(region_list), post_list) + result_array = apply_post_index(_read_box(region_array), post_array) + + np.testing.assert_array_equal(result_list, expected) + np.testing.assert_array_equal(result_array, expected) + + +def test_normalize_orthogonal_wrong_length_list_of_bools_raises_error() -> None: + """Wrong-length list of booleans should raise IndexError like ndarray.""" + wrong_length_bool_list = [True, False, True] # Length 3, but axis 0 has size 10 + + with pytest.raises(IndexError, match="boolean index for axis 0 must be 1-d with length 10"): + normalize_orthogonal((wrong_length_bool_list, slice(None)), SHAPE) diff --git a/tests/engine/test_protocols.py b/tests/engine/test_protocols.py new file mode 100644 index 0000000000..d8badc3a04 --- /dev/null +++ b/tests/engine/test_protocols.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import pytest + +from zarr.abc.engine import ArrayEngine, AsyncArrayEngine, Region +from zarr.errors import UnsupportedEngineError + +if TYPE_CHECKING: + from zarr.core.buffer import BufferPrototype, NDArrayLike, NDBuffer + from zarr.core.metadata import ArrayMetadata + + +def test_region() -> None: + """Region carries start/end_exclusive and derives shape.""" + r = Region(start=(1, 2), end_exclusive=(4, 2)) + assert r.start == (1, 2) + assert r.end_exclusive == (4, 2) + assert r.shape == (3, 0) + + +class _FakeSyncEngine: + def read_selection(self, selection: Region, *, prototype: BufferPrototype) -> NDArrayLike: + return np.zeros(selection.shape) + + def write_selection( + self, selection: Region, value: NDBuffer, *, prototype: BufferPrototype + ) -> None: + return None + + def with_metadata(self, metadata: ArrayMetadata) -> _FakeSyncEngine: + return self + + +def test_runtime_checkable_protocols() -> None: + """isinstance checks verify method presence for the sync protocol.""" + assert isinstance(_FakeSyncEngine(), ArrayEngine) + assert not isinstance(object(), ArrayEngine) + assert ( + not isinstance(_FakeSyncEngine(), AsyncArrayEngine) or True + ) # names match; mypy is authoritative + + +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..8b074b95ef --- /dev/null +++ b/tests/engine/test_sync_path.py @@ -0,0 +1,82 @@ +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: + from zarr.abc.engine import Region + from zarr.core.buffer import BufferPrototype, NDArrayLike, NDBuffer + 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, selection: Region, *, prototype: BufferPrototype) -> NDArrayLike: + self.calls += 1 + with pytest.raises(RuntimeError): + asyncio.get_running_loop() + return self._inner.read_selection(selection, prototype=prototype) + + def write_selection( + self, selection: Region, value: NDBuffer, *, prototype: BufferPrototype + ) -> None: + self.calls += 1 + with pytest.raises(RuntimeError): + asyncio.get_running_loop() + return self._inner.write_selection(selection, value, prototype=prototype) + + 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 0d6d2d5906..109119e315 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 ( @@ -1596,10 +1596,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/test_indexing.py b/tests/test_indexing.py index a9358e4fcf..ee0e6b7580 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -1821,15 +1821,33 @@ async def test_accessed_chunks( z = zarr_array_from_numpy_array(StorePath(store), np.zeros(shape), chunk_shape=chunks) for ii, (optype, slices) in enumerate(ops): - # Resolve the slices into the accessed chunks for each dimension - chunks_per_dim = [] + # The array engine interchanges selections as contiguous, step-1 `Region` + # boxes: a selection is served through its step-1 *bounding box*, + # deliberately replacing the old per-chunk gather/scatter (see + # docs/superpowers/specs/2026-07-22-array-engine-protocol-design.md, + # "Consequence, accepted for now"). + # + # Reads fetch the whole bounding box, so a strided slice reads every chunk + # between its first and last selected index. A write is a bounding-box + # read-modify-write: it reads the box, patches the selected elements, and + # writes the box back -- but `write_empty_chunks=False` (the default) + # drops box chunks that stay all-fill, so only chunks that actually + # contain a selected element are written. Model the two sets separately. + box_chunks_per_dim = [] # every chunk the step-1 bounding box overlaps + selected_chunks_per_dim = [] # chunks containing an actually selected index for N, C, sl in zip(shape, chunks, slices, strict=True): - chunk_ind = np.arange(N, dtype=int)[sl] // C - chunks_per_dim.append(np.unique(chunk_ind)) + selected = np.arange(N, dtype=int)[sl] + selected_chunks_per_dim.append(np.unique(selected // C)) + if selected.size: + bounding_box = np.arange(selected.min(), selected.max() + 1) + else: + bounding_box = selected + box_chunks_per_dim.append(np.unique(bounding_box // C)) - # Combine and generate the cartesian product to determine the chunks keys that - # will be accessed - chunks_accessed = [".".join(map(str, comb)) for comb in itertools.product(*chunks_per_dim)] + box_chunks = [".".join(map(str, comb)) for comb in itertools.product(*box_chunks_per_dim)] + selected_chunks = [ + ".".join(map(str, comb)) for comb in itertools.product(*selected_chunks_per_dim) + ] counts_before = store.counter.copy() @@ -1837,24 +1855,28 @@ async def test_accessed_chunks( if optype == "__getitem__": z[slices] else: + # a non-zero fill value keeps the selected chunks non-empty so they + # are written (ii == 0 only for the very first op, which is a read) z[slices] = ii # Get the change in counts delta_counts = store.counter - counts_before - # Check that the access counts for the operation have increased by one for all - # the chunks we expect to be included - for ci in chunks_accessed: - assert delta_counts.pop((optype, ci)) == 1 - - # If the chunk was partially written to it will also have been read once. We - # don't determine if the chunk was actually partial here, just that the - # counts are consistent that this might have happened - if optype == "__setitem__": - assert ("__getitem__", ci) not in delta_counts or delta_counts.pop( - ("__getitem__", ci) - ) == 1 - # Check that no other chunks were accessed + if optype == "__getitem__": + # a read fetches the whole bounding box, and nothing else + for ci in box_chunks: + assert delta_counts.pop(("__getitem__", ci)) == 1 + else: + # the RMW reads bounding-box chunks (unless the write covers the whole + # box, in which case there is no facade read); an edge chunk only + # partially covered by the box is read again by the codec pipeline's + # partial-chunk merge, so tolerate more than one read per box chunk. + for ci in box_chunks: + delta_counts.pop(("__getitem__", ci), None) + # ...and writes back the chunks that end up non-empty + for ci in selected_chunks: + assert delta_counts.pop(("__setitem__", ci)) >= 1 + # Check that no chunks outside the bounding box / selection were accessed assert len(delta_counts) == 0 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..1a590c6757 --- /dev/null +++ b/tests/zarrista/test_engine.py @@ -0,0 +1,142 @@ +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 + +if TYPE_CHECKING: + from pathlib import Path + + +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, so the + # facade must copy such a result -- both on the full-box identity path and + # on a partial (bounding-box) read. + _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_not_implemented(tmp_path: Path) -> None: + # zarrista decodes vlen dtypes (e.g. strings) to a `VariableArray`, which + # only exposes the Arrow C Data interface -- `np.asarray` on it does not + # raise, it silently produces a wrong-shape 0-d `object` array. The engine + # must reject this itself rather than let a bogus read through. + 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") + with pytest.raises(NotImplementedError, match="VariableArray"): + ze[:] + + +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 test_zarrista_async_engine_read_write_combinations(tmp_path: Path) -> None: + # exercises `ZarristaAsyncEngine` directly (over an obstore-backed + # `ObjectStore`), rather than through the sync `Array`/`ZarristaEngine` + # path the other tests in this module cover. + obstore = pytest.importorskip("obstore") + from zarr.api import asynchronous as async_api + from zarr.storage import ObjectStore + + store = ObjectStore(obstore.store.LocalStore(prefix=str(tmp_path))) + z = await async_api.create_array(store=store, shape=(10, 9), chunks=(3, 4), dtype="float32") + await z.setitem((slice(None), slice(None)), np.arange(90, dtype="float32").reshape(10, 9)) + ze = await async_api.open_array(store=store, engine="zarrista") + + # contiguous read + out = await ze.getitem((slice(2, 7), slice(1, 5))) + expected = await z.getitem((slice(2, 7), slice(1, 5))) + np.testing.assert_array_equal(np.asarray(out), np.asarray(expected)) + # full-chunk-aligned write + await ze.setitem((slice(0, 3), slice(0, 4)), np.zeros((3, 4), dtype="float32")) + written = await z.getitem((slice(0, 3), slice(0, 4))) + np.testing.assert_array_equal(np.asarray(written), np.zeros((3, 4), dtype="float32")) + # partial-chunk RMW write + await ze.setitem((slice(1, 2), slice(1, 2)), np.float32(99.0)) + assert float(np.asarray(await z.getitem((1, 1)))) == 99.0 + assert float(np.asarray(await z.getitem((0, 0)))) == 0.0 # untouched neighbor 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 ab2ff8b1e7..22c5c31f2c 100644 --- a/uv.lock +++ b/uv.lock @@ -1157,6 +1157,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0a/1e/8222edaee03c37350eaa726213614e343a62f1e56396dd000ad9277bfa3d/hypothesis-6.152.7-py3-none-any.whl", hash = "sha256:c0b17dd428fcb6e962f60315f6f4a77816c72fbb281ce9ba73699dabead5ec82", size = 533802, upload-time = "2026-05-13T04:19:30.635Z" }, ] +[[package]] +name = "icechunk" +version = "2.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "zarr" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/89/41/ea6d2ebe3fced2be5ca4e8540b776d9544a31e30f0a90fe7968b28288aff/icechunk-2.1.1.tar.gz", hash = "sha256:a46637e2e079edcfb92a6cc2a7a0cfb990224b0be2fb1931f967361d17498fa0", size = 3528133, upload-time = "2026-07-08T20:30:01.408Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b7/9c/9cad0e3c6f4ab671816e513a10ddde6c13ca5b290ee38ae4f6611b098263/icechunk-2.1.1-cp312-abi3-macosx_10_12_x86_64.whl", hash = "sha256:660e8ad231cddad0f86cd26cc2f46a845d6b5cfed7bf023e9ceefbdb78b7d3e4", size = 17092044, upload-time = "2026-07-08T20:30:09.175Z" }, + { url = "https://files.pythonhosted.org/packages/79/ba/e7af165c67a8707e18f5e97233093a9d13153948334788f168002bfabffd/icechunk-2.1.1-cp312-abi3-macosx_11_0_arm64.whl", hash = "sha256:8fe37a901d7ad6a6bedb074e60a521505cce11d418777b43d974601a53c03ed5", size = 15776756, upload-time = "2026-07-08T20:30:06.79Z" }, + { url = "https://files.pythonhosted.org/packages/63/5b/b02968cd97ef2d910dd46eadf03e078f630ddf14c869c3b0c05363e6739a/icechunk-2.1.1-cp312-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9131213e1d831e63a0f74d20792934f3e1c97d4084680788df0648f4c19cb985", size = 17473456, upload-time = "2026-07-08T20:30:04.128Z" }, + { url = "https://files.pythonhosted.org/packages/c8/59/d1a379c7206011da4d89ab32064ff98398eaf4cc32227189cfb443d252a1/icechunk-2.1.1-cp312-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:34cc925a3d339dde989889d6f703bed3eb0d5d2e41887cf12a485c36122159eb", size = 17117003, upload-time = "2026-07-08T20:29:59.159Z" }, + { url = "https://files.pythonhosted.org/packages/ca/08/dcab15706126a1ac05639d6213ce34fdf72ff5c77a85542e09a9fcab4356/icechunk-2.1.1-cp312-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:2ceef9b8b0d4988b35c33411e5c2da15364611ba8a88292f6d382fe6d216f452", size = 17338644, upload-time = "2026-07-08T20:30:11.397Z" }, + { url = "https://files.pythonhosted.org/packages/ed/93/6ca99877532b9f214ed62ead6c9921b5798d0ce1eed468ab0f71331bcd77/icechunk-2.1.1-cp312-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:f6bbcb216f636fc6b53682297a6438eeb81f577e976b430ba58407d05f34bb31", size = 17886257, upload-time = "2026-07-08T20:30:13.761Z" }, + { url = "https://files.pythonhosted.org/packages/5b/ca/14d4a288343b2e01e06888a86640e8ce3e16b6f2b8bd10616add382f71d4/icechunk-2.1.1-cp312-abi3-win_amd64.whl", hash = "sha256:b87329b7088301777f5dde7f6f4180938cf410768e4cb49d3b9cc064eda2b036", size = 16209937, upload-time = "2026-07-08T20:30:16.41Z" }, +] + [[package]] name = "idna" version = "3.15" @@ -2257,52 +2275,39 @@ wheels = [ [[package]] name = "obstore" -version = "0.9.4" +version = "0.11.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/d3/7a/3a37b0bf0da898478029fcc511a0d2a7252689b1f29e46db7ae74a219c74/obstore-0.9.4.tar.gz", hash = "sha256:e2b93f1372c59da2c7e74122fc6dc4b713d84fd4528b5b500ef7f548425496b5", size = 124167, upload-time = "2026-04-22T19:51:05.261Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ac/25/4449a0066796b91e282d7604a66387bba399b14752598c748ea9557c4c32/obstore-0.9.4-cp311-abi3-macosx_10_12_x86_64.whl", hash = "sha256:0d17cd04e7f22960050a85f8daa6e274d693e8fb3b97b81eeaa293c6f9e62eb4", size = 4090743, upload-time = "2026-04-22T19:49:26.461Z" }, - { url = "https://files.pythonhosted.org/packages/93/91/639fe5f5644593b9f4bea66f8f29c7bfd4de3b3381fb74b4f7df678f505f/obstore-0.9.4-cp311-abi3-macosx_11_0_arm64.whl", hash = "sha256:d4beec92710fb8826fb357baf28fb79a91ee07dcdfe73777207aa762164aaa35", size = 3876313, upload-time = "2026-04-22T19:49:28.107Z" }, - { url = "https://files.pythonhosted.org/packages/ce/71/d6675f845ebe1e3927f2dce6a2a4d5a393359274762ee00c5e6855d5f468/obstore-0.9.4-cp311-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:d523c8c365ab60afb8d232614a00a92bea439a9f5c55b92486c23a47af038a1e", size = 4029950, upload-time = "2026-04-22T19:49:30.279Z" }, - { url = "https://files.pythonhosted.org/packages/0e/3a/5915a173f5c6a95f9ec186a7e29b0ce6a23bd9b04c2b0b29a351dbe2baf6/obstore-0.9.4-cp311-abi3-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ee0483619088337ee365cb344fceee337e2670ec4de2a1da92ac7f6b2220f18e", size = 4129455, upload-time = "2026-04-22T19:49:31.934Z" }, - { url = "https://files.pythonhosted.org/packages/b5/a9/63c31d2d436c06c4d39ed5cb154fe54202b303854532ec09537c4ce0755b/obstore-0.9.4-cp311-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:83da348bf0a7dd84e5839c0cd54d79dcd08e0729c394e566f73a605b93b9e998", size = 4416727, upload-time = "2026-04-22T19:49:34.016Z" }, - { url = "https://files.pythonhosted.org/packages/7f/fa/23c5c6db02be0e13abcbe01c1ca94c5f7876e8c58e74cb9ac2b57b068866/obstore-0.9.4-cp311-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:f282a17200bcc37b8d7a1d02a146ed41812eb6e76fd0a4c9a154f02da1b8031f", size = 4311520, upload-time = "2026-04-22T19:49:35.905Z" }, - { url = "https://files.pythonhosted.org/packages/86/f0/49f6b02dab9c05e3fd79d6129e4d9e7e9874d6e5e05369ca3b3b80a48aaa/obstore-0.9.4-cp311-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d29dcfceaa0a205ded2263d29a2a3aa206819d549e0325c1f2106f79e2658584", size = 4220536, upload-time = "2026-04-22T19:49:38.343Z" }, - { url = "https://files.pythonhosted.org/packages/50/ab/d0bfd6d68422e7d8f2204d91736c7e62767e0576ad749da442a71e7773b2/obstore-0.9.4-cp311-abi3-manylinux_2_24_aarch64.whl", hash = "sha256:caecb912723ab8e9da8da26def249d66da4318959df2bafc0a55af64f3255902", size = 4105099, upload-time = "2026-04-22T19:49:40.384Z" }, - { url = "https://files.pythonhosted.org/packages/66/3b/f595d0ee354f9daa69438991f8818602f34bc59498c8468456a02d45fb27/obstore-0.9.4-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:0c1c06fec8837595a2829b5f7536d0d01e940ce10b07ad2a8594fec1cfd0b7d5", size = 4294206, upload-time = "2026-04-22T19:49:42.016Z" }, - { url = "https://files.pythonhosted.org/packages/60/54/3c5af2d59258aaa9e5bef05320658ea6e9b1f3897a3a977bf7f54a0b6ec1/obstore-0.9.4-cp311-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:c132795a789ec5ade31bf4d5b55ed321fb41d9749e9145520bf19063e1da5f7b", size = 4265047, upload-time = "2026-04-22T19:49:43.983Z" }, - { url = "https://files.pythonhosted.org/packages/fb/af/a8ba1feb81b9833b253147839da40405ec6bfa51feb3abfe909c800208a5/obstore-0.9.4-cp311-abi3-musllinux_1_2_i686.whl", hash = "sha256:c6e342360a5d0ae71486bc5f8311778aa144ec1a905c23593f8ef57b5bceae24", size = 4255361, upload-time = "2026-04-22T19:49:45.864Z" }, - { url = "https://files.pythonhosted.org/packages/15/f7/3ccc0288111e057f8ba3d99bee14f95d9e9bb00acaf6e9700e0eb4cd82c3/obstore-0.9.4-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:aeb6f7e7e862550f5020a10692ef6f02d5ba4912dba08942eb59bb7d73f93fe0", size = 4439378, upload-time = "2026-04-22T19:49:47.581Z" }, - { url = "https://files.pythonhosted.org/packages/c4/b1/3ac8b5772743c60064f3c7e02d27f346dbb58feaa99a49ee09798d1cfb00/obstore-0.9.4-cp311-abi3-win_amd64.whl", hash = "sha256:a58ef942292841f99d69ac11d19d05544c835447c8c09dacbfb7409c6374c4a1", size = 4191594, upload-time = "2026-04-22T19:49:49.308Z" }, - { url = "https://files.pythonhosted.org/packages/9d/81/8f6b6509f8df603261cdb5ddb521c49891457775669c6ad857812bf4a7c1/obstore-0.9.4-cp313-cp313t-macosx_10_12_x86_64.whl", hash = "sha256:fff17f59390ed307afcd1fb18c56076c1f911dd9f5c2636b7d7133c4d07f8c3f", size = 4071300, upload-time = "2026-04-22T19:49:51.386Z" }, - { url = "https://files.pythonhosted.org/packages/ab/fe/0c74ddf3ab9b24ef356925bfb613bc7846f869220361a784b63f754d8563/obstore-0.9.4-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:4527c4c7889f1bd1f1952017d74774870e14e199d6b50b9e72f291f9498d898c", size = 3870593, upload-time = "2026-04-22T19:49:53.481Z" }, - { url = "https://files.pythonhosted.org/packages/73/fa/260ec94f9a7b4f4c8afbdd016710bed0736615488d3ac0c5620f9179bfcd/obstore-0.9.4-cp313-cp313t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:a57c2016e3e569de35050f95c679ffe61813c4e3cb6d6028c4c3f57231021eb4", size = 4023990, upload-time = "2026-04-22T19:49:55.644Z" }, - { url = "https://files.pythonhosted.org/packages/8d/84/5b8e2b9607fb93c96a39a4cfa6d37bd3049ebf7265d0e9f8afa938bf32fe/obstore-0.9.4-cp313-cp313t-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:bd5327cee4fb3578b51beb1c92915cc3a05ffe794be40f50bd68d27e97d78c5c", size = 4119971, upload-time = "2026-04-22T19:49:57.745Z" }, - { url = "https://files.pythonhosted.org/packages/3c/2b/e6c093acb7e62009d5b1678d82839903287c29d4a6e1dfbea8fbf41313d5/obstore-0.9.4-cp313-cp313t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:12b1e6105eafe02d8973dbeb2d274eeac2271c67f1126ffa16f18ddea8dd5443", size = 4407147, upload-time = "2026-04-22T19:49:59.928Z" }, - { url = "https://files.pythonhosted.org/packages/ce/3d/5c93a9adee8f045b89d5f21b337f53667499db770bda129f805723ab14e4/obstore-0.9.4-cp313-cp313t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:0b0d378248fda4e36652808d73eaaeb7e67154427e6c724248c9b0b9b03e70a6", size = 4312215, upload-time = "2026-04-22T19:50:01.534Z" }, - { url = "https://files.pythonhosted.org/packages/9a/de/507f60b4e6a8c0cad9f93a51a7b28132c9db49e20aadbcd542fa2abc57c4/obstore-0.9.4-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a78fb77c346abd2bcdfa071d7166be2bdc38c28573ae5a230746df6158a5593e", size = 4216936, upload-time = "2026-04-22T19:50:03.244Z" }, - { url = "https://files.pythonhosted.org/packages/7b/ff/612bd5f8258349bfe9e8c349d184b5ea3333038d4cce0d003eefafb2160c/obstore-0.9.4-cp313-cp313t-manylinux_2_24_aarch64.whl", hash = "sha256:f4e5a6dfe6877fb599868d560d6fcf4d7416cadbdf3bd947254b53830c2f11c0", size = 4105091, upload-time = "2026-04-22T19:50:05.038Z" }, - { url = "https://files.pythonhosted.org/packages/e3/73/b083b99e7bc0b529bee7b4437cafd7cc7d9f59c10995a48b6c26447fdf7f/obstore-0.9.4-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:f8114a2b84268c991232d89b105d9239299b6afb56e4941a61c09f3a89033022", size = 4292570, upload-time = "2026-04-22T19:50:06.823Z" }, - { url = "https://files.pythonhosted.org/packages/8b/cd/3c4555f98db9a49432bc0afa68bfc33dd47bdfa3699c915b4b0e887577e3/obstore-0.9.4-cp313-cp313t-musllinux_1_2_armv7l.whl", hash = "sha256:9d7b959f5f74532a142fb449c0bef5814dfe3fa5c43c31ac4284a15221a75aaf", size = 4261946, upload-time = "2026-04-22T19:50:08.789Z" }, - { url = "https://files.pythonhosted.org/packages/96/f8/bdc66df3d0dfdcfb3931a585a7fb3b74336619baf6d3540b1425b424232b/obstore-0.9.4-cp313-cp313t-musllinux_1_2_i686.whl", hash = "sha256:a8e9101fc2659dd938e7ae06512075bc0a8f02ab28d2ee438d6fca8b4f3bdfba", size = 4245595, upload-time = "2026-04-22T19:50:10.765Z" }, - { url = "https://files.pythonhosted.org/packages/d7/22/1aa58ea676293e5b888391c8433ff6ab8f66622aae30427287f9daac6d46/obstore-0.9.4-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:538384255545b5c575497fcab26389c8f01707402b6ddcdd73b769b66311635d", size = 4436599, upload-time = "2026-04-22T19:50:12.585Z" }, - { url = "https://files.pythonhosted.org/packages/1d/9e/b52f2c97be27952d488cf1980af0c635f9947003e5744e3e1dc6252f0040/obstore-0.9.4-cp313-cp313t-win_amd64.whl", hash = "sha256:eef1c772657bb1293adad0d671ca1ff1e1dcae84ec4dfbf1a34e47c2a1f134ac", size = 4180463, upload-time = "2026-04-22T19:50:14.288Z" }, - { url = "https://files.pythonhosted.org/packages/19/76/c53583f95c6811057abd3116756dca46785318d564a0e99c207cbb2d8938/obstore-0.9.4-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:e009e7437770c85beae4c32cb79f662f0a9922676ef127e943d107a5c082d38d", size = 4071302, upload-time = "2026-04-22T19:50:15.967Z" }, - { url = "https://files.pythonhosted.org/packages/2f/23/ac3b9c05a09b3d5f178ed6f288c5d6913df8f7386059590194e0fee65d15/obstore-0.9.4-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:ac5f3ad314bd4592fe484b79c229518be7bb5f6218bed33c20742026d5caf860", size = 3870813, upload-time = "2026-04-22T19:50:17.62Z" }, - { url = "https://files.pythonhosted.org/packages/f7/ec/c3458e0f24d2d1a4f185f541905b07e51c91b3fec589b1600c77d511e585/obstore-0.9.4-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:db79d5ebc4177360565ffcec4abd49930cf052cdbeb94e3a3ece2e2d08f087d0", size = 4024237, upload-time = "2026-04-22T19:50:19.81Z" }, - { url = "https://files.pythonhosted.org/packages/a7/eb/6cf468a200e491fdc6c04075e2fbbac1707bbecd243f0f56ae1e75d052ed/obstore-0.9.4-cp314-cp314t-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:05b565d89c3115fb74385852dd628e12f6645a1bba97523dceae016b538a3f33", size = 4119635, upload-time = "2026-04-22T19:50:21.605Z" }, - { url = "https://files.pythonhosted.org/packages/81/fb/b44d002767fa5af95ab4ca8e16c3a9057fc11f13de03f498b99adf0c4e50/obstore-0.9.4-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:7dfc4fc98403d8fbb316eb04257c8122b6f1dda37e80869491fdacf60a815e4c", size = 4406906, upload-time = "2026-04-22T19:50:23.654Z" }, - { url = "https://files.pythonhosted.org/packages/4b/18/9a75ad5082cd581c4a55f0e62bedf4b030a8b53824976fc1f030eff225b3/obstore-0.9.4-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:c69af620fd3d06a8cfb62d25faf1adb6ccc97cc572f47ee04dddcde5a5e5444e", size = 4311826, upload-time = "2026-04-22T19:50:25.458Z" }, - { url = "https://files.pythonhosted.org/packages/8d/03/b0f945b31f40364a7ed4dbc5677abc66331fcf478732f4d643e17e56bb13/obstore-0.9.4-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fa78c0230e0b9d49b25ed18980e1751331ddfe05782d6ce97579a9ccda8229ea", size = 4217086, upload-time = "2026-04-22T19:50:27.266Z" }, - { url = "https://files.pythonhosted.org/packages/1a/26/bdd85264c806802086f21d73cc7c95a5baca5feeeac4bce8acb97142163f/obstore-0.9.4-cp314-cp314t-manylinux_2_24_aarch64.whl", hash = "sha256:c828719f0bb310a9cf0e0f08cb62a0b8cc550138617cb03ac897900aec9d3d47", size = 4105560, upload-time = "2026-04-22T19:50:29.324Z" }, - { url = "https://files.pythonhosted.org/packages/6a/36/4a4a6a398e5f145edd1886388ebe5e6f6bbaf74950a5dea1a6ceae63e6b5/obstore-0.9.4-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:49a0455519f284b6bc2e0694298114926aff1d1f3d5d344e9163e03b446826cc", size = 4292582, upload-time = "2026-04-22T19:50:31.028Z" }, - { url = "https://files.pythonhosted.org/packages/39/4c/9caa197cd2eba726e9a5285db34027049b9527a23e1a7e08479678ad6a4a/obstore-0.9.4-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:cf437309fc0fe852591ae50405300490229f876ea06574651fd753ca3fd23f25", size = 4261613, upload-time = "2026-04-22T19:50:32.868Z" }, - { url = "https://files.pythonhosted.org/packages/b8/94/a3fbe6fb3ee1c57fd4943ddbb21848eea3925b77e0789614c857d86b795e/obstore-0.9.4-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:d83dbd20b6a5d42e35794ef64046de39040854829ec4f1eb2f6dfb54df48cc3d", size = 4245638, upload-time = "2026-04-22T19:50:35.009Z" }, - { url = "https://files.pythonhosted.org/packages/56/a7/d18e168f318327d63512dfa7cf3b5e89ed9bfba6d6a8917ad7d4700b8657/obstore-0.9.4-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:5a0c337f37f30a2d66555d69bf3abd840457a279c57ede93bd02e014721ed364", size = 4437226, upload-time = "2026-04-22T19:50:36.635Z" }, - { url = "https://files.pythonhosted.org/packages/f4/ce/66aadd155db1e273c6ec2236c0fb904666d10c2e3b791b40624c272e586c/obstore-0.9.4-cp314-cp314t-win_amd64.whl", hash = "sha256:24e37a1c713c95a964e119f8ef879415a495432162e74e80ed29d645aeeca114", size = 4180746, upload-time = "2026-04-22T19:50:38.396Z" }, +sdist = { url = "https://files.pythonhosted.org/packages/2e/2f/f83afaab7945509d72245b2b00af0b4834ce78fdd2d9ae9f0ad1a3036a91/obstore-0.11.0.tar.gz", hash = "sha256:a2f55163bcd348b4a60d12e6893eac50eddc742bad8032a1705d49140b992204", size = 130565, upload-time = "2026-06-25T18:29:49.405Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fc/b2/00c213e7e5ca8065f97e37e55294adab836e3f6a88b23e4029069aaecf95/obstore-0.11.0-cp311-abi3-macosx_10_12_x86_64.whl", hash = "sha256:42f36546c7ac44dbab1173d2330a8a1b1a3f0e37950e553b8c904e3dd0744b25", size = 5491935, upload-time = "2026-06-25T18:28:32.029Z" }, + { url = "https://files.pythonhosted.org/packages/ac/37/6a6b9a5e15a8a37c24d14317a87648097c4888593b588510c03c030d2e90/obstore-0.11.0-cp311-abi3-macosx_11_0_arm64.whl", hash = "sha256:687bb9d3962d568b7c439c5d0c6fea19b2749862a8e5c8eebd0c058c4eccde9e", size = 4672619, upload-time = "2026-06-25T18:28:33.852Z" }, + { url = "https://files.pythonhosted.org/packages/28/f9/6745ce8c4f7bfac19dc14a4438b48a2e93a689b92b0cecfc695e41a4e8b1/obstore-0.11.0-cp311-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:010b51578c7514a41719d795cdb7a1e6529be509dac3772e477187a59422bb97", size = 5072806, upload-time = "2026-06-25T18:28:36.127Z" }, + { url = "https://files.pythonhosted.org/packages/6c/18/991d3b3cdd851c0225e55f3dc45b47fd9e249827d188995011469f805132/obstore-0.11.0-cp311-abi3-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:cfaa8129a3f5d8518a3a75184d4b02348db0f6263177cd1f0951f6568243cc9e", size = 5303777, upload-time = "2026-06-25T18:28:37.89Z" }, + { url = "https://files.pythonhosted.org/packages/8d/e9/90e56015a45b5e56a84fc3188c4e5fb088b288d41992c73a629e10df6760/obstore-0.11.0-cp311-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:c790a5cb9ff2970d1f464a6a708d734dce9939e9f668cb6708c5dba5d61589b2", size = 5493871, upload-time = "2026-06-25T18:28:39.981Z" }, + { url = "https://files.pythonhosted.org/packages/66/02/f1744091d59ce71c5523174eb860fbb298275c901e89b9ea6fbf3e654a33/obstore-0.11.0-cp311-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:827113e12fe8088e0281a9d57b90b2b8dbc8a6ffe3b15dadb9baa5feb3d266c1", size = 5361913, upload-time = "2026-06-25T18:28:42.089Z" }, + { url = "https://files.pythonhosted.org/packages/5d/59/3f47822683ee2b6db8685faa25829946d6343a561251ec2704548455d946/obstore-0.11.0-cp311-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a2ff6d3ed553298828fb760b4aef6347fbcc7b5c5e3ce3f8381ce805c370021a", size = 5638724, upload-time = "2026-06-25T18:28:43.897Z" }, + { url = "https://files.pythonhosted.org/packages/23/50/1df335fdf9b527b3933f1e94ab6fc720ad314260fab8591cb0b6668ff192/obstore-0.11.0-cp311-abi3-manylinux_2_24_aarch64.whl", hash = "sha256:39d04b324fcf984e7050734ebda77b81764025b0c011750201a0d8954087f7aa", size = 5413508, upload-time = "2026-06-25T18:28:45.624Z" }, + { url = "https://files.pythonhosted.org/packages/de/dc/a259aba149b841ca7c91fea177df9972a60a636b54077beed1a35b254994/obstore-0.11.0-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:37c0d15d775b1370ef5204ee3919a5ddf7e2592d11815213105f8db031f2ab8d", size = 5619995, upload-time = "2026-06-25T18:28:47.599Z" }, + { url = "https://files.pythonhosted.org/packages/e5/b4/ec25fdb4d6b060bc6eea647fc0e88f75fcc20fe8d16d67fb0dbe999d323b/obstore-0.11.0-cp311-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:7f468caf9b6e0f12ff151e5fe618de5fc9192befa9bd02734b06de4efd2e49f6", size = 5299512, upload-time = "2026-06-25T18:28:49.629Z" }, + { url = "https://files.pythonhosted.org/packages/a8/e5/29be060d06ec13e2af3d1b6cfb77b7c37f8be6c56b77295c945fefad73e4/obstore-0.11.0-cp311-abi3-musllinux_1_2_i686.whl", hash = "sha256:42d8e8fad85be8ee488c1a9a9b7c6a42128abb84e67175da40d3d1165c1846df", size = 5427026, upload-time = "2026-06-25T18:28:51.317Z" }, + { url = "https://files.pythonhosted.org/packages/57/b7/577a965f440e9ea64243518663f9d16be7df8eafc7123818e8e841fa21ce/obstore-0.11.0-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:9c8fd2a544e2e0b926669c47fcfb8d2314e234abc240ea165dae04ee42e1d7ac", size = 5869187, upload-time = "2026-06-25T18:28:53.166Z" }, + { url = "https://files.pythonhosted.org/packages/e2/18/8fdbaee22bfd5b9c44e1fdff8ca0508e2fe60c42bf9fc85f0c9c27b4ecf2/obstore-0.11.0-cp311-abi3-win_amd64.whl", hash = "sha256:6fb3d4678c0f4242d3109362e9b1df5d7b27765f43d5aacb2e81af53a75cb9ef", size = 5329384, upload-time = "2026-06-25T18:28:55.305Z" }, + { url = "https://files.pythonhosted.org/packages/8b/8b/7555e48ec768728fcfc71a051c6b28d6ddaf1bececf492ce5ef995aab5f0/obstore-0.11.0-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:f3132393eff9f3f2b543ecbb3bcc12319a7c433fef06493b4350d6854d505a14", size = 5515763, upload-time = "2026-06-25T18:28:57.527Z" }, + { url = "https://files.pythonhosted.org/packages/23/8f/94d83f3336421cbb5e436ab0ae5695eae72f7c82990d6b1ac090712c8052/obstore-0.11.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:a3a8da19b47af4c14ecc694209b3c18ac6d89f96be5656ee3a19b77947c14155", size = 4649491, upload-time = "2026-06-25T18:28:59.386Z" }, + { url = "https://files.pythonhosted.org/packages/fe/86/11f4e1f51a8c6cf21a5915c018d2357201ad3c5799d418f0c6529fafaab2/obstore-0.11.0-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:e91298b9b6c3a0408c28eece62bca6c5b6cda2f6350351d84e07b4dc8fb2631f", size = 5060659, upload-time = "2026-06-25T18:29:01.139Z" }, + { url = "https://files.pythonhosted.org/packages/49/e7/fd3036b0923d10e878e2073020f1ef692a618ed1cc3980d3e4a468c93713/obstore-0.11.0-cp314-cp314t-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:3422c532486671dfb5e3e739bf15ec9ca2a8da544a3c23b74ae3857dcab1c6a8", size = 5277058, upload-time = "2026-06-25T18:29:02.96Z" }, + { url = "https://files.pythonhosted.org/packages/2f/a7/b016c3ac6857ac856326dc0a292e3871b8500d03f25f3b90e168b05de357/obstore-0.11.0-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:3de027ce46cf0592c2b41654e2c27dbc52a4a726f6fbc511380acc0da3a9f658", size = 5475852, upload-time = "2026-06-25T18:29:05.032Z" }, + { url = "https://files.pythonhosted.org/packages/85/9e/644ffe8db7757de7f71f94a036ab24222bccb4de290d3ca69f76547e812d/obstore-0.11.0-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:1a80a95548678210bc336b866e37139c293565c3b163eb3fef2433d5d6640a33", size = 5363082, upload-time = "2026-06-25T18:29:06.878Z" }, + { url = "https://files.pythonhosted.org/packages/ae/0b/26af6b5fa6ba96af84086f44f78b4e5b0af1729c31402d7b28c68989d174/obstore-0.11.0-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:acaa261dc15efb95bbeca06f8fe9b47ee23d7302a6aa1fa3a9654baab8b23d7c", size = 5629116, upload-time = "2026-06-25T18:29:08.771Z" }, + { url = "https://files.pythonhosted.org/packages/ae/ce/f30d502991c6719b2fbd7b8385ef3e39da07bfca099108bcc5eeed8b9c20/obstore-0.11.0-cp314-cp314t-manylinux_2_24_aarch64.whl", hash = "sha256:a3300cabbc3129670987b3723629c791d83c117ef1b6a0c670c2043648e000a1", size = 5404534, upload-time = "2026-06-25T18:29:11.294Z" }, + { url = "https://files.pythonhosted.org/packages/52/20/d5bf5f816e868717ba647ed9a2109e800deb402d0265d410456b3fcb4376/obstore-0.11.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:f4e6a9480843645cd4ee122c41d5c5a46a56f1e9cdda85638826f2e0e439fe5c", size = 5613159, upload-time = "2026-06-25T18:29:13.414Z" }, + { url = "https://files.pythonhosted.org/packages/db/b4/6d4c1c211e3b06cc8554189e0d4406e8fa1f98ed55f9213b8e398a11599f/obstore-0.11.0-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:9fb2b1814c4314b8903f4e2ebbe8c3365fea6543669615ee9a0288b0d3a2edeb", size = 5286279, upload-time = "2026-06-25T18:29:15.424Z" }, + { url = "https://files.pythonhosted.org/packages/02/5e/d7b5589424a56171b16ab94cf92eb493490c300aaa044913bbdd94cace68/obstore-0.11.0-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:63fb9b072815eafe4705f617f567d896b1c858adcb390795fa1e269367791031", size = 5401780, upload-time = "2026-06-25T18:29:17.514Z" }, + { url = "https://files.pythonhosted.org/packages/c4/18/841baea8936e51a18b0e5d4c51f09c0a7798cb73b027e9794be2362a0f0b/obstore-0.11.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:086fafba314ff98cfab1c4bf7814699e862513e8889720cf6f7462296cb32787", size = 5853618, upload-time = "2026-06-25T18:29:19.354Z" }, + { url = "https://files.pythonhosted.org/packages/83/9a/d6127f5422b78e0222b0a9eadcfd7a5aa8d873a9498da7d4a77d4ac8ce2e/obstore-0.11.0-cp314-cp314t-win_amd64.whl", hash = "sha256:676d1154f6f08721110f9b7d14ee3a3c0293abaf9da135bb90f54e276dca1cac", size = 5314113, upload-time = "2026-06-25T18:29:21.209Z" }, ] [[package]] @@ -4057,6 +4062,23 @@ 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 = "zarrista" }, +] [package.metadata] requires-dist = [ @@ -4162,3 +4184,40 @@ test = [ { name = "tomlkit" }, { name = "uv" }, ] +zarrista = [ + { name = "coverage", specifier = ">=7.10" }, + { name = "hypothesis" }, + { name = "icechunk", specifier = ">=1.1.21" }, + { name = "numpydoc" }, + { name = "obstore", specifier = ">=0.10.1" }, + { 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 = "zarrista", git = "https://github.com/developmentseed/zarrista?rev=95e47ad4c414c5920f0cf15550f923039641da8e" }, +] + +[[package]] +name = "zarr-metadata" +version = "0.3.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f2/9c/cfd5aa02a27c63ecec702a77834b395411518da5c748414d7e6a323638ed/zarr_metadata-0.3.0.tar.gz", hash = "sha256:d8fe02feef43380056ea0429ceb50974b7b5afe6f0386853977506b034e89d53", size = 36398, upload-time = "2026-06-19T13:17:38.276Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/91/96/64137930fb40b96b4d207eb1f1e4e42c5d6c9e682a5a7c3e6feff6eb0e29/zarr_metadata-0.3.0-py3-none-any.whl", hash = "sha256:e6651f418fcc89cc3c6fc11aa852fb6f8dd6f31d62913abc0d4d37ce7302d671", size = 45636, upload-time = "2026-06-19T13:17:37.097Z" }, +] + +[[package]] +name = "zarrista" +version = "0.1.0b7" +source = { git = "https://github.com/developmentseed/zarrista?rev=95e47ad4c414c5920f0cf15550f923039641da8e#95e47ad4c414c5920f0cf15550f923039641da8e" } +dependencies = [ + { name = "zarr-metadata" }, +]