Fix upstream-dev CI job: zarr test overrides, stateful testing, upstream GHA

.github/workflows/python-upstream.yaml

@@ -219,7 +219,7 @@ jobs:
 
       - name: Generate and publish the pytest report
         if: needs.build.outputs.pytest-outcome == 'failure'
-        uses: xarray-contrib/issue-from-pytest-log@87351a8f864e969567cda22a25a2f214cbe2340f # v1.6.0
+        uses: xarray-contrib/issue-from-pytest-log@054799b34bd75a5fd6c86277a4a8a575224e60c6 # v1.6.1
         with:
           log-path: icechunk-python/output-pytest-log.jsonl
 
@@ -294,7 +294,7 @@ jobs:
           path: xarray-backends-logs
 
       - name: Generate and publish the xarray backends report
-        uses: xarray-contrib/issue-from-pytest-log@87351a8f864e969567cda22a25a2f214cbe2340f # v1.6.0
+        uses: xarray-contrib/issue-from-pytest-log@054799b34bd75a5fd6c86277a4a8a575224e60c6 # v1.6.1
         with:
           log-path: xarray-backends-logs/output-pytest-log.jsonl
           issue-title: "Nightly Xarray backends tests failed"

icechunk-python/tests/test_stateful_repo_ops.py

@@ -556,6 +556,7 @@ def garbage_collect(
             if created_at_by_id[k] < older_than
         }
 
+        new_parents: dict[str, str | None] = {}
         if self.spec_version >= 2:
             # GC re-parents a kept snapshot whose parent run was deleted to its
             # nearest surviving ancestor, harvesting the deleted ancestors' tx
