feat(planner): Pass 2 reuse rules (rollup, exact top-k reuse) and nested sum merge in Pass 1 (#680)

[zzylol]

Closes #680.

Why

The legacy search could answer a computation from a cheaper related one (rollup, top-k reuse, nested sum merge). The Stage 1 → Stage 3 pipeline had none of these, so it missed cheaper plans (#676 capability audit). This PR follows the decisions recorded for #680:

• the nested sum merge is a Pass 1 candidate, not a Pass 2 rule;
• rollup reuses only finer computations already in the workload, and only for families the family-capability table marks mergeable;
• top-k reuse is exact-only for now.

What

• Pass 1, nested sum merge. sum by (job) (sum_over_time(m[5m])), or any grouped sum over the value of a sum directly beneath it, gets a new candidate NestedSumMerge next to its own accumulator. The candidate is one exact Sum accumulator over the inner sum's input rows, grouped by the outer keys. It absorbs the inner target, the same way the whole-expression top-k does.
• Pass 2, rollup rule. A coarse exact accumulator (sum, min, max over Float64 values) is answered from a finer one that the candidate already computes over the same rows (for example sum by (job) from sum by (job, instance)). The rollup is an exact accumulator over the finer finalized values, grouped by the coarse keys. The value stays exact. (The table lets counts roll up as the sum of the finer counts, but the executor's exact accumulators take only Float64 updates, so integer counts are refused.)
• Pass 2, exact top-k reuse rule. Exact top-k computations of one input and partition with different k all read one ranking, Limit(k_max) over Sort.
• Stage 1 lists the reuse candidates after the window-composition ones: exact top-k reuse first, then each rollup candidate followed by its own top-k reuse candidates. Each is followed by its identical-expression candidates. Plan selection prices them against the separate computations.

How

• add_nested_sum_candidates and nested_sum_input check that the merge applies:

  • the outer sum groups by listed keys, has no filter or HAVING, and is not itself absorbed;
  • the inner sum has one consumer and no filter or HAVING;
  • the outer sum reads the inner value column;
  • when the inner sum is per series, no outer key is __name__: a per-series sum drops the metric name while the raw samples keep it, so the merge would group under m where the nested sums group under "";
  • neither the outer sum is absorbed nor the inner sum absorbs another target, so merges never stack and an absorbing target is never absorbed (the combinations() invariant).

  The outer keys are moved onto the inner input by column (partitions_over). ComputationCandidate::absorbs_inner replaces the earlier WholeExpression matches in combinations, absorbed_choice_varies and the display name.

• rollup_kind reads the family-capability table. A family qualifies when its states merge over disjoint or overlapping selections with an Exact merged error. Counter states (DisjointInTime), IRate (no merge) and every sketch family return None.

• rollup_sites pairs each coarse build (read only through FinalizeExactAccumulator) with its coarsest finer build. The finer build must have the same family, update and filter, read the same rows, and have group keys that strictly include the coarse keys. "Same rows" means equal nodes, or PromQL scans of one metric without predicates under the same time ranges and shifts, which may project other labels. A global SQL coarse aggregate is refused, because it has a row even over empty input. The finer values must be Float64, the only update type exact accumulators execute, so integer counts do not roll up. The rollup value is renamed to the coarse columns, and its types must match.

• topk_reuse_sites groups plain Aggregate top-k nodes over an exact, unfiltered input by equal input and partition. In a reuse candidate every member reads one Limit(k_max, Sort) ranking. Stage 2 still implements each top-k as Limit(Sort) over that ranking.

• Every variant composes its guarantees again (with_composed_guarantees).

• Cap. Reuse candidates are listed eagerly, like every other Pass 2 rule's. Every rollup subset is combined with every top-k reuse subset, so the cap applies to the product: more than 20 rollup and top-k sites together (MAX_REUSE_SITES) list one Unsupported error instead of 2^n candidates.

Key code interfaces added / modified

crates/asap-aware-mapping/src/per_computation_candidates.rs

/// One Pass 1 candidate for a computation (one measure of an aggregate).
/// The type stays; `NestedSumMerge` is new.
pub enum ComputationCandidate {
    /// Plain exact execution of the original sub-DAG.
    Exact, // unchanged
    /// An exact mergeable accumulator state.
    ExactAccumulator { kind: ExactKind, params: ExactParams }, // unchanged
    /// An unsized summary family, sized later for the measure's accuracy.
    Summary(SketchAlgorithm), // unchanged
    /// `avg` as exact `sum` and `count` accumulator states.
    AvgAsSumCount, // unchanged
    /// A whole-expression top-k heap summary over the input of the per-item sum or
    /// count beneath it; absorbs that target.
    WholeExpression(SketchAlgorithm), // unchanged
    /// A grouped `sum` over a `sum` beneath it as one exact `Sum` accumulator over
    /// the inner sum's input, grouped by the outer keys; absorbs the inner target.
    NestedSumMerge, // new
}

impl ComputationCandidate {
    /// Whether the candidate absorbs the target beneath its own
    /// (`ComputationTarget::absorbs`): a whole-expression top-k or a nested-sum merge.
    pub fn absorbs_inner(&self) -> bool; // new
}

/// One aggregate measure Pass 1 offers candidates for. The type stays; `absorbs`
/// now also serves nested-sum merges.
pub struct ComputationTarget {
    /// The target directly beneath this one that its absorbing candidates absorb;
    /// `None` (the default) when it has none.
    pub absorbs: Option<usize>, // was: absorbed only by WholeExpression candidates
    /// The aggregate node this target computes.
    pub target: Rc<OperatorNode>, // unchanged
    /// Index into the aggregate's measures.
    pub measure: usize, // unchanged
    /// The per-computation candidates offered for this target (Pass 1).
    pub choices: ComputationChoices, // unchanged
    /// The measure's own filter; only rows that pass it update the state; `None` keeps all.
    pub filter: Option<Predicate>, // unchanged
    /// The aggregate's own `HAVING`, then one entry per consumer `Filter` directly above it.
    pub having: Vec<HavingEntry>, // unchanged
    /// The metric whose raw samples (or plain sums of them) are every value the aggregate
    /// reads, for heap-sketch weight checks; `None` otherwise.
    pub samples_of: Option<String>, // unchanged
}

crates/asap-aware-mapping/src/asap_aware_cse.rs

/// One coarse build the rollup rule can answer from a finer one. Produced by
/// `rollup_sites` for a composed candidate.
pub struct RollupSite {                 // new
    /// The coarse exact accumulator build (`SummaryAgg`); every reader of it is a
    /// `FinalizeExactAccumulator`, which the rollup replaces.
    pub coarse: Rc<OperatorNode>,
    /// The finer build's `FinalizeExactAccumulator`, already in the candidate,
    /// whose values the rollup reads.
    pub finer: Rc<OperatorNode>,
    /// The finalized exact accumulator over `finer`'s values, grouped by the
    /// coarse keys, with the coarse value's column names.
    pub rollup: Rc<OperatorNode>,
}

/// A Pass 2 candidate of the rollup rule, produced by `rollup_variants`.
pub struct RollupVariant {              // new
    /// The sites rolled up (indices into `rollup_sites`), ascending.
    pub sites: Vec<usize>,
    /// The roots with those sites rolled up and every guarantee composed again;
    /// `Err` when the rebuilt DAG does not derive or compose (an unsatisfiable
    /// candidate, not bad input), which drops only this candidate.
    pub roots: Result<Vec<QueryRoot>, PerComputationCandidateError>,
}

/// Exact top-k computations of one input and partition with different `k`.
/// Produced by `topk_reuse_sites`.
pub struct TopKReuseSite {              // new
    /// The plain `Aggregate` top-k nodes of one input and partition, in
    /// children-first order; at least two different `k`.
    pub members: Vec<Rc<OperatorNode>>,
    /// `Limit(k_max)` over a descending `Sort` of their input, partitioned as they
    /// are: what every member reads in a reuse candidate.
    pub ranking: Rc<OperatorNode>,
}

/// A Pass 2 candidate of the exact top-k reuse rule, produced by `topk_reuse_variants`.
pub struct TopKReuseVariant {           // new
    /// The sites reused (indices into `topk_reuse_sites`), ascending.
    pub sites: Vec<usize>,
    /// The roots with every member of those sites reading its ranking and every
    /// guarantee composed again; `Err` when the rebuilt DAG does not derive or
    /// compose (an unsatisfiable candidate), which drops only this candidate.
    pub roots: Result<Vec<QueryRoot>, PerComputationCandidateError>,
}

/// The most rollup and top-k reuse sites one candidate may have together (every
/// subset of one rule is combined with every subset of the other); more is
/// `PerComputationCandidateError::Unsupported` (an unsatisfiable request).
pub const MAX_REUSE_SITES: usize = 20;  // new

/// The exact accumulator that merges finalized values of `family` across disjoint
/// groups, by the family-capability table (`Sum` for sums and counts, `Min`, `Max`);
/// `None` for counter states, IRate and every sketch family.
pub fn rollup_kind(family: &FieldDataType) -> Option<ExactKind>;            // new
/// The rollup sites of a composed candidate, each with its coarsest finer build over
/// the same raw scan or time-range rows; a site needs Float64 finer values (the only
/// update type exact accumulators execute), so integer counts give none.
pub fn rollup_sites(roots: &[QueryRoot]) -> Vec<RollupSite>;               // new
/// One rollup candidate per non-empty subset of `sites`, in binary-counter order;
/// `Err(Unsupported)` past `MAX_REUSE_SITES` sites.
pub fn rollup_variants(roots: &[QueryRoot], sites: &[RollupSite])
    -> Result<Vec<RollupVariant>, PerComputationCandidateError>;            // new
/// The exact top-k reuse sites of a composed candidate.
pub fn topk_reuse_sites(roots: &[QueryRoot]) -> Vec<TopKReuseSite>;        // new
/// One top-k reuse candidate per non-empty subset of `sites`; `Err(Unsupported)`
/// past `MAX_REUSE_SITES` sites. Stage 1 caps rollup and top-k sites together.
pub fn topk_reuse_variants(roots: &[QueryRoot], sites: &[TopKReuseSite])
    -> Result<Vec<TopKReuseVariant>, PerComputationCandidateError>;         // new

crates/asap-aware-mapping/src/pass/stage1_to_stage3.rs

/// A Stage 1 candidate Pass 2 added, as reported in `Stage1ToStage3Run::stage1_pass2`.
/// The type stays; `reused` is new.
pub struct Pass2Candidate {
    /// The candidate's id, `L<n>`.
    pub id: String, // unchanged
    /// The id of the Pass 1 candidate it derives from.
    pub from_logical: String, // unchanged
    /// Summary build sets made one build each (summary-capability rule).
    pub summaries: Vec<String>, // unchanged
    /// Builds built as tumbling panes (window-composition rule).
    pub panes: Vec<String>, // unchanged
    /// The computations answered from another (rollup and exact top-k reuse rules),
    /// e.g. "Sum acc for Q1 rolled up from Q2"; empty (the default) for other rules.
    pub reused: Vec<String>, // new
    /// Identical node groups made one node (identical-expression rule).
    pub shared: Vec<String>, // unchanged
}

/// What a Pass 2 candidate shares of a Pass 1 candidate (`Stage1Candidate::pass2_of`).
/// The type stays; `reused` is new.
pub struct Pass2Of {
    /// The Pass 1 candidate's position in the enumeration.
    pub from: usize, // unchanged
    /// As `Pass2Candidate::summaries`.
    pub summaries: Vec<String>, // unchanged
    /// As `Pass2Candidate::panes`.
    pub panes: Vec<String>, // unchanged
    /// As `Pass2Candidate::reused`.
    pub reused: Vec<String>, // new
    /// As `Pass2Candidate::shared`.
    pub shared: Vec<String>, // unchanged
}

Usage, end to end (checked with stage_pipeline at cde2ca3d, built-in models, no --epsilon, so every query is exact):

# What the user writes: four panels over one metric `x`.
# Q1 and Q2 sum the same rows at two groupings (Q1's keys ⊂ Q2's): a rollup site.
# Q3 and Q4 rank the same input by `job` with k = 5 and k = 10: a top-k reuse site.
cargo run -p asap-devtools --bin stage_pipeline -- \
    --promql "sum by (job) (x)" \
    --promql "sum by (job, instance) (x)" \
    --promql "topk by (job) (5, sum by (job, instance) (x))" \
    --promql "topk by (job) (10, sum by (job, instance) (x))" \
    --out run.json

What Stage 1 does with it, for each Pass 1 candidate (here only the all-exact one), after its window-composition candidates:

// `roots` holds that Pass 1 candidate's composed query roots, one per query (Q1–Q4).
// `topk_sites` holds the exact top-k groups that can share one ranking: here one,
// Q3 and Q4 (same input and partition), with `ranking` = Limit(10, Sort).
let topk_sites = topk_reuse_sites(&roots);
// `topk` holds one candidate per non-empty subset of `topk_sites`: here one, listed as
// "top-k for Q3, Q4 from one top-10" (L6), where both top-k read that ranking.
let topk = topk_reuse_variants(&roots, &topk_sites)?;
// `sites` holds each coarse exact accumulator paired with a finer one over the same rows:
// here one, Q1's Sum by (job) answered from Q2's finalized Sum by (job, instance).
let sites = rollup_sites(&roots);
// `rollups` holds one candidate per non-empty subset of `sites`: here one, "Sum acc for Q1
// rolled up from Q2" (L11), still exact. Each rollup candidate is then followed by its
// own top-k reuse candidates (L16: both rules), and every candidate by its
// identical-expression candidates (shared Scan, TimeRange, ...).
let rollups = rollup_variants(&roots, &sites)?;

What the user gets in run.json: 20 Stage 1 candidates (L1–L5 base, L6–L10 top-k reuse, L11–L15 rollup, L16–L20 both), all implemented by Stage 2. Stage 3 selects P1 (all exact, separate computations) at 0.000119 cost_per_second. P6 (top-k reuse) costs 0.000119 and P11 (rollup) 0.000126, because Stage 2 leaves every node NotMaterialized and a NotMaterialized node is charged once per query that reads it (see the pricing limitation below).

The nested sum merge is a Pass 1 candidate, so it appears for a single query: stage_pipeline --promql "sum by (job) (sum_over_time(x[5m]))" lists L1 "Q1 exact (Sum acc, Sum acc)" and L2 "Q1 exact (nested Sum acc)" (ComputationCandidate::NestedSumMerge), and Stage 3 selects P2 at 0.00065003 against P1's 0.00065058.

Before this PR

• sum by (job) (sum_over_time(m[5m])) was always two accumulators: a per-series Sum, finalized, then a Sum by job over those values.
• sum by (job) beside sum by (job, instance), or topk(5, X) beside topk(10, X), was always computed separately from the raw rows.

After this PR

stage_pipeline --promql "sum by (job) (sum_over_time(x[5m]))" (default statistics):

L1 Q1 exact (Sum acc, Sum acc)   SummaryAgg(PerEntity) → Finalize → SummaryAgg(by job) → Finalize   0.00065058/s
L2 Q1 exact (nested Sum acc)     SummaryAgg(by job) over the 5m range → Finalize                    0.00065003/s  ← selected

For sum by (job) (x), sum by (job, instance) (x) and two topk by (job) with k = 5 and k = 10 over the finer sums, Stage 1 now also lists candidates such as "Sum acc for Q1 rolled up from Q2" and "top-k for Q3, Q4 from one top-10". All 20 candidates are implemented by Stage 2.

Limitation (pricing). Stage 2 currently leaves every node NotMaterialized, and Stage 3 charges a NotMaterialized node once per query that reads it. Until G8 (Stage 2 materialization) lands, rollup and top-k reuse are therefore priced but never cheaper. The integration test materializes the highest shared computed nodes as a stand-in for G8, and with that the reuse candidates are selected.

Verification

• asap-aware-mapping unit tests:
  • nested_sum_merges_into_one_exact_accumulator, nested_sum_merge_needs_a_sum_of_sums, nested_sum_merge_never_groups_by_the_dropped_metric_name and nested_sum_merges_do_not_stack (the last two failed before the review fix);
  • rollup_from_a_finer_exact_sum, extrema_roll_up_and_integer_counts_do_not and rollup_needs_a_finer_computation_of_the_same_rows;
  • rollup_is_refused_for_families_without_an_exact_merge_across_groups: Rate, Increase, IRate, CMS+heap and KLL, plus a grouped KLL pair that gives no site;
  • exact_topk_reuses_the_larger_ranking and topk_reuse_needs_exact_topk_of_one_input_and_partition (including heap sketches, which are not reused);
  • reuse_sites_past_20_together_are_an_error (hung enumerating 2^21 candidates before the cap).
• crates/integration-tests/tests/reuse_rules.rs, executing each plan on samples of two jobs of two series:
  • nested_sum_by_the_metric_name_answers_as_the_nested_sums failed before the fix (merged __name__="m", nested ""); every candidate of sum by (__name__) and sum by (job) over sum_over_time now answers alike;
  • the nested-sum merge is selected and answers as the nested sums;
  • rollups of sums (over merged nested sums), minima and maxima, and exact top-k reuse, answer query by query as the candidate they reuse from. This execution check found that integer counts could not be summed again by the executor, hence the Float64 gate;
  • rollup and exact top-k reuse are selected under the docs: propose workload-wide planning, summary sharing, and materialization #509 shared data workload once shared nodes are computed once, and are not selected without that.
• Local runs on the head b9957817e6f1f4b3f3b563e8d3a3d5c7d1d5dcb7 (rebased on f0b07992):
  • cargo nextest run -p asap-aware-mapping -p asap-integration-tests -p asap-devtools: 880 passed;
  • cargo fmt --all -- --check: clean;
  • cargo clippy --workspace --all-targets --all-features --locked -- -D warnings: clean.

Limitations and follow-ups

• Sketch rollup (KLL) is deferred to Sketch rollup: a regrouping SummaryMerge so mergeable sketches can roll up to coarser keys #719. The table marks KLL mergeable, but rolling up sketch states needs a merge that regroups states by coarser keys, and the IR has no such operator (SummaryMerge merges states of one grouping).
• Rollup fires only when both builds read raw scan or time-range rows (equal nodes, or PromQL scans of one metric without predicates under the same time ranges and shifts). Builds over finalized inner values, such as count by (job) (count_over_time(m[5m])) beside count by (job, series) (...), are not matched. Integer-valued finer values (counts) do not roll up.
• Exact top-k reuse needs a plain exact top-k Aggregate. A top-k over sum ... sum_over_time lowers to Sort + Limit, so the rule does not apply there.
• Reuse candidates are not combined with summary-capability or window-composition candidates. A nested-sum merge is not offered beneath a whole-expression top-k, which absorbs it. Rollup across scans with label predicates is not attempted.
• Legacy code kept. The legacy search still calls rollup.rs (RollupStrategy), topk_reuse.rs (TopKLimitReuseStrategy) and replacement.rs's query_time_nested_sum / contains_aggregate. Nothing became unused, so nothing is deleted here. These go with the final legacy-core deletion.

[zzylol]

/// One Pass 1 candidate for a computation (one measure of an aggregate).
pub enum ComputationCandidate {
Exact,
ExactAccumulator { kind: ExactKind, params: ExactParams },
Summary(SketchAlgorithm),
AvgAsSumCount,
WholeExpression(SketchAlgorithm),
/// A grouped sum over a sum beneath it as one exact Sum accumulator over
/// the inner sum's input, grouped by the outer keys; absorbs the inner target.
NestedSumMerge, // new
}
How did you do the category of this? where do they come from? It is confusing

[zzylol]

Where the six variants come from: they were not derived from one scheme. They were added one rule at a time. The design doc's Pass 1 table has two kinds, exact or a summary. The code then grew four more:

Variant	Added by	What it actually is
Exact	Pass 1 table	run the original sub-DAG
Summary(alg)	Pass 1 table	a summary state for this target
ExactAccumulator { kind, params }	exact mergeable states (sum/count/min/max…)	an exact state, so it can be shared, windowed and materialized like a summary
AvgAsSumCount	avg rewrite	the same thing as ExactAccumulator, but with two states (sum and count)
WholeExpression(alg)	whole-expression top-k	a summary (CMS/CountSketch + heap) for a fused target: the top-k plus the per-item sum/count beneath it
NestedSumMerge (#716)	nested-sum rule	an exact Sum accumulator for a fused target: the outer sum plus the inner sum beneath it

So the enum mixes two independent questions:

1. Which target is computed: this one computation, or this one fused with the one beneath it (absorbs_inner()).
2. How it is computed: exact by re-running the sub-DAG, exact state(s), or a summary.

WholeExpression is really "summary × fused target", and NestedSumMerge is really "exact accumulator × fused target". That is why the list reads as inconsistent.

A cleaner shape would separate the two questions:

/// How one Pass 1 target is computed.
pub enum ComputationCandidate {
    Exact,                              // re-run the original sub-DAG
    ExactAccumulator(Vec<ExactState>),  // one or more exact mergeable states; avg = [Sum, Count]
    Summary(SketchAlgorithm),           // an unsized summary family
}

A fused target (top-k over per-item sum; sum over sum) then becomes its own ComputationTarget that absorbs the inner one. A whole-expression top-k is Summary(CmsHeap) on that target, and a nested-sum merge is ExactAccumulator([Sum]) on it. This matches the design doc's "exact or a summary" table. I'll confirm the direction before changing it.