From e59b1beebc384f9298c8b4ad72d8c434a7c8224a Mon Sep 17 00:00:00 2001 From: zz_y Date: Thu, 10 Sep 2026 07:26:52 -0600 Subject: [PATCH] Preserve nested SQL sums at read time --- crates/asap-aware-mapping/src/replacement.rs | 90 ++++++++++++++----- .../tests/sql_to_post_asap.rs | 59 ++++++++++++ 2 files changed, 128 insertions(+), 21 deletions(-) diff --git a/crates/asap-aware-mapping/src/replacement.rs b/crates/asap-aware-mapping/src/replacement.rs index 5ba17f88..712b637c 100644 --- a/crates/asap-aware-mapping/src/replacement.rs +++ b/crates/asap-aware-mapping/src/replacement.rs @@ -350,11 +350,12 @@ use std::collections::{HashMap, HashSet, VecDeque}; use asap_types::post_asap::{ validate_execution_data_states_at, CandidateCompleteness, EntityIdentity, ErrorMetric, - ExactKind, ExactOperationSchemaError, ExactParams, ExecutionDataState, ExecutionDataStateError, - ExecutionTiming, GroupingStrategy, NonNegativeWeightProof, SamplingKind, SamplingParams, - SketchAlgorithm, SketchKind, SketchParams, SketchQuery as PostAsapSketchQuery, StatModelKind, - StatModelParams, SummaryExpr, SummaryFamilyType, SummaryField, SummaryInputExpr, SummaryNode, - SummarySchema, SummaryUpdate, ValueOperation, WaveletKind, WaveletParams, WeightDomain, + ExactKind, ExactOperation, ExactOperationSchemaError, ExactParams, ExecutionDataState, + ExecutionDataStateError, ExecutionTiming, GroupingStrategy, NonNegativeWeightProof, + SamplingKind, SamplingParams, SketchAlgorithm, SketchKind, SketchParams, + SketchQuery as PostAsapSketchQuery, StatModelKind, StatModelParams, SummaryExpr, + SummaryFamilyType, SummaryField, SummaryInputExpr, SummaryNode, SummarySchema, SummaryUpdate, + ValueOperation, WaveletKind, WaveletParams, WeightDomain, }; use asap_types::post_asap::{AccuracyError, CompositionOperator, GuaranteeSource, ResultGuarantee}; use asap_types::pre_asap::agg_intent::{agg_is_mergeable, AggIntent}; @@ -3713,22 +3714,26 @@ impl<'a> GlobalSelection<'a> { if let Some(node) = self.materialized.borrow().get(&ptr) { return Ok(Rc::clone(node)); } - let node = match self - .groups - .get(&ptr) - .and_then(|sel| sel.chosen) - .map(|c| &c.replacement) - { - None => self.materialize_residual(target)?, - Some(Replacement::Rewrite(rewritten)) => keep_pre_asap(rewritten)?, - Some(Replacement::Summary(node)) => self.relink_summary(node, target)?, - Some(Replacement::ExactComposition(_)) => Rc::clone( - &self.groups[&ptr] - .composition - .as_ref() - .expect("selected compositions have a validated decision") - .plan, - ), + let node = if read_time_nested_sum(target) { + self.materialize_residual(target)? + } else { + match self + .groups + .get(&ptr) + .and_then(|sel| sel.chosen) + .map(|c| &c.replacement) + { + None => self.materialize_residual(target)?, + Some(Replacement::Rewrite(rewritten)) => keep_pre_asap(rewritten)?, + Some(Replacement::Summary(node)) => self.relink_summary(node, target)?, + Some(Replacement::ExactComposition(_)) => Rc::clone( + &self.groups[&ptr] + .composition + .as_ref() + .expect("selected compositions have a validated decision") + .plan, + ), + } }; self.materialized.borrow_mut().insert(ptr, Rc::clone(&node)); Ok(node) @@ -3808,6 +3813,21 @@ impl<'a> GlobalSelection<'a> { offset: *offset, }, ), + QueryExpr::Aggregate { + reduction, + measures, + output_names, + having, + child, + } if read_time_nested_sum(target) => ( + child, + ValueOperation::Exact(ExactOperation::Aggregate { + reduction: reduction.clone(), + measures: measures.clone(), + output_names: output_names.clone(), + having: having.clone(), + }), + ), _ => return keep_pre_asap(target), }; let child = self.materialize_inner(child_target)?; @@ -3858,6 +3878,34 @@ impl<'a> GlobalSelection<'a> { } } +/// A mergeable outer SUM over a relationally wrapped aggregate is a read-time +/// reduction of the inner summary values. Maintaining the outer SUM directly +/// would hide that inner temporal aggregate inside `KeepPreAsap` and lose its +/// independently selected summary. +fn read_time_nested_sum(target: &QueryExpr) -> bool { + let QueryExpr::Aggregate { + measures, + having: None, + child, + .. + } = target + else { + return false; + }; + matches!(measures.as_slice(), [AggIntent::Sum { .. }]) && contains_aggregate(child) +} + +fn contains_aggregate(expr: &QueryExpr) -> bool { + match expr { + QueryExpr::Aggregate { .. } => true, + QueryExpr::Project { child, .. } + | QueryExpr::Filter { child, .. } + | QueryExpr::Sort { child, .. } + | QueryExpr::Limit { child, .. } => contains_aggregate(child), + _ => false, + } +} + /// Rebuild `node` (a `SummaryAgg`, possibly under a `SummaryEstimate`) with /// `new_child` as the `SummaryAgg`'s child, if the result still validates /// as maintained state; otherwise return `node` unchanged. diff --git a/crates/integration-tests/tests/sql_to_post_asap.rs b/crates/integration-tests/tests/sql_to_post_asap.rs index f625bfbd..ef1f353a 100644 --- a/crates/integration-tests/tests/sql_to_post_asap.rs +++ b/crates/integration-tests/tests/sql_to_post_asap.rs @@ -144,6 +144,65 @@ async fn clickhouse_temporal_sql_reuses_rate_and_increase_physical_summaries() { } } +#[tokio::test] +async fn clickhouse_outer_sum_recursively_binds_inner_temporal_aggregate() { + for (function, window_ms) in [ + ("asap_rate", 300_000), + ("asap_rate", 3_600_000), + ("asap_increase", 300_000), + ] { + let sql = format!( + "SELECT sum(v) AS value FROM (\ + SELECT service, {function}(latency, ts, {window_ms}) AS v \ + FROM metrics GROUP BY service)" + ); + let pre_asap = Rc::new( + lower_sql_dialect( + &sql, + &catalog(), + SqlDialect::ClickhouseSQL, + AccuracyTarget::Exact, + ) + .await + .expect("nested temporal SQL must lower"), + ); + let space = search_workload(vec![("nested", Rc::clone(&pre_asap))]); + let selection = space.global_selection(&DefaultCostModel); + let root = selection + .materialize(&space.roots[0].1) + .expect("materialization failed") + .expect("root must be discovered"); + + fn has_temporal_summary(node: &SummaryNode) -> bool { + match &node.expr { + SummaryExpr::SummaryAgg { + family: + SummaryFamilyType::ExactAggregate(ExactKind::Rate | ExactKind::Increase, _), + .. + } => true, + SummaryExpr::ValueOperation { child, .. } + | SummaryExpr::SummaryEstimate { + summary_input: child, + .. + } => has_temporal_summary(child), + _ => false, + } + } + assert!( + has_temporal_summary(&root), + "inner {function} was hidden: {root:?}" + ); + let executable = compile_executable_dag(&root).expect("nested SQL DAG must be executable"); + assert!(executable.nodes.iter().any(|node| matches!( + node.payload, + ExecutableOperatorPayload::Value { + operation: ValueOperation::Exact(_), + .. + } + ))); + } +} + /// The `Aggregate` node beneath the identity `Project` DataFusion's planner /// always wraps a top-level aggregate in — see the module docs above. fn inner_aggregate(qe: &QueryExpr) -> &QueryExpr {