Design: memory-bounded adaptive chunking for large raster→H3 hex (finer H3-parent chunks + windowed COG reads + Armada)

[cboettig]

Design: memory-bounded adaptive chunking for large raster→H3 hex

Type: design / architecture (not a bug). Goal: make cng-datasets raster scale to large/fine imports with uniform small memory requests automatically, instead of a few oversized big-instance pods that need bespoke handling.

Surfaced by data-workflows #453 (Annual NLCD, res-10 CONUS × 19 yr) and complementary to #172 (multi-year output structure). Applies to any large raster hex — e.g. the 300 m global res-8 layers already flagged as needing 64Gi.

Problem (measured)

The hex unit of parallelism is one h0 base cell (--h0-index 0–121, the only chunk knob). RAM scales with the largest chunk's H3-cell count, and h0 is too coarse at fine resolutions:

• res-10 densest CONUS h0 ≈ 282 M cells; default exact-extract peaks at ~140 GiB (live kubectl top, #453).
• The build requested 256Gi — ~1.8× the real peak — and NRP 256Gi nodes have only ~250Gi allocatable, so the request barely places → chronic FailedScheduling.
• The profile is bimodal: ~116/122 h0 are empty ocean (finish in seconds, ~GiB), a handful of land h0 cost everything. A flat request oversizes the many to survive the few.

Net: today large imports require either an oversized flat request (wastes RAM, starves scheduling) or bespoke two-tier / node-taint handling. We want neither.

Root cause

RAM ≈ f(cells in the single largest chunk), chunk = h0. Coarse chunk + fine resolution = a few enormous chunks. Fix the chunk granularity and the memory problem dissolves.

What exists today (from cng-datasets raster --help)

• Chunking: only --h0-index (0–121). No sub-h0 chunking.
• --method {exact-extract | warp-centroid} — exact-extract (default) is area-weighted, one row/cell, memory-heavy; warp-centroid is documented "fast and low-memory" (point-samples pixel→cell centroid, one row per warped pixel).
• --target-extent xmin,ymin,xmax,ymax clip bbox (mosaic clipping).
• --local-cache-dir / --no-local-cache (each pod localizes the whole COG today, e.g. 1.5 GiB/pod).

Design levers (best → band-aid)

A. Finer H3-parent chunking (the real fix). Chunk by h1 (7 children) or h2 (49) instead of h0. H3 nests perfectly, so sub-cells tile an h0 with no seam artifacts / no dedup (unlike arbitrary bbox tiling of exact-extract, where a res-10 cell straddling a bbox edge would be split/double-counted). Estimated peak (proportional to cell count — needs empirical confirmation): h1 ≈ 20 GiB, h2 ≈ 3 GiB. Makes RAM uniform and small across all pods; density no longer needs to be known; no two-tier. Needs --chunk-resolution N / --parent-index.

B. --method=warp-centroid (low-mem, exists now). For categorical/mode (e.g. NLCD) centroid-sampling ~15 pixels/cell ≈ area-weighted mode. May be a near-free RAM win. Caveats: semantics differ slightly (not for area-weighted continuous layers where exact fractional coverage matters), and it may emit a large per-pixel intermediate needing a GROUP BY. A/B test required (peak RAM + value agreement vs exact-extract).

C. Windowed COG reads. At finer chunking each small pod should read only its sub-region of the COG (COG range reads), not localize the full file — otherwise per-pod I/O explodes with pod count (h2 → thousands of pods each pulling 1.5 GiB).

D. Execution backend. k8s indexed Jobs are fine at h0 (~2318 pods) but finer chunking blows the ~200-pod namespace guideline (h2 ≈ 6000 chunks/COG). This is the documented Armada regime (no pod cap, external queue, --chunk-size 1 so one slow chunk never blocks). Small uniform pods also schedule on any node. So: fine chunking ⇒ Armada backend.

E. Band-aid (superseded by A): two-tier memory (big request only for populated h0). Helps scheduling, not peak RAM.

Target behavior

raster-workflow picks a chunk resolution so the largest chunk stays under a memory budget (e.g. --max-hex-memory 96Gi → auto-select h1/h2/h3), emits uniform small-request chunks, reads windowed COG, and routes to Armada when the chunk count is large — no per-import custom handling. Target: < 120 GiB, ideally < 32 GiB, uniform.

How we proceed (phased — design + test)

• Phase 0 — empirical de-risk (no tool changes; can run on NRP now):
  1. warp-centroid vs exact-extract on one dense h0: peak RAM + confirm mode values match.
  2. RAM-vs-granularity: measure peak for a sub-region of one h0 (via --target-extent, for the RAM curve only — not production data) to confirm the ~1/7, ~1/49 scaling before committing to the interface.
• Phase 1 — interface design: --chunk-resolution / auto from --max-hex-memory; windowed COG read; Armada auto-route; how sub-chunks compose parents/partitions; preserve the antimeridian/pole handling (Raster Plan A: h0 boundary polyfill double-counts cells, ~16% overcount globally #88/_split_antimeridian (#88 fix) throws GEOSException on polar antimeridian cells #92) under sub-chunking.
• Phase 2 — implement + test: correctness gate — sub-chunked hex must be bit-identical (values + coverage) to the h0-chunk baseline on a fixture; performance (RAM/pod, wall-clock, total I/O); edge cases (seam h0, all-nodata chunks, points/lines).
• Phase 3 — adopt as the default path for large/fine imports in raster-workflow; keep h0-chunk as a fallback.

Open questions

• Auto chunk-resolution from a memory budget, or explicit --chunk-resolution? (auto is the "no custom handling" goal.)
• Does warp-centroid already give correct mode/min/max per cell, or only per-pixel rows needing aggregation? (determines whether B is a method or a pipeline stage.)
• Windowed COG read: GDAL -projwin/range reads per chunk — worth it vs a one-time COG→tiles pre-split?
• Interaction with the multi-year year= pattern (raster-workflow: built-in multi-year / time-series pattern (per-year COGs + year=-partitioned unified hex) #172): chunk res is orthogonal (per year × sub-chunk).

Refs: data-workflows #453 (evidence + measured RAM), boettiger-lab/datasets #172 (multi-year output structure, complementary), AGENTS memory model ("RAM ∝ largest chunk's cell count").

[cboettig]

⛔ Hard requirement (from review): decouple COMPUTE chunk granularity from OUTPUT partition granularity

Finer compute chunking (h1/h2) must not change the final layout. The output must always coalesce to the canonical ≤122 h0 partitions — one data_0.parquet per h0 (per year= for the multi-year case). A blow-up to tens of thousands of sub-h0 shards would wreck global-read performance and ** globs — a non-starter.

So the fine-chunk pipeline is map → reduce, exactly like the existing vector path (chunks/ → repartition → hex/ by h0):

1. map: many small pods hex their h1/h2 sub-chunk → sub-shards in a staging prefix (e.g. hex-chunks/h0=<cell>/<subid>.parquet).
2. reduce/repartition: coalesce each h0's sub-shards into one file per h0 — DuckDB COPY ... PARTITION_BY (h0), FILENAME_PATTERN 'data_0' with threads=1 to avoid re-fragmentation (the gfw-consolidate recipe already does exactly this) → final hex/h0=*/data_0.parquet.
3. purge staging + verify empty (the #409 purge-and-verify discipline).

Result: fine chunks for bounded RAM, canonical ≤122 h0 partitions for fast reads — the two are independent by construction. The repartition step's memory stays bounded (streams per-h0; per-h0 total is what one h0-chunk pod already holds today).

The current NLCD build (#453) sidesteps this because it chunks by h0 (one writer per h0 → already one file each, ~114 mode files, verified). The reduce step is only needed once compute chunking goes finer than the output partition key. Add to Phase 1 (interface) + Phase 2 (correctness gate: final file count == populated-h0 count, byte-identical to the h0-chunk baseline).