From 8f8f13cdc50ffed68cc686a4257aedac54470528 Mon Sep 17 00:00:00 2001 From: eloramirez1356 Date: Tue, 28 Jul 2026 15:46:35 -0500 Subject: [PATCH 1/3] perf: fast-path append-only leading-dimension resize --- py_hamt/sharded_zarr_store.py | 25 ++- tests/test_z16_fast_leading_resize.py | 240 ++++++++++++++++++++++++++ 2 files changed, 264 insertions(+), 1 deletion(-) create mode 100644 tests/test_z16_fast_leading_resize.py 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..ae1d70c --- /dev/null +++ b/tests/test_z16_fast_leading_resize.py @@ -0,0 +1,240 @@ +import json +from unittest.mock import AsyncMock + +import numpy as np +import pytest +import xarray as xr +import zarr.core.buffer +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() + + 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: + 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 From e4b3b439687072dbfb06aa919bfb7cb0a53cf48a Mon Sep 17 00:00:00 2001 From: eloramirez1356 Date: Wed, 29 Jul 2026 16:24:55 -0500 Subject: [PATCH 2/3] test: clean up fast resize coverage --- tests/test_z16_fast_leading_resize.py | 25 +++++++++---------------- 1 file changed, 9 insertions(+), 16 deletions(-) diff --git a/tests/test_z16_fast_leading_resize.py b/tests/test_z16_fast_leading_resize.py index ae1d70c..647d794 100644 --- a/tests/test_z16_fast_leading_resize.py +++ b/tests/test_z16_fast_leading_resize.py @@ -5,6 +5,7 @@ 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 ( @@ -54,9 +55,7 @@ async def test_v1_leading_growth_skips_snapshot_and_preserves_chunks( old_shard_cids = list(writable.array_indices[""].shard_cids) snapshot = AsyncMock( - side_effect=AssertionError( - "leading-dimension growth must not snapshot shards" - ) + side_effect=AssertionError("leading-dimension growth must not snapshot shards") ) monkeypatch.setattr(writable, "_snapshot_shards_for_resize", snapshot) @@ -69,9 +68,7 @@ async def test_v1_leading_growth_skips_snapshot_and_preserves_chunks( 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") - ) + 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, @@ -127,9 +124,7 @@ async def test_v2_leading_growth_is_fast_and_array_local( humidity_manifest = writable.array_indices["humidity"].to_manifest() snapshot = AsyncMock( - side_effect=AssertionError( - "leading-dimension growth must not snapshot shards" - ) + side_effect=AssertionError("leading-dimension growth must not snapshot shards") ) monkeypatch.setattr(writable, "_snapshot_shards_for_resize", snapshot) @@ -139,9 +134,7 @@ async def test_v2_leading_growth_is_fast_and_array_local( 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") - ) + 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, @@ -188,9 +181,7 @@ async def test_v1_xarray_append_persists_visible_shape_without_snapshot( ) snapshot = AsyncMock( - side_effect=AssertionError( - "leading-dimension growth must not snapshot shards" - ) + side_effect=AssertionError("leading-dimension growth must not snapshot shards") ) monkeypatch.setattr(writable, "_snapshot_shards_for_resize", snapshot) @@ -227,7 +218,9 @@ async def test_non_leading_change_and_shrink_use_general_resize( snapshot_calls = 0 original_snapshot = store._snapshot_shards_for_resize - async def count_snapshot(array_index: ArrayIndex) -> dict: + async def count_snapshot( + array_index: ArrayIndex, + ) -> dict[int, list[CID | None]]: nonlocal snapshot_calls snapshot_calls += 1 return await original_snapshot(array_index) From 24aff8e8c7088e5305d75274bcaa8d089dea69e3 Mon Sep 17 00:00:00 2001 From: eloramirez1356 Date: Wed, 29 Jul 2026 16:37:56 -0500 Subject: [PATCH 3/3] test: prevent aliasing in fast resize coverage --- tests/test_z16_fast_leading_resize.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_z16_fast_leading_resize.py b/tests/test_z16_fast_leading_resize.py index 647d794..a2e3f2e 100644 --- a/tests/test_z16_fast_leading_resize.py +++ b/tests/test_z16_fast_leading_resize.py @@ -122,6 +122,7 @@ async def test_v2_leading_growth_is_fast_and_array_local( 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")