Parallelize the spill reducer (intra-block partitions + inter-block k-way) — blocked by #279

[espg]

🤖 from Claude

Summary

Follow-on to #279. Once the t-digest reduce loop is vectorized numpy (GIL-releasing) and the fold is a single order-independent k-way merge, the SpillAggregator reducer can be parallelized for another ~2–4× on heavy shards. This is an engine-level change to src/zagg/processing/spill.py (the reducer thread + the memory budget), not a user-side agg function — at most gated by a worker knob (e.g. streaming.reduce_workers: N). Like #279 it ships in the packaged zagg, so fleet use needs a layer redeploy; it is fully testable locally without one.

Blocked by #279 — specifically:

• The reduce hot loop is a Python-scalar k1 loop today (holds the GIL), so threading it buys nothing on CPU until t-digest agg hot path is unvectorized: ~15× build + ~6× fold available in pure numpy #279 makes it vectorized numpy. Only then does parallelism scale the reduce.
• Cross-block parallelism (Flavor B below) is only correct once t-digest agg hot path is unvectorized: ~15× build + ~6× fold available in pure numpy #279's k-way merge decouples per-block reduction from the final order-independent combine.

Current model

_close_block → _reduce_one on one threading.Thread; _join_reducer before the next close. So reads overlap one block's reduce, but blocks reduce strictly serially in close order, and _fold_block walks its partitions sequentially (spill.py:598-621).

Two flavors of parallelism

