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