diff --git a/crates/executor/src/operators/mod.rs b/crates/executor/src/operators/mod.rs index 1859fbdb..cec05d79 100644 --- a/crates/executor/src/operators/mod.rs +++ b/crates/executor/src/operators/mod.rs @@ -135,9 +135,10 @@ enum Kind { kind: planner_types::ir::operator::JoinKind, predicate: Box, }, + /// `value: None`: every row adds a unit weight (SQL `COUNT(*)`). SummaryBuild { family: SummaryFamilyType, - value: usize, + value: Option, time: Option, groups: Vec, }, diff --git a/crates/executor/src/operators/summary/mod.rs b/crates/executor/src/operators/summary/mod.rs index 18b6dbe5..bb57b343 100644 --- a/crates/executor/src/operators/summary/mod.rs +++ b/crates/executor/src/operators/summary/mod.rs @@ -168,6 +168,32 @@ impl Operator { { return Err(invalid("counter summary requires a timestamp column")); } + Self::grouped_build(input, family, Some(value), time, groups) + } + /// An exact count where every row adds one, whatever its values (SQL + /// `COUNT(*)`); the rows need no value column. + pub fn unit_count_build( + input: SchemaRef, + family: SummaryFamilyType, + groups: Vec, + ) -> Result { + crate::values::validate_family(&family)?; + validate_groups(&input, &groups)?; + if !matches!( + family, + SummaryFamilyType::ExactAggregate(planner_types::ir::schema::ExactKind::Count, _) + ) { + return Err(invalid("a unit-weight build needs an exact count")); + } + Self::grouped_build(input, family, None, None, groups) + } + fn grouped_build( + input: SchemaRef, + family: SummaryFamilyType, + value: Option, + time: Option, + groups: Vec, + ) -> Result { crate::capability::validate_summary_kernel( &family, &SummaryUpdate::column(ColumnRef::SampleValue), @@ -439,7 +465,7 @@ pub(super) fn execute_merge<'a>( async fn build_summary( mut input: Input<'_, Batch>, family: &SummaryFamilyType, - value: usize, + value: Option, time: Option, groups: &[usize], emit_empty_global: bool, @@ -496,7 +522,10 @@ async fn build_summary( states.get_mut(&key).expect("inserted group"); // SQL aggregates ignore NULL samples while retaining the group. // A missing counter sample also contributes no observation. - let value = &row[value]; + let value = match value { + Some(value) => &row[value], + None => &Value::Float64(1.0), + }; if matches!(value, Value::Null) { continue; } diff --git a/crates/executor/src/operators/unchecked.rs b/crates/executor/src/operators/unchecked.rs index 98310430..6061a377 100644 --- a/crates/executor/src/operators/unchecked.rs +++ b/crates/executor/src/operators/unchecked.rs @@ -155,7 +155,10 @@ impl TryFrom for Operator { value, time, groups, - } => Operator::summary_build(input(0)?, family, value, time, groups)?, + } => match value { + Some(value) => Operator::summary_build(input(0)?, family, value, time, groups)?, + None => Operator::unit_count_build(input(0)?, family, groups)?, + }, Kind::KeyedSummaryBuild { family, value, diff --git a/crates/executor/src/physical_planner/mod.rs b/crates/executor/src/physical_planner/mod.rs index 4e0dfb41..0bc73c50 100644 --- a/crates/executor/src/physical_planner/mod.rs +++ b/crates/executor/src/physical_planner/mod.rs @@ -1001,16 +1001,26 @@ fn bind_operation(node: &PhysicalASAPDAGNode, inputs: &[SchemaRef]) -> Result column, + SummaryInputExpr::Constant(weight) if *weight == 1.0 => { + return Operator::unit_count_build( + input.clone(), + family.clone(), + groups(input, keys)?, + ) + } + _ => { + return Err(invalid( + "summary update expression must be projected to a column", + )) + } + }; Operator::summary_build( input.clone(), family.clone(), diff --git a/crates/integration-tests/tests/pass1_sql_coverage.rs b/crates/integration-tests/tests/pass1_sql_coverage.rs new file mode 100644 index 00000000..c692a7ba --- /dev/null +++ b/crates/integration-tests/tests/pass1_sql_coverage.rs @@ -0,0 +1,147 @@ +//! Pass 1 alternatives over SQL row sources compose, compile and execute. +mod physical_common; +use std::collections::BTreeMap; +use std::rc::Rc; + +use asap_executor::values::Value; +use asap_frontend_sql::{lower_sql, SqlCatalog}; +use asap_logical_optimizer::pass1::logical_candidates::{ + compose_logical_candidate, enumerate_choices, enumerate_local_logical_candidates, + LocalLogicalCandidates, +}; +use asap_plan_selection::{plan_stages, PlanningModels}; +use asap_types::ir::schema::{DataType, Field, Schema}; +use asap_types::ir::{ASAPOp, Operator, OperatorNode, QueryRoot}; +use asap_types::types::AccuracyTarget; +use asap_types::workload::{ + DataArrival, DataWorkload, Evidence, EvidenceSource, Predictability, QueryRecurrence, Rate, + RootDemand, +}; + +fn catalog() -> SqlCatalog { + SqlCatalog::new().with_table( + "flows", + Schema::new(vec![ + Field::plain("ts", DataType::Timestamp, false), + Field::plain("src_ip", DataType::Utf8, false), + ]), + ) +} + +fn target() -> AccuracyTarget { + AccuracyTarget::EpsilonDelta { + epsilon: 0.01, + delta: 0.01, + } +} + +fn rows() -> Vec> { + ["a", "a", "b", "a", "c", "b"] + .into_iter() + .map(|ip| vec![Value::Timestamp(0), Value::Utf8(ip.into())]) + .collect() +} + +async fn inventory(sql: &str) -> LocalLogicalCandidates { + let root = lower_sql(sql, &catalog(), target()).await.unwrap(); + enumerate_local_logical_candidates(vec![(0, QueryRoot::Operator(root))], &BTreeMap::new()) + .unwrap() +} + +fn has_estimate(root: &Rc) -> bool { + OperatorNode::reachable(root) + .iter() + .any(|n| matches!(n.operator, Operator::ASAP(ASAPOp::SummaryEstimate { .. }))) +} + +/// `COUNT(*) … GROUP BY src_ip` (#509 Example 2's inner query) reads no +/// sample value. Every candidate composes and compiles; the exact ones, +/// the exact `Count` accumulator included, count rows per `src_ip`. +#[tokio::test] +async fn sql_count_star_group_by_candidates_compose_and_execute() { + let inventory = inventory("SELECT src_ip, COUNT(*) AS c FROM flows GROUP BY src_ip").await; + let choices = enumerate_choices(&inventory, usize::MAX); + assert!(choices.len() > 2, "exact and summary candidates"); + let mut executed = 0; + for choice in choices { + let roots = compose_logical_candidate(&inventory, &choice) + .unwrap_or_else(|e| panic!("{choice:?} composes: {e}")); + let QueryRoot::Operator(root) = &roots[0].1 else { + panic!("operator root") + }; + physical_common::compile_physical_asap_dag(root) + .unwrap_or_else(|e| panic!("{choice:?} compiles: {e}")); + if has_estimate(root) { + continue; + } + let mut rows = physical_common::execute_raw_rows(root, rows()); + rows.sort_by_key(|row| format!("{row:?}")); + let counts: Vec<_> = rows + .iter() + .map(|row| match (&row[0], &row[1]) { + (Value::Utf8(ip), Value::Int64(n)) => (ip.to_string(), *n), + other => panic!("{choice:?}: unexpected row {other:?}"), + }) + .collect(); + assert_eq!( + counts, + [("a".into(), 3), ("b".into(), 2), ("c".into(), 1)], + "{choice:?}" + ); + executed += 1; + } + assert_eq!(executed, 2, "pass-through and the exact Count accumulator"); +} + +fn declared(value: T) -> Evidence { + Evidence { + value: Some(value), + source: EvidenceSource::Declared, + ..Default::default() + } +} + +/// #509 Example 2's design queries (its integer Q3 keeps `COUNT(*) GROUP BY +/// src_ip` as a target): every candidate builds through Stages 1 and 2. +/// Stage 3 may still reject one, e.g. for accuracy. +#[tokio::test] +async fn example2_design_candidates_all_build() { + const QUERIES: [&str; 3] = [ + "SELECT COUNT(DISTINCT src_ip) FROM flows", + "SELECT -SUM(p * LN(p)) FROM (SELECT COUNT(*) * 1.0 / SUM(COUNT(*)) OVER () AS p FROM flows GROUP BY src_ip)", + "SELECT SQRT(SUM(c * c)) FROM (SELECT src_ip, COUNT(*) AS c FROM flows GROUP BY src_ip)", + ]; + let mut roots = vec![]; + for (i, sql) in QUERIES.into_iter().enumerate() { + let root = lower_sql(sql, &catalog(), target()).await.unwrap(); + roots.push((i, QueryRoot::Operator(root))); + } + let demand = vec![ + RootDemand { + accuracy: Some(target()), + recurrence: QueryRecurrence::OneTime { + invocations: 1, + execute_at: None, + }, + predictability: Predictability::default(), + latency_ms: None, + }; + roots.len() + ]; + let data = DataWorkload { + arrival: DataArrival::ContinuouslyIngesting, + ingestion_rate: declared(Rate(100_000.0)), + input_cardinality: declared(10_000_000), + ..Default::default() + }; + let run = plan_stages(roots, &demand, &data, PlanningModels::builtin(), 4096).unwrap(); + let enumeration = run.enumeration.unwrap(); + let unbuilt: Vec<_> = enumeration + .selection + .rejected + .iter() + .filter(|r| r.reason.starts_with("Stage ")) + .collect(); + assert!(enumeration.candidates.len() > 1); + assert!(unbuilt.is_empty(), "{unbuilt:#?}"); +} diff --git a/crates/logical-optimizer/src/pass1/logical_candidates.rs b/crates/logical-optimizer/src/pass1/logical_candidates.rs index a9341abc..778187b0 100644 --- a/crates/logical-optimizer/src/pass1/logical_candidates.rs +++ b/crates/logical-optimizer/src/pass1/logical_candidates.rs @@ -658,6 +658,11 @@ fn summary_update( FieldDataType::Sketch(kind, _) => Some(kind.algorithm()), _ => None, }; + if let AggIntent::Count { .. } = intent { + if child.closed && !child.has_promql_series_identity() { + return sql_row_count_update(algorithm.is_some(), reduction, child); + } + } let weight = crate::pass1::replacement::summarised_input(intent, child) .map_err(|_| LogicalCandidateError::Unsupported("input column outside child schema"))?; Ok(match (intent, algorithm) { @@ -715,6 +720,43 @@ fn summary_update( }) } +/// SQL `COUNT(*)`: rows have no sample value, and every row counts, so each +/// adds a unit weight. A sketch hashes an item per row; its bare count +/// ignores the item's value, so any non-null column serves, a grouping +/// column first. +fn sql_row_count_update( + sketch: bool, + reduction: &Reduction, + child: &Schema, +) -> Result { + let item = if sketch { + let keys = reduction + .group_keys() + .filter(|keys| !keys.is_without()) + .into_iter() + .flat_map(|keys| keys.iter().copied()); + let column = keys + .chain(0..child.fields.len()) + .filter_map(|index| child.fields.get(index)) + .find(|field| field.is_plain() && !field.nullable) + .ok_or(LogicalCandidateError::Unsupported( + "a COUNT(*) sketch needs a non-null item column", + ))?; + Some(SummaryInputExpr::Column( + crate::pass1::replacement::column_ref(column), + )) + } else { + None + }; + Ok(SummaryUpdate { + item, + weight: SummaryInputExpr::Constant(1.0), + weight_domain: WeightDomain::NonNegative { + proof: NonNegativeWeightProof::UnitCount, + }, + }) +} + fn statistic(intent: &AggIntent) -> Result { Ok(match intent { AggIntent::Quantile { q, .. } => SketchStatistic::Quantile { q: *q },