diff --git a/crates/asap-aware-mapping/src/explanation.rs b/crates/asap-aware-mapping/src/explanation.rs index b63b9119..84a5212e 100644 --- a/crates/asap-aware-mapping/src/explanation.rs +++ b/crates/asap-aware-mapping/src/explanation.rs @@ -161,7 +161,7 @@ //! //! | Catalog entry | Status | Where a future `ExplanationKind` would come from | //! |---|---|---| -//! | Semantic-equivalent rewriting (e.g. `avg` → `sum`/`count`) | No `ReplacementStrategy` implementation in this codebase yet (tracked separately, issue #253) | Once wired into `default_strategies()`: any `Replacement::Rewrite` candidate that strategy proposes | +//! | Semantic-equivalent rewriting (e.g. `avg` → `sum`/`count`) | [`AvgToSumOverCountStrategy`](crate::rewrite::AvgToSumOverCountStrategy) exists and is wired into `default_strategies()` (issue #253) — but still no `ExplanationKind` of its own below, since this table is about *direct* findings for a catalog entry, and this strategy's whole point is indirect: its `Replacement::Rewrite` candidate exposes `sum`/`count` as independently bindable discovered targets, which can then earn `CommonSubexpressionReuse` findings when the workload actually reuses them | A dedicated variant would need `findings_from_plan_space` to recognize a `LogicalRewrite`-provenance candidate as a finding in its own right, not just rely on what it exposes downstream | //! | Roll-ups (fine-to-coarse group-by reuse) | [`RollupStrategy`](crate::rollup::RollupStrategy), derived from workload siblings after CSE/target discovery (issue #254) | Any `Replacement::Rewrite` candidate that rolls a coarse aggregate up from a compatible finer aggregate | //! | Wavelets/OMP | Params type exists (`WaveletKind`/`WaveletParams`), reachable only via a deployment `CostModel::realize_extension` (no core `AggIntent` dispatch picks it) | A `ReplacementStrategy` that inspects a deployment's own `CostModel`, once some intent shape actually maps to `Implementation::Wavelet` | //! | Sampling | Same story as Wavelets: `SamplingKind`/`SamplingParams` exist, unreachable from core dispatch | Same hook as Wavelets, for `Implementation::Sample` | diff --git a/crates/asap-aware-mapping/src/replacement.rs b/crates/asap-aware-mapping/src/replacement.rs index fbf4d55a..257455c7 100644 --- a/crates/asap-aware-mapping/src/replacement.rs +++ b/crates/asap-aware-mapping/src/replacement.rs @@ -291,8 +291,12 @@ //! once every ancestor's own selected candidate is accounted for — and, for //! every [`SharedSubtreeStrategy`]-shaped group, re-decides //! [`CostModel::cse_share_decision`] against *that* corrected count instead -//! of the group's raw structural one (see [`multiplier`]'s doc for the -//! exact recurrence): a group that chooses `Share` collapses its own +//! of the group's raw structural one. When that group also contains a +//! non-CSE alternative such as a semantic rewrite, the chosen CSE candidate +//! and the cheapest non-CSE candidate additionally compete through +//! [`CostModel::estimate_cost`]; the CSE pair is no longer allowed to hide an +//! otherwise valid logical alternative. See [`multiplier`]'s doc for the +//! exact recurrence: a group that chooses `Share` collapses its own //! multiplicity to exactly `1` for everything beneath it (one shared //! execution backs every use of it); a group that chooses //! `RecomputeIndependently` — or has no Share/Recompute decision of its own @@ -2024,8 +2028,34 @@ impl PlanSpace { let chosen = if effective >= 2 && cse_candidate_pair(group).is_some() { match decide_with_effective_count(group, effective, cost_model) { Some(decision) => { - chosen_share.insert(*ptr, decision); - pick_shared_subtree_candidate(group, decision) + let cse = pick_shared_subtree_candidate(group, decision); + let effective_target = + TargetSubDAG::with_consumer_count(&group.target, effective); + let logical = group + .candidates + .iter() + .filter(|candidate| !is_cse_candidate(candidate)) + .min_by(|a, b| { + cost_model + .estimate_cost(a, &effective_target) + .total_cmp(&cost_model.estimate_cost(b, &effective_target)) + }); + match (cse, logical) { + (Some(cse), Some(logical)) + if cost_model + .estimate_cost(logical, &effective_target) + .total_cmp(&cost_model.estimate_cost(cse, &effective_target)) + .is_lt() => + { + Some(logical) + } + (cse, _) => { + if cse.is_some() { + chosen_share.insert(*ptr, decision); + } + cse + } + } } // `realize_child` couldn't produce even a logical fallback — // not expected in practice for a target that's already @@ -2393,11 +2423,24 @@ fn topological_order(order: &[*const QueryExpr], graph: &ReferenceGraph) -> Vec< /// explanation-specific list to stay in sync with. Use /// [`default_strategies_with`] to plug in a deployment-specific /// [`CostModel`] instead. +/// +/// [`AvgToSumOverCountStrategy`](crate::rewrite::AvgToSumOverCountStrategy) is +/// included here (issue #253) even though it's a +/// [`Replacement::Rewrite`]-only strategy with no [`CostModel`] of its own to +/// plug in — it's context-free (`matches`/`replacements` need nothing beyond +/// the target itself) exactly like [`SharedSubtreeStrategy`], so it belongs +/// in this list rather than being derived per-workload the way +/// [`RollupStrategy`] is. Rewriting `avg` into `sum`/`count` upfront is what +/// lets [`SketchAlgorithmStrategy`] and [`SharedSubtreeStrategy`] see a +/// mergeable accumulator to sketch or share at all — see that module's own +/// doc comment for why a bare `avg` node otherwise never becomes a +/// [`ReplacementStrategy`] target for anything. pub fn default_strategies() -> Vec> { vec![ Box::new(SketchAlgorithmStrategy::default_cost_model()), Box::new(HydraGroupingStrategy::default_cost_model()), Box::new(SharedSubtreeStrategy), + Box::new(crate::rewrite::AvgToSumOverCountStrategy), ] } @@ -2411,6 +2454,7 @@ pub fn default_strategies_with<'a>( Box::new(SketchAlgorithmStrategy::new(cost_model)), Box::new(HydraGroupingStrategy::new(cost_model)), Box::new(SharedSubtreeStrategy), + Box::new(crate::rewrite::AvgToSumOverCountStrategy), ] } @@ -4593,6 +4637,54 @@ mod tests { ); } + #[test] + fn global_selection_compares_a_logical_rewrite_with_the_cse_choice() { + struct PreferLogicalRewrite; + + impl CostModel for PreferLogicalRewrite { + fn rank_candidates( + &self, + _intent: &AggIntent, + candidates: &[SketchAlgorithm], + ) -> Vec { + candidates.to_vec() + } + + fn estimate_cost( + &self, + candidate: &ReplacementSubDAG, + _target: &TargetSubDAG<'_>, + ) -> f64 { + match candidate.provenance { + ReplacementProvenance::LogicalRewrite => 0.0, + _ => 100.0, + } + } + } + + let a = Rc::new(agg( + vec![2], + AggIntent::Avg { col: None }, + metric_scan(&["job"]), + )); + let b = Rc::new(agg( + vec![2], + AggIntent::Avg { col: None }, + metric_scan(&["job"]), + )); + let space = search_workload(vec![("a", a), ("b", b)]); + let root = &space.roots[0].1; + let selected = space.global_selection(&PreferLogicalRewrite); + + assert_eq!( + selected + .for_target(root) + .and_then(|group| group.chosen) + .map(|candidate| candidate.provenance), + Some(ReplacementProvenance::LogicalRewrite) + ); + } + #[test] fn topological_order_puts_a_later_discovered_parent_before_its_child() { // Mirrors nested_shared_subtree_below_an_unshared_parent_is_still_discovered's diff --git a/crates/asap-aware-mapping/src/rewrite.rs b/crates/asap-aware-mapping/src/rewrite.rs index a728dcc9..92e39d6e 100644 --- a/crates/asap-aware-mapping/src/rewrite.rs +++ b/crates/asap-aware-mapping/src/rewrite.rs @@ -18,12 +18,12 @@ //! `Sum` and `Count` are both ordinary mergeable accumulators //! (`agg_is_mergeable`) — exactly the shape [`SharedSubtreeStrategy`] and a //! future sketch-family search already know how to reuse across a -//! workload. Rewriting `Aggregate{ measures: [Avg{col}], .. }` into -//! `Aggregate{ measures: [Sum{col}, Count], .. }` re-divided back into the -//! original `avg` column by a wrapping `Project` computes exactly the same -//! result, but *reshapes* it into pieces other strategies can now share or -//! sketch. This module only performs that reshaping — see "Non-goals" below -//! for why it does not also decide whether the reshaping is worth it. +//! workload. Rewriting `Aggregate{ measures: [Avg{col}], .. }` into two +//! independent single-measure `Sum` and `Count` aggregates, divided with a +//! `BinaryOp`, computes the same result but *reshapes* it into targets other +//! strategies can bind and share independently. This module only performs +//! that reshaping — see "Non-goals" below for why it does not also decide +//! whether the reshaping is worth it. //! //! ## Scope: `by(...)` grouping only (issue #253's own scope note) //! @@ -68,7 +68,7 @@ use std::rc::Rc; use asap_types::pre_asap::agg_intent::AggIntent; use asap_types::pre_asap::expr_ir::ArithmeticOpKind; -use asap_types::pre_asap::query_expr::{ProjectItem, QueryExpr, Reduction}; +use asap_types::pre_asap::query_expr::{BinaryOpKind, ProjectItem, QueryExpr, Reduction}; use asap_types::pre_asap::schema::{ColumnId, DataType}; use asap_types::types::AccuracyTarget; @@ -113,8 +113,10 @@ fn avg_rewrite_target(node: &QueryExpr) -> Option<(usize, Option)> { Some((by.keys().len(), *col)) } -/// Build the rewritten `Aggregate{ [Sum, Count] } |> Project{ sum/count }` -/// tree for `root`, or `None` if `root` isn't [`avg_rewrite_target`]'s shape. +/// Build the rewritten `Project{ cast(sum) } / Aggregate{ Count }` tree for +/// `root`, or `None` if `root` isn't [`avg_rewrite_target`]'s shape. `Sum` and +/// `Count` deliberately live in separate, single-measure aggregates so the +/// replacement fixpoint discovers each as an independently bindable target. /// /// The `Project`'s leading `by.len()` items are bare `Column(i)` /// pass-throughs of the grouping keys — identical in name/type to the @@ -158,21 +160,24 @@ fn build_rewrite(root: &Rc) -> Option> { .cloned() .unwrap_or_else(|| "avg".to_string()); - let sum_count_agg = QueryExpr::Aggregate { + let sum_agg = Rc::new(QueryExpr::Aggregate { reduction: reduction.clone(), - measures: vec![ - AggIntent::Sum { col }, - AggIntent::Count { - accuracy: AccuracyTarget::Exact, - }, - ], + measures: vec![AggIntent::Sum { col }], output_names: Vec::new(), having: None, child: Rc::clone(child), - }; + }); + let count_agg = Rc::new(QueryExpr::Aggregate { + reduction: reduction.clone(), + measures: vec![AggIntent::Count { + accuracy: AccuracyTarget::Exact, + }], + output_names: Vec::new(), + having: None, + child: Rc::clone(child), + }); let sum_idx = group_count; - let count_idx = group_count + 1; let mut cols: Vec = (0..group_count) .map(|i| ProjectItem { alias: None, @@ -181,30 +186,30 @@ fn build_rewrite(root: &Rc) -> Option> { .collect(); cols.push(ProjectItem { alias: Some(avg_name), - expr: QueryExpr::Arithmetic { - op: ArithmeticOpKind::Div, - left: Rc::new(QueryExpr::Cast { - expr: Rc::new(QueryExpr::Column(sum_idx)), - to: DataType::Float64, - try_cast: false, - }), - right: Rc::new(QueryExpr::Column(count_idx)), + expr: QueryExpr::Cast { + expr: Rc::new(QueryExpr::Column(sum_idx)), + to: DataType::Float64, + try_cast: false, }, }); - Some(Rc::new(QueryExpr::Project { + let float_sum = Rc::new(QueryExpr::Project { cols, qualifier: None, - child: Rc::new(sum_count_agg), + child: sum_agg, + }); + Some(Rc::new(QueryExpr::BinaryOp { + op: BinaryOpKind::Arithmetic(ArithmeticOpKind::Div), + lhs: float_sum, + rhs: count_agg, + vector_match: None, })) } /// Rewrites `Aggregate{ measures: [Avg{col}], .. }` into the semantically -/// equivalent `Aggregate{ measures: [Sum{col}, Count], .. } |> Project{ -/// sum/count re-divided under the original `avg` column name }` — see the -/// module docs for why this reshaping lets `SharedSubtreeStrategy` (and a -/// future sketch-family search) reach further than a bare `avg` node ever -/// could. +/// equivalent pair of single-measure `Sum` and `Count` aggregates divided by +/// a `BinaryOp` — see the module docs for why keeping the accumulators in +/// separate relational nodes lets later strategies reach them independently. /// /// A unit struct: unlike [`SketchAlgorithmStrategy`], this strategy doesn't /// bind anything (its one [`Replacement`] is always [`Replacement::Rewrite`], @@ -377,18 +382,8 @@ mod tests { // ── replacements / schema round-trip ───────────────────────────────── - /// The rewritten tree must be a `Project` over an `Aggregate{[Sum, - /// Count]}`, and — the correctness bug this test exists to catch — its - /// `output_schema()` must equal the original `Avg` aggregate's, exactly. - /// An ungrouped (`by: []`) aggregate is the shape used here specifically - /// because it's the one shape where a `Project`-wrapped rewrite's - /// `unique_keys`/`closed` bookkeeping (always reset by `Project`, see - /// `QueryExpr::output_schema`'s `Project` arm) coincides exactly with - /// what an ungrouped `Aggregate` already reports for those same fields - /// (`unique_keys: []` whenever `by` is empty, `closed: true` either way) - /// — see `does_not_match_a_without_grouped_avg_aggregate`'s sibling test - /// below for the (column-level, not whole-`Schema`) invariant a - /// non-trivial `by(...)` grouping still gets. + /// The rewritten tree must keep `Sum` and `Count` in separate aggregates, + /// and its `output_schema()` must equal the original `Avg` aggregate's. #[test] fn avg_rewrites_and_schema_matches_exactly_when_ungrouped() { let original = avg_agg(vec![], None, metric_scan(&[])); @@ -403,20 +398,22 @@ mod tests { Replacement::Rewrite(rc) => rc, other => panic!("expected a Rewrite replacement, got {other:?}"), }; - assert!( - matches!(rewritten.as_ref(), QueryExpr::Project { .. }), - "expected a Project wrapping the rewritten Aggregate, got {rewritten:?}" - ); - if let QueryExpr::Project { child, .. } = rewritten.as_ref() { - assert!(matches!( - child.as_ref(), - QueryExpr::Aggregate { measures, .. } - if matches!( - measures.as_slice(), - [AggIntent::Sum { col: None }, AggIntent::Count { accuracy: AccuracyTarget::Exact }] - ) - )); - } + let QueryExpr::BinaryOp { lhs, rhs, .. } = rewritten.as_ref() else { + panic!("expected sum/count BinaryOp, got {rewritten:?}"); + }; + let QueryExpr::Project { child: sum, .. } = lhs.as_ref() else { + panic!("expected cast Project above Sum, got {lhs:?}"); + }; + assert!(matches!( + sum.as_ref(), + QueryExpr::Aggregate { measures, .. } + if matches!(measures.as_slice(), [AggIntent::Sum { col: None }]) + )); + assert!(matches!( + rhs.as_ref(), + QueryExpr::Aggregate { measures, .. } + if matches!(measures.as_slice(), [AggIntent::Count { accuracy: AccuracyTarget::Exact }]) + )); let original_schema = original.output_schema().unwrap(); let rewritten_schema = rewritten.output_schema().unwrap(); @@ -449,16 +446,10 @@ mod tests { assert_eq!(rewritten_schema.columns[0].name, "avg_latency"); } - /// For a non-trivial `by(...)` grouping, `Project`'s own schema - /// derivation always resets `unique_keys` to `[]` (see - /// `QueryExpr::output_schema`'s `Project` arm) — a pre-existing - /// property of every `Project`-wrapped rewrite in this IR, not specific - /// to this strategy — so a grouped `Avg` aggregate's `unique_keys = - /// [[0]]` doesn't survive verbatim. What must still hold, and does, is - /// the column-level shape (names, types, nullability, and grouping - /// column count) downstream schema derivation actually reads. + /// A grouped rewrite must preserve the aggregate's grouping-key metadata; + /// CSE and roll-up legality both depend on it. #[test] - fn grouped_avg_rewrite_matches_columns_though_not_the_whole_schema_struct() { + fn grouped_avg_rewrite_preserves_the_whole_schema() { let original = avg_agg(vec![2], None, metric_scan(&["job"])); let original_schema = original.output_schema().unwrap(); let original_rc = Rc::new(original); @@ -471,14 +462,51 @@ mod tests { }; let rewritten_schema = rewritten.output_schema().unwrap(); - assert_eq!(rewritten_schema.columns, original_schema.columns); - assert_eq!(rewritten_schema.closed, original_schema.closed); - assert!( - original_schema.unique_keys == vec![vec![0]] && rewritten_schema.unique_keys.is_empty(), - "documents the known Project-reset gap this test exists to pin down; if this \ - assertion starts failing, the gap has closed and this test (and its comment) \ - should be simplified to a plain schema equality assertion instead" - ); + assert_eq!(rewritten_schema, original_schema); + } + + #[test] + fn default_search_discovers_bindable_sum_and_count_targets() { + let root = Rc::new(avg_agg(vec![2], None, metric_scan(&["job"]))); + let space = crate::replacement::search_workload(vec![("avg", Rc::clone(&root))]); + + let avg_group = space.group_for(&space.roots[0].1).expect("avg group"); + assert!(avg_group.candidates.iter().any(|candidate| { + candidate.provenance == crate::replacement::ReplacementProvenance::LogicalRewrite + })); + + let mut found_sum = false; + let mut found_count = false; + for group in space.groups() { + let QueryExpr::Aggregate { measures, .. } = group.target.as_ref() else { + continue; + }; + let expected = matches!(measures.as_slice(), [AggIntent::Sum { .. }]) + || matches!( + measures.as_slice(), + [AggIntent::Count { + accuracy: AccuracyTarget::Exact + }] + ); + if !expected { + continue; + } + assert!( + group + .candidates + .iter() + .any(|candidate| matches!(candidate.replacement, Replacement::Summary(_))), + "rewritten accumulator must be independently bindable: {measures:?}" + ); + found_sum |= matches!(measures.as_slice(), [AggIntent::Sum { .. }]); + found_count |= matches!( + measures.as_slice(), + [AggIntent::Count { + accuracy: AccuracyTarget::Exact + }] + ); + } + assert!(found_sum && found_count); } #[test] @@ -519,16 +547,18 @@ mod tests { DataType::Float64 ); - let QueryExpr::Project { cols, .. } = rewritten.as_ref() else { - unreachable!(); + let QueryExpr::BinaryOp { lhs, .. } = rewritten.as_ref() else { + panic!("expected sum/count BinaryOp"); + }; + let QueryExpr::Project { cols, .. } = lhs.as_ref() else { + panic!("expected cast Project above Sum"); }; assert!(matches!( &cols.last().unwrap().expr, - QueryExpr::Arithmetic { - op: ArithmeticOpKind::Div, - left, + QueryExpr::Cast { + to: DataType::Float64, .. - } if matches!(left.as_ref(), QueryExpr::Cast { to: DataType::Float64, .. }) + } )); } diff --git a/crates/types/src/pre_asap/query_expr.rs b/crates/types/src/pre_asap/query_expr.rs index dad4ddc4..9f346f71 100644 --- a/crates/types/src/pre_asap/query_expr.rs +++ b/crates/types/src/pre_asap/query_expr.rs @@ -1073,9 +1073,11 @@ impl QueryExpr { // π — one output column per projection item. Each item's type is // inferred from its expression against the child schema; the name - // is the explicit alias or a derived default. Projection may drop - // the grouping/time columns, so unique_keys reset and time_index - // is re-found by name. + // is the explicit alias or a derived default. A child unique key + // survives exactly when every one of its columns is passed through + // as a bare `Column` item (possibly reordered or aliased). Derived + // expressions cannot carry key identity. `time_index` is re-found + // by name. QueryExpr::Project { cols, qualifier, child } => { let in_schema = child.output_schema()?; let columns: Vec = cols @@ -1098,10 +1100,23 @@ impl QueryExpr { }) .collect(); let time_index = columns.iter().position(|c| c.name == "ts"); + let unique_keys = in_schema + .unique_keys + .iter() + .filter_map(|key| { + key.iter() + .map(|input_col| { + cols.iter().position(|item| { + matches!(&item.expr, QueryExpr::Column(col) if col == input_col) + }) + }) + .collect::>>() + }) + .collect(); Ok(Schema { columns, time_index, - unique_keys: Vec::new(), + unique_keys, // Projection enumerates exactly its items → closed. closed: true, }) @@ -1586,6 +1601,68 @@ mod tests { } } + #[test] + fn project_preserves_unique_keys_that_are_passed_through() { + let input = Rc::new(scan( + vec![ + col("tenant", DataType::Utf8, false), + col("region", DataType::Utf8, false), + col("value", DataType::Int64, false), + ], + None, + vec![vec![0, 1]], + )); + let projected = QueryExpr::Project { + cols: vec![ + ProjectItem { + alias: Some("r".into()), + expr: QueryExpr::Column(1), + }, + ProjectItem { + alias: Some("t".into()), + expr: QueryExpr::Column(0), + }, + ProjectItem { + alias: None, + expr: QueryExpr::Arithmetic { + op: ArithmeticOpKind::Add, + left: Rc::new(QueryExpr::Column(2)), + right: Rc::new(QueryExpr::Literal(ScalarValue::Int64(1))), + }, + }, + ], + qualifier: None, + child: input, + }; + + assert_eq!( + projected.output_schema().unwrap().unique_keys, + vec![vec![1, 0]] + ); + } + + #[test] + fn project_drops_a_unique_key_when_a_key_column_is_omitted() { + let input = Rc::new(scan( + vec![ + col("tenant", DataType::Utf8, false), + col("region", DataType::Utf8, false), + ], + None, + vec![vec![0, 1]], + )); + let projected = QueryExpr::Project { + cols: vec![ProjectItem { + alias: None, + expr: QueryExpr::Column(0), + }], + qualifier: None, + child: input, + }; + + assert!(projected.output_schema().unwrap().unique_keys.is_empty()); + } + #[test] fn legacy_window_json_without_frame_deserializes_as_unspecified() { let window = QueryExpr::SQLWindowFunc {