Refactor the per-cell aggregation handoff: sort/hash grouping + Arrow path (additive, benchmarked)

[espg]

Refactor the per-cell aggregation handoff: sort/hash grouping + Arrow path (additive, benchmarked)

Summary

Rework the per-shard aggregation hot path in processing.py to (1) replace the per-child
boolean-mask loop with a single sort/hash-based group split, and (2) introduce an Arrow
table handoff alongside the existing pandas one so we can benchmark representation cost head to
head. Pandas is retained in this issue — the goal is measured comparison, not removal. The
pandas→Arrow cutover (and pandas removal) is a deferred follow-up gated on the benchmark.

This lands ahead of #29 (non-scalar outputs): #29's cell→table handoff is the exact seam under
test here, so a benchmarked answer + a cleaner grouping substrate should exist before #29 commits
its handoff design. Keeping the change additive means it does not block #29.

What process_shard does today (and the costs)

As of main (@ 83b92ab), src/zagg/processing.py:

1. _read_group reads h5coro datasets (numpy) per beam, spatial-filters, wraps in
   pd.DataFrame(data_dict).
2. df_all = pd.concat(all_dataframes, ignore_index=True) across the 6 beams.
3. cell_col = grid.cells_of(df_all["leaf_id"].values) → cell id per observation.
4. Per-child loop: for i, child in enumerate(children): df_cell = df_all[cell_col == child]
   → calculate_cell_statistics(df_cell). This is O(n_children × n_obs) — a HEALPix O6→O12
   shard scans ~4,096 boolean masks over the full obs array, and boolean indexing copies the
   matched rows + builds a fresh DataFrame object per cell.
5. calculate_cell_statistics pulls df_cell[col].values (numpy) and runs numpy reductions or an
   eval'd expression; each variable is float(func(...)) (scalar — see Support non-scalar aggregation outputs (vectors + ragged per-cell payloads) #29).
6. df_out = pd.DataFrame({var: stats_arrays[var]}) → write_dataframe_to_zarr reshapes each
   column to grid.chunk_shape and array.set_block_selection(chunk_idx, values).

Pandas here is a thin columnar wrapper: no groupby, no joins — just concat, boolean-mask
selection, and .values. pyarrow is already a dependency; the catalog layer
(catalog/sources.py) is already Arrow/stac-geoparquet, so the per-cell handoff is the only
remaining pandas stack.

Two separable levers (keep them decoupled)

Lever 1 — grouping algorithm (backend-independent, behavior-preserving)

Replace the per-child mask loop with a single group split: sort cell_col once, find run
boundaries (np.diff/np.flatnonzero), and iterate contiguous slices — O(n_obs log n_obs)
instead of O(n_children × n_obs). (This is exactly agg_benchmark.build_groups.) Pure win, no
representation change, byte-for-byte identical outputs.

Lever 2 — representation (pandas vs Arrow, benchmarked, additive)

Add an Arrow handoff path behind a flag/config (default stays pandas), and measure it against
pandas on a real shard. Two parts:

• Hybrid kernel/loop: route the kernel-able stats (count, min/max, mean, var, and
  tdigest) through pyarrow.compute hash-aggregate kernels (Table.group_by().aggregate()),
  which vectorize the whole grouping in one call; keep the per-group loop only for the genuinely
  custom fields (expression, weighted average).
• Zero-copy extraction rule (implementation note): the custom/expression aggregations are the
  same numpy math regardless of backend — they do not regress iff extraction stays zero-copy:
  combine_chunks() once after concat, group by contiguous slice (not scattered take), then
  column.to_numpy(zero_copy_only=True). The footgun that would make them slower: chunked or
  null-bearing columns, or scattered take, force a per-cell copy (back to pandas-level cost, not
  worse). At our density (median ~16 obs/cell) per-cell Python overhead dominates either way, so
  the win is moving kernel-able stats out of the loop, not the array extraction itself.

Why bother (the case for eventually landing Arrow)

• Consistency: catalog is already Arrow; this removes the second dataframe stack and the
  pandas==2.2.3 hard pin.
• Backend support (i.e., local processing support) #20 (multi-backend): Arrow is the zero-copy currency for vaex/duckdb/ray/dask — the cleanest
  substrate for the local-parallel and BYO-cluster backends.
• Support non-scalar aggregation outputs (vectors + ragged per-cell payloads) #29 (non-scalar outputs): Arrow FixedSizeList (vectors) and List/LargeList (ragged)
  carry the non-scalar per-cell payloads natively, vs awkward pandas object columns.
• Memory: lighter than pandas (no index/object overhead) — relevant to the 2 GB Lambda budget.

Proposed work

• Lever 1: sort/hash group split replacing the per-child mask loop in process_shard.
  Byte-for-byte verified against current scalar outputs.
• Pass numpy slices, not sub-DataFrames, into calculate_cell_statistics (drop per-cell
  DataFrame construction); keep the function's numpy/eval math unchanged.
• Lever 2 (additive): an Arrow handoff path (config/flag; pandas remains default) —
  _read_group returns an Arrow table, concat_tables + combine_chunks, hybrid
  kernel/loop aggregation, zero-copy slice extraction.
• Benchmark harness: one real shard (e.g. an ATL03/ATL06 HEALPix shard), measure
  wall-time and peak RSS for {pandas, arrow} × {old loop, new grouping}, asserting outputs
  are byte-for-byte equal. Report a small table.
• Write path: write_dataframe_to_zarr already only needs name → ndarray; adapt it (or a
  thin sibling) to accept an Arrow table / dict of arrays so it is handoff-agnostic.

Scope / non-goals

