diff --git a/CLAUDE.md b/CLAUDE.md index e1a043d..4f57073 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -43,3 +43,20 @@ exposed to Python via `pyo3`. `from zarrista.codec._array_to_array import ArrayToArrayCodec`, not `from ._array_to_array import ...`. The package root is `zarrista` (maturin's `python-source = "python"`). + +## Documentation conventions + +- **The `.pyi` stubs are the single source of truth for user-facing docs.** The + docs site renders from them only (`allow_inspection: false` in `mkdocs.yml`), + so full prose, examples, and mkdocstrings cross-references + (`` [`Array.store_metadata`][zarrista.Array.store_metadata] ``) belong in the + `.pyi` and nowhere else. +- **A `///` doc comment on a Python-facing item gets the one-line summary only.** + On `#[pymethods]`, `#[pyclass]`, and `#[pyfunction]` items the doc comment + becomes Python's `__doc__` (what `help()` shows), which would otherwise + duplicate the stub. Don't restate the stub's prose there; a single summary line + is enough. Everything else in `src/` — plain `impl` blocks, `FromPyObject` + types, private helpers — is Rust-only and keeps full doc comments. +- **Implementation notes go in `//` comments, not `///`.** Rationale aimed at + Rust readers (why an upstream type was avoided, alignment caveats, safety + reasoning) must not leak into a Python docstring. diff --git a/dev-docs/specs/2026-07-30-array-with-shape-design.md b/dev-docs/specs/2026-07-30-array-with-shape-design.md new file mode 100644 index 0000000..b183991 --- /dev/null +++ b/dev-docs/specs/2026-07-30-array-with-shape-design.md @@ -0,0 +1,162 @@ +# `Array.with_shape` — resizing existing arrays + +Design for [#111](https://github.com/developmentseed/zarrista/issues/111). + +## Problem + +zarrista has no way to change the shape of an existing array. zarrs supplies the +underlying operation, but only the metadata half of it. + +## What zarrs provides + +`Array::set_shape(&mut self, ArrayShape) -> Result<&mut Self, ArrayCreateError>` +rebuilds the chunk grid from the array's existing chunk-grid metadata against the +new shape, then patches `metadata.shape` for V3 or V2. It errors when the new +shape is incompatible with the chunk grid. + +It is purely in-memory. Persisting is a separate `store_metadata()` / +`async_store_metadata()`, both of which zarrista already wraps. + +zarrs does *not* delete chunks that fall outside a shrunken array's bounds, and +has no reload/refresh on `Array` — `Array::open` is the only way to re-read +metadata, and it constructs a fresh object. + +## The model + +zarrista's `Array` is a cached snapshot of `zarr.json` plus the objects derived +from it (chunk grid, codecs, data type, fill value), taken at `open()` and never +revisited. That cache is load-bearing: it is why `arr.shape` is a field read +rather than a network round-trip. + +The design question is therefore not "mutable or immutable" but **what the sync +points are, and whether the API has one in each direction**: + +| Direction | Cause | Answer | +| --- | --- | --- | +| Local ahead of store | `from_metadata()`, `with_shape()` | `store_metadata()` | +| Store ahead of local | another process, another handle, an Icechunk commit | none today | + +The two are not symmetric. Local-ahead is *authored* — deliberate, with a fix to +hand. Store-ahead is *ambient* — nothing local causes it and there is no local +signal it happened. Closing that second gap is out of scope here; see +[Deferred](#deferred). + +## Decision: return a new array + +`with_shape` returns a new `Array` and leaves the receiver untouched. It does not +mutate in place. + +Rebinding is the intended usage, and under it no stale handle exists: + +```python +arr = Array.open(store) +arr = arr.with_shape([8, 8]) +arr.store_metadata() +``` + +This follows established Python convention (`Path.with_suffix`, +`datetime.replace`) and matches `read_only()`, which already derives a new +`Array` from an existing one in this same class. + +### Why not mutate in place + +Mutation would match zarr-python's `Array.resize()`, and would keep a +non-rebound handle correct. It was rejected because: + +- `#[pyclass(frozen)]` would need interior mutability + (`RwLock>>`), touching all 57 `self.inner` sites across + `sync.rs`, `async.rs`, and `shared.rs`. +- `shape()`, `chunk_grid_shape()`, and `dimension_names()` currently return + borrowed data, which is impossible when borrowing from a temporary guard. + They would have to return owned values. +- Locking is per-accessor, so a Python-level sequence such as + `(arr.shape, arr.chunk_grid_shape)` could straddle a mutation and observe a + mix of before and after. There is no way to hold the lock across that from + Python. An immutable object is internally consistent for its whole lifetime. + +Stale handles held elsewhere in a program are possible under either design, so +mutation does not solve that problem — it only shifts which object is correct. + +## Implementation + +`with_shape` performs no I/O and is byte-identical for `PyArray` and +`PyAsyncArray`, so it belongs in the `shared_array_methods!` macro in +`src/array/shared.rs` alongside the metadata accessors — not in `sync.rs` / +`async.rs`. Both classes expose a `pub(crate) fn new(Arc>, store)` and +a `store` field, and the macro expands in each class's own module, so the same +body compiles for both. + +```rust +/// Return a new array with `shape`, leaving this one unchanged. +fn with_shape( + &self, + shape: $crate::array::PyArrayShape, +) -> $crate::error::ZarristaResult { + let mut resized = self.inner.with_storage(self.inner.storage()); + resized.set_shape(shape)?; + Ok(Self::new(::std::sync::Arc::new(resized), self.store.clone())) +} +``` + +`Array` is not `Clone`; `with_storage` supplies the deep copy, and passing the +array's own storage back leaves the storage type parameter unchanged. + +`AsyncArray.with_shape` is a **sync** method — there is nothing to await. The +stub must say so, since every other `AsyncArray` method that touches the store is +a coroutine. + +The macro's module doc comment currently reads "These accessors only read from +`self.inner` and perform no I/O". `with_shape` is not an accessor, so the wording +needs a light revision; the no-I/O invariant still holds. + +## Error handling + +`set_shape` returns `ArrayCreateError` when the shape does not fit the chunk +grid — wrong dimensionality, or a `rectangular` grid whose explicit boundaries do +not sum to the new extent. `ZarristaError` already converts from +`ArrayCreateError` (`Array::open` relies on it), so this needs no new error type. + +There is no manual validation in the function body, per the project's +"no manual validation in function bodies" rule: zarrs is the validator. Because +the error propagates before `Self::new`, a rejected shape never yields an object. + +## Documented semantics + +The `.pyi` prose must be explicit about three points, each of which a user could +reasonably assume otherwise: + +1. Nothing is written. `store_metadata()` is the push, as after + `from_metadata()`. +2. The receiver is unchanged and remains valid. Rebinding is the intended usage. +3. Shrinking leaves chunks outside the new bounds in the store. Reclaiming them + will be a separate vacuum call. + +## Testing + +`tests/array/test_with_shape.py`, following the `tests/array/test_store_metadata.py` +layout — a `_metadata()` helper, sync tests, then an async section. + +- returns a new array at the new shape; the original still reports the old shape +- grow → `store_metadata()` → reopen reports the new shape +- shrink → `store_metadata()` → reopen reports the new shape +- data written before a grow reads back at the same coordinates afterward +- wrong dimensionality raises `ZarristaError` +- write a chunk, shrink so that chunk falls outside the new bounds, persist — + the chunk is still present in the store afterward. Pins current behavior so the + future vacuum call has something to flip. +- async: grow → `await store_metadata()` → reopen reports the new shape + +## Deferred + +Out of scope for #111, each warranting its own issue: + +- **Vacuum.** A call that erases chunks outside the array's current bounds. + zarr-python folds this into `resize()`; here it must be separate, because + `with_shape` does no I/O at all. zarrs supplies `erase_chunks(&dyn + ArraySubsetTraits)` (and `async_erase_chunks`) to build on. +- **A pull sync point.** A `reopen()` returning a fresh handle from a re-read of + `zarr.json`, closing the store-ahead-of-local gap. +- **Optimistic concurrency.** Track an ETag at open and use a conditional PUT in + `store_metadata()`, raising on mismatch rather than silently clobbering a + concurrent writer. This is the real fix for lost updates, and obstore supports + conditional put, but it has its own semantics to design. diff --git a/dev-docs/specs/README.md b/dev-docs/specs/README.md new file mode 100644 index 0000000..d5b8df1 --- /dev/null +++ b/dev-docs/specs/README.md @@ -0,0 +1,3 @@ +The full verbose history of specs are included here for future agent context. + +Most human readers should not read these specs. diff --git a/python/zarrista/_array.pyi b/python/zarrista/_array.pyi index a99ae85..c9094b5 100644 --- a/python/zarrista/_array.pyi +++ b/python/zarrista/_array.pyi @@ -213,6 +213,59 @@ class Array: Reads behave identically, but any write (`store_chunk`, `erase_chunk`, `erase_metadata`, ...) raises at runtime. """ + def with_chunk_grid(self, chunk_grid: ChunkGrid) -> Array: + """Return a new array reference with `chunk_grid`. + + This does not mutate the existing `array`. + + The new array's shape comes from the grid, so this can change the shape + and the chunking together: + + ```py + array = array.with_chunk_grid(ChunkGrid.regular([8, 8], [4, 4])) + array.store_metadata() + ``` + + Nothing is persisted to the store. Call + [`Array.store_metadata`][zarrista.Array.store_metadata] to persist the new + grid to the store. This array is unaffected and remains usable; rebinding, + as above, is the intended usage. + + **Existing chunks are neither migrated nor erased.** If the chunk shape + changes, chunks already in the store sit at keys that no longer describe + the same region of the array, so later reads may fail to decode or return + wrong data. It is the caller's responsibility to ensure the new grid is + compatible with whatever is already stored. The safe uses are setting the + grid before any chunks are written, or erasing and rewriting the existing + chunks yourself. + + If only the array shape is changing, use + [`Array.with_shape`][zarrista.Array.with_shape] instead — it preserves the + chunking, so existing chunks stay valid. + """ + def with_shape(self, shape: list[int]) -> Array: + """Return a new array reference with `shape`, leaving this one unchanged. + + Nothing is persisted to the store. Call + [`Array.store_metadata`][zarrista.Array.store_metadata] to persist the new + shape to the store: + + ```py + array = array.with_shape([8, 8]) + array.store_metadata() + ``` + + This array is unaffected and remains usable; it simply goes on describing + the old shape. Rebinding, as above, is the intended usage. + + Growing an array leaves the new region reading as the fill value. Shrinking + leaves any chunks outside the new bounds in the store, where they are no + longer addressable through this array; reclaiming that space will be a + separate call. + + Raises a `ZarristaError` if `shape` is incompatible with the array's chunk + grid, such as a shape of the wrong dimensionality. + """ @property def is_sharded(self) -> bool: """Whether the array's array-to-bytes codec is `sharding_indexed`.""" @@ -447,6 +500,58 @@ class AsyncArray: Reads behave identically, but any write (`store_chunk`, `erase_chunk`, `erase_metadata`, ...) raises at runtime. """ + def with_chunk_grid(self, chunk_grid: ChunkGrid) -> AsyncArray: + """Return a new array reference with `chunk_grid`, leaving this one unchanged. + + This method is synchronous: it performs no I/O. The new array's shape comes + from the grid, so this can change the shape and the chunking together: + + ```py + array = array.with_chunk_grid(ChunkGrid.regular([8, 8], [4, 4])) + await array.store_metadata() + ``` + + Nothing is persisted to the store. Call + [`AsyncArray.store_metadata`][zarrista.AsyncArray.store_metadata] to persist + the new grid to the store. This array is unaffected and remains usable; + rebinding, as above, is the intended usage. + + **Existing chunks are neither migrated nor erased.** If the chunk shape + changes, chunks already in the store sit at keys that no longer describe + the same region of the array, so later reads may fail to decode or return + wrong data. It is the caller's responsibility to ensure the new grid is + compatible with whatever is already stored. The safe uses are setting the + grid before any chunks are written, or erasing and rewriting the existing + chunks yourself. + + If only the array shape is changing, use + [`AsyncArray.with_shape`][zarrista.AsyncArray.with_shape] instead — it + preserves the chunking, so existing chunks stay valid. + """ + def with_shape(self, shape: list[int]) -> AsyncArray: + """Return a new array reference with `shape`, leaving this one unchanged. + + This method is synchronous: it performs no I/O. Nothing is persisted to the + store. Call + [`AsyncArray.store_metadata`][zarrista.AsyncArray.store_metadata] to persist + the new shape to the store: + + ```py + array = array.with_shape([8, 8]) + await array.store_metadata() + ``` + + This array is unaffected and remains usable; it simply goes on describing + the old shape. Rebinding, as above, is the intended usage. + + Growing an array leaves the new region reading as the fill value. Shrinking + leaves any chunks outside the new bounds in the store, where they are no + longer addressable through this array; reclaiming that space will be a + separate call. + + Raises a `ZarristaError` if `shape` is incompatible with the array's chunk + grid, such as a shape of the wrong dimensionality. + """ @property def is_sharded(self) -> bool: """Whether the array's array-to-bytes codec is `sharding_indexed`.""" diff --git a/src/array/shared.rs b/src/array/shared.rs index eb790d5..8012dfc 100644 --- a/src/array/shared.rs +++ b/src/array/shared.rs @@ -1,6 +1,6 @@ //! Shared `#[pymethods]` for `PyArray` / `PyAsyncArray`. //! -//! These accessors only read from `self.inner` and perform no I/O, so they are +//! These methods only read from `self.inner` and perform no I/O, so they are //! identical between the sync and async variants. The macro emits a separate //! `#[pymethods]` block (requires the `multiple-pymethods` pyo3 feature). macro_rules! shared_array_methods { @@ -164,6 +164,46 @@ macro_rules! shared_array_methods { fn subset_all(&self) -> $crate::array::PyArraySubset { self.inner.subset_all().into() } + + /// Return a new array reference with `chunk_grid`, leaving this one unchanged. + fn with_chunk_grid( + &self, + chunk_grid: $crate::array::PyChunkGrid, + ) -> $crate::error::ZarristaResult { + let chunk_grid = chunk_grid.into_inner(); + // Workaround for missing Clone + let mut regridded = self.inner.with_storage(self.inner.storage()); + + // SAFETY: existing chunks are not checked against the new grid. + // + // This call touches no storage, so it invalidates nothing on its + // own; the hazard is documented on the Python type stub. + unsafe { + regridded.set_shape_and_chunk_grid( + chunk_grid.array_shape().to_vec(), + chunk_grid.metadata(), + )?; + } + + Ok(Self::new( + ::std::sync::Arc::new(regridded), + self.store.clone(), + )) + } + + /// Return a new array reference with `shape`, leaving this one unchanged. + fn with_shape( + &self, + shape: $crate::array::PyArrayShape, + ) -> $crate::error::ZarristaResult { + // Workaround for missing Clone + let mut resized = self.inner.with_storage(self.inner.storage()); + resized.set_shape(shape)?; + Ok(Self::new( + ::std::sync::Arc::new(resized), + self.store.clone(), + )) + } } }; } diff --git a/tests/array/test_with_chunk_grid.py b/tests/array/test_with_chunk_grid.py new file mode 100644 index 0000000..3c16e9e --- /dev/null +++ b/tests/array/test_with_chunk_grid.py @@ -0,0 +1,113 @@ +"""`Array.with_chunk_grid()` returns a new array with a different chunk grid.""" + +from pathlib import Path + +import numpy as np +from obstore.store import LocalStore + +from zarrista import ( + Array, + ArrayBuilder, + ArrayBytes, + AsyncArray, + ChunkGrid, + DataType, + FillValue, +) +from zarrista.store import MemoryStore + + +def _array(store: MemoryStore) -> Array: + """A 4x4 int8 array in `store` at `/a`, chunked 2x2, fill 0. + + `ArrayBuilder.create` writes the metadata, so the array is openable. + """ + return ArrayBuilder( + ChunkGrid.regular([4, 4], [2, 2]), + DataType.from_string("int8"), + FillValue(b"\x00"), + ).create(store, "/a") + + +def test_returns_new_array_and_leaves_original() -> None: + array = _array(MemoryStore()) + + regridded = array.with_chunk_grid(ChunkGrid.regular([8, 8], [4, 4])) + + assert regridded.shape == [8, 8] + assert regridded.chunk_grid.metadata == ChunkGrid.regular([8, 8], [4, 4]).metadata + assert array.shape == [4, 4] + assert array.chunk_grid.metadata == ChunkGrid.regular([4, 4], [2, 2]).metadata + + +def test_does_not_write() -> None: + store = MemoryStore() + array = _array(store) + + array.with_chunk_grid(ChunkGrid.regular([8, 8], [4, 4])) + + # The stored metadata is untouched until `store_metadata` is called. + reopened = Array.open(store, "/a") + assert reopened.shape == [4, 4] + assert reopened.chunk_grid.metadata == ChunkGrid.regular([4, 4], [2, 2]).metadata + + +def test_chunk_shape_change_persists() -> None: + store = MemoryStore() + _array(store).with_chunk_grid(ChunkGrid.regular([4, 4], [4, 4])).store_metadata() + + reopened = Array.open(store, "/a") + assert reopened.shape == [4, 4] + assert reopened.chunk_grid.metadata == ChunkGrid.regular([4, 4], [4, 4]).metadata + + +def test_shape_may_change_too() -> None: + store = MemoryStore() + _array(store).with_chunk_grid(ChunkGrid.regular([8, 8], [2, 2])).store_metadata() + + assert Array.open(store, "/a").shape == [8, 8] + + +def test_rectilinear_grid_accepted() -> None: + array = _array(MemoryStore()) + grid = ChunkGrid.rectilinear([4, 4], [2, 2]) + + regridded = array.with_chunk_grid(grid) + + assert regridded.chunk_grid.metadata == grid.metadata + + +def test_existing_chunks_are_not_migrated() -> None: + # Pins the hazard the docs describe: regridding rewrites metadata only, so + # bytes already in the store stay exactly as they were, under keys that no + # longer describe the same region. + store = MemoryStore() + array = _array(store) + # One 2x2 int8 chunk is 4 bytes. + array.store_chunk([0, 0], ArrayBytes(np.arange(4, dtype="int8").tobytes())) + + # A 4x4 chunk would be 16 bytes. + array.with_chunk_grid(ChunkGrid.regular([4, 4], [4, 4])).store_metadata() + + # `array` still describes the 2x2 grid, so it can still address the old chunk. + # It is untouched: still 4 bytes, not re-encoded to 16. + assert len(array.retrieve_encoded_chunk([0, 0])) == 4 + + +# --- async -------------------------------------------------------------------- + + +async def test_async_with_chunk_grid(tmp_path: Path) -> None: + array = await ArrayBuilder( + ChunkGrid.regular([4, 4], [2, 2]), + DataType.from_string("int8"), + FillValue(b"\x00"), + ).create_async(LocalStore(str(tmp_path)), "/a") + + # `with_chunk_grid` is sync even on `AsyncArray`: it performs no I/O. + regridded = array.with_chunk_grid(ChunkGrid.regular([4, 4], [4, 4])) + await regridded.store_metadata() + + assert array.chunk_grid.metadata == ChunkGrid.regular([4, 4], [2, 2]).metadata + reopened = await AsyncArray.open_async(LocalStore(str(tmp_path)), "/a") + assert reopened.chunk_grid.metadata == ChunkGrid.regular([4, 4], [4, 4]).metadata diff --git a/tests/array/test_with_shape.py b/tests/array/test_with_shape.py new file mode 100644 index 0000000..3ff1df2 --- /dev/null +++ b/tests/array/test_with_shape.py @@ -0,0 +1,115 @@ +"""`Array.with_shape()` returns a new array at a new shape, without writing.""" + +from pathlib import Path + +import numpy as np +import pytest +from obstore.store import LocalStore + +from zarrista import ( + Array, + ArrayBuilder, + ArrayBytes, + AsyncArray, + ChunkGrid, + DataType, + FillValue, +) +from zarrista.exceptions import ArrayCreateError +from zarrista.store import MemoryStore + + +def _array(store: MemoryStore) -> Array: + """A 4x4 int8 array in `store` at `/a`, chunked 2x2, fill 0. + + `ArrayBuilder.create` writes the metadata, so the array is openable. + """ + return ArrayBuilder( + ChunkGrid.regular([4, 4], [2, 2]), + DataType.from_string("int8"), + FillValue(b"\x00"), + ).create(store, "/a") + + +def test_with_shape_returns_new_array_and_leaves_original() -> None: + array = _array(MemoryStore()) + + resized = array.with_shape([8, 8]) + + assert resized.shape == [8, 8] + assert array.shape == [4, 4] + + +def test_with_shape_does_not_write() -> None: + store = MemoryStore() + array = _array(store) + + array.with_shape([8, 8]) + + # The stored metadata is untouched until `store_metadata` is called. + assert Array.open(store, "/a").shape == [4, 4] + + +def test_grow_then_store_metadata_persists() -> None: + store = MemoryStore() + _array(store).with_shape([8, 8]).store_metadata() + + assert Array.open(store, "/a").shape == [8, 8] + + +def test_shrink_then_store_metadata_persists() -> None: + store = MemoryStore() + _array(store).with_shape([2, 2]).store_metadata() + + assert Array.open(store, "/a").shape == [2, 2] + + +def test_data_survives_a_grow() -> None: + store = MemoryStore() + array = _array(store) + chunk = np.arange(4, dtype="int8").reshape(2, 2) + array.store_chunk([0, 0], ArrayBytes(chunk.tobytes())) + + array.with_shape([8, 8]).store_metadata() + + grown = Array.open(store, "/a") + np.testing.assert_array_equal(grown.retrieve_chunk([0, 0]).to_numpy(), chunk) + + +def test_shrink_leaves_out_of_bounds_chunks() -> None: + # Pins current behavior: reclaiming these chunks is a future vacuum call. + store = MemoryStore() + array = _array(store) + array.store_chunk([1, 1], ArrayBytes(np.ones(4, dtype="int8").tobytes())) + + array.with_shape([2, 2]).store_metadata() + + # `array` still describes the 4x4 shape, so it can still address chunk [1, 1]. + assert array.retrieve_encoded_chunk([1, 1]) is not None + + +@pytest.mark.parametrize("shape", [[8, 8, 8], [8]]) +def test_wrong_dimensionality_raises(shape: list[int]) -> None: + array = _array(MemoryStore()) + + with pytest.raises(ArrayCreateError, match="inconsistent dimensionality"): + array.with_shape(shape) + + +# --- async -------------------------------------------------------------------- + + +async def test_async_with_shape_returns_new_array(tmp_path: Path) -> None: + array = await ArrayBuilder( + ChunkGrid.regular([4, 4], [2, 2]), + DataType.from_string("int8"), + FillValue(b"\x00"), + ).create_async(LocalStore(str(tmp_path)), "/a") + + # `with_shape` is sync even on `AsyncArray`: it performs no I/O. + resized = array.with_shape([8, 8]) + await resized.store_metadata() + + assert array.shape == [4, 4] + reopened = await AsyncArray.open_async(LocalStore(str(tmp_path)), "/a") + assert reopened.shape == [8, 8]