diff --git a/py_hamt/sharded_zarr_store.py b/py_hamt/sharded_zarr_store.py index 7d0d072..3f9ed01 100644 --- a/py_hamt/sharded_zarr_store.py +++ b/py_hamt/sharded_zarr_store.py @@ -1223,15 +1223,38 @@ async def _replace_shards_after_resize( async with shard_lock: await self._shard_data_cache.discard(cache_key) + @staticmethod + def _can_fast_resize_leading_dimension( + array_index: ArrayIndex, new_shape: tuple[int, ...] + ) -> bool: + """Return whether resizing only appends to the row-major chunk grid.""" + old_shape = array_index.array_shape + return ( + len(old_shape) > 0 + and len(new_shape) == len(old_shape) + and new_shape[0] >= old_shape[0] + and new_shape[1:] == old_shape[1:] + and array_index.order == "C" + ) + async def _resize_array_index( self, array_index: ArrayIndex, new_shape: tuple[int, ...] ) -> None: + new_shape = tuple(new_shape) + if self._can_fast_resize_leading_dimension(array_index, new_shape): + array_index.resize(new_shape) + if self._primary_array_path == array_index.array_path: + self._set_legacy_geometry_from_index(array_index) + self._sync_arrays_to_root() + self._dirty_root = True + return + old_num_shards = array_index.num_shards old_total_chunks = array_index.total_chunks old_chunks_per_dim = array_index.chunks_per_dim old_shard_cids = list(array_index.shard_cids) old_shards_by_index = await self._snapshot_shards_for_resize(array_index) - array_index.resize(tuple(new_shape)) + array_index.resize(new_shape) new_shards_by_index = self._remap_shards_for_resize( old_shards_by_index, old_chunks_per_dim, diff --git a/tests/test_z16_fast_leading_resize.py b/tests/test_z16_fast_leading_resize.py new file mode 100644 index 0000000..a2e3f2e --- /dev/null +++ b/tests/test_z16_fast_leading_resize.py @@ -0,0 +1,234 @@ +import json +from unittest.mock import AsyncMock + +import numpy as np +import pytest +import xarray as xr +import zarr.core.buffer +from multiformats import CID +from testing_utils import CIDInMemoryCAS + +from py_hamt.sharded_zarr_store import ( + SHARDED_ZARR_V2, + ArrayIndex, + ShardedZarrStore, +) + + +def test_fast_leading_resize_requires_strict_append_only_growth() -> None: + index = ArrayIndex.new( + array_path="temperature", + array_shape=(2, 4, 6), + chunk_shape=(1, 2, 3), + chunks_per_shard=4, + ) + + assert ShardedZarrStore._can_fast_resize_leading_dimension(index, (3, 4, 6)) + assert ShardedZarrStore._can_fast_resize_leading_dimension(index, (2, 4, 6)) + assert not ShardedZarrStore._can_fast_resize_leading_dimension(index, (1, 4, 6)) + assert not ShardedZarrStore._can_fast_resize_leading_dimension(index, (3, 5, 6)) + assert not ShardedZarrStore._can_fast_resize_leading_dimension(index, (3, 4)) + + +@pytest.mark.asyncio +async def test_v1_leading_growth_skips_snapshot_and_preserves_chunks( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cas = CIDInMemoryCAS() + prototype = zarr.core.buffer.default_buffer_prototype() + store = await ShardedZarrStore.open( + cas=cas, + read_only=False, + array_shape=(2, 2), + chunk_shape=(1, 1), + chunks_per_shard=3, + ) + await store.set("temperature/c/0/0", prototype.buffer.from_bytes(b"old-0-0")) + await store.set("temperature/c/1/1", prototype.buffer.from_bytes(b"old-1-1")) + root_cid = await store.flush() + + writable = await ShardedZarrStore.open( + cas=cas, + read_only=False, + root_cid=root_cid, + ) + old_shard_cids = list(writable.array_indices[""].shard_cids) + + snapshot = AsyncMock( + side_effect=AssertionError("leading-dimension growth must not snapshot shards") + ) + monkeypatch.setattr(writable, "_snapshot_shards_for_resize", snapshot) + + await writable.resize_store((3, 2)) + + snapshot.assert_not_awaited() + resized_index = writable.array_indices[""] + assert resized_index.array_shape == (3, 2) + assert resized_index.shard_cids[: len(old_shard_cids)] == old_shard_cids + assert writable._root_obj["chunks"]["array_shape"] == [3, 2] + assert writable._root_obj["chunks"]["shard_cids"] == resized_index.shard_cids + + await writable.set("temperature/c/2/0", prototype.buffer.from_bytes(b"new-2-0")) + resized_root_cid = await writable.flush() + reopened = await ShardedZarrStore.open( + cas=cas, + read_only=True, + root_cid=resized_root_cid, + ) + + old_chunk = await reopened.get("temperature/c/0/0", prototype) + appended_chunk = await reopened.get("temperature/c/2/0", prototype) + assert old_chunk is not None + assert appended_chunk is not None + assert old_chunk.to_bytes() == b"old-0-0" + assert appended_chunk.to_bytes() == b"new-2-0" + + +@pytest.mark.asyncio +async def test_v2_leading_growth_is_fast_and_array_local( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cas = CIDInMemoryCAS() + prototype = zarr.core.buffer.default_buffer_prototype() + store = await ShardedZarrStore.open( + cas=cas, + read_only=False, + chunks_per_shard=3, + manifest_version=SHARDED_ZARR_V2, + ) + for array_path in ("temperature", "humidity"): + metadata = { + "zarr_format": 3, + "node_type": "array", + "shape": [2, 2], + "chunk_grid": { + "name": "regular", + "configuration": {"chunk_shape": [1, 1]}, + }, + } + await store.set( + f"{array_path}/zarr.json", + prototype.buffer.from_bytes(json.dumps(metadata).encode()), + ) + await store.set( + f"{array_path}/c/0/0", + prototype.buffer.from_bytes(array_path.encode()), + ) + root_cid = await store.flush() + + writable = await ShardedZarrStore.open( + cas=cas, + read_only=False, + root_cid=root_cid, + ) + humidity_manifest = writable.array_indices["humidity"].to_manifest() + humidity_manifest["shard_cids"] = list(humidity_manifest["shard_cids"]) + + snapshot = AsyncMock( + side_effect=AssertionError("leading-dimension growth must not snapshot shards") + ) + monkeypatch.setattr(writable, "_snapshot_shards_for_resize", snapshot) + + await writable.resize_store((3, 2), array_path="temperature") + + snapshot.assert_not_awaited() + assert writable.array_indices["temperature"].array_shape == (3, 2) + assert writable.array_indices["humidity"].to_manifest() == humidity_manifest + + await writable.set("temperature/c/2/1", prototype.buffer.from_bytes(b"appended")) + resized_root_cid = await writable.flush() + reopened = await ShardedZarrStore.open( + cas=cas, + read_only=True, + root_cid=resized_root_cid, + ) + + old_chunk = await reopened.get("temperature/c/0/0", prototype) + appended_chunk = await reopened.get("temperature/c/2/1", prototype) + assert old_chunk is not None + assert appended_chunk is not None + assert old_chunk.to_bytes() == b"temperature" + assert appended_chunk.to_bytes() == b"appended" + assert reopened.array_indices["temperature"].array_shape == (3, 2) + assert reopened.array_indices["humidity"].array_shape == (2, 2) + + +@pytest.mark.asyncio +async def test_v1_xarray_append_persists_visible_shape_without_snapshot( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cas = CIDInMemoryCAS() + initial = xr.Dataset( + {"temperature": (("time", "x"), np.arange(4).reshape(2, 2))}, + coords={"time": np.arange(2), "x": np.arange(2)}, + ).chunk({"time": 1, "x": 1}) + appended = xr.Dataset( + {"temperature": (("time", "x"), np.arange(4, 8).reshape(2, 2))}, + coords={"time": np.arange(2, 4), "x": np.arange(2)}, + ).chunk({"time": 1, "x": 1}) + store = await ShardedZarrStore.open( + cas=cas, + read_only=False, + array_shape=(2, 2), + chunk_shape=(1, 1), + chunks_per_shard=3, + ) + initial.to_zarr(store=store, mode="w") + root_cid = await store.flush() + writable = await ShardedZarrStore.open( + cas=cas, + read_only=False, + root_cid=root_cid, + ) + + snapshot = AsyncMock( + side_effect=AssertionError("leading-dimension growth must not snapshot shards") + ) + monkeypatch.setattr(writable, "_snapshot_shards_for_resize", snapshot) + + appended.to_zarr(store=writable, append_dim="time") + + snapshot.assert_not_awaited() + resized_root_cid = await writable.flush() + reopened = await ShardedZarrStore.open( + cas=cas, + read_only=True, + root_cid=resized_root_cid, + ) + actual = xr.open_zarr(store=reopened).compute() + expected = xr.concat([initial, appended], dim="time").compute() + + assert actual.sizes == {"time": 4, "x": 2} + xr.testing.assert_identical(actual, expected) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("new_shape", [(2, 3), (1, 2)]) +async def test_non_leading_change_and_shrink_use_general_resize( + new_shape: tuple[int, int], + monkeypatch: pytest.MonkeyPatch, +) -> None: + cas = CIDInMemoryCAS() + store = await ShardedZarrStore.open( + cas=cas, + read_only=False, + array_shape=(2, 2), + chunk_shape=(1, 1), + chunks_per_shard=3, + ) + snapshot_calls = 0 + original_snapshot = store._snapshot_shards_for_resize + + async def count_snapshot( + array_index: ArrayIndex, + ) -> dict[int, list[CID | None]]: + nonlocal snapshot_calls + snapshot_calls += 1 + return await original_snapshot(array_index) + + monkeypatch.setattr(store, "_snapshot_shards_for_resize", count_snapshot) + + await store.resize_store(new_shape) + + assert snapshot_calls == 1 + assert store.array_indices[""].array_shape == new_shape