diff --git a/.github/scripts/run_benchmark.py b/.github/scripts/run_benchmark.py index 0daa6e9a..09efd03d 100644 --- a/.github/scripts/run_benchmark.py +++ b/.github/scripts/run_benchmark.py @@ -169,7 +169,9 @@ def run_target( catalog=str(shardmap_path), store=store, backend="lambda", - morton_cell=str(shard_key), + # morton_cell takes the external shard label (the decimal morton + # string for HEALPix — issue #199), not the raw packed-word digits. + morton_cell=grid.shard_label(shard_key), region=region, function_name=function_name, overwrite=True, diff --git a/deployment/aws/lambda_handler.py b/deployment/aws/lambda_handler.py index b30693fe..b38867d6 100644 --- a/deployment/aws/lambda_handler.py +++ b/deployment/aws/lambda_handler.py @@ -31,15 +31,19 @@ allocated), keeping the flag-off event and outputs byte-identical. "result_url": str (optional, issue #151) -- where to ALSO write this invocation's response envelope as JSON (e.g. - "s3://bucket/out.zarr.status//.json"). Set by the - orchestrator's async dispatch (InvocationType="Event", which discards + "s3://bucket/out.zarr.status//.json", where the + label is the decimal morton string for HEALPix -- issue #199). Set by + the orchestrator's async dispatch (InvocationType="Event", which discards the return value); the orchestrator polls this object instead of holding a synchronous connection open while the shard runs. Written with the output-store credentials. Absent -> no write, and the event and behavior are byte-identical to the synchronous path. } -Setup mode (creates the zarr template once before per-cell fan-out): +Setup mode (creates the zarr template once before per-cell fan-out; for a +hive-layout config -- output.store_layout: hive, issue #199 -- it writes the +morton_hive.json manifest instead, and each process-mode worker emits its own +leaf template): { "mode": "setup", "store_path": str, @@ -47,7 +51,10 @@ "n_parent_cells": int, # OPTIONAL -- dense layout only (populated count) "overwrite": bool, "config": dict, # single source of truth: child_order, chunk_inner, - # layout, and grid type all come from here + # layout, store_layout, and grid type all come from here + "dataset": dict (optional, hive only) -- {"short_name", "version"} identity + block for the manifest, sourced from the ShardMap metadata by the + orchestrator (matching the local dispatcher). Absent on flat runs. "output_credentials": dict (optional, same shape as process mode), } @@ -126,7 +133,8 @@ from zarr.errors import GroupNotFoundError # Import cloud-agnostic processing -from zagg.config import get_handoff, load_config_from_dict +from zagg.config import get_handoff, get_store_layout, load_config_from_dict +from zagg.grids.base import shard_label from zagg.processing import ( write_dataframe_to_zarr, write_ragged_to_zarr, @@ -590,12 +598,38 @@ def _handle_extract(event: Dict[str, Any], context: Any) -> Dict[str, Any]: def _handle_setup(event: Dict[str, Any]) -> Dict[str, Any]: - """Create the zarr template at ``event['store_path']``.""" + """Create the zarr template at ``event['store_path']``. + + For a hive-layout config (issue #199 phase 3) template time writes ONLY + the ``morton_hive.json`` manifest — no global zarr template exists (zero + metadata above the leaves, D5); each worker emits its own leaf template. + The optional ``dataset`` event key carries the manifest's identity block + (the orchestrator sources it from the ShardMap metadata, same as the local + path). The flat path below is byte-identical to before, bar one addition: + the success body now ECHOES the layout it acted on (``"layout"``) — a + stale deployment without the hive branch returns the old echo-less body, + which the dispatcher rejects for hive runs instead of silently letting old + workers write a flat store at the hive root (review finding, PR #205). + """ from zagg.grids import from_config logger.info(f"Setup mode: creating template at {event.get('store_path')}") try: config = load_config_from_dict(event["config"]) + if get_store_layout(config) == "hive": + from zagg.hive import build_manifest, ensure_manifest + + grid = from_config(config, parent_order=event.get("parent_order")) + ensure_manifest( + event["store_path"], + build_manifest(grid, dataset=event.get("dataset")), + overwrite=event.get("overwrite", False), + **_output_store_kwargs(event), + ) + return { + "statusCode": 200, + "body": json.dumps({"ok": True, "mode": "setup", "layout": "hive"}), + } store = open_store(event["store_path"], **_output_store_kwargs(event)) # Build the grid exactly as the worker does (from_config), so the # template's chunk structure can't drift from what workers write. The @@ -616,7 +650,10 @@ def _handle_setup(event: Dict[str, Any]) -> Dict[str, Any]: populated_shards=populated, ) grid.emit_template(store, overwrite=event.get("overwrite", False)) - return {"statusCode": 200, "body": json.dumps({"ok": True, "mode": "setup"})} + return { + "statusCode": 200, + "body": json.dumps({"ok": True, "mode": "setup", "layout": "flat"}), + } except Exception as e: logger.exception(e) return {"statusCode": 500, "body": json.dumps({"error": str(e), "mode": "setup"})} @@ -871,92 +908,131 @@ def _handle_process(event: Dict[str, Any], context: Any) -> Dict[str, Any]: # chunk is keyed by its own block index. single_chunk = int(getattr(grid, "chunks_per_shard", 1)) == 1 - # Lazy store + one-time template check, opened on the FIRST chunk write so a - # no-data shard (zero chunks) never touches the store, exactly as before. A - # missing template or a failed write is RECORDED (not raised) so ``metadata`` - # from ``process_shard`` survives — the buffered path returned its 500 with - # that metadata; folding the error in after the stream preserves that body. store_box: dict = {} write_error: dict = {} _write_elapsed = 0.0 - - def _get_store(): - """Open + template-check once; returns the store, or None if the template - is missing (recording the error so the write is skipped).""" - if "store" in store_box: - return store_box["store"] - if write_error: - return None - store = open_store(store_path, **_output_store_kwargs(event)) - # Validate the Zarr template exists before writing. ``store`` is a zarr v3 - # ``Store`` whose ``exists()`` is async, so open the group via the high-level - # sync API and catch the missing-node error instead (issue #118), in the same - # open-and-catch spirit as ``readers/tdigest_tensor.py``. - # ``GroupNotFoundError`` is raised identically on LocalStore and obstore (S3); - # a present-but-wrong-type node surfaces as a real error, not "missing". - try: - open_group(store, path=grid.group_path, mode="r", zarr_format=3) - except GroupNotFoundError: - msg = f"Zarr template not found at {store_path}/{grid.group_path}" - logger.error(msg) - write_error["msg"] = msg - return None - logger.info(f" Writing data to {store_path}...") - store_box["store"] = store - return store - - def _write_chunk(block_index, carrier, ragged): - nonlocal _write_elapsed - if write_error: - return # a prior chunk failed (or template missing) — skip the rest - store = _get_store() - if store is None: - return # template missing — recorded in write_error, skip the rest - _t0 = time.time() if profile else None - try: - # write_dataframe_to_zarr no-ops on an empty carrier, so no per-chunk - # emptiness check is needed. Use each chunk's own block_index. - write_dataframe_to_zarr(carrier, store, grid=grid, chunk_idx=block_index) - ragged_key = int(shard_key) if single_chunk else _block_index_key(block_index, grid) - write_ragged_to_zarr(ragged, store, grid=grid, shard_key=ragged_key) - except Exception as e: - # Mirror the buffered path's ``except``: record the failure, stop - # writing, and let the run surface a 500 after process_shard returns. - logger.error(f"Failed to write zarr to {store_path}: {e}") - write_error["msg"] = f"Failed to write zarr: {e}" - return - if profile: - _write_elapsed += time.time() - _t0 - - chunk_results: list | None = [] if sharded else None - _df_out, metadata = process_shard( - grid, - shard_key, - event["granule_urls"], - s3_credentials=s3_creds, - config=config, - chunk_results=chunk_results, - write_chunk=None if sharded else _write_chunk, - handoff=handoff, - profile=profile, - aoi_payload=aoi_payload, - ) - - # Sharded output (issue #108): bundle the shard's K inner chunks into one - # ShardingCodec shard object — one block selection per dense array (a per- - # inner-chunk loop would read-modify-write the same shard object). This path - # accumulated all K, so it opens + validates + writes here (same recording). - if sharded and chunk_results: - store = _get_store() - if store is not None: - _write_t0 = time.time() if profile else None + chunk_results: list | None = None + _df_out = None + + if get_store_layout(config) == "hive": + # Hive layout (issue #199 phase 3): the worker owns its WHOLE leaf + # zarr — it derives the leaf path from shard_key + the event's + # config orders, emits its own leaf template (lazily, on the first + # chunk), writes its data, and stamps completion as its FINAL PUT + # (D4), on error-free shards only. process_and_write_hive is the + # same code path the local dispatcher runs, so leaf semantics + # cannot drift between backends. A write failure raises out to the + # handler's exception envelope: the leaf is then unstamped debris, + # overwritten wholesale on retry — the same recovery model as the + # local path (no per-chunk error recording needed). + from zagg.hive import process_and_write_hive + + metadata = process_and_write_hive( + shard_key, + event["granule_urls"], + grid, + s3_creds, + store_path, + config, + store_kwargs=_output_store_kwargs(event), + handoff=handoff, + aoi_payload=aoi_payload, + profile=profile, + ) + else: + # Flat layout: lazy store + one-time template check, opened on the + # FIRST chunk write so a no-data shard (zero chunks) never touches + # the store, exactly as before. A missing template or a failed + # write is RECORDED (not raised) so ``metadata`` from + # ``process_shard`` survives — the buffered path returned its 500 + # with that metadata; folding the error in after the stream + # preserves that body. + def _get_store(): + """Open + template-check once; returns the store, or None if the template + is missing (recording the error so the write is skipped).""" + if "store" in store_box: + return store_box["store"] + if write_error: + return None + store = open_store(store_path, **_output_store_kwargs(event)) + # Validate the Zarr template exists before writing. ``store`` is a zarr v3 + # ``Store`` whose ``exists()`` is async, so open the group via the high-level + # sync API and catch the missing-node error instead (issue #118), in the same + # open-and-catch spirit as ``readers/tdigest_tensor.py``. + # ``GroupNotFoundError`` is raised identically on LocalStore and obstore (S3); + # a present-but-wrong-type node surfaces as a real error, not "missing". + try: + open_group(store, path=grid.group_path, mode="r", zarr_format=3) + except GroupNotFoundError: + msg = f"Zarr template not found at {store_path}/{grid.group_path}" + logger.error(msg) + write_error["msg"] = msg + return None + logger.info(f" Writing data to {store_path}...") + store_box["store"] = store + return store + + def _write_chunk(block_index, carrier, ragged): + nonlocal _write_elapsed + if write_error: + return # a prior chunk failed (or template missing) — skip the rest + store = _get_store() + if store is None: + return # template missing — recorded in write_error, skip the rest + _t0 = time.time() if profile else None try: - write_shard_to_zarr(chunk_results, store, grid=grid, shard_key=int(shard_key)) - if profile: - _write_elapsed += time.time() - _write_t0 + # write_dataframe_to_zarr no-ops on an empty carrier, so no per-chunk + # emptiness check is needed. Use each chunk's own block_index. + write_dataframe_to_zarr(carrier, store, grid=grid, chunk_idx=block_index) + # At K==1 the CSR subgroup is named by the grid's shard label (the + # decimal morton string for HEALPix — issue #199); K>1 keeps the + # flattened block-index int. Mirrors runner._process_and_write. + ragged_key = ( + shard_label(grid, shard_key) + if single_chunk + else _block_index_key(block_index, grid) + ) + write_ragged_to_zarr(ragged, store, grid=grid, shard_key=ragged_key) except Exception as e: + # Mirror the buffered path's ``except``: record the failure, stop + # writing, and let the run surface a 500 after process_shard returns. logger.error(f"Failed to write zarr to {store_path}: {e}") write_error["msg"] = f"Failed to write zarr: {e}" + return + if profile: + _write_elapsed += time.time() - _t0 + + chunk_results = [] if sharded else None + _df_out, metadata = process_shard( + grid, + shard_key, + event["granule_urls"], + s3_credentials=s3_creds, + config=config, + chunk_results=chunk_results, + write_chunk=None if sharded else _write_chunk, + handoff=handoff, + profile=profile, + aoi_payload=aoi_payload, + ) + + # Sharded output (issue #108): bundle the shard's K inner chunks into one + # ShardingCodec shard object — one block selection per dense array (a per- + # inner-chunk loop would read-modify-write the same shard object). This path + # accumulated all K, so it opens + validates + writes here (same recording). + if sharded and chunk_results: + store = _get_store() + if store is not None: + _write_t0 = time.time() if profile else None + try: + write_shard_to_zarr( + chunk_results, store, grid=grid, shard_key=int(shard_key) + ) + if profile: + _write_elapsed += time.time() - _write_t0 + except Exception as e: + logger.error(f"Failed to write zarr to {store_path}: {e}") + write_error["msg"] = f"Failed to write zarr: {e}" # A recorded template-missing / write failure folds into ``metadata`` so the # response surfaces a 500 with the structured log, exactly as the buffered diff --git a/docs/hive_layout.md b/docs/hive_layout.md new file mode 100644 index 00000000..4504625f --- /dev/null +++ b/docs/hive_layout.md @@ -0,0 +1,152 @@ +# Hive store layout (morton-hive/1) + +zagg can write each dispatch shard as its **own self-describing leaf zarr** +under a morton digit tree, instead of into one shared flat store. The layout is +the write-side half of the sparse-coverage design record +([`docs/design/sparse_coverage.md`](design/sparse_coverage.md) §2–§3, decisions +D1–D6); the convention itself is owned by the mortie spec and versioned as +`morton-hive/1`. + +It is opt-in (default `flat`, today's single shared store) and — in this first +round — wired to the **local backend only** (see [Status](#status)). + +## Layout + +``` +{store_root}/ + morton_hive.json <- static manifest (root-only exception) + {sign+base}/{d1}/.../{d_n}/ <- one decimal digit per level (D2) + {full_id}.zarr/ <- vanilla zarr v3 leaf, one per shard (D3) +``` + +- **Ids are morton decimal strings** (D1): sign + base digit (`1..6` / + `-1..-6`), then one digit per order, digits `1..4`, never `0`. A string + prefix *is* the spatial ancestor, so cross-resolution containment is + `fine_id.startswith(coarse_id)` — arithmetic, not I/O. +- **One digit per path component** (D2), so shards at mixed orders nest + naturally: every order is a legal node. +- **Full id at the leaf** (D3): `.../-5/1/1/2/3/3/3/-5112333.zarr` is + self-describing without parsing its directory chain — greppable in + inventories, unambiguous if moved. Each leaf is a completely vanilla zarr v3 + store: the same group/array template as the flat layout, sized to one shard + (dense arrays hold `cells_per_shard` cells; `resolution: chunk` companions + hold the shard's K inner chunks; CSR ragged subgroups sit under the same + group path, named by the shard label at K==1). +- **Node invariant** (D5): below the root, a node contains *only* digit + children (`[1-4]/`) and `*.zarr` objects — zero zarr metadata above the + leaf, no shared mutable state across workers. The root alone also carries + the manifest (and, in a follow-on, `coverage.moc`). `zagg.hive` + re-checks every computed leaf path against this invariant before writing. + +## Config + +```yaml +output: + store: s3://bucket/product # becomes the hive root + store_layout: hive # default: flat + grid: + type: healpix # hive is HEALPix-only (morton digit tree) + parent_order: 9 # shard order -> tree depth + child_order: 13 # cell order +``` + +Validation rejects `hive` with a rectilinear grid (node names are morton +digits), with `sharded: true` (a leaf is a vanilla zarr v3 store), and with +`consolidate_metadata: true` (there is no store-root zarr hierarchy to +consolidate — D5/D12). + +## The manifest (`morton_hive.json`) + +Written **once at template time**, before any shard dispatches; never touched +during a run (D6). With it, every shard path is computable arithmetically with +zero requests: + +```json +{ + "spec": "morton-hive/1", + "dataset": {"short_name": "ATL03", "version": "007"}, + "cell_order": 13, + "shard_order": 9, + "split_schedule": [1, 1, 1, 1, 1, 1, 1, 1, 1], + "pyramid": {"orders": [], "aggregation": {}}, + "generated_at": "2026-07-10T12:00:00+00:00" +} +``` + +`split_schedule` is implicit under D2 (one digit per level down to the shard +order) but recorded explicitly for forward compatibility. `pyramid` is +declared-only in round one: overview zarrs are generated by a later +post-process sweep (D11), never at fan-out time. + +A rerun into an existing root verifies the manifest's **frozen keys** match +the run's own configuration (`spec`, `dataset`, `cell_order`, `shard_order`, +`split_schedule`) and fails loudly on a mismatch — the hive analogue of the +flat layout's shard-map signature guard. The sweep-mutable `pyramid` block and +the `generated_at` timestamp are deliberately excluded, so a swept store still +resumes (and the sweep's pyramid declaration is preserved, not clobbered). + +**Re-templating does not remove existing leaves.** Overwriting the manifest +replaces *only* the JSON — committed leaves written under the old +configuration would survive, stamped and walker-discoverable, and (because +mixed shard orders are legal under D2) indistinguishable from intentional +data. The writer therefore refuses an overwrite that changes the frozen keys +while the digit tree has any `{sign+base}` children (one delimiter-LIST): +clear the store root, or pick a new one, before writing with a different +configuration. + +## The commit stamp + +S3 has no empty directories and LIST is strongly consistent, so **absence is +trustworthy**: a delimiter-LIST with no digit children means nothing finer +exists. **Presence is not** — a worker that dies mid-shard has already created +the `.zarr/` prefix. So the shard's *final* write is a root +`group.attrs.update(...)` recording completion (D4): + +```json +"morton_hive_commit": { + "spec": "morton-hive/1", + "complete": true, + "cells_with_data": 412, + "granule_count": 17, + "written_at": "2026-07-10T12:03:41+00:00" +} +``` + +A leaf whose root metadata lacks the stamp is **debris**: incomplete, +ignorable, safe to overwrite on retry (the writer re-emits the leaf template +with `overwrite=True`, so retries are idempotent). This is *not* consolidated +metadata — one small PUT rewriting the root `zarr.json`, which the leaf +template creates anyway. A shard that errors, or streams no chunks (no data), +leaves no stamp; a fully empty shard leaves no `.zarr/` prefix at all (the +leaf is created lazily on the first chunk write). + +## Reading a hive store + +There is no store-root `zarr.open()` (deliberately — D12; a root hierarchy can +be added later by the sweep as a derived artifact). Readers: + +1. GET `morton_hive.json` (once, cacheable) → `shard_order`, `cell_order`. +2. Compute a shard's leaf path by string arithmetic on its decimal id + (`zagg.hive.shard_leaf_path`), open the leaf zarr, and **check the commit + stamp** (`zagg.hive.read_commit`) before trusting the contents. +3. Discovery without a shard list is a delimiter-LIST walk: recurse on + `[1-4]/` children; a `*.zarr` entry is data at that node; no digit children + ⇒ nothing finer. Never LIST per observation in a join loop (D10). + +The `coverage.moc` domain declaration (§4 of the design record) is a follow-on +issue and will remove the walk from the discovery path too. + +## Status + +- **Both backends** write hive stores end-to-end through the same + `zagg.hive.process_and_write_hive` code path. On **Lambda** + ([issue #199](https://github.com/englacial/zagg/issues/199) phase 3) + `mode: "setup"` writes only the manifest — the orchestrator still needs no + S3 access — and each worker derives its leaf path from its `shard_key` + + the event config's orders, emits its own leaf template, and stamps + completion as its final PUT. The async status channel stays at the flat + sibling prefix (`{store_root}.status//…`), outside the digit tree. +- The store-level `coverage.moc` (and any per-shard MOC stamping) waits on + the design in [issue #200](https://github.com/englacial/zagg/issues/200). +- Write-throughput validation at fleet scale is tracked with the benchmark + machinery in [issue #202](https://github.com/englacial/zagg/issues/202). diff --git a/mkdocs.yml b/mkdocs.yml index dfe3ae86..e7a4dde5 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -26,6 +26,7 @@ nav: - "Architecture": "design/architecture.md" - "Schema": "design/schema.md" - "Sharded output": "sharding.md" + - "Hive store layout": "hive_layout.md" - "Typed morton boundary": "morton_arrow.md" - "Deployment": - "Standing Up the Backend": "deployment/standup.md" diff --git a/pyproject.toml b/pyproject.toml index b45f3686..d1cff6dc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,7 +35,10 @@ dependencies = [ # register from it). 0.2.0 is not on PyPI yet — resolution fails until # upstream publishes it. "h5coro-hidefix>=0.3.1", - "mortie>=0.8.5", + # 0.9.0 ships the decimal-emit surface shard ids route through (issue #199): + # decimal_repr / to_decimal, the _decimal_to_word parse-back, hive_path, and + # the MortonIndexScalar decimal display. + "mortie>=0.9.0", "earthaccess", "boto3", "fastparquet", diff --git a/src/zagg/__main__.py b/src/zagg/__main__.py index 3ba877c1..02f4f270 100644 --- a/src/zagg/__main__.py +++ b/src/zagg/__main__.py @@ -29,29 +29,50 @@ def main(): """, ) parser.add_argument("--config", required=True, help="Path to pipeline config YAML") - parser.add_argument("--catalog", default=None, help="Path to granule catalog JSON (overrides config)") + parser.add_argument( + "--catalog", default=None, help="Path to granule catalog JSON (overrides config)" + ) parser.add_argument("--store", default=None, help="Output store path (overrides config)") - parser.add_argument("--backend", default="local", choices=["local", "lambda"], - help="Execution backend (default: local)") - parser.add_argument("--driver", default=None, choices=["s3", "https"], - help="Data access driver (default: from config, or s3)") - parser.add_argument("--max-cells", type=int, default=None, help="Limit number of cells (for testing)") - parser.add_argument("--morton-cell", type=str, default=None, help="Process a specific morton cell") + parser.add_argument( + "--backend", + default="local", + choices=["local", "lambda"], + help="Execution backend (default: local)", + ) + parser.add_argument( + "--driver", + default=None, + choices=["s3", "https"], + help="Data access driver (default: from config, or s3)", + ) + parser.add_argument( + "--max-cells", type=int, default=None, help="Limit number of cells (for testing)" + ) + parser.add_argument( + "--morton-cell", + type=str, + default=None, + help="Process a single shard: its decimal morton id (e.g. -31123) for " + "HEALPix, or the shard-key int for other grids", + ) parser.add_argument("--max-workers", type=int, default=None, help="Max concurrent workers") parser.add_argument("--overwrite", action="store_true", help="Overwrite existing Zarr template") parser.add_argument("--dry-run", action="store_true", help="Show what would be processed") parser.add_argument( - "--profile", action="store_true", + "--profile", + action="store_true", help="Emit per-phase worker timings (read/index/aggregate) on the lambda " - "backend. Off by default to avoid the per-worker probe tax (issue #100).", + "backend. Off by default to avoid the per-worker probe tax (issue #100).", ) parser.add_argument("--region", default="us-west-2", help="AWS region (default: us-west-2)") parser.add_argument( - "--output-creds", default=None, metavar="PATH", + "--output-creds", + default=None, + metavar="PATH", help="Path to a JSON file with credentials for writing the output store " - "(keys: accessKeyId, secretAccessKey, optional sessionToken/" - "endpointUrl/region; camelCase, snake_case, or STS PascalCase " - "spellings accepted). Omit to use the ambient/execution-role creds.", + "(keys: accessKeyId, secretAccessKey, optional sessionToken/" + "endpointUrl/region; camelCase, snake_case, or STS PascalCase " + "spellings accepted). Omit to use the ambient/execution-role creds.", ) parser.add_argument( "--function-name", @@ -94,17 +115,23 @@ def main(): if args.dry_run: print(f"\n[DRY RUN] Would process {results['total_cells']} cells") - print(f" Granules per cell: min={results['granules_per_cell_min']}, " - f"max={results['granules_per_cell_max']}, " - f"avg={results['granules_per_cell_avg']:.1f}") + print( + f" Granules per cell: min={results['granules_per_cell_min']}, " + f"max={results['granules_per_cell_max']}, " + f"avg={results['granules_per_cell_avg']:.1f}" + ) print(f" Output: {results['store_path']}") else: - print(f"\nDone: {results['cells_with_data']} cells with data, " - f"{results['total_obs']:,} obs, {results['cells_error']} errors, " - f"{results['wall_time_s']:.1f}s") + print( + f"\nDone: {results['cells_with_data']} cells with data, " + f"{results['total_obs']:,} obs, {results['cells_error']} errors, " + f"{results['wall_time_s']:.1f}s" + ) if "estimated_cost_usd" in results: - print(f"Lambda compute: {results['lambda_time_s']:.0f}s total, " - f"{results['gb_seconds']:.0f} GB-s, ~${results['estimated_cost_usd']:.2f}") + print( + f"Lambda compute: {results['lambda_time_s']:.0f}s total, " + f"{results['gb_seconds']:.0f} GB-s, ~${results['estimated_cost_usd']:.2f}" + ) if results.get("worker_phase_max"): breakdown = ", ".join( f"{phase} {secs:.0f}s" for phase, secs in results["worker_phase_max"].items() diff --git a/src/zagg/config.py b/src/zagg/config.py index 0912037b..f75127df 100644 --- a/src/zagg/config.py +++ b/src/zagg/config.py @@ -318,6 +318,33 @@ def validate_config(config: PipelineConfig) -> None: f"output.consolidate_metadata must be a boolean (got {consolidate_metadata!r})" ) + # Optional store layout (issue #199 phase 2): "flat" (default, today's single + # shared store) or "hive" (one leaf zarr per shard under a morton digit tree — + # docs/design/sparse_coverage.md D1-D6). Hive ids are morton decimal strings, + # so the layout is HEALPix-only; the sharded (ShardingCodec) write path and + # metadata consolidation both assume the single shared store, so they are + # rejected with hive rather than silently mis-writing. + store_layout = config.output.get("store_layout") + if store_layout is not None and store_layout not in ("flat", "hive"): + raise ValueError(f"output.store_layout must be 'flat' or 'hive' (got {store_layout!r})") + if store_layout == "hive": + if (grid or {}).get("type", "healpix") != "healpix": + raise ValueError( + "output.store_layout: hive requires a healpix grid (hive node names " + f"are morton decimal digits; grid type is {(grid or {}).get('type')!r})" + ) + if (grid or {}).get("sharded"): + raise ValueError( + "output.store_layout: hive does not support sharded (ShardingCodec) " + "output yet — each hive leaf is a vanilla zarr v3 store (D3); drop " + "sharded or use the flat layout" + ) + if consolidate_metadata: + raise ValueError( + "output.store_layout: hive has no store-root zarr hierarchy to " + "consolidate (D5/D12) — drop output.consolidate_metadata" + ) + # Validate bounds structure (optional) if config.bounds is not None: allowed_keys = {"temporal", "spatial"} @@ -1507,6 +1534,26 @@ def get_store_path(config: PipelineConfig) -> str | None: return config.output.get("store") +def get_store_layout(config: PipelineConfig) -> str: + """Return the STORE layout — flat single store vs morton hive (issue #199). + + Distinct from ``output.grid.layout`` (the HEALPix *array* layout inside one + store): ``output.store_layout`` selects how shard output is arranged under + ``output.store``. ``"flat"`` (default) is today's single shared zarr store. + ``"hive"`` writes one self-describing leaf zarr per shard at + ``{store}/{sign+base}/{d1}/.../{full_id}.zarr`` with a root + ``morton_hive.json`` manifest and a per-leaf commit stamp — see + ``docs/design/sparse_coverage.md`` (D1-D6) and :mod:`zagg.hive`. A + present-but-null key falls back to the default, like ``layout``. + + Returns + ------- + str + ``"flat"`` (default) or ``"hive"``. + """ + return config.output.get("store_layout") or "flat" + + def get_aoi_mask(config: PipelineConfig) -> bool: """Whether the optional strict-AOI cell mask is enabled (issue #101). diff --git a/src/zagg/grids/base.py b/src/zagg/grids/base.py index 313b6a54..a0d9a1cd 100644 --- a/src/zagg/grids/base.py +++ b/src/zagg/grids/base.py @@ -258,6 +258,17 @@ def block_index(self, shard_key: ShardKey) -> tuple[int, ...]: """ ... + def shard_label(self, shard_key: ShardKey) -> str: + """External string form of a shard key (issue #199). + + Used wherever a shard id surfaces outside the process — CSR subgroup + names, async ``.status`` object keys, log lines. HEALPix renders the + packed word as its decimal morton string (D1 in + ``docs/design/sparse_coverage.md``); rectilinear keeps the packed tile + int's decimal digits. + """ + ... + def shard_footprint(self, shard_key: ShardKey): """Shard polygon in WGS84 (shapely.Geometry).""" ... @@ -302,10 +313,23 @@ def nests_with(self, other: "OutputGrid") -> bool: ... +def shard_label(grid, shard_key) -> str: + """External string form of ``shard_key`` under ``grid`` (issue #199). + + Dispatches to ``grid.shard_label`` (decimal morton string for HEALPix, + plain int digits for rectilinear) and falls back to ``str(int(...))`` for + minimal grid stand-ins that don't implement the method (the same tolerance + the worker extends to ``iter_chunks``). + """ + fn = getattr(grid, "shard_label", None) + return fn(shard_key) if fn is not None else str(int(shard_key)) + + __all__ = [ "OutputGrid", "ShardKey", "InconsistentShardError", + "shard_label", "vector_array_spec", "chunk_array_spec", "sharded_array_spec", diff --git a/src/zagg/grids/healpix.py b/src/zagg/grids/healpix.py index eeaedcf4..427065d4 100644 --- a/src/zagg/grids/healpix.py +++ b/src/zagg/grids/healpix.py @@ -24,7 +24,7 @@ sharded_array_spec, vector_array_spec, ) -from zagg.grids.morton import to_morton_array +from zagg.grids.morton import morton_decimal, to_morton_array HEALPIX_BASE_CELLS: int = 12 # Reference order at which ``assign`` resolves points before ``cells_of`` / @@ -458,6 +458,16 @@ def block_index(self, shard_key) -> tuple[int, ...]: raise RuntimeError("block_index requires set_populated_shards() for dense layout") return (self._position_map[int(shard_key)],) + def shard_label(self, shard_key) -> str: + """Decimal morton string for this shard's packed word (issue #199). + + The external form of a HEALPix shard id (D1 in + ``docs/design/sparse_coverage.md``): CSR subgroup names, ``.status`` + object keys, and log lines all carry e.g. ``-31123``, never the raw + packed-word integer. + """ + return morton_decimal(shard_key) + def shard_footprint(self, shard_key): """Parent-cell polygon in WGS84 (lon, lat).""" from mortie.tools import mort2polygon @@ -567,10 +577,39 @@ def emit_template(self, store: Store, *, overwrite: bool = False) -> Store: spec.to_zarr(store, self.group_path, overwrite=overwrite) return store + def emit_shard_template(self, store: Store, *, overwrite: bool = False) -> Store: + """Write ONE shard's leaf-zarr template to ``store`` (issue #199 phase 2). + + The hive layout (D3 in ``docs/design/sparse_coverage.md``) gives every + shard its own self-describing leaf zarr: the same group structure as + :meth:`emit_template` but with every dense array sized to a single + shard (``cells_per_shard`` cells, K inner chunks) and a ROOT group so + the D4 commit stamp is one attrs update on an object that exists + anyway. Writes go at leaf-LOCAL block indices (0..K-1). + """ + if self.sharded: + # Validated at config load too; re-checked here because a leaf is a + # vanilla zarr v3 store (D3) — the ShardingCodec write path assumes + # the single shared store. + raise ValueError("hive leaf templates do not support sharded output") + spec = GroupSpec(members={self.group_path: self.shard_spec()}, attributes={}) + with zarr_config.set({"async.concurrency": 128}): + spec.to_zarr(store, "", overwrite=overwrite) + return store + def spec(self) -> GroupSpec: """Return the pydantic-zarr GroupSpec for this grid's template.""" return self._spec() + def shard_spec(self) -> GroupSpec: + """GroupSpec for ONE shard's hive leaf (issue #199 phase 2). + + Identical member set to :meth:`spec` — same dtypes, fills, chunking — + with the cells axis sized to one shard and the ``resolution: chunk`` + companions sized to the shard's K inner chunks. + """ + return self._group_spec(self.cells_per_shard, (self.chunks_per_shard,)) + # ── internals ──────────────────────────────────────────────────────── def _spec(self) -> GroupSpec: @@ -578,7 +617,9 @@ def _spec(self) -> GroupSpec: n_pixels = HEALPIX_BASE_CELLS * (4**self.child_order) else: n_pixels = self.n_children * self.n_shards + return self._group_spec(n_pixels, self.chunk_grid_shape) + def _group_spec(self, n_pixels: int, chunk_grid_shape: tuple[int, ...]) -> GroupSpec: base = ArraySpec( attributes={}, shape=(n_pixels,), @@ -651,7 +692,7 @@ def _shard(arr): members[name] = vector_array_spec( chunk_array_spec( spec, - chunk_grid_shape=self.chunk_grid_shape, + chunk_grid_shape=chunk_grid_shape, chunk_dims=("chunks",), ), sig, diff --git a/src/zagg/grids/morton.py b/src/zagg/grids/morton.py index 0b2b629c..d464104f 100644 --- a/src/zagg/grids/morton.py +++ b/src/zagg/grids/morton.py @@ -67,6 +67,54 @@ def to_morton_array(words): return MortonIndexArray.from_words(np.asarray(words, dtype=np.uint64)) +def morton_decimal(word) -> str: + """Decimal morton string for one packed shard-key word (issue #199). + + The external/path form of a shard id per the sparse-coverage design record + (``docs/design/sparse_coverage.md`` D1): the packed ``uint64`` word stays + the canonical in-memory/wire form, and every externally visible string — + CSR subgroup names, ``.status`` object keys, log lines — renders through + mortie's decode-through-kernel decimal repr (e.g. ``-31123``). Raises + ``ValueError`` on an empty, invalid, or negative word (a path component + must never be silently wrong) — a NEGATIVE int here is usually a *legacy + signed decimal id* handed in where the packed word belongs; parse it with + ``morton_word(str(id))`` instead. + """ + word = int(word) + if word < 0: + # np.uint64 coercion would raise an opaque OverflowError; normalize to + # the documented ValueError with the likely cause spelled out. + raise ValueError( + f"packed morton word must be non-negative (got {word}); a signed " + f"decimal id like '-4211322' is the external form — parse it with " + f"morton_word(str(id))" + ) + from mortie import MortonIndexArray + + return MortonIndexArray.from_words(np.asarray([word], dtype=np.uint64)).decimal_repr()[0] + + +def morton_word(label: str) -> int: + """Parse a decimal morton string back to its packed word (issue #199). + + The inverse of :func:`morton_decimal` at the zagg boundary — used where an + external decimal id re-enters (``--morton-cell``, CSR subgroup names on the + read path). Raises ``ValueError`` on a malformed id. + + Implementation note: this rides mortie's private-but-documented + ``_decimal_to_word`` (the issue-104 parse-back) rather than the public + ``MortonIndexArray.from_hive_path(label, suffix="")`` because the array + classes are built lazily and require pandas — the private function is + numpy-only, keeping the reader path light. The upstream ask (a public + numpy-only export) stands. Same non-injectivity caveat mortie documents: an + order-29 *point* id parses back to the *area* word (irrelevant for shard + keys at order <= 11; noted since this is a general boundary helper). + """ + from mortie.morton_index import _decimal_to_word + + return _decimal_to_word(str(label)) + + def morton_to_arrow(values): """Export ``values`` as a typed ``arro3.core.Array`` (issue #135). @@ -116,8 +164,10 @@ def is_morton_arrow(col) -> bool: "MORTON_EXTENSION_NAME", "is_morton_array", "is_morton_arrow", + "morton_decimal", "morton_from_arrow", "morton_to_arrow", + "morton_word", "morton_words", "to_morton_array", ] diff --git a/src/zagg/grids/rectilinear.py b/src/zagg/grids/rectilinear.py index d60fabb8..89be176d 100644 --- a/src/zagg/grids/rectilinear.py +++ b/src/zagg/grids/rectilinear.py @@ -274,6 +274,15 @@ def shard_local_region(self, block_index, shard_key) -> tuple: slice(lic * self.inner_w, (lic + 1) * self.inner_w), ) + def shard_label(self, shard_key) -> str: + """External string form of a rect shard key (issue #199). + + Rect shard keys are packed ``(rb, cb)`` tile ints, not morton words, so + the label is just the int's decimal digits (unlike HEALPix, which + renders the decimal morton string). + """ + return str(int(shard_key)) + @property def group_path(self) -> str: return "rectilinear" diff --git a/src/zagg/hive.py b/src/zagg/hive.py new file mode 100644 index 00000000..de91d407 --- /dev/null +++ b/src/zagg/hive.py @@ -0,0 +1,352 @@ +"""Morton-hive store layout: leaf paths, manifest, and commit stamp (issue #199). + +Phase 2 of the layout migration (``docs/design/sparse_coverage.md`` §2-§3). +Under ``output.store_layout: hive`` each shard is its own **self-describing +leaf zarr** under a morton digit tree:: + + {store_root}/ + morton_hive.json <- static manifest (§3); root-only exception + {sign+base}/{d1}/.../{d_n}/ <- one decimal digit per level (D2) + {full_id}.zarr/ <- vanilla zarr v3 leaf (D3) + +- Ids are morton decimal strings (D1); the leaf path is computed by mortie's + ``hive_path`` (the convention is owned by the mortie spec) and re-checked + here against the node invariant. +- **Node invariant (D5)**: below the root a node contains only digit children + (``[1-4]/``, or the ``{sign+base}`` component at the first level) and + ``*.zarr`` objects — zero zarr metadata above the leaf, so 2,000 workers + share no mutable state and a delimiter-LIST with no digit children is a + definitive "nothing finer exists". +- **Manifest (D6)**: ``morton_hive.json`` is written once at template time and + never touched during a run; with it every shard path is computable with zero + requests. The convention is versioned (``morton-hive/1``) from day one. +- **Commit stamp (D4)**: the shard's FINAL write is a root + ``group.attrs.update(...)`` marking completion (plus cell count, timestamp, + granule count). A ``.zarr/`` prefix whose root metadata lacks the stamp is + debris — incomplete, ignorable, safe to overwrite on retry. This is NOT + consolidated metadata: one small PUT on an object that must exist anyway. + +The coverage MOC (§4) and pyramid sweep (§7) are follow-on issues; the +manifest's ``pyramid`` block is declared-only in round one (D11/D12). +""" + +from __future__ import annotations + +import json +from datetime import datetime, timezone + +import numpy as np +import zarr +from zarr.errors import GroupNotFoundError + +from zagg.store import open_object_store + +#: Convention version recorded in the manifest and the commit stamp (D6). +HIVE_SPEC = "morton-hive/1" +#: Root manifest object name (the root-only exception to the node invariant). +MANIFEST_NAME = "morton_hive.json" +#: Root-group attrs key carrying the commit stamp (D4). +COMMIT_ATTR = "morton_hive_commit" + + +def shard_leaf_path(store_root: str, shard_key) -> str: + """Absolute path of a shard's leaf zarr under ``store_root`` (D2/D3). + + Computed by mortie's ``hive_path`` — the layout convention is owned by the + mortie spec — and re-checked against the node invariant (D5) so a future + drift in either side fails loudly instead of writing a stray prefix. + Raises ``ValueError`` on an invalid shard key. + """ + from mortie import MortonIndexArray + + word = int(shard_key) + if word < 0: + raise ValueError( + f"shard key must be a packed morton word (got {word}); parse a decimal " + f"id with zagg.grids.morton.morton_word first" + ) + rel = MortonIndexArray.from_words(np.asarray([word], dtype=np.uint64)).hive_path()[0] + check_node_invariant(rel) + return f"{store_root.rstrip('/')}/{rel}" + + +def check_node_invariant(rel_path: str) -> None: + """Raise unless ``rel_path`` is a legal hive leaf path (D5). + + Below the root only digit components are allowed — ``{sign+base}`` + (optional ``-``, one digit ``1..6``) at the first level, one ``1..4`` digit + per level after — terminating in ``{full_id}.zarr`` whose id equals the + concatenated components. This is the walker's contract: any other name + under the root (bar the manifest and the future ``coverage.moc``) breaks + child classification. + """ + parts = rel_path.strip("/").split("/") + leaf = parts[-1] + ok = len(parts) >= 2 and leaf.endswith(".zarr") + if ok: + head, digits = parts[0], parts[1:-1] + ok = _is_base_component(head) + ok = ok and all(len(d) == 1 and d in "1234" for d in digits) + ok = ok and leaf[: -len(".zarr")] == head + "".join(digits) + if not ok: + raise ValueError(f"path {rel_path!r} violates the hive node invariant (D5)") + + +def build_manifest(grid, dataset: dict | None = None) -> dict: + """Build the static ``morton_hive.json`` payload for one store (§3, D6). + + ``grid`` supplies the orders; ``dataset`` (typically the ShardMap's + ``metadata``) supplies identity — only ``short_name`` and ``version`` are + recorded. The split schedule is implicit under D2 (one digit per level down + to the shard order) but recorded explicitly for forward compatibility; the + ``pyramid`` block is declared-only in round one (D11: overviews are a + second-pass sweep, never written at fan-out time). + """ + dataset = dataset or {} + return { + "spec": HIVE_SPEC, + "dataset": { + "short_name": dataset.get("short_name"), + "version": dataset.get("version"), + }, + "cell_order": int(grid.child_order), + "shard_order": int(grid.parent_order), + "split_schedule": [1] * int(grid.parent_order), + "pyramid": {"orders": [], "aggregation": {}}, + "generated_at": _utcnow(), + } + + +def ensure_manifest(store_root: str, manifest: dict, *, overwrite: bool = False, **store_kwargs): + """Write the root manifest once at template time; verify it on reruns. + + A retry into an existing hive store must be able to proceed (that is the + D4 debris/retry model), so an existing manifest is accepted — but only if + its FROZEN keys match the run's own (:data:`_FROZEN_MANIFEST_KEYS`: orders + + identity + schedule — the flat path's ``_check_signature`` analogue). + ``generated_at`` and ``pyramid`` are excluded: the pyramid block is + populated/updated by the §7 sweep by design (D11), so comparing it would + brick every resume after the first sweep. + + ``overwrite=True`` replaces the MANIFEST ONLY — it never clears data. To + guard against the silent-corruption footgun (committed leaves from the old + orders would survive a "re-template" and be indistinguishable from legal + mixed-order data, D2), an overwrite that CHANGES the frozen keys refuses + when the digit tree already has children (one delimiter-LIST); clear the + store root first. Returns the manifest now in effect. + """ + import obstore + + store = open_object_store(store_root, **store_kwargs) + existing = _read_json(store, MANIFEST_NAME) + frozen_matches = existing is not None and _frozen(existing) == _frozen(manifest) + if existing is not None and not overwrite: + if not frozen_matches: + raise ValueError( + f"{MANIFEST_NAME} at {store_root} does not match this run " + f"(existing {existing!r} vs {manifest!r}); this store was templated " + f"for different orders/identity — clear the store root (or pick a " + f"new one) before writing with this configuration" + ) + return existing + if overwrite and existing is not None and not frozen_matches: + # One delimiter-LIST: a {sign+base}-shaped child means shards were + # already written under the OLD configuration. Their leaves are + # stamped and walker-discoverable, so replacing just the manifest + # would leave them masquerading as legal mixed-order data (D2). + listing = obstore.list_with_delimiter(store) + children = [p.rstrip("/").split("/")[-1] for p in listing["common_prefixes"]] + if any(_is_base_component(c) for c in children): + raise ValueError( + f"refusing to overwrite {MANIFEST_NAME} at {store_root} with " + f"different orders/identity: the digit tree already has shard " + f"data (e.g. {children[0]!r}/), and overwrite replaces the " + f"manifest only — clear the store root first" + ) + obstore.put(store, MANIFEST_NAME, json.dumps(manifest, indent=1).encode()) + return manifest + + +#: Manifest keys the resume match-check compares (orders + identity + schedule). +#: ``generated_at`` (a timestamp) and ``pyramid`` (populated by the §7 sweep, +#: D11) are mutable by design and excluded. +_FROZEN_MANIFEST_KEYS = ("spec", "dataset", "cell_order", "shard_order", "split_schedule") + + +def _frozen(manifest: dict) -> dict: + """The frozen-key projection of a manifest (resume/overwrite match-check).""" + return {k: manifest.get(k) for k in _FROZEN_MANIFEST_KEYS} + + +def _is_base_component(name: str) -> bool: + """Whether ``name`` is a ``{sign+base}``-shaped hive root child (D5).""" + base = name[1:] if name.startswith("-") else name + return len(base) == 1 and base in "123456" + + +def read_manifest(store_root: str, **store_kwargs) -> dict | None: + """Read ``morton_hive.json`` from a store root; ``None`` when absent.""" + return _read_json(open_object_store(store_root, **store_kwargs), MANIFEST_NAME) + + +def stamp_commit(leaf_store, *, cells_with_data: int, granule_count: int) -> None: + """Stamp a shard leaf complete — the shard's FINAL write (D4). + + One small PUT rewriting the leaf's root ``zarr.json`` (which the template + already created), not consolidation. Until this lands, the leaf prefix is + debris: a worker that dies mid-shard leaves no stamp, and a retry may + overwrite the prefix wholesale. + """ + group = zarr.open_group(leaf_store, path="", mode="r+", zarr_format=3) + group.attrs[COMMIT_ATTR] = { + "spec": HIVE_SPEC, + "complete": True, + "cells_with_data": int(cells_with_data), + "granule_count": int(granule_count), + "written_at": _utcnow(), + } + + +def read_commit(leaf_store) -> dict | None: + """The leaf's commit stamp, or ``None`` for debris / absent leaves (D4). + + Absence (no root group at all) and an unstamped root are the same answer: + the shard is not complete. Presence requires the stamp — never infer + completeness from the ``.zarr/`` prefix existing. + """ + try: + group = zarr.open_group(leaf_store, path="", mode="r", zarr_format=3) + except (FileNotFoundError, GroupNotFoundError): + return None + stamp = group.attrs.get(COMMIT_ATTR) + # A malformed (non-mapping) stamp is debris too — never half-trusted. + return dict(stamp) if isinstance(stamp, dict) else None + + +def leaf_block_index(grid, block_index, shard_key) -> tuple: + """Leaf-LOCAL storage block for a chunk in a hive leaf (issue #199 phase 2). + + The hive leaf's arrays are sized to one shard, so a chunk's block index is + its position WITHIN the shard, not the global block ``iter_chunks`` yields. + Derived from the existing ``shard_local_region`` seam (the sharded path's + within-shard placement): the region's start divided by the chunk extent. + At K==1 this is always ``(0,)``. + """ + region = grid.shard_local_region(block_index, shard_key) + return tuple(int(s.start) // int(c) for s, c in zip(region, grid.chunk_shape)) + + +def process_and_write_hive( + shard_key, + granule_urls, + grid, + s3_creds, + store_root, + config, + *, + store_kwargs, + driver=None, + handoff="arrow", + aoi_payload=None, + profile=False, +): + """Process one shard into its own hive leaf store (issue #199 phase 2). + + The SHARED per-shard write path for both backends (phase 3): the local + runner's ``_cell_work`` and the Lambda handler's hive branch both call + this, so leaf templating, chunk placement, CSR naming, and stamp ordering + cannot drift between dispatchers. The shard's output is a self-describing + leaf zarr at :func:`shard_leaf_path` ``(store_root, shard_key)`` (D3), with + dense chunks written at leaf-LOCAL block indices and — as the shard's + FINAL write — the D4 commit stamp on the leaf's root group. The leaf + template is emitted lazily on the first chunk write (mirroring the Lambda + handler's lazy store open), so a no-data shard never creates the + ``.zarr/`` prefix; a worker that dies mid-shard leaves an UNSTAMPED prefix + — debris, overwritten wholesale on retry (``overwrite=True`` on the leaf + template makes the retry idempotent). ``profile`` forwards to + ``process_shard`` (issue #100); the write phase is interleaved with the + stream on this path, so no separate ``write`` timing is recorded. + """ + from zagg.grids.base import shard_label + from zagg.processing import process_shard, write_dataframe_to_zarr, write_ragged_to_zarr + from zagg.store import open_store + + leaf_path = shard_leaf_path(store_root, shard_key) + box: dict = {} + + def _leaf(): + if "store" not in box: + store = open_store(leaf_path, **store_kwargs) + # overwrite=True: any existing prefix here is either debris from a + # torn run (D4) or a prior committed write being redone — both are + # replaced wholesale; per-leaf state never blocks a retry. + grid.emit_shard_template(store, overwrite=True) + box["store"] = store + return box["store"] + + single_chunk = int(getattr(grid, "chunks_per_shard", 1)) == 1 + + def _write_chunk(block_index, carrier, ragged): + store = _leaf() + local = leaf_block_index(grid, block_index, shard_key) + write_dataframe_to_zarr(carrier, store, grid=grid, chunk_idx=local) + # CSR subgroup naming inside a leaf mirrors the flat layout: the shard + # label at K==1; the LOCAL chunk ordinal at K>1 (leaf arrays are + # 1-D — hive is HEALPix-only, validated at config load). + ragged_key = shard_label(grid, shard_key) if single_chunk else int(local[0]) + write_ragged_to_zarr(ragged, store, grid=grid, shard_key=ragged_key) + + _df_out, metadata = process_shard( + grid, + int(shard_key), + granule_urls, + s3_credentials=s3_creds, + config=config, + driver=driver, + handoff=handoff, + aoi_payload=aoi_payload, + write_chunk=_write_chunk, + profile=profile, + ) + # Stamp ONLY a fully-written leaf: an errored shard (or one that streamed + # no chunks) stays unstamped — debris by definition (D4). The stamp is the + # last write, so its presence certifies everything before it landed. + if "store" in box and not metadata.get("error"): + stamp_commit( + box["store"], + cells_with_data=metadata.get("cells_with_data", 0), + granule_count=metadata.get("granule_count", 0), + ) + return metadata + + +def _read_json(obj_store, key: str) -> dict | None: + """GET+parse one small JSON object; ``None`` when it does not exist.""" + import obstore + from obstore.exceptions import NotFoundError + + try: + data = obstore.get(obj_store, key).bytes() + except (FileNotFoundError, NotFoundError): + return None + return json.loads(bytes(data)) + + +def _utcnow() -> str: + return datetime.now(timezone.utc).isoformat(timespec="seconds") + + +__all__ = [ + "COMMIT_ATTR", + "HIVE_SPEC", + "MANIFEST_NAME", + "build_manifest", + "check_node_invariant", + "ensure_manifest", + "leaf_block_index", + "process_and_write_hive", + "read_commit", + "read_manifest", + "shard_leaf_path", + "stamp_commit", +] diff --git a/src/zagg/processing/worker.py b/src/zagg/processing/worker.py index 0665227f..3957c239 100644 --- a/src/zagg/processing/worker.py +++ b/src/zagg/processing/worker.py @@ -35,6 +35,7 @@ get_data_vars, get_output_signature, ) +from zagg.grids.base import shard_label from zagg.index import index_from_config from zagg.processing.aggregate import ( _aggregate_chunk_cells, @@ -205,7 +206,11 @@ def process_shard( index_backend = index_from_config(config) shard_key = int(shard_key) - logger.info(f"Processing shard: {shard_key}") + # Log lines carry the external shard label (decimal morton string for + # HEALPix — issue #199); ``shard_key`` itself stays the packed int (the + # canonical wire/metadata form). + label = shard_label(grid, shard_key) + logger.info(f"Processing shard: {label}") start_time = datetime.now() # Resolve driver @@ -234,7 +239,7 @@ def process_shard( # Check for granules if not granule_urls: - logger.info(f" No granules provided for shard {shard_key} - skipping") + logger.info(f" No granules provided for shard {label} - skipping") metadata["error"] = "No granules found" metadata["duration_s"] = (datetime.now() - start_time).total_seconds() return pd.DataFrame(), metadata @@ -458,12 +463,12 @@ def _iter_granule_reads(): # message is "no data AND N raised", not "all groups raised". if read_errors: logger.warning( - f" No data after filtering for shard {shard_key} and " + f" No data after filtering for shard {label} and " f"{read_errors} group read(s) raised - skipping" ) metadata["error"] = f"No data after filtering ({read_errors} group reads raised)" else: - logger.info(f" No data after filtering for shard {shard_key} - skipping") + logger.info(f" No data after filtering for shard {label} - skipping") metadata["error"] = "No data after filtering" metadata["duration_s"] = (datetime.now() - start_time).total_seconds() if profile: @@ -628,7 +633,7 @@ def _iter_granule_reads(): metadata["phase_timings"] = phase_timings duration = (datetime.now() - start_time).total_seconds() - logger.info(f"Completed shard {shard_key} in {duration:.1f}s") + logger.info(f"Completed shard {label} in {duration:.1f}s") metadata["cells_with_data"] = cells_with_data metadata["total_obs"] = n_obs_total diff --git a/src/zagg/processing/write.py b/src/zagg/processing/write.py index 20abd358..95b6f3b3 100644 --- a/src/zagg/processing/write.py +++ b/src/zagg/processing/write.py @@ -496,7 +496,7 @@ def write_ragged_to_zarr( store: Store, *, grid, - shard_key: int, + shard_key: int | str, ) -> Store: """Write a shard's ``kind: ragged`` (CSR) fields to the Zarr store (issue #48). @@ -514,9 +514,10 @@ def write_ragged_to_zarr( At **cell resolution** (default) ``cell_ids[k]`` is each populated cell's position in the chunk's ``children`` block (the index collected by - ``process_shard``); the per-shard subgroup name is the ``shard_key`` (the - coverage cell's morton id for HEALPix), recovered by the reader directly from - the store. + ``process_shard``); the per-shard subgroup name is the ``shard_key`` label + (the coverage cell's **decimal morton string** for HEALPix — issue #199, D1 + in ``docs/design/sparse_coverage.md``), recovered by the reader directly + from the store. At **chunk resolution** (``resolution: chunk``, issue #82) a ragged field stores ONE variable-length payload per chunk, not per cell. The populated @@ -540,9 +541,10 @@ def write_ragged_to_zarr( grid : OutputGrid Provides ``group_path`` for routing the write (and ``config`` for the per-field dtype + resolution). - shard_key : int - Shard identifier; the CSR subgroup name (one chunk per shard at cell - resolution). + shard_key : int or str + CSR subgroup name, used verbatim: the grid's shard label at cell + resolution (``grid.shard_label`` — decimal morton string for HEALPix, + issue #199), or the flattened block-index int at K>1. Returns ------- @@ -553,7 +555,7 @@ def write_ragged_to_zarr( return store agg_fields = get_agg_fields(grid.config) if getattr(grid, "config", None) else {} chunk_res_fields = _chunk_resolution_fields(getattr(grid, "config", None)) - shard_key = int(shard_key) + shard_key = str(shard_key) for name, entry in ragged.items(): # Located fields (issue #87) deliver (values_list, cell_ids, locations_list); # unlocated fields keep the 2-tuple. diff --git a/src/zagg/readers/tdigest_tensor.py b/src/zagg/readers/tdigest_tensor.py index 8f11caea..4f8dee37 100644 --- a/src/zagg/readers/tdigest_tensor.py +++ b/src/zagg/readers/tdigest_tensor.py @@ -19,9 +19,11 @@ ``cell_ids[k]`` is a cell's position in the chunk's row-major ``(64, 64)`` children block, and ``values[offsets[k]:offsets[k+1]]`` is that cell's ``(k_centroids, 2)`` ``(mean, weight)`` digest. The subgroup name is the -chunk's morton index, recovered directly from the store (issue #79 design -decision (3)); when a sibling ``{field}/morton`` ``uint64`` coordinate array is -present it is used in preference, mapping chunk order → morton id. +chunk's morton index — its **decimal morton string** since issue #199 (D1 in +``docs/design/sparse_coverage.md``) — recovered directly from the store (issue +#79 design decision (3)); when a sibling ``{field}/morton`` ``uint64`` +coordinate array is present it is used in preference, mapping chunk order → +morton id. The reader (generator) ---------------------- @@ -43,6 +45,7 @@ from zarr.abc.store import Store from zagg.csr import iter_csr_cells, read_csr +from zagg.grids.morton import morton_word from zagg.stats.tdigest import cdf_from_tdigest, quantile_from_tdigest __all__ = [ @@ -209,6 +212,47 @@ def chunk_z_range( raise ValueError(f"unknown fit mode {fit!r}") +def _is_decimal_shard_name(name: str) -> bool: + """Whether a CSR subgroup name fits the decimal morton *shard-id* grammar. + + Grammar: optional sign, base digit ``1..6``, then one ``1..4`` digit per + order — with at least ONE order digit. zagg shard orders are >= 1, so a + real shard id always has >= 2 digits; that floor is what disambiguates the + lone digits ``"1".."6"``, which the raw grammar would accept as order-0 ids + but which in practice are small block-index keys (review finding, PR #205). + """ + body = name[1:] if name.startswith("-") else name + return len(body) >= 2 and body[:1] in "123456" and all(d in "1234" for d in body[1:]) + + +def _subgroup_words(shard_keys: list[str]) -> dict[str, int]: + """Parse CSR subgroup names to their numeric keys (issue #199). + + Shard-keyed subgroups are named by the decimal morton string (D1), so the + numeric key is the parsed packed word; block-index-keyed subgroups (the K>1 + layout) keep plain non-negative ints. The flavor is decided per store: + + - every name fits the shard-id grammar (:func:`_is_decimal_shard_name`) + -> decimal morton store; keys are the parsed packed words; + - otherwise -> block-keyed store; keys are ``int(name)`` — and a leading + ``-`` in this branch raises, because a signed name is provably NOT a + block index (block keys are non-negative flattened indices): a mixed + store (e.g. pre-#199 debris alongside decimal-named subgroups) must fail + loudly rather than silently mis-key every shard. + """ + if shard_keys and all(_is_decimal_shard_name(k) for k in shard_keys): + return {k: morton_word(k) for k in shard_keys} + signed = [k for k in shard_keys if k.startswith("-")] + if signed: + raise ValueError( + f"CSR subgroup names mix decimal morton ids with non-id names " + f"(signed name(s) {signed[:4]} cannot be block-index keys); the " + f"store layout is ambiguous — likely pre-issue-199 debris next to " + f"decimal-named subgroups. Rewrite the store." + ) + return {k: int(k) for k in shard_keys} + + def _resolve_chunk_morton( store: Store, field: str, shard_keys: list[str], zarr_format: Literal[2, 3] ) -> dict[str, int]: @@ -216,22 +260,24 @@ def _resolve_chunk_morton( Prefers a sibling ``{field}/morton`` ``uint64`` coordinate array (chunk order → morton id); falls back to parsing the subgroup name as the parent - morton id (the shard key is the parent morton — see :func:`process_shard`). + morton id (the shard key is the parent morton — see :func:`process_shard`), + a decimal morton string since issue #199. The coordinate array is aligned against the subgroup names sorted - **numerically** (the canonical ascending-morton chunk order), not - lexicographically — string sorting would mis-align names of differing digit - counts (e.g. ``"1000"`` before ``"99"``). + **numerically** by parsed key (the canonical ascending-morton chunk order), + not lexicographically — string sorting would mis-align names of differing + digit counts (e.g. ``"1000"`` before ``"99"``). """ + words = _subgroup_words(shard_keys) try: arr = zarr.open_array(store, path=f"{field}/morton", mode="r", zarr_format=zarr_format) morton = np.asarray(arr[...]) except (FileNotFoundError, KeyError, ValueError): - return {k: int(k) for k in shard_keys} + return words if len(morton) != len(shard_keys): # Coordinate present but not 1:1 with subgroups — fall back to the names. - return {k: int(k) for k in shard_keys} - return {k: int(m) for k, m in zip(sorted(shard_keys, key=int), morton)} + return words + return {k: int(m) for k, m in zip(sorted(shard_keys, key=words.__getitem__), morton)} def _shard_groups( diff --git a/src/zagg/runner.py b/src/zagg/runner.py index 512f8ad9..c7e7a82c 100644 --- a/src/zagg/runner.py +++ b/src/zagg/runner.py @@ -42,6 +42,7 @@ get_output_region, get_parent_order, get_pipeline_type, + get_store_layout, get_store_path, ) from zagg.dispatch import ( @@ -54,6 +55,8 @@ PreflightReport, dispatch, ) +from zagg.grids.base import shard_label +from zagg.grids.morton import morton_word from zagg.processing import ( process_shard, write_dataframe_to_zarr, @@ -171,7 +174,8 @@ def agg( Lambda-only (issue #151). ``"async"`` (default) dispatches each cell with ``InvocationType="Event"`` and polls a per-shard result object the worker writes next to the output store - (``.status//.json``), so no synchronous + (``.status//.json``, where the label is + the decimal morton string for HEALPix — issue #199), so no synchronous connection sits idle while the shard runs -- shards longer than a NAT idle window (~4 min on GitHub-hosted runners) complete reliably, and the 6 MB synchronous response cap no longer applies. Requires (a) a @@ -907,7 +911,9 @@ def _select_cells( catalog_data : dict Loaded catalog (shard_keys/granules format). morton_cell : str, optional - Process a single shard, identified by stringified key. + Process a single shard. For a HEALPix catalog this is the shard's + decimal morton string (e.g. ``-31123`` — issue #199); for other grids + it is the stringified shard-key int. max_cells : int, optional Truncate to the first N shards. @@ -922,10 +928,33 @@ def _select_cells( """ pairs = list(zip(catalog_data["shard_keys"], catalog_data["granules"])) if morton_cell: - target = int(morton_cell) + grid_type = (catalog_data.get("grid_signature") or {}).get("type") + if grid_type == "healpix": + try: + target = morton_word(morton_cell) + except ValueError as e: + raise ValueError( + f"--morton-cell {morton_cell!r} is not a decimal morton id " + f"(shard ids are decimal morton strings since issue #199): {e}" + ) from e + else: + target = int(morton_cell) matches = [(k, urls) for k, urls in pairs if k == target] if not matches: - raise ValueError(f"shard '{morton_cell}' not in catalog") + msg = f"shard '{morton_cell}' not in catalog" + # A well-formed decimal id can still miss because the catalog itself + # predates the packed-word form: legacy shard_keys are signed i64 + # decimal ids (small or negative), while packed words carry a 1..12 + # prefix in the top 4 bits (always >= 2^60). Hard break — no shim — + # but say why the lookup likely failed (review finding, PR #205). + keys = catalog_data["shard_keys"] + if grid_type == "healpix" and any(int(k) < (1 << 60) for k in keys): + msg += ( + " (catalog shard_keys look like legacy signed decimal ids, not " + "packed morton words — a pre-issue-199 shard map; regenerate it " + "with `python -m zagg.catalog`)" + ) + raise ValueError(msg) return matches if max_cells: pairs = pairs[:max_cells] @@ -937,6 +966,22 @@ def _select_cells( return pairs +def _safe_label(grid, shard_key) -> str: + """Render a shard key for log/report lines; NEVER raises (issue #199). + + ``shard_label`` -> ``morton_decimal`` raises on an invalid word — right at + path-construction sites (a path component must never be silently wrong), + wrong in error-*reporting* paths: re-rendering the same malformed key that + made a cell fail would abort the accumulation loop precisely while + reporting that failure (review finding, PR #205; mortie's own scalar repr + is non-raising for the same reason). Falls back to the raw digits. + """ + try: + return shard_label(grid, shard_key) + except Exception: + return str(shard_key) + + def _lambda_dispatch_order(cells: list[tuple]) -> list[tuple]: """Order cells for lambda fan-out: biggest work first, in coarse buckets. @@ -1069,8 +1114,12 @@ def _write_chunk(block_index, carrier, ragged): # table), so no carrier-specific emptiness check is needed here. write_dataframe_to_zarr(carrier, zarr_store, grid=grid, chunk_idx=block_index) # Persist this chunk's ragged (CSR) fields — one CSR group per field per - # chunk (issue #48). No-ops when ``ragged`` is empty. - ragged_key = int(shard_key) if single_chunk else _block_index_key(block_index, grid) + # chunk (issue #48). No-ops when ``ragged`` is empty. At K==1 the CSR + # subgroup is named by the grid's shard label (the decimal morton string + # for HEALPix — issue #199); K>1 keeps the flattened block-index int. + ragged_key = ( + shard_label(grid, shard_key) if single_chunk else _block_index_key(block_index, grid) + ) write_ragged_to_zarr(ragged, zarr_store, grid=grid, shard_key=ragged_key) # Sharded output (issue #108): the shard's K inner chunks bundle into one @@ -1155,13 +1204,30 @@ def _run_local( else: grid = from_config(config) _check_signature(grid, catalog_data) - zarr_store = open_store( - store_path, - region=region, - credentials=output_credentials, - endpoint_url=output_endpoint_url, - ) - zarr_store = grid.emit_template(zarr_store, overwrite=overwrite) + store_layout = get_store_layout(config) + store_kwargs = { + "region": region, + "credentials": output_credentials, + "endpoint_url": output_endpoint_url, + } + if store_layout == "hive": + # Hive layout (issue #199 phase 2): no shared zarr template — zero + # metadata above the leaves (D5). Template time writes only the root + # manifest (D6); each shard emits its own leaf template lazily inside + # hive.process_and_write_hive (the leaf write path lives in zagg.hive, + # next to the manifest/stamp machinery it exercises). + from zagg.hive import build_manifest, ensure_manifest, process_and_write_hive + + ensure_manifest( + store_path, + build_manifest(grid, dataset=catalog_data.get("metadata")), + overwrite=overwrite, + **store_kwargs, + ) + zarr_store = None + else: + zarr_store = open_store(store_path, **store_kwargs) + zarr_store = grid.emit_template(zarr_store, overwrite=overwrite) # Per-cell work, catching its own exceptions so one bad cell counts as an # error and the run continues (the old loop's ``except`` branch). The @@ -1186,18 +1252,32 @@ def _cell_work(payload): ds = _clamped_data_source(config.data_source, len(_resolve_urls(records, driver))) cell_config = replace(config, data_source=ds) if ds is not None else config try: - meta = _process_and_write( - shard_key, - grid.block_index(int(shard_key)), - records, - grid, - s3_creds, - zarr_store, - cell_config, - driver=driver, - handoff=handoff, - **extra, - ) + if store_layout == "hive": + meta = process_and_write_hive( + shard_key, + _resolve_urls(records, driver), + grid, + s3_creds, + store_path, + cell_config, + store_kwargs=store_kwargs, + driver=driver, + handoff=handoff, + **extra, + ) + else: + meta = _process_and_write( + shard_key, + grid.block_index(int(shard_key)), + records, + grid, + s3_creds, + zarr_store, + cell_config, + driver=driver, + handoff=handoff, + **extra, + ) return {"shard_key": shard_key, "ok": True, "meta": meta} except Exception as e: return {"shard_key": shard_key, "ok": False, "error": e} @@ -1212,21 +1292,23 @@ def _cell_work(payload): n = len(cells) def _accumulate(report, i, outcome): - shard_key = outcome["shard_key"] + # _safe_label, not shard_label: a malformed key that failed the cell + # must not also kill the loop reporting that failure (issue #199). + label = _safe_label(grid, outcome["shard_key"]) if not outcome["ok"]: report.cells_error += 1 - logger.warning(f" [{i}/{n}] {shard_key}: ERROR {outcome['error']}") + logger.warning(f" [{i}/{n}] {label}: ERROR {outcome['error']}") return meta = outcome["meta"] report.results.append(meta) if meta.get("error"): - logger.info(f" [{i}/{n}] {shard_key}: {meta['error']}") + logger.info(f" [{i}/{n}] {label}: {meta['error']}") else: obs = meta.get("total_obs", 0) report.total_obs += obs report.cells_with_data += 1 if i % 10 == 0 or n <= 20: - logger.info(f" [{i}/{n}] {shard_key}: {obs:,} obs") + logger.info(f" [{i}/{n}] {label}: {obs:,} obs") start_time = time.time() try: @@ -1346,7 +1428,8 @@ def _run_lambda( # Async result channel (issue #151): a per-run unique status prefix next to # the output store. Each worker mirrors its response envelope to - # /.json and the dispatch threads poll for it instead of + # /.json (decimal morton string for HEALPix, issue + # #199) and the dispatch threads poll for it instead of # holding a synchronous invoke connection open -- GitHub-hosted runners sit # behind a ~4 min NAT idle timeout that severed every >250 s benchmark # target. The run_id keeps reruns into the same store from reading stale @@ -1414,6 +1497,17 @@ def _preflight(n): # inline ``executor.submit(_invoke_lambda_cell, ...)`` passed. def _cell_work(payload): shard_key, records = payload + # Rendered once per cell: the status-object name (below) and the + # payload-cap error message in _invoke_lambda_cell both carry it + # (issue #199). On ASYNC runs the label becomes a path component (the + # status key), so it must raise on a malformed key; on SYNC runs it is + # purely cosmetic and a raise out of _cell_work would be RUN-fatal + # (dispatch() re-raises), so fall back to the raw digits instead + # (review finding, PR #205). + if result_prefix is not None: + label = shard_label(grid, shard_key) + else: + label = _safe_label(grid, shard_key) # Only thread aoi_payload when the manifest carries a mask (flag on); # otherwise omit the kwarg so the event payload is byte-identical to the # pre-feature path (issue #101). Mirrors the local runner's _cell_work. @@ -1425,7 +1519,9 @@ def _cell_work(payload): # timeout + queue/write margin). Sync runs pass none of these, keeping # the invoke byte-identical to the legacy path. if result_prefix is not None: - key = f"{int(shard_key)}.json" + # Status objects are named by the shard label — the decimal morton + # string for HEALPix (issue #199) — not the raw packed word. + key = f"{label}.json" extra["result_url"] = f"{result_prefix}/{key}" extra["result_fetch"] = _result_fetcher( result_box, result_prefix, output_creds_event, region, key @@ -1456,6 +1552,7 @@ def _cell_work(payload): max_workers=state["workers"], handoff=handoff, profile=profile, + label=label, **extra, ) @@ -1496,6 +1593,15 @@ def _finalize_fn(): # calls that already happen, so no worker probe tax -- issue #100). They # decompose wall time into setup invoke / fan-out / finalize invoke so # "where did wall time go" is answerable from the summary. + # Hive layout (issue #199 phase 3): setup mode writes morton_hive.json + # instead of a global template. The worker builds the manifest itself from + # the event's config; only the dataset identity (from the ShardMap + # metadata, matching the local path's source) rides in the event — flat + # runs omit the key, keeping their setup event byte-identical. + dataset = None + if get_store_layout(config) == "hive": + md = catalog_data.get("metadata") or {} + dataset = {"short_name": md.get("short_name"), "version": md.get("version")} setup_start = time.time() _invoke_lambda_setup( state["lambda_client"], @@ -1507,6 +1613,7 @@ def _finalize_fn(): overwrite=overwrite, config_dict=config_dict, output_creds_event=output_creds_event, + dataset=dataset, ) setup_s = time.time() - setup_start @@ -1521,7 +1628,10 @@ def _accumulate(report, i, result): report.cells_with_data += 1 elif error not in ("No granules found", "No data after filtering"): report.cells_error += 1 - logger.warning(f" [{i}/{n}] shard {result.get('shard_key')}: {error}") + key = result.get("shard_key") + # _safe_label: error reporting must not raise on the bad key itself. + label = _safe_label(grid, key) if key is not None else key + logger.warning(f" [{i}/{n}] shard {label}: {error}") report.results.append(result) if i % 50 == 0: @@ -2196,8 +2306,25 @@ def _invoke_lambda_setup( overwrite, config_dict, output_creds_event=None, + dataset=None, ): - """Invoke Lambda in setup mode to create the zarr template.""" + """Invoke Lambda in setup mode to create the zarr template. + + For a hive-layout config (issue #199 phase 3) the worker writes + ``morton_hive.json`` instead; ``dataset`` (``{"short_name", "version"}`` + from the ShardMap metadata) is threaded through for the manifest's + identity block. The key is added only when set, so flat setup events stay + byte-identical. + + Stale-deployment guard (review finding, PR #205): a function deployed + before phase 3 has no hive branch — it would emit the flat GLOBAL template + at the hive root, return the same 200 body, and the whole run would then + silently write a flat store. The phase-3 handler echoes the layout it + acted on (``"layout"`` in the response body), and a hive dispatch REQUIRES + that echo — failing fast at setup with a redeploy message instead of after + a full fan-out. Flat dispatch does not require the echo, so a new + dispatcher keeps working against old deployed functions for flat runs. + """ event = { "mode": "setup", "store_path": store_path, @@ -2207,6 +2334,8 @@ def _invoke_lambda_setup( "overwrite": overwrite, "config": config_dict, } + if dataset is not None: + event["dataset"] = dataset if output_creds_event is not None: event["output_credentials"] = output_creds_event response = lambda_client.invoke( @@ -2220,6 +2349,19 @@ def _invoke_lambda_setup( result = json.loads(payload) if result.get("statusCode") != 200: raise RuntimeError(f"Lambda setup error: {result.get('body')}") + requested = ((config_dict or {}).get("output") or {}).get("store_layout") or "flat" + if requested == "hive": + try: + echoed = json.loads(result.get("body") or "{}").get("layout") + except (TypeError, ValueError): + echoed = None + if echoed != "hive": + raise RuntimeError( + f"Lambda setup did not confirm the hive store layout (response body " + f"{result.get('body')!r}): the deployed function predates issue #199 " + f"phase 3 and would write a FLAT store at the hive root — redeploy " + f"the function before dispatching hive runs" + ) def _invoke_lambda_finalize(lambda_client, function_name, store_path, output_creds_event=None): @@ -2261,9 +2403,14 @@ def _invoke_lambda_cell( result_url=None, result_fetch=None, poll_timeout_s=None, + label=None, ): """Invoke Lambda for a single cell with retry logic. + ``label`` (issue #199) is the shard's rendered external id (decimal morton + string for HEALPix) for user-facing messages; ``None`` falls back to the + raw ``shard_key`` digits. The event payload always carries the int. + ``max_workers`` is used only for the file-descriptor-exhaustion message (#28); it does not affect dispatch. ``profile`` (issue #100) forwards a ``"profile": true`` event key so the worker emits ``phase_timings``; when @@ -2330,7 +2477,7 @@ def _invoke_lambda_cell( payload = json.dumps(event) if invocation_type == "Event" and len(payload) > _ASYNC_PAYLOAD_CAP_BYTES: raise ValueError( - f"cell {shard_key} event payload is {len(payload):,} bytes, over the " + f"cell {label or shard_key} event payload is {len(payload):,} bytes, over the " f"{_ASYNC_PAYLOAD_CAP_BYTES:,}-byte async dispatch budget (Lambda caps " 'Event invokes at 256 KB): pass invocation="sync" for this run, or ' "shrink the per-cell payload (e.g. the strict-AOI aoi_payload)" diff --git a/tests/test_benchmark.py b/tests/test_benchmark.py index 61210f3a..27d089c3 100644 --- a/tests/test_benchmark.py +++ b/tests/test_benchmark.py @@ -745,8 +745,10 @@ def fake_agg(config, **kwargs): } }, } - # The o10 shardmap meta has no shard_key here; run_target reads it, so add one. - manifest["shardmaps"]["healpix_o10"]["shard_key"] = 0 + # The o10 shardmap meta has no shard_key here; run_target reads it (and + # renders it through grid.shard_label — issue #199), so add a valid + # packed morton word. + manifest["shardmaps"]["healpix_o10"]["shard_key"] = 5347395636851376137 run_benchmark.run_target( "t", manifest, diff --git a/tests/test_grids.py b/tests/test_grids.py index e5fd8f6f..a299a37a 100644 --- a/tests/test_grids.py +++ b/tests/test_grids.py @@ -768,6 +768,80 @@ def test_is_morton_arrow_false_for_plain_columns(self): assert not is_morton_arrow(np.arange(3)) # no Arrow field at all +class TestShardLabel: + """Issue #199: shard ids surface externally as decimal morton strings (D1 in + ``docs/design/sparse_coverage.md``) — ``morton_decimal``/``morton_word`` are + the boundary pair, and ``grid.shard_label`` is the grid seam every external + string (CSR subgroup names, status keys, log lines) routes through.""" + + def test_decimal_round_trip_both_hemispheres(self): + from mortie import geo2mort + + from zagg.grids.morton import morton_decimal, morton_word + + # Southern points render signed, northern unsigned; cover both plus a + # spread of orders (order-0 base cells through fine cells). + for lat, lon in [(-78.5, -132.0), (-72.1, 25.4), (78.3, 12.0), (0.1, 0.1)]: + for order in (0, 6, 9, 18): + word = int(geo2mort(np.array([lat]), np.array([lon]), order=order)[0]) + s = morton_decimal(word) + # Grammar: optional sign, base 1..6, then one 1..4 digit per order. + body = s.lstrip("-") + assert body[0] in "123456" and all(d in "1234" for d in body[1:]) + assert len(body) == order + 1 + # Packed word -> decimal string -> packed word is lossless. + assert morton_word(s) == word + + def test_order_29_round_trip_full_grammar_width(self): + # Orders 19-29 extend the decimal form past the legacy i64; one order-29 + # AREA case pins the full grammar width (review nit, PR #205). + from mortie import geo2mort + + from zagg.grids.morton import morton_decimal, morton_word + + word = int(geo2mort(np.array([-78.5]), np.array([-132.0]), order=29)[0]) + s = morton_decimal(word) + assert len(s.lstrip("-")) == 30 + assert morton_word(s) == word + + def test_morton_word_rejects_malformed(self): + from zagg.grids.morton import morton_word + + for bad in ("", "0", "7", "150", "abc", "11827859996358475782"): + with pytest.raises(ValueError): + morton_word(bad) + + def test_morton_decimal_raises_valueerror_on_invalid(self): + # The emit direction's advertised contract (review finding, PR #205): + # the empty sentinel raises, and a NEGATIVE input — a legacy signed + # decimal id where the packed word belongs — raises ValueError (not the + # opaque uint64-coercion OverflowError) with the remedy named. + from zagg.grids.morton import morton_decimal + + with pytest.raises(ValueError): + morton_decimal(0) + with pytest.raises(ValueError, match="morton_word"): + morton_decimal(-4211322) + + def test_healpix_shard_label_is_decimal(self, cfg): + from zagg.grids.morton import morton_word + + g = HealpixGrid(parent_order=6, child_order=12, config=cfg) + for parent in _valid_parents(3): + # The round trip through the parse-back is what carries this test + # (label == morton_decimal(parent) would be tautological). + assert morton_word(g.shard_label(parent)) == parent + + def test_base_helper_dispatches_and_falls_back(self, cfg): + from zagg.grids.base import shard_label + + g = HealpixGrid(parent_order=6, child_order=12, config=cfg) + parent = _valid_parents(1)[0] + assert shard_label(g, parent) == g.shard_label(parent) + # Minimal grid stand-ins without the method keep plain int digits. + assert shard_label(object(), 42) == "42" + + class TestCellIdsEncoding: """Issue #135: ``output.grid.cell_ids_encoding`` — ``cell_ids`` stays NESTED HEALPix uint64 by default; ``morton`` emits the packed morton words instead diff --git a/tests/test_hive.py b/tests/test_hive.py new file mode 100644 index 00000000..d1876387 --- /dev/null +++ b/tests/test_hive.py @@ -0,0 +1,750 @@ +"""Tests for the morton-hive store layout — issue #199 phase 2. + +Covers the config flag, leaf-path computation + node invariant (D2/D3/D5), +the ``morton_hive.json`` manifest (D6), the commit stamp / debris / torn-write +retry semantics (D4), and the local runner's hive write path. +""" + +import json +import os +from dataclasses import asdict + +import numpy as np +import pandas as pd +import pytest +import zarr +from zarr.storage import MemoryStore + +from zagg import hive +from zagg.config import default_config, get_data_vars, validate_config +from zagg.grids import HealpixGrid +from zagg.grids.morton import morton_decimal, morton_word + + +@pytest.fixture +def cfg(): + return default_config("atl06") + + +def _shard_word(order=6): + """A real southern packed shard word (decimal form ``-5112333`` at order 6).""" + from mortie import geo2mort + + return int(geo2mort(np.array([-78.5]), np.array([-132.0]), order=order)[0]) + + +# ── config flag ────────────────────────────────────────────────────────────── + + +class TestStoreLayoutConfig: + def test_default_is_flat(self, cfg): + from zagg.config import get_store_layout + + assert get_store_layout(cfg) == "flat" + validate_config(cfg) # flat default validates unchanged + + def test_hive_accepted_for_healpix(self, cfg): + cfg.output["store_layout"] = "hive" + validate_config(cfg) + + def test_null_key_falls_back_to_flat(self, cfg): + from zagg.config import get_store_layout + + cfg.output["store_layout"] = None + assert get_store_layout(cfg) == "flat" + validate_config(cfg) + + def test_unknown_value_rejected(self, cfg): + cfg.output["store_layout"] = "tree" + with pytest.raises(ValueError, match="store_layout"): + validate_config(cfg) + + def test_hive_rejects_rectilinear(self, cfg): + cfg.output["store_layout"] = "hive" + cfg.output["grid"] = { + "type": "rectilinear", + "crs": "EPSG:3031", + "resolution": 100, + "bounds": [0, 0, 1000, 1000], + } + with pytest.raises(ValueError, match="healpix"): + validate_config(cfg) + + def test_hive_rejects_sharded(self, cfg): + cfg.output["store_layout"] = "hive" + cfg.output.setdefault("grid", {})["sharded"] = True + with pytest.raises(ValueError, match="sharded"): + validate_config(cfg) + + def test_hive_rejects_consolidate_metadata(self, cfg): + cfg.output["store_layout"] = "hive" + cfg.output["consolidate_metadata"] = True + with pytest.raises(ValueError, match="consolidate"): + validate_config(cfg) + + +# ── leaf paths + node invariant ────────────────────────────────────────────── + + +class TestLeafPath: + def test_matches_mortie_hive_path(self): + # The convention is owned by the mortie spec: zagg's leaf path must be + # exactly mortie's hive_path under the store root. + from mortie import MortonIndexArray + + word = _shard_word() + expected = MortonIndexArray.from_words(np.asarray([word], dtype=np.uint64)).hive_path( + root="s3://b/root" + )[0] + assert hive.shard_leaf_path("s3://b/root", word) == expected + + def test_one_digit_per_level_full_id_leaf(self): + # D2/D3: sign+base, one digit per order, full decimal id at the leaf. + word = _shard_word() + assert morton_decimal(word) == "-5112333" + assert hive.shard_leaf_path("root", word) == "root/-5/1/1/2/3/3/3/-5112333.zarr" + + def test_trailing_slash_root_normalized(self): + word = _shard_word() + assert hive.shard_leaf_path("root/", word) == hive.shard_leaf_path("root", word) + + def test_negative_key_rejected(self): + # A signed legacy id is the DECIMAL form, not a packed word. + with pytest.raises(ValueError, match="packed morton word"): + hive.shard_leaf_path("root", -4211322) + + def test_node_invariant_accepts_computed_paths(self): + for order in (1, 6, 11): + word = _shard_word(order) + s = morton_decimal(word) + head = 2 if s.startswith("-") else 1 + rel = "/".join([s[:head], *s[head:]]) + f"/{s}.zarr" + hive.check_node_invariant(rel) + + @pytest.mark.parametrize( + "bad", + [ + "-4211322.zarr", # bare leaf: no digit chain at all + "0/1/01.zarr", # base digit 0 + "-4/5/-45.zarr", # order digit outside 1..4 + "-4/2/-43.zarr", # leaf id does not match the chain + "-4/2/-42", # not a .zarr leaf + "-4/21/-421.zarr", # grouped digits (one digit per level, D2) + ], + ) + def test_node_invariant_rejects(self, bad): + with pytest.raises(ValueError, match="node invariant"): + hive.check_node_invariant(bad) + + +# ── manifest (D6) ──────────────────────────────────────────────────────────── + + +class TestManifest: + def _grid(self, cfg): + return HealpixGrid(parent_order=6, child_order=8, layout="fullsphere", config=cfg) + + def test_build_contents(self, cfg): + m = hive.build_manifest(self._grid(cfg), dataset={"short_name": "ATL06", "version": "007"}) + assert m["spec"] == "morton-hive/1" + assert m["dataset"] == {"short_name": "ATL06", "version": "007"} + assert m["cell_order"] == 8 + assert m["shard_order"] == 6 + # Explicit split schedule: one digit per level down to the shard order. + assert m["split_schedule"] == [1] * 6 + # Declared-only in round one (populated by the pyramid sweep, D11). + assert m["pyramid"] == {"orders": [], "aggregation": {}} + assert m["generated_at"] + + def test_ensure_write_read_round_trip(self, cfg, tmp_path): + root = str(tmp_path / "store") + m = hive.build_manifest(self._grid(cfg)) + assert hive.ensure_manifest(root, m) == m + assert hive.read_manifest(root) == m + # The object is the root-only exception: it lives at the root, as JSON. + assert json.loads((tmp_path / "store" / hive.MANIFEST_NAME).read_text()) == m + + def test_rerun_with_matching_manifest_is_accepted(self, cfg, tmp_path): + # Retry semantics (D4): a rerun into the same root must proceed. + root = str(tmp_path / "store") + grid = self._grid(cfg) + hive.ensure_manifest(root, hive.build_manifest(grid)) + again = hive.build_manifest(grid) # fresh generated_at + assert hive.ensure_manifest(root, again)["spec"] == "morton-hive/1" + + def test_rerun_ignores_sweep_mutated_pyramid(self, cfg, tmp_path): + # The pyramid block is populated/updated by the §7 sweep BY DESIGN + # (D11), so the resume match-check must not compare it — else the + # first sweep would brick every later resume (review finding, PR #205). + root = str(tmp_path / "store") + grid = self._grid(cfg) + swept = hive.build_manifest(grid) + swept["pyramid"] = {"orders": [4, 5], "aggregation": {"count": "sum"}} + hive.ensure_manifest(root, swept) + # A later run's fresh (declared-only) manifest still resumes, and the + # sweep's pyramid declaration is preserved, not clobbered. + resumed = hive.ensure_manifest(root, hive.build_manifest(grid)) + assert resumed["pyramid"] == swept["pyramid"] + + def test_mismatched_manifest_says_clear_the_root(self, cfg, tmp_path): + # overwrite=True replaces the manifest ONLY; the remedy must not + # suggest it for an orders change (review finding, PR #205). + root = str(tmp_path / "store") + hive.ensure_manifest(root, hive.build_manifest(self._grid(cfg))) + other = HealpixGrid(parent_order=5, child_order=8, layout="fullsphere", config=cfg) + with pytest.raises(ValueError, match="clear the store root"): + hive.ensure_manifest(root, hive.build_manifest(other)) + + def test_overwrite_replaces_when_tree_is_empty(self, cfg, tmp_path): + root = str(tmp_path / "store") + hive.ensure_manifest(root, hive.build_manifest(self._grid(cfg))) + other = HealpixGrid(parent_order=5, child_order=8, layout="fullsphere", config=cfg) + hive.ensure_manifest(root, hive.build_manifest(other), overwrite=True) + assert hive.read_manifest(root)["shard_order"] == 5 + + def test_overwrite_with_changed_orders_refuses_over_existing_shards(self, cfg, tmp_path): + # Committed leaves from the old orders would survive a manifest-only + # "re-template" as walker-discoverable, stamped, seemingly-legal + # mixed-order data (D2) — refuse via one delimiter-LIST (review + # finding, PR #205). + root = tmp_path / "store" + hive.ensure_manifest(str(root), hive.build_manifest(self._grid(cfg))) + (root / "-5" / "1").mkdir(parents=True) # a {sign+base} child exists + (root / "-5" / "1" / "obj").write_text("x") + other = HealpixGrid(parent_order=5, child_order=8, layout="fullsphere", config=cfg) + with pytest.raises(ValueError, match="clear the store root first"): + hive.ensure_manifest(str(root), hive.build_manifest(other), overwrite=True) + + def test_overwrite_with_same_orders_allowed_over_existing_shards(self, cfg, tmp_path): + # Same frozen keys -> replacing the manifest is safe even with data. + root = tmp_path / "store" + grid = self._grid(cfg) + hive.ensure_manifest(str(root), hive.build_manifest(grid)) + (root / "-5" / "1").mkdir(parents=True) + (root / "-5" / "1" / "obj").write_text("x") + hive.ensure_manifest(str(root), hive.build_manifest(grid), overwrite=True) + + def test_read_absent_returns_none(self, tmp_path): + assert hive.read_manifest(str(tmp_path / "empty")) is None + + +# ── leaf template + commit stamp (D3/D4) ───────────────────────────────────── + + +class TestLeafTemplateAndStamp: + def _grid(self, cfg): + return HealpixGrid(parent_order=6, child_order=8, layout="fullsphere", config=cfg) + + def test_leaf_template_is_shard_sized(self, cfg): + g = self._grid(cfg) + store = MemoryStore() + g.emit_shard_template(store, overwrite=True) + grp = zarr.open_group(store, path=g.group_path, mode="r", zarr_format=3) + for name in ("morton", "cell_ids", *get_data_vars(cfg)): + assert grp[name].shape == (g.cells_per_shard,) + assert grp[name].chunks == (g.cells_per_chunk,) + + def test_leaf_has_root_group_for_the_stamp(self, cfg): + # D4: the stamp is one attrs update on an object that exists anyway. + store = MemoryStore() + self._grid(cfg).emit_shard_template(store, overwrite=True) + root = zarr.open_group(store, path="", mode="r", zarr_format=3) + assert hive.COMMIT_ATTR not in root.attrs # fresh leaf is unstamped + + def test_emit_is_idempotent_with_overwrite(self, cfg): + store = MemoryStore() + g = self._grid(cfg) + g.emit_shard_template(store, overwrite=True) + g.emit_shard_template(store, overwrite=True) # retry over debris + + def test_sharded_grid_rejected(self, cfg): + g = HealpixGrid(6, 10, layout="fullsphere", config=cfg, chunk_inner=8, sharded=True) + with pytest.raises(ValueError, match="sharded"): + g.emit_shard_template(MemoryStore()) + + def test_stamp_round_trip_and_debris_semantics(self, cfg): + store = MemoryStore() + self._grid(cfg).emit_shard_template(store, overwrite=True) + # An unstamped prefix is debris: present, but not complete. + assert hive.read_commit(store) is None + hive.stamp_commit(store, cells_with_data=5, granule_count=2) + stamp = hive.read_commit(store) + assert stamp["complete"] is True + assert stamp["spec"] == hive.HIVE_SPEC + assert stamp["cells_with_data"] == 5 + assert stamp["granule_count"] == 2 + assert stamp["written_at"] + + def test_read_commit_absent_leaf_is_none(self): + # Walker termination: no leaf at all is the same answer as debris. + assert hive.read_commit(MemoryStore()) is None + + +# ── local write path (runner) ──────────────────────────────────────────────── + + +def _rec(n): + return {"id": f"g{n}", "s3": f"s3://bucket/granule{n}.h5", "https": f"https://h/g{n}.h5"} + + +class TestProcessAndWriteHive: + """Drive ``hive.process_and_write_hive`` with a fake ``process_shard`` that + streams REAL carriers, so the leaf template, dense write, CSR naming, and + stamp ordering are all exercised against real zarr stores.""" + + def _grid(self, cfg): + return HealpixGrid(parent_order=6, child_order=8, layout="fullsphere", config=cfg) + + def _carrier(self, grid, shard): + coords = grid.chunk_coords(shard) + n = len(coords["cell_ids"]) + df = pd.DataFrame( + { + var: np.zeros(n, dtype=np.int32 if var == "count" else np.float32) + for var in get_data_vars(grid.config) + } + ) + for name, vals in coords.items(): + df[name] = vals + return df + + def _meta(self, shard, error=None): + return { + "shard_key": int(shard), + "cells_with_data": 5, + "total_obs": 7, + "granule_count": 1, + "files_processed": 1, + "duration_s": 0.0, + "error": error, + } + + def _run(self, monkeypatch, cfg, tmp_path, fake): + import zagg.processing as processing + + monkeypatch.setattr(processing, "process_shard", fake) + grid = self._grid(cfg) + shard = _shard_word() + root = str(tmp_path / "store") + meta = hive.process_and_write_hive( + shard, ["s3://bucket/granule1.h5"], grid, {}, root, cfg, store_kwargs={} + ) + return grid, shard, root, meta + + def _streaming_fake(self, grid, ragged=None): + def fake(g, shard_key, urls, **kwargs): + carrier = self._carrier(grid, shard_key) + kwargs["write_chunk"](grid.block_index(int(shard_key)), carrier, ragged or {}) + return pd.DataFrame(), self._meta(shard_key) + + return fake + + def test_leaf_written_and_stamped(self, monkeypatch, cfg, tmp_path): + grid_probe = self._grid(cfg) + fake = self._streaming_fake(grid_probe, ragged={"h": ([np.array([1.0, 2.0])], [0])}) + grid, shard, root, meta = self._run(monkeypatch, cfg, tmp_path, fake) + + leaf = hive.shard_leaf_path(root, shard) + from zagg.store import open_store + + leaf_store = open_store(leaf) + # Dense data landed at the leaf-LOCAL block 0. + grp = zarr.open_group(leaf_store, path=grid.group_path, mode="r", zarr_format=3) + np.testing.assert_array_equal( + np.asarray(grp["cell_ids"][:]), + np.asarray(grid.chunk_coords(shard)["cell_ids"]), + ) + # CSR subgroup named by the shard label (decimal morton string). + label = morton_decimal(shard) + assert morton_word(label) == shard + sub = zarr.open_group(leaf_store, path=f"{grid.group_path}/h/{label}", mode="r") + assert "values" in sub.array_keys() + # The commit stamp is present and carries the worker's counters (D4). + stamp = hive.read_commit(leaf_store) + assert stamp["complete"] is True + assert stamp["cells_with_data"] == meta["cells_with_data"] + assert stamp["granule_count"] == meta["granule_count"] + + def test_no_data_shard_leaves_no_prefix(self, monkeypatch, cfg, tmp_path): + # The leaf is created lazily on the first chunk write, so a no-data + # shard leaves NO .zarr/ prefix (absence stays trustworthy). + def fake(g, shard_key, urls, **kwargs): + return pd.DataFrame(), self._meta(shard_key, error="No granules found") + + grid, shard, root, meta = self._run(monkeypatch, cfg, tmp_path, fake) + leaf = hive.shard_leaf_path(root, shard) + assert not os.path.exists(leaf) + + def test_torn_write_leaves_debris_then_retry_succeeds(self, monkeypatch, cfg, tmp_path): + # Torn-write simulation: the worker dies after the dense write, before + # the stamp. The prefix exists (debris), read_commit says incomplete, + # and a clean retry overwrites it WHOLESALE and stamps. The torn + # attempt also writes a CSR subgroup the retry does NOT rewrite, so + # the wholesale claim is pinned against upstream drift: if the leaf + # re-template merely re-put metadata instead of delete_dir-ing the + # prefix, the stale subgroup would survive inside a leaf whose stamp + # certifies it complete (review finding, PR #205). + import zagg.processing as processing + from zagg.store import open_store + + grid = self._grid(cfg) + shard = _shard_word() + root = str(tmp_path / "store") + leaf = hive.shard_leaf_path(root, shard) + + def torn(g, shard_key, urls, **kwargs): + carrier = self._carrier(grid, shard_key) + stale_ragged = {"h": ([np.array([1.0])], [0])} + kwargs["write_chunk"](grid.block_index(int(shard_key)), carrier, stale_ragged) + raise RuntimeError("worker died mid-shard") + + monkeypatch.setattr(processing, "process_shard", torn) + with pytest.raises(RuntimeError, match="died mid-shard"): + hive.process_and_write_hive( + shard, ["s3://bucket/g1.h5"], grid, {}, root, cfg, store_kwargs={} + ) + assert os.path.exists(leaf) # the prefix exists... + assert hive.read_commit(open_store(leaf)) is None # ...but is debris + stale = os.path.join(leaf, grid.group_path, "h", morton_decimal(shard)) + assert os.path.exists(stale) # the torn attempt's CSR subgroup + + # Retry (no ragged this time): same leaf, overwritten wholesale — + # the stale subgroup is GONE — and stamped at the end. + monkeypatch.setattr(processing, "process_shard", self._streaming_fake(grid)) + hive.process_and_write_hive( + shard, ["s3://bucket/g1.h5"], grid, {}, root, cfg, store_kwargs={} + ) + assert hive.read_commit(open_store(leaf))["complete"] is True + assert not os.path.exists(stale), "stale torn-write object survived the re-template" + + def test_errored_shard_is_not_stamped(self, monkeypatch, cfg, tmp_path): + # A shard that wrote chunks but ended in error stays unstamped debris. + from zagg.store import open_store + + grid_probe = self._grid(cfg) + + def fake(g, shard_key, urls, **kwargs): + carrier = self._carrier(grid_probe, shard_key) + kwargs["write_chunk"](grid_probe.block_index(int(shard_key)), carrier, {}) + return pd.DataFrame(), self._meta(shard_key, error="No data after filtering (1 ...)") + + grid, shard, root, _meta = self._run(monkeypatch, cfg, tmp_path, fake) + leaf = hive.shard_leaf_path(root, shard) + assert os.path.exists(leaf) + assert hive.read_commit(open_store(leaf)) is None + + def test_tree_walk_node_invariant(self, monkeypatch, cfg, tmp_path): + # Walker semantics (D5): below the root only digit dirs and *.zarr + # nodes; no zarr metadata above the leaf; the root additionally holds + # only the manifest. A LIST with no digit children is thus a + # definitive "nothing finer exists". + grid_probe = self._grid(cfg) + fake = self._streaming_fake(grid_probe) + grid, shard, root, _meta = self._run(monkeypatch, cfg, tmp_path, fake) + hive.ensure_manifest(root, hive.build_manifest(grid)) + + for dirpath, dirnames, filenames in os.walk(root): + if dirpath == root: + assert filenames == [hive.MANIFEST_NAME] + base = [d[1:] if d.startswith("-") else d for d in dirnames] + assert all(len(b) == 1 and b in "123456" for b in base) + continue + if dirpath.endswith(".zarr") or ".zarr" + os.sep in dirpath: + continue # inside a leaf: vanilla zarr v3, its own business + # An intermediate digit node: no objects (zarr.json or otherwise), + # only digit children and leaf dirs. + assert filenames == [], f"object above the leaf at {dirpath}: {filenames}" + for d in dirnames: + assert d.endswith(".zarr") or (len(d) == 1 and d in "1234"), ( + f"non-hive child {d!r} at {dirpath}" + ) + + def test_stamp_is_the_final_write(self, monkeypatch, cfg, tmp_path): + """D4 ordering pin (review finding, PR #205): the commit stamp is the + shard's LAST write — presence certifies everything before it landed. + ONE test covers BOTH backends: the local dispatcher and the Lambda + handler execute this same ``process_and_write_hive`` function, so the + op ordering cannot diverge between them.""" + import zagg.processing as processing + + ops: list = [] + + def rec(name, fn): + def wrapped(*a, **k): + ops.append(name) + return fn(*a, **k) + + return wrapped + + grid = self._grid(cfg) + fake = self._streaming_fake(grid, ragged={"h": ([np.array([1.0])], [0])}) + monkeypatch.setattr(processing, "process_shard", fake) + monkeypatch.setattr( + processing, "write_dataframe_to_zarr", rec("dense", processing.write_dataframe_to_zarr) + ) + monkeypatch.setattr( + processing, "write_ragged_to_zarr", rec("ragged", processing.write_ragged_to_zarr) + ) + monkeypatch.setattr(hive, "stamp_commit", rec("stamp", hive.stamp_commit)) + hive.process_and_write_hive( + _shard_word(), ["s3://b/g1.h5"], grid, {}, str(tmp_path / "store"), cfg, store_kwargs={} + ) + assert ops == ["dense", "ragged", "stamp"] + + +class TestLeafBlockIndex: + def test_k1_maps_to_zero(self, cfg): + g = HealpixGrid(6, 8, layout="fullsphere", config=cfg) + shard = _shard_word() + (block,) = [b for b, _ in g.iter_chunks(shard)] + assert hive.leaf_block_index(g, block, shard) == (0,) + + def test_k_gt_1_enumerates_local_ordinals(self, cfg): + g = HealpixGrid(6, 10, layout="fullsphere", config=cfg, chunk_inner=8) + assert g.chunks_per_shard == 16 + shard = _shard_word() + locals_ = [hive.leaf_block_index(g, b, shard) for b, _ in g.iter_chunks(shard)] + assert locals_ == [(i,) for i in range(16)] + + +class TestRunnerWiring: + """The local backend writes the manifest (no shared template) under hive; + the lambda backend dispatches hive runs (issue #199 phase 3), threading + the manifest's dataset identity through the setup invoke.""" + + def _catalog(self, tmp_path): + shard = _shard_word() + catalog = { + "metadata": {"short_name": "ATL06", "version": "007"}, + "grid_signature": { + "type": "healpix", + "indexing_scheme": "nested", + "parent_order": 6, + "child_order": 12, + "layout": "fullsphere", + }, + "shard_keys": [shard], + "granules": [[_rec(1)]], + } + p = tmp_path / "catalog.json" + p.write_text(json.dumps(catalog)) + return str(p), shard + + def test_local_hive_writes_manifest_not_template(self, monkeypatch, cfg, tmp_path): + from zagg import runner + from zagg.runner import agg + + cfg.output["store_layout"] = "hive" + catalog_path, shard = self._catalog(tmp_path) + root = str(tmp_path / "out") + calls = [] + + monkeypatch.setattr(runner, "get_nsidc_s3_credentials", lambda: {"accessKeyId": "a"}) + + def fake_hive_write(shard_key, granule_urls, grid, s3_creds, store_root, config, **kw): + calls.append((int(shard_key), store_root)) + return {"shard_key": int(shard_key), "error": None, "total_obs": 1} + + monkeypatch.setattr(hive, "process_and_write_hive", fake_hive_write) + agg(cfg, catalog=catalog_path, store=root, backend="local") + + assert calls == [(shard, root)] + # Template time wrote ONLY the manifest — no shared zarr template (D5). + assert sorted(os.listdir(root)) == [hive.MANIFEST_NAME] + assert hive.read_manifest(root)["shard_order"] == 6 + + def test_lambda_hive_dispatches_with_manifest_dataset(self, monkeypatch, cfg, tmp_path): + # Issue #199 phase 3: hive is wired to the lambda backend. The setup + # invoke carries the manifest's dataset identity (from the ShardMap + # metadata, same source as the local path) and the per-cell events need + # NO new keys — the worker derives everything from the config dict. + from unittest.mock import MagicMock + + import boto3 + + from zagg import runner + from zagg.concurrency import ConcurrencyReport + from zagg.runner import agg + + cfg.output["store_layout"] = "hive" + catalog_path, shard = self._catalog(tmp_path) + captured: dict = {} + + monkeypatch.setattr( + runner, + "get_nsidc_s3_credentials", + lambda: {"accessKeyId": "a", "secretAccessKey": "s", "sessionToken": "t"}, + ) + monkeypatch.setattr(boto3, "Session", lambda *a, **k: MagicMock()) + monkeypatch.setattr(runner, "_get_function_timeout_s", lambda *a, **k: 720) + monkeypatch.setattr( + runner, + "compute_available_workers", + lambda requested, *a, **k: ( + 1, + ConcurrencyReport( + account_limit=1000, + current_concurrent=0, + padding=100, + available=900, + function_reserved=None, + ), + ), + ) + monkeypatch.setattr( + runner, "_invoke_lambda_setup", lambda *a, **kw: captured.update(setup=kw) + ) + monkeypatch.setattr(runner, "_invoke_lambda_finalize", lambda *a, **k: None) + + def fake_cell(client, chunk_idx, shard_key, *a, **k): + captured["cell_shard_key"] = shard_key + return { + "status_code": 200, + "body": {"total_obs": 1}, + "error": None, + "lambda_duration": 1.0, + "shard_key": shard_key, + } + + monkeypatch.setattr(runner, "_invoke_lambda_cell", fake_cell) + agg(cfg, catalog=catalog_path, store="s3://out/product", backend="lambda") + + assert captured["setup"]["dataset"] == {"short_name": "ATL06", "version": "007"} + # store_layout rides in the config dict already serialized into events. + assert captured["setup"]["config_dict"]["output"]["store_layout"] == "hive" + # The per-cell event schema is unchanged: shard_key stays the packed int. + assert captured["cell_shard_key"] == shard + + def test_lambda_flat_setup_omits_dataset(self, monkeypatch, cfg, tmp_path): + # Flat runs keep their setup call byte-identical: dataset stays None. + from unittest.mock import MagicMock + + import boto3 + + from zagg import runner + from zagg.concurrency import ConcurrencyReport + from zagg.runner import agg + + catalog_path, shard = self._catalog(tmp_path) + captured: dict = {} + + monkeypatch.setattr( + runner, + "get_nsidc_s3_credentials", + lambda: {"accessKeyId": "a", "secretAccessKey": "s", "sessionToken": "t"}, + ) + monkeypatch.setattr(boto3, "Session", lambda *a, **k: MagicMock()) + monkeypatch.setattr(runner, "_get_function_timeout_s", lambda *a, **k: 720) + monkeypatch.setattr( + runner, + "compute_available_workers", + lambda requested, *a, **k: ( + 1, + ConcurrencyReport( + account_limit=1000, + current_concurrent=0, + padding=100, + available=900, + function_reserved=None, + ), + ), + ) + monkeypatch.setattr( + runner, "_invoke_lambda_setup", lambda *a, **kw: captured.update(setup=kw) + ) + monkeypatch.setattr(runner, "_invoke_lambda_finalize", lambda *a, **k: None) + monkeypatch.setattr( + runner, + "_invoke_lambda_cell", + lambda *a, **k: { + "status_code": 200, + "body": {"total_obs": 1}, + "error": None, + "lambda_duration": 1.0, + "shard_key": shard, + }, + ) + agg(cfg, catalog=catalog_path, store="s3://out/x.zarr", backend="lambda") + assert captured["setup"]["dataset"] is None + + +class TestInvokeLambdaSetupEvent: + """Pin the ACTUAL setup event on the wire and the stale-deployment guard + (review findings, PR #205). The dispatch tests above monkeypatch + ``_invoke_lambda_setup`` at kwarg level, so the event-shaping conditional + (``dataset`` added only when set) and the layout-echo assertion are + exercised here directly, with a mocked boto3 client capturing ``Payload``.""" + + @staticmethod + def _client(body: dict): + from unittest.mock import MagicMock + + payload = MagicMock() + payload.read.return_value = json.dumps( + {"statusCode": 200, "body": json.dumps(body)} + ).encode() + client = MagicMock() + client.invoke.return_value = {"Payload": payload, "FunctionError": None} + return client + + @staticmethod + def _invoke(client, config_dict, dataset=None): + from zagg.runner import _invoke_lambda_setup + + _invoke_lambda_setup( + client, + "process-shard", + "s3://out/product", + parent_order=6, + child_order=12, + n_parent_cells=None, + overwrite=False, + config_dict=config_dict, + dataset=dataset, + ) + return json.loads(client.invoke.call_args.kwargs["Payload"]) + + def test_hive_event_carries_dataset(self, cfg): + cfg.output["store_layout"] = "hive" + client = self._client({"ok": True, "mode": "setup", "layout": "hive"}) + event = self._invoke(client, asdict(cfg), dataset={"short_name": "ATL06", "version": "007"}) + assert event["dataset"] == {"short_name": "ATL06", "version": "007"} + + def test_flat_event_omits_dataset_and_matches_baseline(self, cfg): + # The byte-identity claim, pinned on the wire: no "dataset" key, and + # the event is exactly the pre-phase-3 flat setup event. + config_dict = asdict(cfg) + client = self._client({"ok": True, "mode": "setup", "layout": "flat"}) + event = self._invoke(client, config_dict) + assert "dataset" not in event + assert event == { + "mode": "setup", + "store_path": "s3://out/product", + "parent_order": 6, + "child_order": 12, + "n_parent_cells": None, + "overwrite": False, + "config": config_dict, + } + + def test_flat_without_layout_echo_unaffected(self, cfg): + # Old deployed functions return the echo-less body: flat dispatch must + # keep working against them. + self._invoke(self._client({"ok": True, "mode": "setup"}), asdict(cfg)) + + @pytest.mark.parametrize( + "body", + [ + {"ok": True, "mode": "setup"}, # pre-phase-3 function: no echo + {"ok": True, "mode": "setup", "layout": "flat"}, # wrong layout acted on + ], + ) + def test_hive_without_hive_echo_fails_fast(self, cfg, body): + # Stale-deployment guard: an old function would emit the flat GLOBAL + # template at the hive root and return a 200 the dispatcher couldn't + # tell apart — the layout echo makes that fail at setup, pre-fan-out. + cfg.output["store_layout"] = "hive" + with pytest.raises(RuntimeError, match="redeploy"): + self._invoke( + self._client(body), asdict(cfg), dataset={"short_name": "A", "version": "1"} + ) diff --git a/tests/test_lambda_handler.py b/tests/test_lambda_handler.py index 9f86b8a8..1232f6a0 100644 --- a/tests/test_lambda_handler.py +++ b/tests/test_lambda_handler.py @@ -377,6 +377,10 @@ def fake_write_ragged(ragged, st, *, grid, shard_key): # Regular (non-sharded) write path — a MagicMock attr is truthy by default, # so pin it off explicitly (issue #108 routes sharded grids elsewhere). grid_stub.sharded = False + # K==1 ragged subgroups are keyed by the grid's shard label (issue #199); + # a deterministic stub (decimal-string shaped) pins that the handler + # routes through the seam rather than stringifying the raw key itself. + grid_stub.shard_label = lambda key: f"-{int(key)}" monkeypatch.setattr(processing, "process_shard", fake_process_shard) monkeypatch.setattr(grids, "from_config", lambda *a, **k: grid_stub) @@ -454,9 +458,10 @@ def fake_process_shard(grid, shard_key, granule_urls, **kwargs): assert seen["chunk_results"] is None # no accumulation sink assert written == [(0,), (1,)] # each chunk written as it streamed - def test_k_eq_1_ragged_keyed_by_shard_key(self, handler_mod, monkeypatch): - """K=1: the lone chunk's ragged CSR is persisted (the gap this phase closes: - the old handler never called write_ragged_to_zarr), keyed by shard_key.""" + def test_k_eq_1_ragged_keyed_by_shard_label(self, handler_mod, monkeypatch): + """K=1: the lone chunk's ragged CSR is persisted, keyed by the grid's + shard label (the decimal morton string for HEALPix — issue #199), not + the raw shard-key int.""" chunks = [((0,), pd.DataFrame(), {"h_tdigest": ([], [])})] cap = self._patch(handler_mod, monkeypatch, chunks) event = _base_event(_healpix_config_dict()) @@ -464,9 +469,9 @@ def test_k_eq_1_ragged_keyed_by_shard_key(self, handler_mod, monkeypatch): resp = handler_mod._handle_process(event, _context()) assert resp["statusCode"] == 200 assert cap["dense"] == [(0,)] - # Single chunk -> ragged keyed by shard_key (cell-resolution contract). + # Single chunk -> ragged keyed by the stub grid's label of shard_key. assert len(cap["ragged"]) == 1 - assert cap["ragged"][0][0] == event["shard_key"] + assert cap["ragged"][0][0] == f"-{event['shard_key']}" def test_streaming_later_chunk_write_failure_partial_writes_and_500( self, handler_mod, monkeypatch @@ -823,6 +828,220 @@ def test_present_but_not_a_group_is_not_masked_as_missing(self, handler_mod, mon assert "Zarr template not found" not in resp["body"] +class TestProcessHive: + """Issue #199 phase 3: under ``output.store_layout: hive`` the worker owns + its WHOLE leaf — it derives the leaf path from ``shard_key`` + the event + config's orders, emits its own leaf template, writes its data, and stamps + completion as its FINAL PUT — through ``zagg.hive.process_and_write_hive``, + the same code path the local dispatcher runs.""" + + # Order-6 southern shard; decimal morton string -4211322. + _WORD = 11827859996358475782 + + @staticmethod + def _hive_config_dict(): + cfg = default_config("atl06") + cfg.output["store_layout"] = "hive" + return asdict(cfg) + + def _event(self, tmp_path): + ev = _base_event(self._hive_config_dict()) + ev["shard_key"] = self._WORD + ev["child_order"] = 12 + ev["store_path"] = str(tmp_path / "hive-out") + return ev + + @staticmethod + def _grid(): + from zagg.config import load_config_from_dict + + cfg = load_config_from_dict(TestProcessHive._hive_config_dict()) + return from_config(cfg) + + @staticmethod + def _carrier(grid, shard): + import numpy as np + + from zagg.config import get_data_vars + + coords = grid.chunk_coords(shard) + n = len(coords["cell_ids"]) + df = pd.DataFrame( + { + var: np.zeros(n, dtype=np.int32 if var == "count" else np.float32) + for var in get_data_vars(grid.config) + } + ) + for name, vals in coords.items(): + df[name] = vals + return df + + def _streaming_fake(self, grid, ragged=None, die=False): + def fake(g, shard_key, urls, **kwargs): + carrier = self._carrier(grid, shard_key) + kwargs["write_chunk"](grid.block_index(int(shard_key)), carrier, ragged or {}) + if die: + raise RuntimeError("worker died mid-shard") + return pd.DataFrame(), { + "shard_key": int(shard_key), + "cells_with_data": 5, + "total_obs": 7, + "granule_count": 1, + "files_processed": 1, + "duration_s": 0.0, + "error": None, + } + + return fake + + def test_hive_event_writes_leaf_data_and_stamp(self, handler_mod, monkeypatch, tmp_path): + import numpy as np + import zarr + from mortie import MortonIndexArray + + import zagg.processing as processing + from zagg import hive + from zagg.store import open_store + + grid = self._grid() + ragged = {"h": ([np.array([1.0, 2.0])], [0])} + monkeypatch.setattr(processing, "process_shard", self._streaming_fake(grid, ragged)) + event = self._event(tmp_path) + resp = handler_mod._handle_process(event, _context()) + assert resp["statusCode"] == 200, resp["body"] + body = json.loads(resp["body"]) + assert body["shard_key"] == self._WORD + + # The leaf sits at EXACTLY mortie's hive_path under the store root. + words = MortonIndexArray.from_words(np.asarray([self._WORD], dtype="uint64")) + expected = words.hive_path(root=event["store_path"])[0] + assert hive.shard_leaf_path(event["store_path"], self._WORD) == expected + leaf_store = open_store(expected) + + # Dense data landed at the leaf-local block; CSR is label-named. + grp = zarr.open_group(leaf_store, path=grid.group_path, mode="r", zarr_format=3) + np.testing.assert_array_equal( + np.asarray(grp["cell_ids"][:]), + np.asarray(grid.chunk_coords(self._WORD)["cell_ids"]), + ) + sub = zarr.open_group(leaf_store, path=f"{grid.group_path}/h/-4211322", mode="r") + assert "values" in sub.array_keys() + + # The commit stamp is present and carries the worker's counters (D4). + stamp = hive.read_commit(leaf_store) + assert stamp["complete"] is True + assert stamp["cells_with_data"] == 5 + assert stamp["granule_count"] == 1 + + def test_hive_worker_death_leaves_debris_then_retry_cleans( + self, handler_mod, monkeypatch, tmp_path + ): + import os + + import numpy as np + + import zagg.processing as processing + from zagg import hive + from zagg.store import open_store + + grid = self._grid() + event = self._event(tmp_path) + leaf = hive.shard_leaf_path(event["store_path"], self._WORD) + stale = os.path.join(leaf, grid.group_path, "h", "-4211322") + + # Torn worker: writes a chunk + a CSR subgroup, then dies. The handler + # surfaces a 500 envelope; the leaf prefix exists but is UNSTAMPED. + torn = self._streaming_fake(grid, ragged={"h": ([np.array([1.0])], [0])}, die=True) + monkeypatch.setattr(processing, "process_shard", torn) + resp = handler_mod._handle_process(event, _context()) + assert resp["statusCode"] == 500 + assert "died mid-shard" in resp["body"] + assert os.path.exists(leaf) + assert hive.read_commit(open_store(leaf)) is None # debris + assert os.path.exists(stale) + + # Retry (no ragged this time): the leaf is overwritten WHOLESALE — the + # torn attempt's CSR subgroup is gone — and stamped at the end. + monkeypatch.setattr(processing, "process_shard", self._streaming_fake(grid)) + resp = handler_mod._handle_process(event, _context()) + assert resp["statusCode"] == 200, resp["body"] + assert hive.read_commit(open_store(leaf))["complete"] is True + assert not os.path.exists(stale) + + def test_hive_no_data_shard_leaves_no_prefix(self, handler_mod, monkeypatch, tmp_path): + import os + + import zagg.processing as processing + from zagg import hive + + def fake(g, shard_key, urls, **kwargs): + return pd.DataFrame(), { + "shard_key": int(shard_key), + "cells_with_data": 0, + "total_obs": 0, + "granule_count": 1, + "files_processed": 0, + "duration_s": 0.0, + "error": "No data after filtering", + } + + monkeypatch.setattr(processing, "process_shard", fake) + event = self._event(tmp_path) + resp = handler_mod._handle_process(event, _context()) + assert resp["statusCode"] == 500 # error envelope, as on flat + assert not os.path.exists(hive.shard_leaf_path(event["store_path"], self._WORD)) + + +class TestSetupHive: + """Issue #199 phase 3: for a hive config, setup mode writes ONLY the + ``morton_hive.json`` manifest — no global zarr template (D5) — with the + same frozen-key resume/overwrite semantics as the local path.""" + + def _event(self, tmp_path, config_dict, **extra): + return { + "mode": "setup", + "store_path": str(tmp_path / "hive-out"), + "parent_order": 6, + "overwrite": False, + "config": config_dict, + "dataset": {"short_name": "ATL06", "version": "007"}, + **extra, + } + + def test_setup_writes_manifest_only(self, handler_mod, tmp_path): + import os + + from zagg import hive + + event = self._event(tmp_path, TestProcessHive._hive_config_dict()) + resp = handler_mod._handle_setup(event) + assert resp["statusCode"] == 200, resp["body"] + + manifest = hive.read_manifest(event["store_path"]) + assert manifest["spec"] == "morton-hive/1" + assert manifest["dataset"] == {"short_name": "ATL06", "version": "007"} + assert manifest["shard_order"] == 6 + assert manifest["cell_order"] == 12 + # The manifest is the ONLY object: no template arrays, no zarr.json. + assert sorted(os.listdir(event["store_path"])) == [hive.MANIFEST_NAME] + + def test_setup_frozen_key_mismatch_is_500_clear_root(self, handler_mod, tmp_path): + cfg = TestProcessHive._hive_config_dict() + assert handler_mod._handle_setup(self._event(tmp_path, cfg))["statusCode"] == 200 + # A re-setup with different orders must fail with the pointed remedy, + # matching the local dispatcher's ensure_manifest semantics. + other = TestProcessHive._hive_config_dict() + other["output"]["grid"]["parent_order"] = 5 + resp = handler_mod._handle_setup(self._event(tmp_path, other)) + assert resp["statusCode"] == 500 + assert "clear the store root" in json.loads(resp["body"])["error"] + + def test_setup_rerun_with_matching_config_resumes(self, handler_mod, tmp_path): + cfg = TestProcessHive._hive_config_dict() + assert handler_mod._handle_setup(self._event(tmp_path, cfg))["statusCode"] == 200 + assert handler_mod._handle_setup(self._event(tmp_path, cfg))["statusCode"] == 200 + + class TestSetupTemplate: """Issue #99: the setup handler used to hand-build the HEALPix grid and drop ``chunk_inner``, so the template was chunked at ``parent_order`` while workers diff --git a/tests/test_readers.py b/tests/test_readers.py index 9b474ad7..06cab91a 100644 --- a/tests/test_readers.py +++ b/tests/test_readers.py @@ -9,6 +9,7 @@ from zarr.storage import MemoryStore from zagg.csr import write_csr +from zagg.grids.morton import morton_word from zagg.readers.tdigest_tensor import ( _resolve_chunk_morton, chunk_z_range, @@ -20,7 +21,12 @@ def _write_chunk(store, field, morton_key, cell_to_values, *, delta=512): - """Write one shard subgroup of per-cell t-digests under {field}/{morton_key}.""" + """Write one shard subgroup of per-cell t-digests under {field}/{morton_key}. + + ``morton_key`` is the subgroup name: a decimal morton string for the + shard-keyed layout (issue #199), or a plain int for the K>1 block-keyed + layout. + """ cell_ids = sorted(cell_to_values) payloads = [build_tdigest(np.asarray(cell_to_values[c]), delta=delta) for c in cell_ids] write_csr(store, f"{field}/{morton_key}", payloads, cell_ids) @@ -164,14 +170,18 @@ def test_unknown_fit_raises(self): class TestReadTensors: + # Two coverage chunks named by their decimal morton strings (issue #199); + # the readers report the parsed packed words. + _KEY_A, _KEY_B = "112", "243" + def _store(self): store = MemoryStore() rng = np.random.default_rng(10) - # Two chunks (parent mortons 100 and 250), a few populated cells each. + # Two chunks (decimal morton names), a few populated cells each. _write_chunk( store, "h_tdigest", - 100, + self._KEY_A, { 0: rng.uniform(10.0, 30.0, 3_000), 5: rng.uniform(12.0, 28.0, 2_000), @@ -181,25 +191,26 @@ def _store(self): _write_chunk( store, "h_tdigest", - 250, + self._KEY_B, {7: rng.uniform(40.0, 60.0, 2_500), 63: rng.uniform(42.0, 58.0, 2_000)}, ) return store def test_shape_and_dtype_default(self): out = dict((m, t) for t, m in read_tensors(self._store(), "h_tdigest")) - assert set(out) == {100, 250} + assert set(out) == {morton_word(self._KEY_A), morton_word(self._KEY_B)} for t in out.values(): assert t.shape == (64, 64, 128) assert t.dtype == np.uint32 def test_morton_recovered_from_subgroup_name(self): + # Round trip at the read boundary: decimal subgroup name -> packed word. mortons = sorted(m for _, m in read_tensors(self._store(), "h_tdigest")) - assert mortons == [100, 250] + assert mortons == sorted(morton_word(k) for k in (self._KEY_A, self._KEY_B)) def test_populated_cell_placement_rowmajor(self): out = dict((m, t) for t, m in read_tensors(self._store(), "h_tdigest")) - t = out[100] + t = out[morton_word(self._KEY_A)] # cell 5 → row 0, col 5; cell 4095 → row 63, col 63. assert t[0, 5].sum() > 0 assert t[63, 63].sum() > 0 @@ -210,9 +221,9 @@ def test_counts_match_population(self): store = MemoryStore() rng = np.random.default_rng(11) n = 5_000 - _write_chunk(store, "f", 1, {0: rng.uniform(0.0, 40.0, n)}) + _write_chunk(store, "f", "11", {0: rng.uniform(0.0, 40.0, n)}) t, m = next(read_tensors(store, "f", n_bins=128, resolution=0.5)) - assert m == 1 + assert m == morton_word("11") # Most of the population should land in-window (uniform [0,40] in a 64 m # window anchored at floor of the 5th pct). assert 0.8 * n <= t[0, 0].sum() <= n @@ -229,6 +240,8 @@ def test_dtype_flag(self, dtype, np_dtype): assert t.dtype == np_dtype def test_morton_coord_array_preferred(self): + # Block-keyed (K>1) subgroup names — "100"/"250" are not decimal morton + # ids, so the reader falls back to the int parse for the whole store. store = MemoryStore() rng = np.random.default_rng(13) _write_chunk(store, "f", 100, {0: rng.uniform(0.0, 30.0, 2_000)}) @@ -257,6 +270,39 @@ def test_morton_coord_array_mixed_digit_keys(self): mapping = _resolve_chunk_morton(store, "f", ["99", "100", "1000"], 3) assert mapping == {"99": 900, "100": 901, "1000": 902} + def test_morton_coord_array_decimal_names_word_order(self): + # Decimal morton names align with the coord array in packed-word order + # (issue #199): northern-base "31123" sorts before southern "-31123", + # even though int/lexicographic order would put the negative first. + store = MemoryStore() + rng = np.random.default_rng(34) + for key in ("-31123", "31123"): + _write_chunk(store, "f", key, {0: rng.uniform(0.0, 30.0, 1_000)}) + assert morton_word("31123") < morton_word("-31123") + arr = zarr.open_array( + store, path="f/morton", mode="w", shape=(2,), chunks=(2,), dtype="uint64" + ) + arr[...] = np.array([900, 901], dtype=np.uint64) + mapping = _resolve_chunk_morton(store, "f", ["-31123", "31123"], 3) + assert mapping == {"31123": 900, "-31123": 901} + + def test_lone_digit_names_are_block_keys(self): + # Review finding (PR #205): "5"/"6" fit the raw decimal grammar as + # order-0 ids, but real shard ids always carry >= 2 digits (shard + # orders are >= 1) — lone digits are block-index keys, parsed as ints. + from zagg.readers.tdigest_tensor import _subgroup_words + + assert _subgroup_words(["5", "6", "1"]) == {"5": 5, "6": 6, "1": 1} + + def test_signed_name_in_block_keyed_store_raises(self): + # Review finding (PR #205): a signed name can never be a block index, + # so a store mixing decimal-id and non-id names (pre-#199 debris) must + # fail loudly instead of silently degrading every id to a bogus int. + from zagg.readers.tdigest_tensor import _subgroup_words + + with pytest.raises(ValueError, match="ambiguous"): + _subgroup_words(["-31123", "100"]) + def test_raise_when_chunk_too_wide(self): store = MemoryStore() rng = np.random.default_rng(14) @@ -284,11 +330,11 @@ def test_recovers_unmerged_samples_exactly(self): store = MemoryStore() # Few enough values (< delta) that build_tdigest performs no merges. vals = np.array([3.0, 1.0, 2.0, 5.0, 4.0]) - _write_chunk(store, "f", 42, {7: vals}, delta=512) + _write_chunk(store, "f", "42", {7: vals}, delta=512) out = list(read_raw_values(store, "f")) assert len(out) == 1 morton, cell_id, recovered = out[0] - assert morton == 42 + assert morton == morton_word("42") assert cell_id == 7 # Digest stores centroids sorted by mean → sorted samples. np.testing.assert_allclose(recovered, np.sort(vals)) @@ -332,7 +378,8 @@ def test_yields_per_cell_uint64_vectors(self): store, vals, locs_in = self._located_store() out = {(m, c): locs for m, c, locs in read_locations(store, "f")} - assert set(out) == {(42, 3), (42, 9)} + w = morton_word("42") + assert set(out) == {(w, 3), (w, 9)} for (_, cid), locs in out.items(): assert locs.dtype == np.uint64 # Loss-free regime: locations are the cell's point words co-sorted diff --git a/tests/test_rectilinear.py b/tests/test_rectilinear.py index cfa60630..a105211f 100644 --- a/tests/test_rectilinear.py +++ b/tests/test_rectilinear.py @@ -180,6 +180,12 @@ def test_block_index_in_range(self, grid): assert 0 <= rb < grid.n_row_blocks assert 0 <= cb < grid.n_col_blocks + def test_shard_label_is_int_digits(self, grid): + # Rect shard keys are packed tile ints, not morton words: the external + # label (issue #199) stays the int's decimal digits. + packed = grid._pack(2, 3) + assert grid.shard_label(packed) == str(packed) + class TestChildren: def test_children_count(self, grid): diff --git a/tests/test_runner.py b/tests/test_runner.py index 9e5b33cd..254ac0d0 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -34,13 +34,19 @@ def _rec(n): } +# Packed morton words (the canonical catalog form) whose decimal morton +# strings — the external form ``--morton-cell`` takes since issue #199 — are +# -4211324 / -4211323 / -4211322. +_CATALOG_WORDS = [11828422946311897094, 11828141471335186438, 11827859996358475782] + + @pytest.fixture def catalog_file(tmp_path): """A minimal Phase-5 ShardMap JSON for testing.""" catalog = { "metadata": {"short_name": "ATL06", "total_shards": 3, "total_granules": 6}, "grid_signature": _ATL06_SIG, - "shard_keys": [-4211324, -4211323, -4211322], + "shard_keys": _CATALOG_WORDS, "granules": [[_rec(4), _rec(5), _rec(6)], [_rec(3)], [_rec(1), _rec(2)]], } p = tmp_path / "catalog.json" @@ -89,14 +95,27 @@ def test_dry_run_morton_cell(self, atl06_config, catalog_file): ) assert result["total_cells"] == 1 - def test_dry_run_invalid_morton_cell(self, atl06_config, catalog_file): + def test_dry_run_absent_morton_cell(self, atl06_config, catalog_file): + # A well-formed decimal morton id that isn't in the catalog. with pytest.raises(ValueError, match="not in catalog"): agg( atl06_config, catalog=catalog_file, store="./out.zarr", dry_run=True, - morton_cell="99999999", + morton_cell="-4211321", + ) + + def test_dry_run_malformed_morton_cell(self, atl06_config, catalog_file): + # HEALPix catalogs take decimal morton ids (issue #199): a raw packed + # word (digits outside 1..4) is rejected with a pointed message. + with pytest.raises(ValueError, match="not a decimal morton id"): + agg( + atl06_config, + catalog=catalog_file, + store="./out.zarr", + dry_run=True, + morton_cell="11827859996358475782", ) @@ -135,6 +154,7 @@ def test_shuffle_after_max_cells_truncation(self): assert sorted(k for k, _ in pairs) == list(range(10)) def test_morton_cell(self): + # A non-healpix signature keeps the stringified-int parse. pairs = _select_cells(self._data(3), morton_cell="1") assert [k for k, _ in pairs] == [1] @@ -142,6 +162,72 @@ def test_invalid_morton_cell(self): with pytest.raises(ValueError, match="not in catalog"): _select_cells(self._data(2), morton_cell="99") + def test_morton_cell_healpix_decimal_round_trip(self): + # HEALPix catalogs store packed words; --morton-cell takes the decimal + # morton string and matches after the parse-back (issue #199). + data = { + "metadata": {}, + "grid_signature": {"type": "healpix"}, + "shard_keys": [11827859996358475782, 11828141471335186438], + "granules": [[_rec(1)], [_rec(2)]], + } + pairs = _select_cells(data, morton_cell="-4211322") + assert [k for k, _ in pairs] == [11827859996358475782] + + def test_morton_cell_healpix_rejects_packed_word(self): + data = { + "metadata": {}, + "grid_signature": {"type": "healpix"}, + "shard_keys": [11827859996358475782], + "granules": [[_rec(1)]], + } + with pytest.raises(ValueError, match="not a decimal morton id"): + _select_cells(data, morton_cell="11827859996358475782") + + def test_morton_cell_miss_on_legacy_catalog_hints_regenerate(self): + # A pre-#199 catalog stores legacy signed-i64 keys; a well-formed + # decimal id then misses every packed-word comparison. Hard break — no + # shim — but the error says why (review finding, PR #205). + data = { + "metadata": {}, + "grid_signature": {"type": "healpix"}, + "shard_keys": [-4211322, -4211323], + "granules": [[_rec(1)], [_rec(2)]], + } + with pytest.raises(ValueError, match="legacy signed decimal ids"): + _select_cells(data, morton_cell="-4211322") + + def test_morton_cell_miss_on_packed_catalog_stays_bare(self): + # A genuinely absent id in a packed-word catalog keeps the plain + # message — the legacy hint fires only when the keys look legacy. + data = { + "metadata": {}, + "grid_signature": {"type": "healpix"}, + "shard_keys": [11827859996358475782], + "granules": [[_rec(1)]], + } + with pytest.raises(ValueError, match="not in catalog$"): + _select_cells(data, morton_cell="-4211321") + + +class TestSafeLabel: + """Issue #199 (review finding, PR #205): error-reporting label rendering + must never raise — re-rendering the malformed key that failed a cell would + abort the accumulation loop precisely while reporting that failure.""" + + def test_falls_back_to_digits_on_invalid_key(self): + from zagg.runner import _safe_label + + g = HealpixGrid(parent_order=6, child_order=12, config=default_config("atl06")) + # 0 is the empty sentinel: shard_label raises; _safe_label renders digits. + assert _safe_label(g, 0) == "0" + + def test_renders_decimal_for_valid_key(self): + from zagg.runner import _safe_label + + g = HealpixGrid(parent_order=6, child_order=12, config=default_config("atl06")) + assert _safe_label(g, 11827859996358475782) == "-4211322" + class TestLambdaDispatchOrder: def _data(self, n=50): @@ -776,7 +862,7 @@ def test_oversized_async_payload_raises_before_dispatch(self): client = MagicMock() big_aoi = list(range(_ASYNC_PAYLOAD_CAP_BYTES // 4)) # >250 KB serialized - with pytest.raises(ValueError, match='pass invocation="sync"'): + with pytest.raises(ValueError, match='pass invocation="sync"') as excinfo: _invoke_lambda_cell( client, (0,), @@ -790,10 +876,14 @@ def test_oversized_async_payload_raises_before_dispatch(self): config_dict=None, max_workers=4, aoi_payload=big_aoi, - result_url="s3://out/x.zarr.status/run1/12345.json", + result_url="s3://out/x.zarr.status/run1/-12345.json", result_fetch=lambda: None, poll_timeout_s=10.0, + label="-12345", ) + # The message names the cell by its rendered label (issue #199), not + # the raw packed-word digits. + assert "cell -12345 " in str(excinfo.value) client.invoke.assert_not_called() def test_oversized_payload_allowed_on_sync_path(self): @@ -942,7 +1032,7 @@ class TestInvocationPassthrough: (issue #151); async threads the per-shard result channel into the cell invoke. Same mocked drive as ``TestMaxRetriesPassthrough``.""" - def _drive(self, monkeypatch, atl06_config, catalog_file, **agg_kwargs): + def _drive(self, monkeypatch, atl06_config, catalog_file, *, grid_factory=None, **agg_kwargs): from unittest.mock import MagicMock import boto3 @@ -952,13 +1042,14 @@ def _drive(self, monkeypatch, atl06_config, catalog_file, **agg_kwargs): from zagg.concurrency import ConcurrencyReport captured = {} + grid_factory = grid_factory or _stub_grid monkeypatch.setattr( runner, "get_nsidc_s3_credentials", lambda: {"accessKeyId": "a", "secretAccessKey": "s", "sessionToken": "t"}, ) - monkeypatch.setattr(grids_mod, "from_config", lambda *a, **k: _stub_grid()) + monkeypatch.setattr(grids_mod, "from_config", lambda *a, **k: grid_factory()) monkeypatch.setattr(runner, "_invoke_lambda_setup", lambda *a, **k: None) monkeypatch.setattr(runner, "_invoke_lambda_finalize", lambda *a, **k: None) monkeypatch.setattr(runner, "_get_function_timeout_s", lambda *a, **k: 720) @@ -983,6 +1074,7 @@ def _fake_cell(*a, **k): result_url=k.get("result_url"), result_fetch=k.get("result_fetch"), poll_timeout_s=k.get("poll_timeout_s"), + label=k.get("label"), ) return { "status_code": 200, @@ -1004,9 +1096,13 @@ def _fake_cell(*a, **k): def test_default_async_threads_result_channel(self, monkeypatch, atl06_config, catalog_file): captured = self._drive(monkeypatch, atl06_config, catalog_file) - # .status//.json, run_id unique per run. + # .status//.json, run_id unique per run. assert captured["result_url"].startswith("s3://out/x.zarr.status/") assert captured["result_url"].endswith(".json") + # The object is named by the grid's shard label (issue #199) — the stub + # labels key k as "-{k}" — never the raw packed-word digits alone. + name = captured["result_url"].rsplit("/", 1)[1] + assert name in {f"-{w}.json" for w in _CATALOG_WORDS} assert callable(captured["result_fetch"]) # _drive pins the function timeout at 720 -> deadline 720 + margin. from zagg.runner import _ASYNC_POLL_MARGIN_S @@ -1019,6 +1115,37 @@ def test_sync_invocation_omits_result_channel(self, monkeypatch, atl06_config, c assert captured["result_fetch"] is None assert captured["poll_timeout_s"] is None + @staticmethod + def _raising_label_grid(): + g = _stub_grid() + g.shard_label.side_effect = ValueError("invalid word") + return g + + def test_sync_malformed_key_label_falls_back(self, monkeypatch, atl06_config, catalog_file): + # Review finding (PR #205): on SYNC runs the label never becomes a + # path (no status key), so a malformed key must not raise out of + # _cell_work — an exception there is RUN-fatal on the Lambda path. + # It falls back to the raw digits instead. + captured = self._drive( + monkeypatch, + atl06_config, + catalog_file, + invocation="sync", + grid_factory=self._raising_label_grid, + ) + assert captured["label"] in {str(w) for w in _CATALOG_WORDS} + + def test_async_malformed_key_is_fail_loud(self, monkeypatch, atl06_config, catalog_file): + # On ASYNC runs the label names the status object — a path component — + # so the same malformed key fails loudly before dispatch. + with pytest.raises(ValueError, match="invalid word"): + self._drive( + monkeypatch, + atl06_config, + catalog_file, + grid_factory=self._raising_label_grid, + ) + def test_unknown_invocation_raises(self, atl06_config, catalog_file): with pytest.raises(ValueError, match="Unknown invocation"): agg( @@ -1089,8 +1216,9 @@ class TestProcessAndWriteStreaming: """Issue #91: the non-sharded ``_process_and_write`` streams each chunk through a ``write_chunk`` callback (no ``chunk_results`` accumulation). Drive a fake ``process_shard`` that streams 1 and K>1 chunks through the callback and assert - the dense ``chunk_idx`` sequence + ragged keying (shard_key at K=1, block-index - key at K>1) — the runner-level analogue of the lambda streaming test.""" + the dense ``chunk_idx`` sequence + ragged keying (the grid's shard label at K=1 + — issue #199 — block-index key at K>1) — the runner-level analogue of the + lambda streaming test.""" def _run(self, monkeypatch, atl06_config, *, chunks_per_shard, chunks): from unittest.mock import MagicMock @@ -1112,6 +1240,9 @@ def fake_process_shard(grid, shard_key, urls, **kwargs): grid.sharded = False grid.chunks_per_shard = chunks_per_shard grid.chunk_grid_shape = (4,) + # K==1 ragged subgroups are keyed by the grid's shard label (issue + # #199); a deterministic decimal-string-shaped stub pins the seam. + grid.shard_label = lambda key: f"-{int(key)}" monkeypatch.setattr(runner, "process_shard", fake_process_shard) monkeypatch.setattr( @@ -1138,7 +1269,7 @@ def fake_process_shard(grid, shard_key, urls, **kwargs): ) return cap - def test_k1_streams_ragged_keyed_by_shard_key(self, monkeypatch, atl06_config): + def test_k1_streams_ragged_keyed_by_shard_label(self, monkeypatch, atl06_config): import pandas as pd cap = self._run( @@ -1150,7 +1281,7 @@ def test_k1_streams_ragged_keyed_by_shard_key(self, monkeypatch, atl06_config): # Streaming seam wired: callback passed, no accumulation sink. assert callable(cap["write_chunk"]) and cap["chunk_results"] is None assert cap["dense"] == [(5,)] - assert cap["ragged"] == [5] # K=1 -> keyed by shard_key + assert cap["ragged"] == ["-5"] # K=1 -> keyed by the grid's shard label (#199) def test_k_gt_1_streams_ragged_keyed_by_block_index(self, monkeypatch, atl06_config): import pandas as pd @@ -1176,6 +1307,9 @@ def _stub_grid(): grid.signature.return_value = {} grid.spatial_signature.return_value = {} grid.block_index.side_effect = lambda k: (k,) + # Deterministic decimal-string-shaped labels (issue #199) so status keys + # and CSR subgroup names built off the stub are real strings. + grid.shard_label.side_effect = lambda k: f"-{int(k)}" grid.emit_template.side_effect = lambda store, overwrite=False: store return grid