Optimize append-only leading-dimension resize performance

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,

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