Support non-scalar aggregation outputs (vectors + ragged per-cell payloads)

[espg]

Support non-scalar aggregation outputs (fixed-length vectors + variable-length/ragged per-cell payloads)

Summary

zagg currently reduces every aggregation variable to one scalar per cell. We want to emit
vectors (fixed-length per cell) and variable-length / sparse encodings (per-cell payload
whose size varies by cell) so a cell can hold a mergeable quantile sketch (t-digest centroids,
KLL/REQ) or a raw value tier instead of a lone float. The only hard requirement is that the
per-cell output is representable in Zarr v3.

Motivation (empirical, from the waveform_benchmarks prototype)

Prototyping on real ATL03 (RGT 1336, region 05, E. Greenland, 29-cycle time series; see
waveform_benchmarks/agg_benchmark.py + aggregation_review.ipynb) established:

• Scalar quantiles are not mergeable. Stored per-(season,year) q02/q50/q98 scalars cannot
  be correctly re-pooled across the temporal/spatial merge axis (averaging percentiles is wrong).
  Mergeability is the whole reason for a per-cell summary, and it requires a sketch payload.
• A t-digest is two short vectors (means, weights) shape (k, 2). On real data it
  reconstructs every quantile within ~1 cm of empirical, and two independently-built digests
  combine() into one whose quantiles match the pooled raw — the mergeable object to store per cell.
• Density regime: per 10 m cell per overpass, signal photons are p50 16 / max 138; pooled
  across cycles the busiest cell ≈ 670. So leaf cells are small (store raw or near-raw) and the
  t-digest pays off at the merge level. Either way the per-cell payload is non-scalar and its
  length varies by cell.
• A dense fixed-pad layout is infeasible (~115 GB for one box). The workable layouts are a
  fixed trailing-dimension vector or a ragged/CSR encoding (values + offsets +
  cell_ids), both Zarr-v3-native. (zarr v3 dropped the v2 object_codec path, so CSR is the
  portable ragged representation.)

Relationship to recently merged work

• Design doc for broadening input/output data structures #11 (generalize output grids — merged, docs/design/generalized_grids.md) generalizes the
  grid (HEALPix/H3/rectilinear) but explicitly keeps the aggregation variables grid-agnostic and
  one scalar array per variable. This issue is the orthogonal axis: generalize the per-cell
  value (scalar → vector/ragged). It plugs into the surface Design doc for broadening input/output data structures #11 established — the shared
  aggregation-variable block consumed by every OutputGrid.emit_template(schema).