• In: grouping refactor, an additive Arrow handoff path, a benchmark, byte-for-byte parity.
• Out: removing pandas (deferred follow-up, gated on the benchmark), changing any aggregation
  semantics, non-scalar outputs (Support non-scalar aggregation outputs (vectors + ragged per-cell payloads) #29), new grid types (Design doc for broadening input/output data structures #11). HEALPix/ATL06 outputs must stay
  byte-for-byte identical throughout.

Acceptance criteria

• The grouping is O(n_obs log n_obs), not O(n_children × n_obs); scalar outputs unchanged
  byte-for-byte.
• Both pandas and Arrow paths produce byte-for-byte identical scalar outputs.
• A benchmark table reports wall-time + peak memory for pandas vs Arrow on a real shard.
• Pandas remains the default; the Arrow path is opt-in. No change to HEALPix/ATL06 results.

Phasing relative to #29

Land this ahead of #29, additive. #29 then builds non-scalar outputs on the cleaner grouping
and consumes this benchmark's conclusion for its handoff (Arrow nested types vs pandas object
columns). The pandas removal is a separate later cleanup, not part of either issue.

References

• Hot path: src/zagg/processing.py (_read_group, process_shard per-child loop,
  calculate_cell_statistics, write_dataframe_to_zarr).
• Already-Arrow: src/zagg/catalog/sources.py. Deps: pyproject.toml (pyarrow present,
  pandas==2.2.3 pinned).
• Grouping reference: waveform_benchmarks/agg_benchmark.py build_groups (sort + boundary split).
• Related: non-scalar outputs (Support non-scalar aggregation outputs (vectors + ragged per-cell payloads) #29), multi-backend support (Backend support (i.e., local processing support) #20).

[espg]

claude, this has started to be implemented-- we should discuss an addition on the testing / benchmarking side so we can have robust understanding of how this impacts performance (or if it doesn't).

A testing script(s) would make sense. Three required areas -- NEON site in Maryland (10km by 10km) , similar 10 km by 10km patch outside of kangerlusaaq in greenland (russell glacier, on the ice), and a bathymetric example (Bahamas, more details coming), also at 10km by 10km. All of these would hit an ATL03 aggregation, and scale in the time dimension of the catalog (1d, 5d, 15d, 30d, 91d, 1y).

Additionally, we'd define an optional 'hard' test ; 10 km by 10km at 88 degrees south (max overlaps). This last one is probably differed; the others are likely cheap enough to run within cod-speed or a another git hub runner (hassle to setup with aws though). First pass would be a script that runs these and records timing information outside of CI/CD.

[espg]

🤖 from Claude

Good call — a real-data timing script is the missing half of this issue's benchmark deliverable. Here's how I'd structure it and the few things I need pinned before writing it.

Two complementary harnesses:

1. Synthetic, CI-runnable (representation parity) — I'm adding this as part of PR sort/hash grouping refactor ( #30 ) #33 phase 2: it builds an in-memory shard of N obs across M cells and times {pandas, arrow} × {old loop, new grouping} while asserting byte-for-byte parity. It isolates grouping + extraction + per-cell overhead (no I/O), so it runs anywhere with no creds. This is the "both paths produce identical outputs" acceptance criterion.
2. Real-data, out-of-CI (the script you're describing) — exercises real h5coro/S3 reads and the actual density/overlap regimes the synthetic one can't. Lands as phase 3 of PR sort/hash grouping refactor ( #30 ) #33.

Proposed shape of the real-data script (benchmarks/region_timing.py, not under tests/ since it needs earthaccess/S3 creds and is slow):

• Region table → bbox: NEON Maryland, Russell Glacier (Kangerlussuaq), Bahamas bathymetry, each a 10×10 km patch; optional --hard adds 88°S (max RGT convergence → biggest per-cell n, exactly where the O(n log n) grouping beats the mask loop).
• Time-window sweep {1d, 5d, 15d, 30d, 91d, 1y} driving the catalog query.
• Per (region × window) record: wall-time, peak RSS (via tracemalloc/resource), total_obs, cells_with_data, output bytes — for each {pandas, arrow} path, asserting scalar parity.
• Wired to the existing runner.agg() / catalog path; the actual run left to an AWS-credentialed session (CMR-STAC + NSIDC S3).

The three regions give a good density spread — Maryland (land/veg, moderate), Russell (ice, high photon return), Bahamas (bathymetry, sparse/attenuated subsurface) — which is the point: the representation cost should be measured across regimes, not one.

Questions before I write it (phase 3):

1. Signal-confidence filter per region — ATL03 signal_conf_ph is per-surface-type (land/sea-ice/land-ice/inland-water). Do you want one fixed confidence threshold across all three, or per-region (e.g. land-ice index for Russell, a bathymetry/sea index for the Bahamas)? This materially changes obs counts.
2. Fixed grid resolution for the 10×10 km patch — a single HEALPix child_order (which ~10 m?) and/or rectilinear resolution so density is comparable region-to-region?
3. Where do the numbers go — a committed benchmarks/results/*.md table (so runs are diffable in git), or stdout-only for now?

Once those are set I'll add the script as phase 3 (and keep --hard/88°S behind a flag, deferred as you suggested). The synthetic harness lands now with phase 2 regardless.

---

Generated by Claude Code

[espg]

1. realistically we want everything except for noise; i.e., we should filter signal_conf_ph = 0 and keep the other values (see https://nsidc.org/sites/default/files/icesat2_atl03_data_dict_v003_0.pdf for flag details).
2. we'll start with rectilinear for now; we're waiting on mortie docs: mark SSO execution-role path out of date #35 before we can properly represent a order-19 cell, which is the closest match to 10-meter resolution in healpix.
3. script that outputs to stout / output file. We're far from being to run this in CI/CD given the aws and nasa creds that would need to be configured (potential issue for later in the month).