[codex] Upgrade ShardedZarrStore for pyramid zarrs - #86
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThe PR adds v2 per-array shard indexing to ChangesShardedZarrStore v2 and HAMT Converter
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 12
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@py_hamt/hamt_to_sharded_converter.py`:
- Around line 12-15: Handle classic Zarr v2 chunk keys in _is_zarr_chunk_key and
the conversion flow that reaches set_pointer(), since only c/... keys are
currently recognized. Update the logic in hamt_to_sharded_converter.py to either
reject .zarray-based inputs early or translate v2 chunk keys like 0.0 and
var/0.0 into the sharded chunk-key form before pointer assignment, so they are
not treated as root metadata.
In `@py_hamt/sharded_zarr_store.py`:
- Around line 873-881: Validate the decoded shard length before storing it in
the shard cache: in sharded_zarr_store.py, after dag_cbor.decode in the shard
decode path, add a length check on decoded_shard against
array_index.chunks_per_shard before building shard_data and calling
_shard_data_cache.put. Keep the existing CID/type validation in place, but
reject malformed shards with the wrong number of entries early so later indexed
reads do not fail.
- Around line 150-157: Resizing in sharded_zarr_store.py currently only
truncates shard CID slots in the resize flow, which leaves stale chunk data in
the retained tail shard when the array shrinks. Update the resize logic around
the shard CID handling to also clear any chunk entries whose offsets are beyond
total_chunks in the final retained shard, and make sure any cached shard/chunk
entries are evicted or refreshed so old data cannot reappear after a later grow.
Apply the same cleanup in the related resize paths tied to the shard CID updates
and cache maintenance.
- Around line 1158-1196: The flush path in ShardedZarrStore currently snapshots
dirty shard data and calls mark_clean() without holding the corresponding
_shard_locks, so a concurrent set() or delete() can mutate the shard while
cas.save() is in progress and then get incorrectly cleared. Update the flush
logic around the dirty_shards loop to acquire the shard lock for each cache_key
before reading from _shard_data_cache, saving via cas.save, and marking the
shard clean, and keep the lock held until the cache entry is safely updated; use
the existing _shard_locks, _shard_data_cache, and mark_clean flow as the place
to apply the fix.
- Around line 1315-1316: The public set_pointer mutator currently bypasses the
same read-only write guard used by set(), delete(), resize_store(), and
graft_store(). Update set_pointer in ShardedZarrStore to enforce the read-only
check before calling _set_pointer, using the same guard pattern already applied
in the other mutating methods so pointer updates are blocked when the store is
read-only.
- Around line 1058-1062: The timeout handling in sharded_zarr_store.py should
explicitly chain the caught asyncio.TimeoutError when re-raising the
RuntimeError in this load path. Update the except asyncio.TimeoutError block in
the shard loading logic to preserve the original exception context while keeping
the existing pending-load cleanup for _pending_shard_loads and the timeout
message that references shard_idx.
- Around line 683-688: Reject chunk-shape updates for an existing array index in
sharded_zarr_store.py: when the existing index is reused in the branch that
updates existing.array_shape and existing.chunk_shape, validate that the
incoming chunk_shape matches the stored one and fail instead of overwriting it.
Keep the fix localized to the existing array index handling path in the array
index creation/update logic so coordinate-to-slot mapping cannot change for
already-written shards.
- Around line 297-305: The _V1_COORDINATE_ARRAY_PREFIXES class attribute in
ShardedZarrStore is currently mutable and triggers Ruff; update the symbol to be
an immutable ClassVar backed by a frozenset so it cannot be modified as shared
class state. Keep the same coordinate prefixes, but change the declaration and
any related type annotation in the class definition to use ClassVar and
frozenset to satisfy linting.
- Around line 1324-1335: Invalidate any cached entry in _metadata_read_cache
when set_pointer() updates a metadata pointer so later get() calls do not return
stale bytes. Update the metadata-pointer branch in
ShardedZarrStore.set_pointer() to remove or refresh the cache entry for that key
before returning, and make sure the fix covers keys handled through
_array_path_from_metadata_key() and the existing
_register_array_metadata_from_bytes() flow.
In `@tests/test_sharded_zarr_store_v2.py`:
- Around line 69-906: The v2 path-aware chunk/index behavior in ShardedZarrStore
is only covered by example-based tests, so add at least one Hypothesis invariant
test in test_v2_sharded_zarr_store_v2.py. Use the existing helpers and symbols
such as ShardedZarrStore.open, list_dir, list_prefix, get, exists, and
ArrayIndex to generate varied array/group paths and chunk layouts, then assert
the v2 invariants still hold across those combinations. Keep the test focused on
the new v2 routing/metadata behavior rather than duplicating the current fixed
cases.
- Around line 41-47: The helper method signature in load uses the built-in name
id, which will trigger Ruff A002 under the tests lint rules. Rename that
parameter to a non-builtin identifier in the load method of the test helper
class, and update any internal references to match so the test file continues to
pass ruff check.
In `@tests/test_sharded_zarr_store.py`:
- Around line 785-786: The assertion in the directory listing check is too weak
because it only verifies zarr.json, which would miss regressions where chunk
entries are no longer returned. Update the check around prefixed_dir_keys in the
test to also assert that the expected chunk directory entry c is present, so
list_dir(prefix) is validated for both metadata and chunk contents.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: e52f2759-26cf-43ea-9d50-adb7b40e41b0
📒 Files selected for processing (5)
py_hamt/hamt_to_sharded_converter.pypy_hamt/sharded_zarr_store.pytests/test_sharded_store_grafting.pytests/test_sharded_zarr_store.pytests/test_sharded_zarr_store_v2.py
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/test_sharded_zarr_store_v2.py`:
- Line 310: The test helpers using monkeypatch are missing an explicit type
annotation, which violates the fully typed test rule. Update the affected test
functions in test_sharded_zarr_store_v2, including
test_v2_resize_blocks_interleaved_chunk_mutators and the other monkeypatch-based
tests in that file, so the monkeypatch fixture parameter is annotated as
pytest.MonkeyPatch.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: f05c5aaa-d3dd-4afd-b77d-30350f61861f
📒 Files selected for processing (3)
py_hamt/hamt_to_sharded_converter.pypy_hamt/sharded_zarr_store.pytests/test_sharded_zarr_store_v2.py
🚧 Files skipped from review as they are similar to previous changes (2)
- py_hamt/hamt_to_sharded_converter.py
- py_hamt/sharded_zarr_store.py
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #86 +/- ##
==========================================
Coverage 100.00% 100.00%
==========================================
Files 8 8
Lines 1617 2541 +924
==========================================
+ Hits 1617 2541 +924 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
| if not isinstance(cid, CID): # pragma: no cover | ||
| raise TypeError(f"Expected CID pointer for key {key!r}.") | ||
| cid_base32_str = str(cid.encode("base32")) | ||
| await dest_store.set_pointer(chunk_key, cid_base32_str) |
There was a problem hiding this comment.
Finding: converted Zarr v2 chunks are written under v3-style .../c/... keys, but the .zarray / .zgroup metadata is copied unchanged. Zarr v2 readers still request keys like var/0 or 0.0, so converted stores can read back fill values instead of the original chunk data.
Suggested solution: either preserve the original v2 chunk keys when the source metadata is v2, or fully migrate both metadata and consolidated metadata to a valid v3 layout. Please add an xr.open_zarr round-trip test for a zarr_format=2 HAMT source.
There was a problem hiding this comment.
Handled in 564f0a2. Converted Zarr v2 HAMT sources now keep v2 group metadata (.zgroup) instead of injecting v3 zarr.json, and ShardedZarrStore resolves classic v2 chunk requests such as var/0, var/0.1, and root numeric chunk keys through the path-local shard index. Added an xr.open_zarr round-trip test for a zarr_format=2 HAMT source.
Validation: uv run pytest tests/test_sharded_zarr_store_v2.py, uv run pre-commit run --all-files --show-diff-on-failure, and uv run pytest --cov=py_hamt tests/ --ignore=tests/test_public_gateway.py.
| if self._manifest_version != SHARDED_ZARR_V2: | ||
| return keys | ||
|
|
||
| for array_path in self.array_indices: |
There was a problem hiding this comment.
Finding: v2 directory listing derives entries from array_indices, but delete/clear paths only remove metadata and chunk pointers. After delete_dir("a"), stale array_indices["a"] keeps a visible in listings and can block overwrites with new geometry.
Suggested solution: override delete_dir / clear for v2, or prune matching array_indices, shard cache keys, pending loads, and root arrays entries when deleting an array subtree. Please add tests for delete_dir(""), delete_dir("a"), and to_zarr(mode="w") overwrite with changed shape or chunk shape.
There was a problem hiding this comment.
Handled in 564f0a2. v2 delete_dir and clear now prune matching array_indices, root arrays entries, shard cache entries, and pending shard-load events. delete_dir("a"), delete_dir(""), clear(), and to_zarr(mode="w") overwrites with changed geometry are covered.
Validation: uv run pytest tests/test_sharded_zarr_store_v2.py, uv run pre-commit run --all-files --show-diff-on-failure, and uv run pytest --cov=py_hamt tests/ --ignore=tests/test_public_gateway.py.
| metadata_key = ( | ||
| "zarr.json" if group_path == "" else f"{group_path}/zarr.json" | ||
| ) | ||
| metadata[metadata_key] = await self.cas.save(group_metadata, codec="raw") |
There was a problem hiding this comment.
Finding: this writes default parent group metadata unconditionally. During v1-to-v2 migration, any existing root or group zarr.json metadata is replaced with an empty group document, so group attributes are lost.
Suggested solution: only insert the default group metadata when the metadata key is absent; preserve existing CIDs/data otherwise. Please add a migration regression test with root/group attributes to prove they survive.
There was a problem hiding this comment.
Handled in 564f0a2. Migration now preserves existing group metadata CIDs instead of overwriting them with empty defaults, and strips stale consolidated metadata from preserved v3 group docs so relocated groups can still be discovered. Added a migration regression test for preserved root and group attributes.
Validation: uv run pytest tests/test_sharded_zarr_store_v2.py, uv run pre-commit run --all-files --show-diff-on-failure, and uv run pytest --cov=py_hamt tests/ --ignore=tests/test_public_gateway.py.
| ) | ||
| self._dirty_root = True | ||
|
|
||
| async def migrate_v1_to_v2(self, primary_array_path: str) -> str: |
There was a problem hiding this comment.
Finding: migrate_v1_to_v2() mutates the store and flushes a new root even when the instance was opened read-only. This differs from set, delete, grafting, and resize APIs, which reject writes on read-only stores.
Suggested solution: add the same PermissionError guard before migration does any mutation or flush, and add a read-only migration regression test.
There was a problem hiding this comment.
Handled in 564f0a2. migrate_v1_to_v2() now raises PermissionError immediately on read-only stores before acquiring the write path or flushing. Added a read-only migration regression test.
Validation: uv run pytest tests/test_sharded_zarr_store_v2.py, uv run pre-commit run --all-files --show-diff-on-failure, and uv run pytest --cov=py_hamt tests/ --ignore=tests/test_public_gateway.py.
| if parsed_chunk is None: | ||
| if register_metadata: | ||
| await self._ensure_v2_parent_group_metadata(key) | ||
| self._root_obj["metadata"][key] = pointer_cid_obj |
There was a problem hiding this comment.
Finding: metadata is written into _root_obj before the new metadata is validated and registered against the array index. If _register_array_metadata_from_bytes raises, for example on a chunk-shape change, the store is left with new metadata but the old ArrayIndex geometry.
Suggested solution: validate/register the metadata before mutating _root_obj, or roll back both _root_obj and caches on failure. Metadata and ArrayIndex updates should be atomic for shape/chunk-shape changes.
There was a problem hiding this comment.
Handled in 564f0a2. Metadata pointer updates now load and register array metadata before mutating _root_obj["metadata"]; failed chunk-shape validation leaves the old metadata pointer and ArrayIndex intact. Added a regression test for failed set_pointer() metadata registration.
Validation: uv run pytest tests/test_sharded_zarr_store_v2.py, uv run pre-commit run --all-files --show-diff-on-failure, and uv run pytest --cov=py_hamt tests/ --ignore=tests/test_public_gateway.py.
|
Codex Generated PR Doc: What PR 86 Does PR 86 upgrades Before this PR, because each array can have different shape/chunk geometry, but old sharding assumed one shared geometry. The PR adds a new manifest format: SHARDED_ZARR_V1 = "sharded_zarr_v1"
SHARDED_ZARR_V2 = "sharded_zarr_v2"In v2, the root object has: So each array path gets its own Main Implementation Changes
It adds:
The public
Behavior Changes Reads/writes/deletes now route chunk keys through the relevant array index: Those two keys can now point to different shard lists and different geometries. Listing is also expanded. For v2, Resize is made array-local. Grafting is updated so source and target arrays can be selected by path: graft_store(..., source_array_path="0/FPAR", target_array_path="1/FPAR")Migration And Conversion The PR adds
Tests Added The PR adds
It also tweaks existing sharded store tests for the new behavior. Current PR State Current head is |
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
py_hamt/sharded_zarr_store.py (1)
1091-1130: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winUnhandled
ValueErrorfrom_parse_chunk_keycan crashget()/_delete_unlocked()/_set_pointer().The new fallback at Line 1117-1123 re-raises the
int()ValueErrorwhen a key matches the/c/-or-c/-prefix shape but has non-numeric coordinates and doesn't match any classic v2 key pattern.exists()explicitly catches(ValueError, IndexError, KeyError)around this call, butget()(Line 1493),_delete_unlocked()(Line 1677), and_set_pointer()(Line 1595) call_parse_chunk_keyunguarded. A key such as"c/FPAR/not-a-chunk"— confirmed in this PR's own test (assert not await classic_store.exists("c/FPAR/not-a-chunk"), only exercised viaexists()) — would raise instead of being treated as "not found" if the same key reachedget().🐛 Proposed fix
kind = "metadata" shard_idx_for_trace: int | None = None - parsed_chunk = self._parse_chunk_key(key) + try: + parsed_chunk = self._parse_chunk_key(key) + except (ValueError, IndexError): + return None try:Apply the same guard at the start of
_delete_unlockedand_set_pointer.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@py_hamt/sharded_zarr_store.py` around lines 1091 - 1130, _handle invalid chunk-key parsing in the same way across all callers of _parse_chunk_key. The new ValueError path for malformed /c/ or c/ keys is only tolerated by exists(), but get(), _delete_unlocked(), and _set_pointer() still call _parse_chunk_key directly and can crash on non-numeric coordinates. Add the same defensive handling used by exists() around _parse_chunk_key in _delete_unlocked and _set_pointer, and make sure get() also treats these malformed keys as missing rather than propagating ValueError.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@py_hamt/sharded_zarr_store.py`:
- Around line 1711-1730: `delete_dir` is doing an expensive full-store chunk
enumeration by calling `list_prefix(match_prefix)`, which ultimately scans every
v2 shard across all arrays instead of only the target subtree. Update
`delete_dir` to enumerate only `_root_obj["metadata"]` for individual
`_delete_unlocked` cleanup, and remove the chunk-key loop for shard contents
under the prefix. Keep the existing
`_prune_v2_array_indices_for_prefix(normalized_prefix)` call to handle shard
cache and array-index cleanup efficiently, and preserve the read-only and
root-prefix behavior in `delete_dir`.
---
Outside diff comments:
In `@py_hamt/sharded_zarr_store.py`:
- Around line 1091-1130: _handle invalid chunk-key parsing in the same way
across all callers of _parse_chunk_key. The new ValueError path for malformed
/c/ or c/ keys is only tolerated by exists(), but get(), _delete_unlocked(), and
_set_pointer() still call _parse_chunk_key directly and can crash on non-numeric
coordinates. Add the same defensive handling used by exists() around
_parse_chunk_key in _delete_unlocked and _set_pointer, and make sure get() also
treats these malformed keys as missing rather than propagating ValueError.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: a6961e7f-df26-4c1a-b3d1-03d3734236e2
📒 Files selected for processing (2)
py_hamt/sharded_zarr_store.pytests/test_sharded_zarr_store_v2.py
|
Handled CodeRabbit outside-diff malformed chunk-key feedback in 415e1f2. For v2 stores, malformed chunk-looking keys such as c/FPAR/not-a-chunk now behave as missing across get(), delete(), set_pointer(), and set() instead of propagating parser ValueError/IndexError; v1 keeps its previous strict exceptions for backward compatibility. Added regression coverage for both the new v2 behavior and preserved v1 raise behavior. Validation: uv run pytest tests/test_sharded_zarr_store_v2.py; uv run pre-commit run --all-files --show-diff-on-failure; uv run pytest --cov=py_hamt tests/ --ignore=tests/test_public_gateway.py. |
3 participants
Summary
Implements FPAR pyramid Ticket 1 by upgrading
ShardedZarrStoreto support path-awaresharded_zarr_v2manifests with one shard index per Zarr array path.Changes
ArrayIndex/ChunkKeyhandling and path-aware chunk parsing, get/set/exists/delete/list/resize/graft behavior.c/...metadata chunks and top-levelcgroup paths.migrate_v1_to_v2(primary_array_path)support that reuses existing v1 shard blocks and rewrites metadata under the requested primary path.c-named paths, chunk-prefix listing, grafting, and malformed shard loads.Validation
uv run pytest tests/test_sharded_zarr_store_v2.py tests/test_converter.py -q: 23 passeduv run pre-commit run --all-files --show-diff-on-failure: passedbash run-checks.sh: 157 passed, 2 skipped, 100% coverage, hooks passedNotes
FPAR_PYRAMID_TICKETS.mdwas intentionally not included in this branch.Summary by CodeRabbit
zarr.json), plus more targeted grafting/resizing by array path.