• Signature-enforced runs (Reconcile catalog API + validate grid redesign: Catalog fetch/shard-map split, odc.geo grids, signature-enforced runs #23/Add bring-your-own-role path for IAM-constrained deploys; creds handling for external s3 bucket writes #27 — merged): OutputGrid now exposes signature() /
  nests_with(); a ShardMap records the grid signature and run-time refuses a mismatch. A
  non-scalar payload changes array structure, so the field's output kind/shape must be folded into
  the signature (else a scalar shardmap could pair with a vector run).
• Catalog moved to Arrow (Reconcile catalog API + validate grid redesign: Catalog fetch/shard-map split, odc.geo grids, signature-enforced runs #23): catalog/sources.py is pyarrow/stac-geoparquet, but the
  per-cell aggregation→Zarr handoff in processing.py is still pandas. No impact today; an
  eventual Arrow handoff is synergistic (FixedSizeList for vectors, List/LargeList for ragged).
• Multi-backend (Backend support (i.e., local processing support) #20): see design question 6.

Current behavior — the scalar-per-cell contract

As of main (@ 83b92ab), the scalar assumption is baked into the whole cell→Zarr path:

1. Aggregation eval — processing.calculate_cell_statistics: result[name] = float(func(values, **resolved_params)) hard-casts each variable to a single float per cell.
2. Per-cell accumulator — processing.process_shard: stats_arrays[name] = np.full(n_cells, ...) is 1-D, filled with stats_arrays[key][i] = value.
3. Cell→Zarr handoff — df_out = pd.DataFrame({var: stats_arrays[var]}), one row per cell,
   scalar per cell. (pandas today; catalog layer already Arrow — see above.)
4. Write path — processing.write_dataframe_to_zarr asserts
   len(df_out) == prod(grid.chunk_shape), reshapes each column to grid.chunk_shape,
   array.set_block_selection(chunk_idx, values). No room for a payload axis or ragged data.
5. Template — per-grid OutputGrid.emit_template(schema), built from a shared
   aggregation-variable block via pydantic-zarr (HEALPix path = schema.xdggs_spec /
   xdggs_zarr_template; rectilinear = grids/rectilinear.py _spec). Each agg field becomes one
   ArraySpec of shape = array_shape (grid shape), scalar data_type. No payload dimension.
6. Schema/config — get_agg_fields (config.py) describes a field with
   function/expression/source/params/dtype/fill_value; no notion of output rank or variable length.

Proposed capability (two tiers)

Tier 1 — fixed-length vector per cell. Field declares width C (and optional inner dim, e.g.
(C,2) for t-digest (mean,weight)). Zarr array array_shape + (C,...), chunks
chunk_shape + (C,...), sentinel for unused slots. Covers fixed quantile vectors and a
padded/bounded t-digest. Rectangular but pads to worst case.

Tier 2 — variable-length / ragged (CSR) per cell. Per-cell length varies (true t-digest
centroid count, or a raw tier). Represent as a small group of arrays — values (flat) + offsets

• cell_ids — instead of one grid-shaped array (validated in agg_benchmark.py write_csr).
  Alternative to evaluate: Zarr v3 VLenBytes (simpler API; verify reader round-trip first).

Where the changes land (touch-point checklist)

• config/schema (config.py): declare output kind — scalar (default), vector (with
  length/inner shape), ragged. Scalar stays default so existing configs are unchanged.
• calculate_cell_statistics: drop the unconditional float(...); allow ndarray/bytes
  returns, dispatched on field kind.
• process_shard accumulation: build the right container per kind (N×C array vs
  list-of-payloads + lengths).
• cell→table handoff: the row-per-cell scalar model holds for scalar/fixed-vector but
  breaks for ragged. Generalize the handoff (currently pd.DataFrame) — an Arrow table with
  FixedSizeList/List columns carries vector/ragged natively; otherwise object-column or a
  dedicated non-DataFrame writer.
• write_dataframe_to_zarr / new writer: trailing payload dim (vector) + CSR/vlen writer
  keyed by cell_id (ragged). The len == prod(chunk_shape) invariant only applies to
  dense/scalar/vector.
• template emission — extend the shared aggregation-variable block (pydantic-zarr
  GroupSpec) consumed by every OutputGrid.emit_template(schema), so all grids inherit
  non-scalar support; touches schema.xdggs_spec (HEALPix) and grids/rectilinear.py _spec.
• OutputGrid.signature(): include the per-field output kind/shape so a non-scalar run
  can't be silently paired with a scalar ShardMap (Reconcile catalog API + validate grid redesign: Catalog fetch/shard-map split, odc.geo grids, signature-enforced runs #23/Add bring-your-own-role path for IAM-constrained deploys; creds handling for external s3 bucket writes #27).
• reader/round-trip: helper to read back a cell vector/sketch and (for sketches)
  reconstruct/merge — needed for the merge step and tests.

Design questions to settle first

1. Scope: Tier 1 first + Tier 2 follow-up, or design both together? (Lean: design the
   field-kind abstraction once; implement Tier 1, then Tier 2.)
2. Ragged mechanism: CSR triple vs Zarr v3 VLenBytes — quick spike.
3. Where sketches are built: prototype concludes raw at leaf, t-digest at merge. Leaf-emit
   only here; merge in the temporal work?
4. Sentinel/fill for vectors: weight-0 padding + NaN/inf means so reducers don't read padding.
5. Signature coverage: exactly what of the output kind/shape goes into signature() /
   nests_with() so scalar and vector runs are correctly judged incompatible.
6. Backend portability (Backend support (i.e., local processing support) #20): the output contract must hold across backends (Lambda, local
   numpy, vaex, BYO cluster). A scalar reduction maps onto any groupby-agg; vector/ragged/custom
   sketch does not map cleanly onto vaex built-ins and may need per-group apply. Define the
   field-kind contract backend-agnostically; validate on local-numpy first (the
   waveform_benchmarks prototype is that reference). An Arrow handoff also eases this — Arrow is
   the common currency across numpy/vaex/cluster backends.

Scope / non-goals

• In: per-cell non-scalar output representation (config → eval → template → write → read) in
  the shared aggregation-variable + emit_template surface, so all OutputGrid impls inherit
  it; validated on the rectilinear grid.
• Out (separate): choosing the final sketch (t-digest vs KLL/REQ), tuning CMAX/compression,
  new grid types (Design doc for broadening input/output data structures #11), crossover-cell density studies, and the full pandas→Arrow handoff migration
  (this issue should not depend on it, only stay compatible).

Acceptance criteria

• A config can declare a vector aggregation → correctly-shaped Zarr array, round-tripped in a test.
• A config can declare a ragged aggregation → CSR (or vlen) payload, round-tripped, with a reader
  reconstructing a cell.
• The grid signature() distinguishes scalar/vector/ragged runs (mismatch is rejected).
• Existing scalar configs unchanged (scalar default path); HEALPix/ATL06 path byte-for-byte stable.
• Works end-to-end on the rectilinear grid for the waveform_benchmarks ATL03 case.

References

• Prototype + results: waveform_benchmarks/agg_benchmark.py, aggregation_review.ipynb,
  DESIGN_NOTES.md (CSR layout, t-digest merge demo, density).
• Current scalar path: src/zagg/processing.py (calculate_cell_statistics, process_shard,
  write_dataframe_to_zarr), src/zagg/schema.py (xdggs_spec/xdggs_zarr_template),
  src/zagg/grids/rectilinear.py (_spec/emit_template), src/zagg/grids/base.py
  (OutputGrid, signature, nests_with), src/zagg/config.py (get_agg_fields).
• Related: generalize-output-grids (Design doc for broadening input/output data structures #11, docs/design/generalized_grids.md), signature-enforced
  runs (Reconcile catalog API + validate grid redesign: Catalog fetch/shard-map split, odc.geo grids, signature-enforced runs #23/Add bring-your-own-role path for IAM-constrained deploys; creds handling for external s3 bucket writes #27), multi-backend support (Backend support (i.e., local processing support) #20 — output contract must be backend-agnostic),
  catalog Arrow migration (catalog/sources.py).

[espg]

Prerequisite/phasing: #30 reworks the per-cell aggregation handoff (sort/hash grouping + an additive, benchmarked Arrow path) and is slated to land ahead of this issue. The cell→table handoff #30 settles is the seam this issue extends for non-scalar payloads (Arrow FixedSizeList/List vs pandas object columns), so #30's benchmark conclusion feeds the handoff design here. #30 keeps pandas as default; pandas removal is deferred behind both.

[espg]

Design questions to settle first (answers)

1. Scope: design the field-kind abstraction once; implement Tier 1, then Tier 2.
2. Ragged mechanism: CSR triple vs Zarr v3 VLenBytes — What's the question here? provide more context
3. Where sketches are built: prototype concludes raw at leaf, t-digest at merge. Leaf-emit
   only here; merge is for future (temporal work)
4. Sentinel/fill for vectors: weight-0 padding + NaN/inf means so reducers don't read padding. Sure? What's the question here?
5. Signature coverage: exactly what of the output kind/shape goes into signature() /
   nests_with() so scalar and vector runs are correctly judged incompatible. What are the options?
6. Backend portability (Backend support (i.e., local processing support) #20): the output contract must hold across backends-- JUST Lambda and local
   numpy for now. The vaex, BYO cluster isn't built yet, so don't hamper design to maintain compatibility

[espg]

🤖 from Claude

Thanks — your answers settle the scope. Recording the decisions and then giving the context you asked for on Q2/Q4/Q5.

Settled from your reply: (1) design the field-kind abstraction once, implement Tier 1 (fixed vector) first, Tier 2 (ragged) as a follow-up; (3) leaf-emit only here — t-digest-at-merge is deferred to the temporal work; (6) target Lambda + local-numpy only — don't hamstring the design for vaex/BYO-cluster that don't exist yet. Good — that keeps Tier 1 small.

Q2 — ragged mechanism (the missing context)

The question is purely how the variable-length payload physically lands in Zarr v3, and it's a real fork because Zarr v3 dropped the v2 object_codec path. Two candidates:

• CSR triple — three dense Zarr arrays: values (flat concat of every cell's payload), offsets (n_cells+1, where each cell starts), cell_ids (segment→cell, or implied by grid order). All dense numeric → fully v3-native, chunkable, compressible, no special codec, readable by any Zarr reader with a plain numpy slice. This is what agg_benchmark.write_csr prototyped.
• Zarr v3 VLenBytes — one array whose elements are variable-length byte strings; you'd serialize each cell's payload to bytes. Simpler logical model and per-cell random access, but it's a newer/less-portable v3 codec, costs a serialize/deserialize per cell, and loses the columnar numeric view (you can't slice across cells without decoding).

Recommendation: CSR for numeric ragged payloads (t-digest centroids, raw tiers). It's portable, columnar, and mergeable, and it's already prototyped. Reserve VLenBytes only if a future field needs genuinely heterogeneous per-cell blobs rather than numeric vectors — and the "quick spike" you flagged is exactly: confirm a stock Zarr v3 reader round-trips VLenBytes before betting on it. CSR needs no such verification.

Q4 — sentinel/fill (why it was a question)

Narrow point: for Tier 1, the vector is padded to a worst-case width C, so a reader/merger must distinguish a real slot from padding. Your convention — weight-0 padding + NaN/inf means — is the answer; there's no open question once it's stated, I was asking you to pick it. The one thing to nail: make the sentinel part of the field's declared schema (e.g. pad_weight=0), not an implicit constant, so a t-digest reader and a plain quantile-vector reader each apply the right "skip padding" rule from metadata rather than a hard-coded assumption.

Q5 — signature coverage (the options you asked for)

What of the output kind/shape goes into signature()/nests_with() so a scalar ShardMap can't pair with a vector run:

• Option A — kind only: signature carries per-field {scalar|vector|ragged}. Cheapest; rejects scalar-vs-vector, but treats all vector widths as interchangeable.
• Option B — kind + trailing shape + dtype: also carries the vector width/inner dims (and, for ragged, the CSR component dtypes). Rejects a width-8 run reusing a width-16 shardmap — which matters because a width mismatch silently corrupts set_block_selection. Recommended.
• Option C — full field-spec hash: hash kind+shape+dtype+fill+function-id. Strongest (also catches dtype/sentinel drift) but brittle — any benign metadata edit invalidates an existing store.

Recommendation: Option B. Extend each field entry in signature() to {name, kind, trailing_shape, dtype}; keep nests_with() grid-geometry-only and add a "same field-kind set" check. It's the minimal set that makes a scalar/vector/ragged-or-width mismatch un-pairable without over-coupling to function identity.

With Tier 1 + leaf-emit + Lambda/local scoped, this is ready to move to plan (phased: config field-kind → calculate_cell_statistics ndarray return → process_shard container → template/signature → vector writer → reader/round-trip test). Want me to draft that plan? Note the cell→table handoff it rides on is exactly what #30/PR #33 is reworking, so Tier 1 should land on top of that.

---

Generated by Claude Code