Flavor A — intra-block, across partitions (clean, no correctness change).
_fold_block loops for key in block.partition_keys(). Partitions are disjoint morton cell-groups — within a block a cell lives in exactly one partition — so folding partitions concurrently touches disjoint cells. No associativity/order issue. Each partition worker returns its own {cell: digest} and the disjoint results are dict.update-merged single-threaded (avoid mutating shared self._digests/self._counts from N threads). Parallelizes both the partition read-back I/O and the (post-#279) numpy CPU. Realistic degree = Lambda vCPUs (~2 near 1.8 GB, up to 6 at 10 GB).

Flavor B — inter-block, parallel block reduces (rides on #279 Phase 3).
The same cell recurs in every block, so folding two blocks concurrently into shared self._digests[cell] is racy and order-dependent. With #279's k-way merge, each block reduces independently to its own per-block-per-cell digest (no shared state), then a single k-way merge per cell combines across blocks at the end — parallel-safe and block-order-invariant.

Central design question — the memory budget

_default_block_bytes / _BUILD_MULT (spill.py:344-367) are derived assuming one partition reduced at a time, at ~3× its spilled bytes, alongside the 25%-of-memory read buffer. P-way parallelism multiplies that peak working set by P — it would blow the Lambda memory ceiling unless the threshold is re-derived (shrink block_bytes by ~P, or gate the parallelism degree on measured free memory / vCPU count). This re-derivation + re-measuring the replay peaks (the #217 phase-3 slabs) is the substantive design work here, and it couples to the deployment memory/vCPU tier. This is why it is its own PR, not part of #279.

Proposed phasing (once unblocked)

1. Flavor A — parallelize _fold_block across partitions; per-partition workers return disjoint {cell: digest}; merge single-threaded. Memory-budget re-derivation as the gating design decision (measure peak vs _BUILD_MULT × P).
2. Flavor B — parallel block reduces on top of t-digest agg hot path is unvectorized: ~15× build + ~6× fold available in pure numpy #279's k-way merge; independent per-block digests + final k-way combine; assert block-order-invariance in tests.

Open questions

1. Parallelism degree: fixed knob (reduce_workers), or auto from os.cpu_count() / memory budget? Recommendation: auto-derive, capped, with an override knob.
2. Whether Flavor A alone (simplest, no t-digest agg hot path is unvectorized: ~15× build + ~6× fold available in pure numpy #279 Phase 3 dependency beyond the vectorized loop) is enough on the real heavy-shard slabs before taking on Flavor B — decide after measuring Flavor A on the Chunk-scoped digest streaming: bound worker memory to one inner chunk (unlock coarse HEALPix orders) #217 replays.

[espg]

🤖 from Claude

Updated plan — Flavor A only (Flavor B declined)

Revised after the reducer-coverage discussion. The deciding principle: spill applies to all reducers (single-block is byte-identical to pooled, so non-mergeable reducers work; only multi-block is mergeable-only). The parallel reducer should preserve that "works for anything" property. That principle selects Flavor A and rules out Flavor B.

Flavor A — intra-block, across partitions (the plan)

_fold_block walks block.partition_keys() sequentially. Partitions are disjoint morton cell-groups — within a block a cell lives in exactly one partition — so parallel partition-reduces never combine anything across the parallel units. Each worker just runs whatever the per-partition reduce is (summation, k-way tdigest, pairwise tdigest, or the exact single-block pooled build for a non-mergeable reducer). No merge law, no associativity, no ordering enters.

Why A is the right (and sufficient) mechanism:

• Reducer-agnostic by construction — applies to every reducer, including the non-mergeable ones spill uniquely supports via the exact single-block fallback. This is the property that matches spill's ethos.
• Already saturates the cores — a block has up to 4**_GROUP_LEVELS = 64 partition groups; Lambda gives ~2 vCPUs near 1.8 GB, up to 6 at 10 GB. 64 ≫ 6, so partition-level parallelism alone keeps every core busy.
• Memory is tunable independent of block sizing — throttle to a worker pool of min(cores, budget/partition_peak) concurrent partitions. No need to change block_bytes.

Design work (the real content):

• Have each partition worker return its own {cell: digest} / count dict and merge the disjoint results single-threaded (no shared-state mutation from N threads — partitions are disjoint so this is a plain dict.update).
• Memory-budget re-derivation (the central question): _default_block_bytes / _BUILD_MULT assume one partition reduced at a time. P concurrent partitions multiply the peak working set by P. Either cap P by available_memory / partition_peak, or shrink the reservation. Re-measure the Chunk-scoped digest streaming: bound worker memory to one inner chunk (unlock coarse HEALPix orders) #217 phase-3 replay peaks with P workers.
• Apply to both _fold_block (multi-block) and the single-block _chunk_outputs_exact reduce — both are partition-disjoint.
• Assert output is unchanged vs the sequential reduce (A doesn't alter results, only scheduling).

Flavor B — inter-block, parallel block reduces (DECLINED)

Problem — B cannot be universal. The same cell recurs across blocks, so B must combine per-block results per cell, and that combine must be order-independent to be correct under parallelism. That holds only for reducers with an order-independent cross-block law:

• summation (len/count) ✓, k-way tdigest ✓
• pairwise tdigest ✗ (order-dependent by design)
• non-mergeable reducers — never reach multi-block anyway (N/A)

So B is structurally limited to the order-independent-combinable subset — it can't deliver the "all reducers" property that motivates this issue. It also adds multi-block-in-flight memory cost (several blocks' working sets live at once) on top of A's per-partition cost, for parallelism A already provides (B parallelizes across a handful of blocks; A across up to 64 partitions).

Future potential (not now): if a workload ever has very few partitions per block (K=1 configs → a single partition) and many blocks, A's degree collapses to 1 and B's cross-block parallelism would be the only lever. In that narrow case B could return as an opt-in for the order-independent reducers only (summation + k-way tdigest), layered on A. Until such a workload is measured, B stays declined — it's complexity and a reducer-coverage regression for parallelism that isn't needed.

Acceptance

• _fold_block (and the single-block exact reduce) parallelized across partitions with a bounded worker pool; output byte-unchanged vs sequential.
• Memory-budget re-derivation with re-measured replay peaks; P gated by available memory.
• Still blocked by Vectorize t-digest build/merge + order-independent k-way fold (issue #279) #282 (the reduce loop must be GIL-releasing vectorized numpy first).

[espg]

🤖 from Claude

Status now that the blocker cleared, plus one decision this needs before code.

#279 merged (PR #282), so the two things that gated this — a GIL-releasing vectorized reduce and the order-independent k-way fold — are both on main. Flavor A's correctness precondition (disjoint per-partition cells) and Flavor B's (independent per-block digests + a final k-way combine) are now satisfied.

What still gates the first commit is the central design question this issue already names — the memory budget. _default_block_bytes / _BUILD_MULT (spill.py) are derived for one partition reduced at a time (~3× spilled bytes alongside the 25%-of-mem read buffer); P-way parallelism multiplies that peak by P and would breach the Lambda ceiling unless block_bytes shrinks ~P or the degree is gated on measured free memory / vCPU. Re-deriving that threshold and re-measuring the replay peaks is the substantive work here, and it can't be done blind — it wants the reduce-vs-read split from the sibling measurement PR (#285, the 158-granule CONUS slab) to know whether the reduce is even CPU-bound enough to be worth parallelizing.

So two questions before I open a PR:

(1) Sequencing — proceed with Flavor A now on a conservative budget (shrink block_bytes by the parallelism degree, degree = min(vCPU, cap)), or wait for #285's reduce/read number so the budget is tuned rather than worst-cased? A conservative budget lets A land without the measurement but leaves throughput on the table on read-bound shards.

(2) Parallelism degree — the issue leans auto-derive (from os.cpu_count() / memory) capped, with a streaming.reduce_workers override. Confirm that's the shape, vs. a plain fixed knob defaulting to 1 (opt-in). Auto-derive is the better default but couples the tests to the runner's vCPU tier.

My recommendation for the first PR is Flavor A alone (decide-on-B after measuring, per the issue). Marking waiting on these two answers.

---

Generated by Claude Code