@@ -576,18 +577,22 @@ def garbage_collect(
                     harvested.extend(self.pruned_ancestor_tx_logs.get(x, []))
                     harvested.append(x)
                 self.pruned_ancestor_tx_logs[cid] = harvested
+                new_parents[cid] = anc
 
         for k in deleted:
             self.commits.pop(k, None)
             self.ondisk_snaps.pop(k, None)
             self.pruned_ancestor_tx_logs.pop(k, None)
 
         if self.spec_version >= 2:
-            # V2's delete_snapshots_from_repo_info rewrites parent pointers
-            # for kept snapshots whose parent was GC'd.
-            for c in self.commits.values():
-                if c.parent_id is not None and c.parent_id in deleted:
-                    c.parent_id = self.initial_snapshot_id
+            # V2's delete_snapshots_from_repo_info re-parents a kept snapshot
+            # whose parent was GC'd to its nearest surviving ancestor (the
+            # harvest walk's endpoint), falling back to the initial snapshot
+            # when the whole chain was deleted.
+            for cid, anc in new_parents.items():
+                self.commits[cid].parent_id = (
+                    anc if anc is not None else self.initial_snapshot_id
+                )
         else:
             # V1 doesn't rewrite parent pointers, so commits whose parents
             # were GC'd have broken ancestry and are effectively unusable.
@@ -925,13 +930,29 @@ def checkout_branch(self, ref: str) -> None:
                 note(f"Expecting error when checking out branch {ref!r}")
                 self.repo.writable_session(ref)
 
+    def _v1_chain_fully_modelled(self, commit: str) -> bool:
+        """True when a commit's full on-disk ancestry is still in the model.
+
+        V1 expiration rewrites only ref-reachable chains, so an unreachable
+        commit's on-disk ancestry can still pass through expired snapshots the
+        model dropped — or, after GC deletes their files, become unreadable
+        entirely. Refs must not be created or reset onto such chains.
+        """
+        try:
+            chain = {s.id for s in self.repo.ancestry(snapshot_id=commit)}
+        except IcechunkError:
+            return False
+        return chain <= set(self.model.commits)
+
     @rule(name=ref_name_text, commit=commits, target=branches)
     def create_branch(self, name: str, commit: str) -> str:
         note(f"Creating branch {name!r} for commit {commit!r}")
 
         # V1 expired snapshots stay on disk so create_branch succeeds, but
         # the model no longer tracks their store contents.
-        assume(not (self.model.spec_version == 1 and commit not in self.model.commits))
+        if self.model.spec_version == 1:
+            assume(commit in self.model.commits)
+            assume(self._v1_chain_fully_modelled(commit))
 
         # we can create a tag and branch with the same name
         if name not in self.model.branch_heads and commit in self.model.commits:
@@ -951,7 +972,9 @@ def create_tag(self, name: str, commit_id: str) -> str:
         note(f"Creating tag {name!r} for commit {commit_id!r}")
         # V1 expired snapshots stay on disk so create_tag succeeds, but
         # the model no longer tracks their store contents.
-        assume(not (self.model.spec_version == 1 and commit_id not in self.model.commits))
+        if self.model.spec_version == 1:
+            assume(commit_id in self.model.commits)
+            assume(self._v1_chain_fully_modelled(commit_id))
         if (
             name in self.model.created_tags
             or name in self.model.tags
@@ -989,14 +1012,10 @@ def discard_changes(self) -> None:
     @rule(branch=branches, commit=commits)
     def reset_branch(self, branch: str, commit: str) -> None:
         # V1 expired snapshots stay on disk so reset_branch would succeed,
-        # but modelling that divergence isn't worthwhile — just skip. This also
-        # covers resetting to a commit whose ancestry *reaches* an expired-but-
-        # on-disk snapshot: the real ancestry resurrects it while the model has
-        # popped it, so skip unless the whole chain is still modelled.
+        # but modelling that divergence isn't worthwhile — just skip.
         if self.model.spec_version == 1:
             assume(commit in self.model.commits)
-            ancestry = {s.id for s in self.repo.ancestry(snapshot_id=commit)}
-            assume(ancestry <= set(self.model.commits))
+            assume(self._v1_chain_fully_modelled(commit))
         if branch not in self.model.branch_heads or commit not in self.model.commits:
             note(f"resetting branch {branch}, expecting error.")
             with pytest.raises(IcechunkError):

icechunk-python/tests/test_zarr/test_stateful.py

@@ -4,6 +4,7 @@
 from typing import Any, TypeVar
 
 import hypothesis.strategies as st
+import numpy as np
 import pytest
 from hypothesis import note
 from hypothesis.stateful import (
@@ -34,6 +35,28 @@
 
 Frequency = TypeVar("Frequency", bound=Callable[..., Any])
 
+
+def storage_chunk_sizes(arr: "Array[Any]") -> tuple[tuple[int, ...], ...]:
+    """Per-dimension sizes of the storage-key chunk grid (shards when sharded).
+
+    Store keys are one object per *write* chunk: with sharding, the whole
+    shard is a single ``c/<i>`` object and read (inner) chunks are byte
+    ranges inside it, never separate keys. Shifts move store keys, so
+    offsets and the grid must be in write-chunk units;
+    ``cdata_shape``/``read_chunk_sizes`` count inner chunks and give the
+    wrong grid for sharded arrays. Older zarr lacks write_chunk_sizes but
+    also lacks rectilinear grids, so cells there are regular and can be
+    computed directly.
+    """
+    if hasattr(arr, "write_chunk_sizes"):
+        return arr.write_chunk_sizes  # type: ignore[no-any-return]
+    cell = arr.shards or arr.chunks
+    return tuple(
+        tuple(min(c, s - i * c) for i in range(-(-s // c)))
+        for s, c in zip(arr.shape, cell, strict=True)
+    )
+
+
 # pytestmark = [
 #     pytest.mark.filterwarnings(
 #         "ignore::zarr.core.dtype.common.UnstableSpecificationWarning"
@@ -266,7 +289,8 @@ def shift_array(self, data: st.DataObject) -> None:
 
         arr_model = zarr.open_array(self.model, path=array_path)
         arr_store = zarr.open_array(self.store, path=array_path)
-        num_chunks = arr_model.cdata_shape
+        grid_sizes = storage_chunk_sizes(arr_model)
+        num_chunks = tuple(len(sizes) for sizes in grid_sizes)
 
         # Draw offset: negative shifts left, positive shifts right
         offset = data.draw(
@@ -290,26 +314,18 @@ def shift_array(self, data: st.DataObject) -> None:
         # - Without resize: data shifting beyond bounds is lost
         should_resize = data.draw(st.booleans())
         if should_resize and any(o > 0 for o in offset):
-            # read_chunk_sizes was added alongside rectilinear support; on older
-            # zarr it doesn't exist but grids are always regular, so synthesize
-            # the same tuple-of-tuples shape from the declared chunk size.
-            if hasattr(arr_model, "read_chunk_sizes"):
-                read_sizes = arr_model.read_chunk_sizes
-            else:
-                read_sizes = tuple(
-                    (c,) * n for c, n in zip(arr_model.chunks, num_chunks, strict=True)
-                )
             # Shifting right by offset[i] pushes the last offset[i] chunks
             # past the original extent; grow by their sizes so they fit.
             new_shape = tuple(
                 arr_model.shape[i]
-                + (sum(read_sizes[i][-offset[i] :]) if offset[i] > 0 else 0)
-                for i in range(len(read_sizes))
+                + (sum(grid_sizes[i][-offset[i] :]) if offset[i] > 0 else 0)
+                for i in range(len(grid_sizes))
             )
             note(f"resizing array '{array_path}' from {arr_model.shape} to {new_shape}")
             arr_model.resize(new_shape)
             arr_store.resize(new_shape)
-            num_chunks = arr_model.cdata_shape
+            grid_sizes = storage_chunk_sizes(arr_model)
+            num_chunks = tuple(len(sizes) for sizes in grid_sizes)
 
         note(f"shifting array '{array_path}' by {offset}")
         self.store.session.shift_array(f"/{array_path}", offset)
@@ -398,3 +414,59 @@ def test_zarr_store() -> None:
     # run_state_machine_as_test(
     #     mk_test_instance_sync, settings=settings(report_multiple_bugs=False)
     # )
+
+
+def test_storage_chunk_sizes_granularity() -> None:
+    model = ModelStore()
+    sharded = zarr.create_array(
+        model, name="s", shape=(4,), shards=(4,), chunks=(2,), dtype="i1", fill_value=1
+    )
+    plain = zarr.create_array(
+        model, name="p", shape=(100, 80), chunks=(30, 40), dtype="i1"
+    )
+    assert storage_chunk_sizes(sharded) == ((4,),)
+    # cdata_shape counts inner chunks for sharded arrays; using it as the
+    # storage-key grid is the bug storage_chunk_sizes exists to avoid.
+    assert sharded.cdata_shape == (2,)
+    assert storage_chunk_sizes(plain) == ((30, 30, 30, 10), (40, 40))
+
+
+async def test_shift_sharded_model_vs_store() -> None:
+    """Shifting a sharded array keeps ModelStore and IcechunkStore keys in sync.
+
+    Mirrors the shift_array rule: the array is generated sharded on the model
+    and recreated on the store with only the outer chunk grid, then both are
+    shifted on the storage-key grid.
+    """
+    repo = Repository.create(in_memory_storage(), spec_version=2)
+    session = repo.writable_session("main")
+    store = session.store
+
+    model = ModelStore()
+    zarr.group(store=model)
+
+    arr_model = zarr.create_array(
+        model, name="0", shape=(4,), shards=(4,), chunks=(2,), dtype="i1", fill_value=1
+    )
+    arr_store = zarr.create_array(
+        store, name="0", shape=(4,), chunks=(4,), dtype="i1", fill_value=1
+    )
+    data = np.zeros((4,), dtype="i1")
+    arr_model[:] = data
+    arr_store[:] = data
+
+    model_keys = sorted([k async for k in model.list_prefix("")])
+    store_keys = sorted([k async for k in store.list_prefix("")])
+    assert model_keys == store_keys, (model_keys, store_keys)
+    assert "0/c/0" in model_keys
+
+    grid_sizes = storage_chunk_sizes(arr_model)
+    num_chunks = tuple(len(sizes) for sizes in grid_sizes)
+    assert num_chunks == (1,)
+
+    session.shift_array("/0", (1,))
+    await model.shift_array("0", (1,), num_chunks)
+
+    model_keys = sorted([k async for k in model.list_prefix("")])
+    store_keys = sorted([k async for k in store.list_prefix("")])
+    assert model_keys == store_keys, (model_keys, store_keys)

icechunk-python/tests/test_zarr/test_store/test_icechunk_store.py

@@ -7,7 +7,6 @@
 
 import pytest
 
-import zarr
 from icechunk import IcechunkError, IcechunkStore, local_filesystem_storage
 from icechunk.repository import Repository
 from zarr.abc.store import OffsetByteRequest, RangeByteRequest, Store, SuffixByteRequest
@@ -546,7 +545,8 @@ async def test_getsize_raises(self, store: IcechunkStore) -> None:
             await store.getsize("not-a-real-key")
 
     @pytest.mark.skipif(
-        zarr.__version__ < "3.1.6", reason="Store._get_bytes added in zarr 3.1.6"
+        not hasattr(Store, "_get_bytes"),
+        reason="Store._get_bytes only exists in zarr >=3.1.6,<3.3",
     )
     async def test_get_bytes(self, store: IcechunkStore) -> None:
         # Override: icechunk validates metadata on zarr.json keys,
@@ -561,7 +561,8 @@ async def test_get_bytes(self, store: IcechunkStore) -> None:
             )
 
     @pytest.mark.skipif(
-        zarr.__version__ < "3.1.6", reason="Store._get_bytes_sync added in zarr 3.1.6"
+        not hasattr(Store, "_get_bytes_sync"),
+        reason="Store._get_bytes_sync only exists in zarr >=3.1.6,<3.3",
     )
     def test_get_bytes_sync(self, store: IcechunkStore) -> None:
         # Override: icechunk validates metadata on zarr.json keys
@@ -571,7 +572,8 @@ def test_get_bytes_sync(self, store: IcechunkStore) -> None:
         assert store._get_bytes_sync(key, prototype=default_buffer_prototype()) == data
 
     @pytest.mark.skipif(
-        zarr.__version__ < "3.1.6", reason="Store._get_json added in zarr 3.1.6"
+        not hasattr(Store, "_get_json"),
+        reason="Store._get_json only exists in zarr >=3.1.6,<3.3",
     )
     async def test_get_json(self, store: IcechunkStore) -> None:
         # Override: icechunk validates metadata on zarr.json keys,
@@ -583,7 +585,8 @@ async def test_get_json(self, store: IcechunkStore) -> None:
         assert await store._get_json(key, prototype=default_buffer_prototype()) == data
 
     @pytest.mark.skipif(
-        zarr.__version__ < "3.1.6", reason="Store._get_json_sync added in zarr 3.1.6"
+        not hasattr(Store, "_get_json_sync"),
+        reason="Store._get_json_sync only exists in zarr >=3.1.6,<3.3",
     )
     def test_get_json_sync(self, store: IcechunkStore) -> None:
         # Override: icechunk validates metadata on zarr.json keys