From afad4202985f590492807053800a52069d925bff Mon Sep 17 00:00:00 2001 From: zz_y Date: Tue, 15 Sep 2026 07:26:19 -0600 Subject: [PATCH 1/4] feat(types): add binary aggregate intent with correlation --- .../src/query_physical_lowering.rs | 58 +++++++++++++++ crates/asap-aware-mapping/src/replacement.rs | 20 +++++- .../src/post_asap/execution_data_state.rs | 32 +++++++-- crates/types/src/pre_asap/agg_intent.rs | 70 ++++++++++++++++--- crates/types/src/pre_asap/binder.rs | 30 ++++++-- crates/types/src/pre_asap/mod.rs | 2 +- crates/types/src/pre_asap/resolve.rs | 43 ++++++++++++ 7 files changed, 235 insertions(+), 20 deletions(-) diff --git a/crates/asap-aware-mapping/src/query_physical_lowering.rs b/crates/asap-aware-mapping/src/query_physical_lowering.rs index 39d04ced..6f43aac0 100644 --- a/crates/asap-aware-mapping/src/query_physical_lowering.rs +++ b/crates/asap-aware-mapping/src/query_physical_lowering.rs @@ -1187,6 +1187,10 @@ fn supports_hash_aggregate( | AggIntent::Avg { .. } | AggIntent::StdDev { .. } | AggIntent::Variance { .. } + | AggIntent::Binary { + op: asap_types::pre_asap::BinaryAggOp::Correlation, + .. + } | AggIntent::Group | AggIntent::CountValues { .. } ) @@ -1443,6 +1447,60 @@ mod tests { )); } + // Correlation can be costed as an exact hash aggregate using provider-supplied state size. + #[test] + fn correlation_lowers_to_physical_hash_aggregate() { + use asap_types::pre_asap::{ + AggIntent, BinaryAggOp, Column, DataType, QueryExpr, Reduction, Schema, Source, + }; + let source = Source::Table { + table_ref: "pairs".into(), + }; + let root = Rc::new(QueryExpr::Aggregate { + reduction: Reduction::by(vec![]), + measures: vec![AggIntent::Binary { + op: BinaryAggOp::Correlation, + left: 0, + right: 1, + }], + output_names: vec!["r".into()], + having: None, + child: Rc::new(QueryExpr::Scan { + source: source.clone(), + predicates: vec![], + schema: Schema::new(vec![ + Column::new("x", DataType::Float64, true), + Column::new("y", DataType::Float64, true), + ]), + }), + }); + let scope = scope(vec![coverage(source, vec![])]); + let provided = HashMap::from([ + ( + "query-1".into(), + evidence(scan_stats(edge(100, 1600), 1600)), + ), + ( + "query-0".into(), + evidence(OperatorStatistics::HashAggregate { + edges: unary_edges(edge(100, 1600), edge(1, 8)), + group_count: 1, + key_bytes: 0, + accumulator_bytes_per_group: 48, + }), + ), + ]); + let dag = lower_query_physical_dag(&root, &scope, &scripted(&provided)).unwrap(); + assert!(matches!( + dag.nodes.last().unwrap().operator, + PhysicalOperator::HashAggregate { + grouping_key_count: 0, + accumulator_count: 1, + } + )); + assert!(estimate_physical_dag(&dag.nodes, &dag.root, &scope, &dag.evidence).is_ok()); + } + #[test] fn query_lowering_recurses_and_fuses_global_sort_limit() { use asap_types::pre_asap::{AggIntent, GroupKeys, QueryExpr, Reduction, SortKey, Source}; diff --git a/crates/asap-aware-mapping/src/replacement.rs b/crates/asap-aware-mapping/src/replacement.rs index 075f13a1..a5bd2086 100644 --- a/crates/asap-aware-mapping/src/replacement.rs +++ b/crates/asap-aware-mapping/src/replacement.rs @@ -819,7 +819,10 @@ pub(crate) fn implementations_for_with( // ── Exact, non-mergeable reducers — richer partial state than a // single value (see `agg_is_mergeable`), so no accumulator form. - AggIntent::Avg { .. } | AggIntent::StdDev { .. } | AggIntent::Variance { .. } => { + AggIntent::Avg { .. } + | AggIntent::StdDev { .. } + | AggIntent::Variance { .. } + | AggIntent::Binary { .. } => { vec![Implementation::PassThrough] } @@ -6062,6 +6065,21 @@ mod tests { } } + // Correlation must never acquire a single-input sketch or scalar accumulator. + #[test] + fn binary_aggregate_keeps_exact_paired_input() { + let intent = AggIntent::Binary { + op: asap_types::pre_asap::BinaryAggOp::Correlation, + left: 0, + right: 1, + }; + assert!(matches!( + implementations_for_with(&intent, &crate::cost_model::DefaultCostModel).as_slice(), + [Implementation::PassThrough] + )); + assert!(summary_candidates(&intent).is_empty()); + } + #[test] fn accuracy_target_drives_the_boundary() { // Same intent, three targets → three different decisions. diff --git a/crates/types/src/post_asap/execution_data_state.rs b/crates/types/src/post_asap/execution_data_state.rs index 2dcd15b3..3017821b 100644 --- a/crates/types/src/post_asap/execution_data_state.rs +++ b/crates/types/src/post_asap/execution_data_state.rs @@ -647,13 +647,11 @@ fn check_plain_operands( .map(|keys| keys.keys().to_vec()) .unwrap_or_default(); for m in measures { - if let Some(col) = m.input_col() { - referenced.push(col); - } + referenced.extend(m.input_cols()); } // With no explicit input column (the PromQL sample-value convention) // the operator reads every non-key column, so all must be plain. - let implicit = measures.iter().any(|m| m.input_col().is_none()); + let implicit = measures.iter().any(|m| m.input_cols().is_empty()); for (i, field) in input.fields.iter().enumerate() { if !(implicit || referenced.contains(&i)) { continue; @@ -1069,6 +1067,32 @@ mod tests { assert!(validate_execution_data_states(&root).is_err()); } + // Both paired operands must be plain; an unrelated state column is not an input. + #[test] + fn binary_aggregate_checks_both_operand_states() { + let operation = ValueOperation::Exact(ExactOperation::Aggregate { + reduction: Reduction::by(vec![]), + measures: vec![AggIntent::Binary { + op: crate::pre_asap::BinaryAggOp::Correlation, + left: 0, + right: 1, + }], + output_names: vec![], + having: None, + }); + for operand in [0, 1] { + let mut input = plain(&["x", "y", "unused"]); + input.fields[operand].dtype = kll(); + assert!(matches!( + check_plain_operands(&operation, &input), + Err(ExecutionDataStateError::NonPlainOperand { .. }) + )); + } + let mut input = plain(&["x", "y", "unused"]); + input.fields[2].dtype = kll(); + check_plain_operands(&operation, &input).unwrap(); + } + #[test] fn exact_operator_schema_matches_pre_asap_aggregate_derivation() { let child_schema = lift_plain(&scan().output_schema().unwrap()); diff --git a/crates/types/src/pre_asap/agg_intent.rs b/crates/types/src/pre_asap/agg_intent.rs index a3ae5ceb..2051bef0 100644 --- a/crates/types/src/pre_asap/agg_intent.rs +++ b/crates/types/src/pre_asap/agg_intent.rs @@ -40,6 +40,14 @@ use crate::types::AccuracyTarget; #[serde(tag = "kind", rename_all = "snake_case")] #[serde(bound(serialize = "C: Serialize", deserialize = "C: Deserialize<'de>"))] pub enum AggIntent { + /// A reduction over paired values from two explicit input columns. + /// Operation semantics live in `BinaryAggOp`; both references participate + /// in binding and dependency tracking, unlike an opaque extension payload. + Binary { + op: BinaryAggOp, + left: C, + right: C, + }, // ── Data-model-agnostic ────────────────────────────────────────────── Count { accuracy: AccuracyTarget, @@ -280,6 +288,15 @@ pub enum AggIntent { }, } +/// Operations supported by the two-input aggregate shape. New operations +/// must define their output type and exact/summary realization explicitly. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum BinaryAggOp { + /// Pearson correlation over pairs where both inputs are non-null. + Correlation, +} + /// Time / calendar accessor functions (issue #46), evaluated over a timestamp. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(tag = "fn", rename_all = "snake_case")] @@ -468,11 +485,19 @@ impl AggIntent { } impl AggIntent { - /// The input column this intent reduces, if it carries one. `None` = the - /// synthetic time-series sample value (PromQL) or an argument-less - /// aggregate (`Count` / `TopK`). Used by schema derivation to resolve each - /// reducer's input column, and by `plan::bind` to pick the column a - /// summary is built over. + /// All explicit value-column dependencies, in argument order. + /// An empty list retains the existing implicit sample/row-count convention. + pub fn input_cols(&self) -> Vec { + match self { + Self::Binary { left, right, .. } => vec![left.clone(), right.clone()], + _ => self.input_col().into_iter().collect(), + } + } + + /// The explicit input of a single-column reducer. Binary aggregates, + /// argument-less aggregates, and implicit PromQL sample inputs return `None`. + /// Use `input_cols` for dependency tracking; this accessor is for consumers + /// that have already selected a single-column implementation. pub fn input_col(&self) -> Option { match self { AggIntent::Sum { col } @@ -504,6 +529,9 @@ impl AggIntent { AggIntent::Avg { .. } => col("avg", DataType::Float64, false), AggIntent::StdDev { .. } => col("stddev", DataType::Float64, false), AggIntent::Variance { .. } => col("variance", DataType::Float64, false), + AggIntent::Binary { op, .. } => match op { + BinaryAggOp::Correlation => col("corr", DataType::Float64, true), + }, AggIntent::Quantile { q, .. } => col( &format!("quantile_{}", quantile_suffix(*q)), DataType::Float64, @@ -628,16 +656,20 @@ pub mod topk { /// Two instances of this aggregation can be merged /// (`agg(A ∪ B) = combine(agg(A), agg(B))`). `Avg` / `StdDev` / `Variance` -/// need richer partial state than a single value, so they are not mergeable. +/// and binary correlation need richer partial state than a single value, so +/// their finalized values are not mergeable. pub fn agg_is_mergeable(op: &AggIntent) -> bool { !matches!( op, - AggIntent::Avg { .. } | AggIntent::StdDev { .. } | AggIntent::Variance { .. } + AggIntent::Avg { .. } + | AggIntent::StdDev { .. } + | AggIntent::Variance { .. } + | AggIntent::Binary { .. } ) } /// Whether this op implies `exact_required` — no sketch benefit. The exact -/// intents are `Sum / Count / Avg / Min / Max`. +/// intents include `Sum / Count / Avg / Min / Max` and binary correlation. pub fn agg_is_exact(op: &AggIntent) -> bool { matches!( op, @@ -646,6 +678,7 @@ pub fn agg_is_exact(op: &AggIntent) -> bool { | AggIntent::Avg { .. } | AggIntent::Min { .. } | AggIntent::Max { .. } + | AggIntent::Binary { .. } | AggIntent::Group | AggIntent::CountValues { .. } ) @@ -700,6 +733,27 @@ mod tests { Column::new(name, dtype, false) } + // Paired aggregates expose both dependencies but cannot merge final scalar results. + #[test] + fn binary_aggregate_contract() { + let intent = AggIntent::Binary { + op: BinaryAggOp::Correlation, + left: 2, + right: 5, + }; + assert_eq!(intent.input_cols(), vec![2, 5]); + assert_eq!(intent.input_col(), None); + assert!(agg_is_exact(&intent)); + assert!(!agg_is_mergeable(&intent)); + let output = intent.output_column(&c("x", DataType::Float64)); + assert_eq!(output.dtype, DataType::Float64); + assert!(output.nullable); + let value = serde_json::to_value(&intent).unwrap(); + assert_eq!(value["kind"], "binary"); + assert_eq!(value["op"], "correlation"); + assert_eq!(serde_json::from_value::(value).unwrap(), intent); + } + #[test] fn output_column_names_are_intent_keyed() { let v = c("value", DataType::Float64); diff --git a/crates/types/src/pre_asap/binder.rs b/crates/types/src/pre_asap/binder.rs index 221aebdb..a3cb6e64 100644 --- a/crates/types/src/pre_asap/binder.rs +++ b/crates/types/src/pre_asap/binder.rs @@ -203,17 +203,14 @@ pub(crate) fn collect_referenced_columns(tree: &UnresolvedQueryExpr) -> Vec, out: &mut Vec) { - if let Some(c) = c { - push_ref_name(c, out); - } - } fn group_keys(g: &super::query_expr::GroupKeys, out: &mut Vec) { g.keys().iter().for_each(|k| push_ref_name(k, out)); } fn measure_cols(measures: &[super::agg_intent::AggIntent], out: &mut Vec) { for m in measures { - opt_ref(&m.input_col(), out); + for c in m.input_cols() { + push_ref_name(&c, out); + } } } fn walk(node: &UnresolvedQueryExpr, out: &mut Vec) { @@ -381,6 +378,27 @@ mod tests { } } + // Both binary inputs must seed a usage-derived schema before positional resolution. + #[test] + fn binary_inputs_seed_usage_derived_schema() { + use crate::pre_asap::{AggIntent, BinaryAggOp, Reduction}; + let tree = UnresolvedQueryExpr::Aggregate { + reduction: Reduction::by(vec![]), + measures: vec![AggIntent::Binary { + op: BinaryAggOp::Correlation, + left: ColumnRef::Named("x".into()), + right: ColumnRef::Named("y".into()), + }], + output_names: vec![], + having: None, + child: Rc::new(src("m")), + }; + assert_eq!(collect_referenced_columns(&tree), vec!["x", "y"]); + let schema = Binder::new().bind(&tree); + assert!(schema.column_id("x").is_some()); + assert!(schema.column_id("y").is_some()); + } + #[test] fn bare_source_yields_ts_value_floor() { let schema = Binder::new().bind(&src("m")); diff --git a/crates/types/src/pre_asap/mod.rs b/crates/types/src/pre_asap/mod.rs index d6c47998..93aae7db 100644 --- a/crates/types/src/pre_asap/mod.rs +++ b/crates/types/src/pre_asap/mod.rs @@ -43,7 +43,7 @@ pub mod schema; pub use agg_intent::{ agg_accuracy, agg_is_exact, agg_is_mergeable, default_cardinality, default_quantile, AggIntent, - MathFunc, TimeFunc, + BinaryAggOp, MathFunc, TimeFunc, }; pub use binder::{Binder, SchemaCatalog, UsageDerivedCatalog}; pub use canonicalize::canonicalize; diff --git a/crates/types/src/pre_asap/resolve.rs b/crates/types/src/pre_asap/resolve.rs index bfd0f1bf..ee44a141 100644 --- a/crates/types/src/pre_asap/resolve.rs +++ b/crates/types/src/pre_asap/resolve.rs @@ -518,6 +518,11 @@ fn resolve_agg_intent( AggIntent::Count { accuracy } => AggIntent::Count { accuracy: accuracy.clone(), }, + AggIntent::Binary { op, left, right } => AggIntent::Binary { + op: *op, + left: resolve_column_ref(left, schema)?, + right: resolve_column_ref(right, schema)?, + }, AggIntent::Sum { col: c } => AggIntent::Sum { col: col(c)? }, AggIntent::Min { col: c } => AggIntent::Min { col: col(c)? }, AggIntent::Max { col: c } => AggIntent::Max { col: col(c)? }, @@ -609,6 +614,44 @@ mod tests { BinaryOpKind, QueryExpr, Source, VectorMatch, VectorMatchKind, }; + // Both sides resolve with qualifiers; an unknown right input is an error. + #[test] + fn resolve_binary_aggregate_inputs() { + use crate::pre_asap::{BinaryAggOp, Column, DataType}; + let schema = Schema::new(vec![ + Column::new("x", DataType::Float64, true).with_table("a"), + Column::new("x", DataType::Float64, true).with_table("b"), + ]); + let intent = AggIntent::Binary { + op: BinaryAggOp::Correlation, + left: ColumnRef::Qualified { + table: "a".into(), + name: "x".into(), + }, + right: ColumnRef::Qualified { + table: "b".into(), + name: "x".into(), + }, + }; + assert_eq!( + resolve_agg_intent(&intent, &schema).unwrap(), + AggIntent::Binary { + op: BinaryAggOp::Correlation, + left: 0, + right: 1, + } + ); + let missing = AggIntent::Binary { + op: BinaryAggOp::Correlation, + left: ColumnRef::Qualified { + table: "a".into(), + name: "x".into(), + }, + right: ColumnRef::Named("missing".into()), + }; + assert!(resolve_agg_intent(&missing, &schema).is_err()); + } + /// `resolve_root` over a `BinaryOp { , PromqlScalarBridge, vector_match }` /// (issue #220): the bridged scalar operand resolves through the same /// generic walk as every other node (its `Literal` child has no From f3718d2a5cf4a3f17061d94a2b6cbe53683b1a97 Mon Sep 17 00:00:00 2001 From: zz_y Date: Tue, 15 Sep 2026 07:26:36 -0600 Subject: [PATCH 2/4] feat(sql): lower corr through binary aggregates --- crates/frontend-sql/src/sql/mod.rs | 35 +++++ .../frontend-sql/tests/binary_aggregates.rs | 143 ++++++++++++++++++ crates/frontend-sql/tests/sql_lowering.rs | 10 ++ crates/sql-function-catalog/src/lib.rs | 15 +- docs/design_docs/pre-asap-ir.md | 17 ++- 5 files changed, 214 insertions(+), 6 deletions(-) create mode 100644 crates/frontend-sql/tests/binary_aggregates.rs diff --git a/crates/frontend-sql/src/sql/mod.rs b/crates/frontend-sql/src/sql/mod.rs index dc9116eb..0d06fc53 100644 --- a/crates/frontend-sql/src/sql/mod.rs +++ b/crates/frontend-sql/src/sql/mod.rs @@ -1623,6 +1623,26 @@ fn lower_agg_intent(expr: &Expr) -> Result, LoweringError> reducer_col(&name, args).map(Some) }; Ok(match semantic { + AggSemantic::Correlation => { + if agg_fn.filter.is_some() + || agg_fn.order_by.is_some() + || agg_fn.null_treatment.is_some() + { + return Err(LoweringError::UnsupportedAggregate( + "corr with FILTER, ORDER BY, or explicit null treatment".into(), + )); + } + let [left, right] = agg_fn.args.as_slice() else { + return Err(LoweringError::UnsupportedAggregate( + "corr requires two arguments".into(), + )); + }; + AggIntent::Binary { + op: asap_types::pre_asap::BinaryAggOp::Correlation, + left: expr_to_group_ref(left)?, + right: expr_to_group_ref(right)?, + } + } AggSemantic::Count if agg_fn.distinct => AggIntent::Cardinality { col: col(&agg_fn.args)?, accuracy: current_accuracy(), @@ -2021,6 +2041,21 @@ impl DerivedCols { let Expr::AggregateFunction(agg_fn) = unalias(expr) else { return Ok(expr.clone()); }; + if matches!( + asap_sql_function_catalog::lookup_native(&agg_fn.func.name().to_lowercase()), + Some(AggSemantic::Correlation) + ) { + // Give each value argument its own projected name, including casts + // and qualified columns. This retains both inputs and avoids losing + // relation qualifiers when the projection becomes an unqualified schema. + let mut rewritten = agg_fn.clone(); + for arg in &mut rewritten.args { + let alias = unalias(arg).to_string(); + self.materialize(alias.clone(), df_expr_to_unresolved(arg)?)?; + *arg = Expr::Column(DfColumn::new_unqualified(alias)); + } + return Ok(Expr::AggregateFunction(rewritten)); + } // `COUNT(*)` reduces no column; `agg_col_name` covers bare/aliased/cast // columns, so `None` here means the argument really is an expression. let counts_rows = agg_fn.func.name().eq_ignore_ascii_case("count") && !agg_fn.distinct; diff --git a/crates/frontend-sql/tests/binary_aggregates.rs b/crates/frontend-sql/tests/binary_aggregates.rs new file mode 100644 index 00000000..09ad7d42 --- /dev/null +++ b/crates/frontend-sql/tests/binary_aggregates.rs @@ -0,0 +1,143 @@ +//! Two-input aggregates retain paired arguments through SQL lowering and exact planning. +use std::rc::Rc; + +use asap_frontend_sql::{lower_sql, SqlCatalog}; +use asap_types::pre_asap::{AggIntent, BinaryAggOp, Column, DataType, QueryExpr, Schema}; +use asap_types::types::AccuracyTarget; + +fn catalog() -> SqlCatalog { + let schema = Schema::new(vec![ + Column::new("x", DataType::Float64, true), + Column::new("y", DataType::Float64, true), + Column::new("g", DataType::Int64, false), + ]); + SqlCatalog::new() + .with_table("a", schema.clone()) + .with_table("b", schema) +} + +async fn lower(sql: &str) -> QueryExpr { + lower_sql(sql, &catalog(), AccuracyTarget::Exact) + .await + .unwrap() +} + +fn aggregate(query: &QueryExpr) -> (&[AggIntent], &QueryExpr) { + match query { + QueryExpr::Aggregate { + measures, child, .. + } => (measures, child), + QueryExpr::Project { child, .. } + | QueryExpr::Filter { child, .. } + | QueryExpr::Sort { child, .. } => aggregate(child), + other => panic!("expected aggregate, got {other:?}"), + } +} + +// Expressions in either argument, including casts, survive as projected values. +#[tokio::test] +async fn corr_materializes_both_arguments() { + for sql in [ + "SELECT corr(x + 1, y * 2) FROM a", + "SELECT corr(x, y * 2) FROM a", + "SELECT corr(CAST(g AS DOUBLE), y) FROM a", + "SELECT corr(x, 1.0) FROM a", + ] { + let query = lower(sql).await; + let (measures, child) = aggregate(&query); + assert_eq!( + measures, + &[AggIntent::Binary { + op: BinaryAggOp::Correlation, + left: 0, + right: 1, + }] + ); + let QueryExpr::Project { cols, .. } = child else { + panic!("derived inputs") + }; + assert_eq!(cols.len(), 2); + assert!(cols + .iter() + .any(|col| !matches!(col.expr, QueryExpr::Column(_)))); + assert_eq!( + query.output_schema().unwrap().columns[0].dtype, + DataType::Float64 + ); + } +} + +// Identically named join columns remain distinct even when projected beneath the aggregate. +#[tokio::test] +async fn corr_preserves_qualified_join_inputs() { + let query = lower("SELECT corr(a.x, b.x) FROM a JOIN b ON a.g = b.g").await; + let (measures, child) = aggregate(&query); + assert_eq!(measures[0].input_cols(), vec![0, 1]); + let QueryExpr::Project { cols, .. } = child else { + panic!("paired projection") + }; + assert_eq!(cols[0].expr, QueryExpr::Column(0)); + assert_eq!(cols[1].expr, QueryExpr::Column(3)); +} + +// Grouping and sibling reducers cannot drop either correlation argument. +#[tokio::test] +async fn corr_coexists_with_grouping_having_and_other_measures() { + let query = lower("SELECT g, corr(x, y) AS r, sum(x * y) AS s FROM a GROUP BY g HAVING corr(x, y) > 0 ORDER BY r").await; + let (measures, child) = aggregate(&query); + let pair = measures + .iter() + .find(|m| matches!(m, AggIntent::Binary { .. })) + .unwrap(); + let schema = child.output_schema().unwrap(); + for id in pair.input_cols() { + assert!(id < schema.columns.len()); + } + assert!(measures.iter().any(|m| matches!(m, AggIntent::Sum { .. }))); + let output = query.output_schema().unwrap(); + assert_eq!(output.columns[1].name, "r"); + assert_eq!(output.columns[1].dtype, DataType::Float64); + assert!(output.columns[1].nullable); +} + +// Repeated inputs reuse their value while retaining two argument positions. +#[tokio::test] +async fn corr_repeated_input_and_serialization() { + let query = lower("SELECT corr(x, x) FROM a").await; + assert_eq!(aggregate(&query).0[0].input_cols(), vec![0, 0]); + let encoded = serde_json::to_string(&query).unwrap(); + let decoded: QueryExpr = serde_json::from_str(&encoded).unwrap(); + assert_eq!(query, decoded); +} + +// Unsupported modifiers and window calls fail instead of silently changing semantics. +#[tokio::test] +async fn corr_rejects_unrepresented_forms() { + for sql in [ + "SELECT corr(DISTINCT x, y) FROM a", + "SELECT corr(x, y) FILTER (WHERE g > 0) FROM a", + "SELECT corr(x, y ORDER BY g) FROM a", + "SELECT corr(x, y) OVER () FROM a", + "SELECT corr(x) FROM a", + ] { + assert!( + lower_sql(sql, &catalog(), AccuracyTarget::Exact) + .await + .is_err(), + "{sql}" + ); + } +} + +// Exact fallback retains the complete typed query and compiles to an executable DAG. +#[tokio::test] +async fn corr_survives_exact_plan_compilation() { + let query = Rc::new(lower("SELECT corr(x, y) AS r FROM a").await); + let plan = asap_aware_mapping::replacement::keep_pre_asap(&query).unwrap(); + assert!(plan.guarantee.as_ref().unwrap().is_exact()); + let asap_types::post_asap::SummaryExpr::KeepPreAsap(retained) = &plan.expr else { + panic!("expected exact fallback"); + }; + assert_eq!(aggregate(retained).0, aggregate(&query).0); + asap_types::post_asap::compile_executable_dag(&plan).unwrap(); +} diff --git a/crates/frontend-sql/tests/sql_lowering.rs b/crates/frontend-sql/tests/sql_lowering.rs index 4ce39b55..7f7ced90 100644 --- a/crates/frontend-sql/tests/sql_lowering.rs +++ b/crates/frontend-sql/tests/sql_lowering.rs @@ -2454,3 +2454,13 @@ async fn clickhouse_tuple_element_preserves_declared_field_metadata() { ); } } + +/// Correlation lowers to a nullable numeric result instead of UnsupportedAggregate. +#[tokio::test] +async fn corr_result_is_nullable_float() { + let query = lower("SELECT corr(latency, bytes) AS correlation FROM metrics").await; + let schema = query.output_schema().unwrap(); + assert_eq!(schema.columns[0].name, "correlation"); + assert_eq!(schema.columns[0].dtype, DataType::Float64); + assert!(schema.columns[0].nullable); +} diff --git a/crates/sql-function-catalog/src/lib.rs b/crates/sql-function-catalog/src/lib.rs index 8991a8a6..a636fd0b 100644 --- a/crates/sql-function-catalog/src/lib.rs +++ b/crates/sql-function-catalog/src/lib.rs @@ -91,6 +91,8 @@ pub enum AggSemantic { Min, Max, Avg, + /// Pearson correlation of two numeric value arguments. + Correlation, /// Sample stddev unless `population`. StdDev { population: bool, @@ -126,6 +128,11 @@ pub struct NativeFunction { /// when DataFusion (or this front end) accepts more than one spelling for /// the same semantic (`avg`/`mean`, `stddev`/`stddev_samp`, ...). pub const NATIVE_FUNCTIONS: &[NativeFunction] = &[ + NativeFunction { + name: "corr", + arity: Arity::Exact(2), + semantic: AggSemantic::Correlation, + }, NativeFunction { name: "count", arity: Arity::Range { min: 0, max: 1 }, @@ -311,7 +318,7 @@ pub const CLICKHOUSE_BUILTINS: &[ClickHouseBuiltin] = &[ /// Aggregate function names DataFusion's own planner resolves out of the box /// that this catalog deliberately does *not* map to a canonical /// [`AggSemantic`] -- either because `AggIntent` has no shape for them -/// (a multi-column correlation/regression aggregate, a bitwise/boolean +/// (an unimplemented covariance/regression aggregate, a bitwise/boolean /// aggregate, a string concatenation aggregate, ...) or because this front /// end already rejects them explicitly elsewhere (`array_agg`, `grouping`). /// @@ -345,10 +352,8 @@ pub const KNOWN_UNMAPPED_NATIVE_FUNCTIONS: &[&str] = &[ "bit_xor", "bool_and", "bool_or", - // Two-column correlation / linear-regression aggregates -- every - // `AggIntent` value reducer takes one input column (`reducer_col` in - // `asap-frontend-sql`), so these have no home yet. - "corr", + // Two-column covariance / regression operations not yet implemented by + // `AggIntent::Binary`. Correlation is the first supported operation. "covar", "covar_pop", "covar_samp", diff --git a/docs/design_docs/pre-asap-ir.md b/docs/design_docs/pre-asap-ir.md index 67b325c0..065d6190 100644 --- a/docs/design_docs/pre-asap-ir.md +++ b/docs/design_docs/pre-asap-ir.md @@ -104,7 +104,7 @@ native-histogram accessors — this list is representative, not exhaustive: ```text Count, Sum(col), Min(col), Max(col), Avg(col), StdDev(col), Variance(col), -Quantile(col, q), TopK(k), Cardinality(col) // data-model-agnostic +Quantile(col, q), TopK(k), Cardinality(col), Binary(op, left, right) // data-model-agnostic Rate, Increase // counter derivatives Changes, Delta, IDelta, Deriv, Resets, PredictLinear(seconds), DoubleExpSmoothing(sf, tf) // range-vector functions @@ -113,6 +113,21 @@ HistogramStdVar, HistogramFraction(lo, hi), HistogramQuantile(q) // native-hist Math(func) // element-wise transform ``` +`Binary { op, left, right }` represents an aggregate over two value columns; +`BinaryAggOp::Correlation` is its first operation. Both references resolve to +positional column IDs, and `input_cols()` exposes both dependencies. SQL lowering +projects both arguments, preserving expressions, casts, and qualified join columns. +The result of correlation is nullable `Float64`, with pairwise null handling owned +by the executing engine. It remains exact: finalized correlation coefficients +cannot be combined as scalar rollups, and no sketch or maintained correlation +accumulator is selected. Physical costing accepts it as a hash aggregate with +provider-supplied accumulator size. + +Further two-input operations can extend `BinaryAggOp` without adding another +argument-carrying aggregate variant. Each needs explicit output and realization +semantics. SQL `corr` currently rejects `DISTINCT`, aggregate `FILTER`, aggregate +`ORDER BY`, explicit null treatment, and window usage (`OVER`). + Additional measures can be added when there is a stable semantic distinction and a meaningful summary implementation. From 1cddccb08ba1d1db687dc9efdc482f163137e7bc Mon Sep 17 00:00:00 2001 From: zz_y Date: Tue, 15 Sep 2026 07:32:49 -0600 Subject: [PATCH 3/4] refactor: name paired aggregates bivariate --- .../src/query_physical_lowering.rs | 10 +++--- crates/asap-aware-mapping/src/replacement.rs | 8 ++--- crates/frontend-sql/src/sql/mod.rs | 4 +-- ..._aggregates.rs => bivariate_aggregates.rs} | 8 ++--- crates/sql-function-catalog/src/lib.rs | 2 +- .../src/post_asap/execution_data_state.rs | 6 ++-- crates/types/src/pre_asap/agg_intent.rs | 32 +++++++++---------- crates/types/src/pre_asap/binder.rs | 8 ++--- crates/types/src/pre_asap/mod.rs | 2 +- crates/types/src/pre_asap/resolve.rs | 18 +++++------ docs/design_docs/pre-asap-ir.md | 8 ++--- 11 files changed, 53 insertions(+), 53 deletions(-) rename crates/frontend-sql/tests/{binary_aggregates.rs => bivariate_aggregates.rs} (95%) diff --git a/crates/asap-aware-mapping/src/query_physical_lowering.rs b/crates/asap-aware-mapping/src/query_physical_lowering.rs index 6f43aac0..75c1f2ca 100644 --- a/crates/asap-aware-mapping/src/query_physical_lowering.rs +++ b/crates/asap-aware-mapping/src/query_physical_lowering.rs @@ -1187,8 +1187,8 @@ fn supports_hash_aggregate( | AggIntent::Avg { .. } | AggIntent::StdDev { .. } | AggIntent::Variance { .. } - | AggIntent::Binary { - op: asap_types::pre_asap::BinaryAggOp::Correlation, + | AggIntent::Bivariate { + op: asap_types::pre_asap::BivariateAggOp::Correlation, .. } | AggIntent::Group @@ -1451,15 +1451,15 @@ mod tests { #[test] fn correlation_lowers_to_physical_hash_aggregate() { use asap_types::pre_asap::{ - AggIntent, BinaryAggOp, Column, DataType, QueryExpr, Reduction, Schema, Source, + AggIntent, BivariateAggOp, Column, DataType, QueryExpr, Reduction, Schema, Source, }; let source = Source::Table { table_ref: "pairs".into(), }; let root = Rc::new(QueryExpr::Aggregate { reduction: Reduction::by(vec![]), - measures: vec![AggIntent::Binary { - op: BinaryAggOp::Correlation, + measures: vec![AggIntent::Bivariate { + op: BivariateAggOp::Correlation, left: 0, right: 1, }], diff --git a/crates/asap-aware-mapping/src/replacement.rs b/crates/asap-aware-mapping/src/replacement.rs index a5bd2086..1eeea0dd 100644 --- a/crates/asap-aware-mapping/src/replacement.rs +++ b/crates/asap-aware-mapping/src/replacement.rs @@ -822,7 +822,7 @@ pub(crate) fn implementations_for_with( AggIntent::Avg { .. } | AggIntent::StdDev { .. } | AggIntent::Variance { .. } - | AggIntent::Binary { .. } => { + | AggIntent::Bivariate { .. } => { vec![Implementation::PassThrough] } @@ -6067,9 +6067,9 @@ mod tests { // Correlation must never acquire a single-input sketch or scalar accumulator. #[test] - fn binary_aggregate_keeps_exact_paired_input() { - let intent = AggIntent::Binary { - op: asap_types::pre_asap::BinaryAggOp::Correlation, + fn bivariate_aggregate_keeps_exact_paired_input() { + let intent = AggIntent::Bivariate { + op: asap_types::pre_asap::BivariateAggOp::Correlation, left: 0, right: 1, }; diff --git a/crates/frontend-sql/src/sql/mod.rs b/crates/frontend-sql/src/sql/mod.rs index 0d06fc53..7c370cc7 100644 --- a/crates/frontend-sql/src/sql/mod.rs +++ b/crates/frontend-sql/src/sql/mod.rs @@ -1637,8 +1637,8 @@ fn lower_agg_intent(expr: &Expr) -> Result, LoweringError> "corr requires two arguments".into(), )); }; - AggIntent::Binary { - op: asap_types::pre_asap::BinaryAggOp::Correlation, + AggIntent::Bivariate { + op: asap_types::pre_asap::BivariateAggOp::Correlation, left: expr_to_group_ref(left)?, right: expr_to_group_ref(right)?, } diff --git a/crates/frontend-sql/tests/binary_aggregates.rs b/crates/frontend-sql/tests/bivariate_aggregates.rs similarity index 95% rename from crates/frontend-sql/tests/binary_aggregates.rs rename to crates/frontend-sql/tests/bivariate_aggregates.rs index 09ad7d42..2c231e45 100644 --- a/crates/frontend-sql/tests/binary_aggregates.rs +++ b/crates/frontend-sql/tests/bivariate_aggregates.rs @@ -2,7 +2,7 @@ use std::rc::Rc; use asap_frontend_sql::{lower_sql, SqlCatalog}; -use asap_types::pre_asap::{AggIntent, BinaryAggOp, Column, DataType, QueryExpr, Schema}; +use asap_types::pre_asap::{AggIntent, BivariateAggOp, Column, DataType, QueryExpr, Schema}; use asap_types::types::AccuracyTarget; fn catalog() -> SqlCatalog { @@ -47,8 +47,8 @@ async fn corr_materializes_both_arguments() { let (measures, child) = aggregate(&query); assert_eq!( measures, - &[AggIntent::Binary { - op: BinaryAggOp::Correlation, + &[AggIntent::Bivariate { + op: BivariateAggOp::Correlation, left: 0, right: 1, }] @@ -87,7 +87,7 @@ async fn corr_coexists_with_grouping_having_and_other_measures() { let (measures, child) = aggregate(&query); let pair = measures .iter() - .find(|m| matches!(m, AggIntent::Binary { .. })) + .find(|m| matches!(m, AggIntent::Bivariate { .. })) .unwrap(); let schema = child.output_schema().unwrap(); for id in pair.input_cols() { diff --git a/crates/sql-function-catalog/src/lib.rs b/crates/sql-function-catalog/src/lib.rs index a636fd0b..34c0a7e1 100644 --- a/crates/sql-function-catalog/src/lib.rs +++ b/crates/sql-function-catalog/src/lib.rs @@ -353,7 +353,7 @@ pub const KNOWN_UNMAPPED_NATIVE_FUNCTIONS: &[&str] = &[ "bool_and", "bool_or", // Two-column covariance / regression operations not yet implemented by - // `AggIntent::Binary`. Correlation is the first supported operation. + // `AggIntent::Bivariate`. Correlation is the first supported operation. "covar", "covar_pop", "covar_samp", diff --git a/crates/types/src/post_asap/execution_data_state.rs b/crates/types/src/post_asap/execution_data_state.rs index 3017821b..aaa99c2a 100644 --- a/crates/types/src/post_asap/execution_data_state.rs +++ b/crates/types/src/post_asap/execution_data_state.rs @@ -1069,11 +1069,11 @@ mod tests { // Both paired operands must be plain; an unrelated state column is not an input. #[test] - fn binary_aggregate_checks_both_operand_states() { + fn bivariate_aggregate_checks_both_operand_states() { let operation = ValueOperation::Exact(ExactOperation::Aggregate { reduction: Reduction::by(vec![]), - measures: vec![AggIntent::Binary { - op: crate::pre_asap::BinaryAggOp::Correlation, + measures: vec![AggIntent::Bivariate { + op: crate::pre_asap::BivariateAggOp::Correlation, left: 0, right: 1, }], diff --git a/crates/types/src/pre_asap/agg_intent.rs b/crates/types/src/pre_asap/agg_intent.rs index 2051bef0..d2907f6c 100644 --- a/crates/types/src/pre_asap/agg_intent.rs +++ b/crates/types/src/pre_asap/agg_intent.rs @@ -41,10 +41,10 @@ use crate::types::AccuracyTarget; #[serde(bound(serialize = "C: Serialize", deserialize = "C: Deserialize<'de>"))] pub enum AggIntent { /// A reduction over paired values from two explicit input columns. - /// Operation semantics live in `BinaryAggOp`; both references participate + /// Operation semantics live in `BivariateAggOp`; both references participate /// in binding and dependency tracking, unlike an opaque extension payload. - Binary { - op: BinaryAggOp, + Bivariate { + op: BivariateAggOp, left: C, right: C, }, @@ -292,7 +292,7 @@ pub enum AggIntent { /// must define their output type and exact/summary realization explicitly. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] -pub enum BinaryAggOp { +pub enum BivariateAggOp { /// Pearson correlation over pairs where both inputs are non-null. Correlation, } @@ -489,12 +489,12 @@ impl AggIntent { /// An empty list retains the existing implicit sample/row-count convention. pub fn input_cols(&self) -> Vec { match self { - Self::Binary { left, right, .. } => vec![left.clone(), right.clone()], + Self::Bivariate { left, right, .. } => vec![left.clone(), right.clone()], _ => self.input_col().into_iter().collect(), } } - /// The explicit input of a single-column reducer. Binary aggregates, + /// The explicit input of a single-column reducer. Bivariate aggregates, /// argument-less aggregates, and implicit PromQL sample inputs return `None`. /// Use `input_cols` for dependency tracking; this accessor is for consumers /// that have already selected a single-column implementation. @@ -529,8 +529,8 @@ impl AggIntent { AggIntent::Avg { .. } => col("avg", DataType::Float64, false), AggIntent::StdDev { .. } => col("stddev", DataType::Float64, false), AggIntent::Variance { .. } => col("variance", DataType::Float64, false), - AggIntent::Binary { op, .. } => match op { - BinaryAggOp::Correlation => col("corr", DataType::Float64, true), + AggIntent::Bivariate { op, .. } => match op { + BivariateAggOp::Correlation => col("corr", DataType::Float64, true), }, AggIntent::Quantile { q, .. } => col( &format!("quantile_{}", quantile_suffix(*q)), @@ -656,7 +656,7 @@ pub mod topk { /// Two instances of this aggregation can be merged /// (`agg(A ∪ B) = combine(agg(A), agg(B))`). `Avg` / `StdDev` / `Variance` -/// and binary correlation need richer partial state than a single value, so +/// and correlation need richer partial state than a single value, so /// their finalized values are not mergeable. pub fn agg_is_mergeable(op: &AggIntent) -> bool { !matches!( @@ -664,12 +664,12 @@ pub fn agg_is_mergeable(op: &AggIntent) -> bool { AggIntent::Avg { .. } | AggIntent::StdDev { .. } | AggIntent::Variance { .. } - | AggIntent::Binary { .. } + | AggIntent::Bivariate { .. } ) } /// Whether this op implies `exact_required` — no sketch benefit. The exact -/// intents include `Sum / Count / Avg / Min / Max` and binary correlation. +/// intents include `Sum / Count / Avg / Min / Max` and correlation. pub fn agg_is_exact(op: &AggIntent) -> bool { matches!( op, @@ -678,7 +678,7 @@ pub fn agg_is_exact(op: &AggIntent) -> bool { | AggIntent::Avg { .. } | AggIntent::Min { .. } | AggIntent::Max { .. } - | AggIntent::Binary { .. } + | AggIntent::Bivariate { .. } | AggIntent::Group | AggIntent::CountValues { .. } ) @@ -735,9 +735,9 @@ mod tests { // Paired aggregates expose both dependencies but cannot merge final scalar results. #[test] - fn binary_aggregate_contract() { - let intent = AggIntent::Binary { - op: BinaryAggOp::Correlation, + fn bivariate_aggregate_contract() { + let intent = AggIntent::Bivariate { + op: BivariateAggOp::Correlation, left: 2, right: 5, }; @@ -749,7 +749,7 @@ mod tests { assert_eq!(output.dtype, DataType::Float64); assert!(output.nullable); let value = serde_json::to_value(&intent).unwrap(); - assert_eq!(value["kind"], "binary"); + assert_eq!(value["kind"], "bivariate"); assert_eq!(value["op"], "correlation"); assert_eq!(serde_json::from_value::(value).unwrap(), intent); } diff --git a/crates/types/src/pre_asap/binder.rs b/crates/types/src/pre_asap/binder.rs index a3cb6e64..5351e494 100644 --- a/crates/types/src/pre_asap/binder.rs +++ b/crates/types/src/pre_asap/binder.rs @@ -380,12 +380,12 @@ mod tests { // Both binary inputs must seed a usage-derived schema before positional resolution. #[test] - fn binary_inputs_seed_usage_derived_schema() { - use crate::pre_asap::{AggIntent, BinaryAggOp, Reduction}; + fn bivariate_inputs_seed_usage_derived_schema() { + use crate::pre_asap::{AggIntent, BivariateAggOp, Reduction}; let tree = UnresolvedQueryExpr::Aggregate { reduction: Reduction::by(vec![]), - measures: vec![AggIntent::Binary { - op: BinaryAggOp::Correlation, + measures: vec![AggIntent::Bivariate { + op: BivariateAggOp::Correlation, left: ColumnRef::Named("x".into()), right: ColumnRef::Named("y".into()), }], diff --git a/crates/types/src/pre_asap/mod.rs b/crates/types/src/pre_asap/mod.rs index 93aae7db..da2a2d8f 100644 --- a/crates/types/src/pre_asap/mod.rs +++ b/crates/types/src/pre_asap/mod.rs @@ -43,7 +43,7 @@ pub mod schema; pub use agg_intent::{ agg_accuracy, agg_is_exact, agg_is_mergeable, default_cardinality, default_quantile, AggIntent, - BinaryAggOp, MathFunc, TimeFunc, + BivariateAggOp, MathFunc, TimeFunc, }; pub use binder::{Binder, SchemaCatalog, UsageDerivedCatalog}; pub use canonicalize::canonicalize; diff --git a/crates/types/src/pre_asap/resolve.rs b/crates/types/src/pre_asap/resolve.rs index ee44a141..0977235f 100644 --- a/crates/types/src/pre_asap/resolve.rs +++ b/crates/types/src/pre_asap/resolve.rs @@ -518,7 +518,7 @@ fn resolve_agg_intent( AggIntent::Count { accuracy } => AggIntent::Count { accuracy: accuracy.clone(), }, - AggIntent::Binary { op, left, right } => AggIntent::Binary { + AggIntent::Bivariate { op, left, right } => AggIntent::Bivariate { op: *op, left: resolve_column_ref(left, schema)?, right: resolve_column_ref(right, schema)?, @@ -616,14 +616,14 @@ mod tests { // Both sides resolve with qualifiers; an unknown right input is an error. #[test] - fn resolve_binary_aggregate_inputs() { - use crate::pre_asap::{BinaryAggOp, Column, DataType}; + fn resolve_bivariate_aggregate_inputs() { + use crate::pre_asap::{BivariateAggOp, Column, DataType}; let schema = Schema::new(vec![ Column::new("x", DataType::Float64, true).with_table("a"), Column::new("x", DataType::Float64, true).with_table("b"), ]); - let intent = AggIntent::Binary { - op: BinaryAggOp::Correlation, + let intent = AggIntent::Bivariate { + op: BivariateAggOp::Correlation, left: ColumnRef::Qualified { table: "a".into(), name: "x".into(), @@ -635,14 +635,14 @@ mod tests { }; assert_eq!( resolve_agg_intent(&intent, &schema).unwrap(), - AggIntent::Binary { - op: BinaryAggOp::Correlation, + AggIntent::Bivariate { + op: BivariateAggOp::Correlation, left: 0, right: 1, } ); - let missing = AggIntent::Binary { - op: BinaryAggOp::Correlation, + let missing = AggIntent::Bivariate { + op: BivariateAggOp::Correlation, left: ColumnRef::Qualified { table: "a".into(), name: "x".into(), diff --git a/docs/design_docs/pre-asap-ir.md b/docs/design_docs/pre-asap-ir.md index 065d6190..9f00e874 100644 --- a/docs/design_docs/pre-asap-ir.md +++ b/docs/design_docs/pre-asap-ir.md @@ -104,7 +104,7 @@ native-histogram accessors — this list is representative, not exhaustive: ```text Count, Sum(col), Min(col), Max(col), Avg(col), StdDev(col), Variance(col), -Quantile(col, q), TopK(k), Cardinality(col), Binary(op, left, right) // data-model-agnostic +Quantile(col, q), TopK(k), Cardinality(col), Bivariate(op, left, right) // data-model-agnostic Rate, Increase // counter derivatives Changes, Delta, IDelta, Deriv, Resets, PredictLinear(seconds), DoubleExpSmoothing(sf, tf) // range-vector functions @@ -113,8 +113,8 @@ HistogramStdVar, HistogramFraction(lo, hi), HistogramQuantile(q) // native-hist Math(func) // element-wise transform ``` -`Binary { op, left, right }` represents an aggregate over two value columns; -`BinaryAggOp::Correlation` is its first operation. Both references resolve to +`Bivariate { op, left, right }` represents an aggregate over two value columns; +`BivariateAggOp::Correlation` is its first operation. Both references resolve to positional column IDs, and `input_cols()` exposes both dependencies. SQL lowering projects both arguments, preserving expressions, casts, and qualified join columns. The result of correlation is nullable `Float64`, with pairwise null handling owned @@ -123,7 +123,7 @@ cannot be combined as scalar rollups, and no sketch or maintained correlation accumulator is selected. Physical costing accepts it as a hash aggregate with provider-supplied accumulator size. -Further two-input operations can extend `BinaryAggOp` without adding another +Further two-input operations can extend `BivariateAggOp` without adding another argument-carrying aggregate variant. Each needs explicit output and realization semantics. SQL `corr` currently rejects `DISTINCT`, aggregate `FILTER`, aggregate `ORDER BY`, explicit null treatment, and window usage (`OVER`). From c53861b4f4a6ed8baff6ee03ddc74e3477a66219 Mon Sep 17 00:00:00 2001 From: Selvomega Date: Tue, 15 Sep 2026 17:30:51 +0000 Subject: [PATCH 4/4] refactor: name the paired aggregate PearsonCorr `AggIntent::Bivariate { op, left, right }` placed a category alongside specific instances such as `Avg` and `Variance`, and `BivariateAggOp` carried exactly one operation. Correlation is the only two-input aggregate, so name the variant after it and drop the operation enum. `covar` and the `regr_*` family can each add a variant when implemented, following the `StdDev` / `Variance` precedent. The variant now sits next to `Variance` rather than ahead of the data-model-agnostic banner. Serialized `kind` becomes `pearson_corr` and the `op` field is gone. `input_cols()` is unchanged: dependency walkers still need a generic accessor, and `input_col()` still returns `None` so single-column consumers cannot pick up half of the pair. Co-Authored-By: Claude Opus 5 (1M context) --- .../src/query_physical_lowering.rs | 13 ++--- crates/asap-aware-mapping/src/replacement.rs | 10 ++-- crates/frontend-sql/src/sql/mod.rs | 3 +- ...ivariate_aggregates.rs => pearson_corr.rs} | 15 ++---- crates/sql-function-catalog/src/lib.rs | 4 +- .../src/post_asap/execution_data_state.rs | 8 +-- crates/types/src/pre_asap/agg_intent.rs | 51 +++++++------------ crates/types/src/pre_asap/binder.rs | 9 ++-- crates/types/src/pre_asap/mod.rs | 2 +- crates/types/src/pre_asap/resolve.rs | 19 +++---- docs/design_docs/pre-asap-ir.md | 31 +++++------ 11 files changed, 60 insertions(+), 105 deletions(-) rename crates/frontend-sql/tests/{bivariate_aggregates.rs => pearson_corr.rs} (91%) diff --git a/crates/asap-aware-mapping/src/query_physical_lowering.rs b/crates/asap-aware-mapping/src/query_physical_lowering.rs index 75c1f2ca..bd2aa675 100644 --- a/crates/asap-aware-mapping/src/query_physical_lowering.rs +++ b/crates/asap-aware-mapping/src/query_physical_lowering.rs @@ -1187,10 +1187,7 @@ fn supports_hash_aggregate( | AggIntent::Avg { .. } | AggIntent::StdDev { .. } | AggIntent::Variance { .. } - | AggIntent::Bivariate { - op: asap_types::pre_asap::BivariateAggOp::Correlation, - .. - } + | AggIntent::PearsonCorr { .. } | AggIntent::Group | AggIntent::CountValues { .. } ) @@ -1451,18 +1448,14 @@ mod tests { #[test] fn correlation_lowers_to_physical_hash_aggregate() { use asap_types::pre_asap::{ - AggIntent, BivariateAggOp, Column, DataType, QueryExpr, Reduction, Schema, Source, + AggIntent, Column, DataType, QueryExpr, Reduction, Schema, Source, }; let source = Source::Table { table_ref: "pairs".into(), }; let root = Rc::new(QueryExpr::Aggregate { reduction: Reduction::by(vec![]), - measures: vec![AggIntent::Bivariate { - op: BivariateAggOp::Correlation, - left: 0, - right: 1, - }], + measures: vec![AggIntent::PearsonCorr { left: 0, right: 1 }], output_names: vec!["r".into()], having: None, child: Rc::new(QueryExpr::Scan { diff --git a/crates/asap-aware-mapping/src/replacement.rs b/crates/asap-aware-mapping/src/replacement.rs index 1eeea0dd..13cffa35 100644 --- a/crates/asap-aware-mapping/src/replacement.rs +++ b/crates/asap-aware-mapping/src/replacement.rs @@ -822,7 +822,7 @@ pub(crate) fn implementations_for_with( AggIntent::Avg { .. } | AggIntent::StdDev { .. } | AggIntent::Variance { .. } - | AggIntent::Bivariate { .. } => { + | AggIntent::PearsonCorr { .. } => { vec![Implementation::PassThrough] } @@ -6067,12 +6067,8 @@ mod tests { // Correlation must never acquire a single-input sketch or scalar accumulator. #[test] - fn bivariate_aggregate_keeps_exact_paired_input() { - let intent = AggIntent::Bivariate { - op: asap_types::pre_asap::BivariateAggOp::Correlation, - left: 0, - right: 1, - }; + fn pearson_corr_keeps_exact_paired_input() { + let intent = AggIntent::PearsonCorr { left: 0, right: 1 }; assert!(matches!( implementations_for_with(&intent, &crate::cost_model::DefaultCostModel).as_slice(), [Implementation::PassThrough] diff --git a/crates/frontend-sql/src/sql/mod.rs b/crates/frontend-sql/src/sql/mod.rs index 7c370cc7..5e235800 100644 --- a/crates/frontend-sql/src/sql/mod.rs +++ b/crates/frontend-sql/src/sql/mod.rs @@ -1637,8 +1637,7 @@ fn lower_agg_intent(expr: &Expr) -> Result, LoweringError> "corr requires two arguments".into(), )); }; - AggIntent::Bivariate { - op: asap_types::pre_asap::BivariateAggOp::Correlation, + AggIntent::PearsonCorr { left: expr_to_group_ref(left)?, right: expr_to_group_ref(right)?, } diff --git a/crates/frontend-sql/tests/bivariate_aggregates.rs b/crates/frontend-sql/tests/pearson_corr.rs similarity index 91% rename from crates/frontend-sql/tests/bivariate_aggregates.rs rename to crates/frontend-sql/tests/pearson_corr.rs index 2c231e45..82e9d60d 100644 --- a/crates/frontend-sql/tests/bivariate_aggregates.rs +++ b/crates/frontend-sql/tests/pearson_corr.rs @@ -1,8 +1,8 @@ -//! Two-input aggregates retain paired arguments through SQL lowering and exact planning. +//! Correlation retains both arguments through SQL lowering and exact planning. use std::rc::Rc; use asap_frontend_sql::{lower_sql, SqlCatalog}; -use asap_types::pre_asap::{AggIntent, BivariateAggOp, Column, DataType, QueryExpr, Schema}; +use asap_types::pre_asap::{AggIntent, Column, DataType, QueryExpr, Schema}; use asap_types::types::AccuracyTarget; fn catalog() -> SqlCatalog { @@ -45,14 +45,7 @@ async fn corr_materializes_both_arguments() { ] { let query = lower(sql).await; let (measures, child) = aggregate(&query); - assert_eq!( - measures, - &[AggIntent::Bivariate { - op: BivariateAggOp::Correlation, - left: 0, - right: 1, - }] - ); + assert_eq!(measures, &[AggIntent::PearsonCorr { left: 0, right: 1 }]); let QueryExpr::Project { cols, .. } = child else { panic!("derived inputs") }; @@ -87,7 +80,7 @@ async fn corr_coexists_with_grouping_having_and_other_measures() { let (measures, child) = aggregate(&query); let pair = measures .iter() - .find(|m| matches!(m, AggIntent::Bivariate { .. })) + .find(|m| matches!(m, AggIntent::PearsonCorr { .. })) .unwrap(); let schema = child.output_schema().unwrap(); for id in pair.input_cols() { diff --git a/crates/sql-function-catalog/src/lib.rs b/crates/sql-function-catalog/src/lib.rs index 34c0a7e1..df2c0926 100644 --- a/crates/sql-function-catalog/src/lib.rs +++ b/crates/sql-function-catalog/src/lib.rs @@ -352,8 +352,8 @@ pub const KNOWN_UNMAPPED_NATIVE_FUNCTIONS: &[&str] = &[ "bit_xor", "bool_and", "bool_or", - // Two-column covariance / regression operations not yet implemented by - // `AggIntent::Bivariate`. Correlation is the first supported operation. + // Two-column covariance / regression operations with no `AggIntent` + // variant; `AggIntent::PearsonCorr` covers `corr` only. "covar", "covar_pop", "covar_samp", diff --git a/crates/types/src/post_asap/execution_data_state.rs b/crates/types/src/post_asap/execution_data_state.rs index aaa99c2a..26b30e63 100644 --- a/crates/types/src/post_asap/execution_data_state.rs +++ b/crates/types/src/post_asap/execution_data_state.rs @@ -1069,14 +1069,10 @@ mod tests { // Both paired operands must be plain; an unrelated state column is not an input. #[test] - fn bivariate_aggregate_checks_both_operand_states() { + fn pearson_corr_checks_both_operand_states() { let operation = ValueOperation::Exact(ExactOperation::Aggregate { reduction: Reduction::by(vec![]), - measures: vec![AggIntent::Bivariate { - op: crate::pre_asap::BivariateAggOp::Correlation, - left: 0, - right: 1, - }], + measures: vec![AggIntent::PearsonCorr { left: 0, right: 1 }], output_names: vec![], having: None, }); diff --git a/crates/types/src/pre_asap/agg_intent.rs b/crates/types/src/pre_asap/agg_intent.rs index d2907f6c..6b020c74 100644 --- a/crates/types/src/pre_asap/agg_intent.rs +++ b/crates/types/src/pre_asap/agg_intent.rs @@ -40,14 +40,6 @@ use crate::types::AccuracyTarget; #[serde(tag = "kind", rename_all = "snake_case")] #[serde(bound(serialize = "C: Serialize", deserialize = "C: Deserialize<'de>"))] pub enum AggIntent { - /// A reduction over paired values from two explicit input columns. - /// Operation semantics live in `BivariateAggOp`; both references participate - /// in binding and dependency tracking, unlike an opaque extension payload. - Bivariate { - op: BivariateAggOp, - left: C, - right: C, - }, // ── Data-model-agnostic ────────────────────────────────────────────── Count { accuracy: AccuracyTarget, @@ -81,6 +73,13 @@ pub enum AggIntent { col: Option, population: bool, }, + /// Pearson correlation — SQL `CORR(left, right)`. The only aggregate with + /// two value inputs; both references take part in binding and dependency + /// tracking, so neither is reachable through `input_col`. + PearsonCorr { + left: C, + right: C, + }, /// φ-quantile of `col`. SQL `approx_percentile_cont(col, φ)` and /// `median(col)` (φ=0.5); PromQL `quantile(φ, …)` leaves `col` as `None` /// (the sample value). @@ -288,15 +287,6 @@ pub enum AggIntent { }, } -/// Operations supported by the two-input aggregate shape. New operations -/// must define their output type and exact/summary realization explicitly. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum BivariateAggOp { - /// Pearson correlation over pairs where both inputs are non-null. - Correlation, -} - /// Time / calendar accessor functions (issue #46), evaluated over a timestamp. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(tag = "fn", rename_all = "snake_case")] @@ -489,12 +479,12 @@ impl AggIntent { /// An empty list retains the existing implicit sample/row-count convention. pub fn input_cols(&self) -> Vec { match self { - Self::Bivariate { left, right, .. } => vec![left.clone(), right.clone()], + Self::PearsonCorr { left, right } => vec![left.clone(), right.clone()], _ => self.input_col().into_iter().collect(), } } - /// The explicit input of a single-column reducer. Bivariate aggregates, + /// The explicit input of a single-column reducer. `PearsonCorr`, /// argument-less aggregates, and implicit PromQL sample inputs return `None`. /// Use `input_cols` for dependency tracking; this accessor is for consumers /// that have already selected a single-column implementation. @@ -529,9 +519,9 @@ impl AggIntent { AggIntent::Avg { .. } => col("avg", DataType::Float64, false), AggIntent::StdDev { .. } => col("stddev", DataType::Float64, false), AggIntent::Variance { .. } => col("variance", DataType::Float64, false), - AggIntent::Bivariate { op, .. } => match op { - BivariateAggOp::Correlation => col("corr", DataType::Float64, true), - }, + // Nullable: correlation is undefined for fewer than two pairs where + // both inputs are non-null, and for a zero-variance input. + AggIntent::PearsonCorr { .. } => col("corr", DataType::Float64, true), AggIntent::Quantile { q, .. } => col( &format!("quantile_{}", quantile_suffix(*q)), DataType::Float64, @@ -664,7 +654,7 @@ pub fn agg_is_mergeable(op: &AggIntent) -> bool { AggIntent::Avg { .. } | AggIntent::StdDev { .. } | AggIntent::Variance { .. } - | AggIntent::Bivariate { .. } + | AggIntent::PearsonCorr { .. } ) } @@ -678,7 +668,7 @@ pub fn agg_is_exact(op: &AggIntent) -> bool { | AggIntent::Avg { .. } | AggIntent::Min { .. } | AggIntent::Max { .. } - | AggIntent::Bivariate { .. } + | AggIntent::PearsonCorr { .. } | AggIntent::Group | AggIntent::CountValues { .. } ) @@ -733,14 +723,10 @@ mod tests { Column::new(name, dtype, false) } - // Paired aggregates expose both dependencies but cannot merge final scalar results. + // Correlation exposes both dependencies but cannot merge final scalar results. #[test] - fn bivariate_aggregate_contract() { - let intent = AggIntent::Bivariate { - op: BivariateAggOp::Correlation, - left: 2, - right: 5, - }; + fn pearson_corr_contract() { + let intent = AggIntent::PearsonCorr { left: 2, right: 5 }; assert_eq!(intent.input_cols(), vec![2, 5]); assert_eq!(intent.input_col(), None); assert!(agg_is_exact(&intent)); @@ -749,8 +735,7 @@ mod tests { assert_eq!(output.dtype, DataType::Float64); assert!(output.nullable); let value = serde_json::to_value(&intent).unwrap(); - assert_eq!(value["kind"], "bivariate"); - assert_eq!(value["op"], "correlation"); + assert_eq!(value["kind"], "pearson_corr"); assert_eq!(serde_json::from_value::(value).unwrap(), intent); } diff --git a/crates/types/src/pre_asap/binder.rs b/crates/types/src/pre_asap/binder.rs index 5351e494..55483841 100644 --- a/crates/types/src/pre_asap/binder.rs +++ b/crates/types/src/pre_asap/binder.rs @@ -378,14 +378,13 @@ mod tests { } } - // Both binary inputs must seed a usage-derived schema before positional resolution. + // Both correlation inputs must seed a usage-derived schema before positional resolution. #[test] - fn bivariate_inputs_seed_usage_derived_schema() { - use crate::pre_asap::{AggIntent, BivariateAggOp, Reduction}; + fn pearson_corr_inputs_seed_usage_derived_schema() { + use crate::pre_asap::{AggIntent, Reduction}; let tree = UnresolvedQueryExpr::Aggregate { reduction: Reduction::by(vec![]), - measures: vec![AggIntent::Bivariate { - op: BivariateAggOp::Correlation, + measures: vec![AggIntent::PearsonCorr { left: ColumnRef::Named("x".into()), right: ColumnRef::Named("y".into()), }], diff --git a/crates/types/src/pre_asap/mod.rs b/crates/types/src/pre_asap/mod.rs index da2a2d8f..d6c47998 100644 --- a/crates/types/src/pre_asap/mod.rs +++ b/crates/types/src/pre_asap/mod.rs @@ -43,7 +43,7 @@ pub mod schema; pub use agg_intent::{ agg_accuracy, agg_is_exact, agg_is_mergeable, default_cardinality, default_quantile, AggIntent, - BivariateAggOp, MathFunc, TimeFunc, + MathFunc, TimeFunc, }; pub use binder::{Binder, SchemaCatalog, UsageDerivedCatalog}; pub use canonicalize::canonicalize; diff --git a/crates/types/src/pre_asap/resolve.rs b/crates/types/src/pre_asap/resolve.rs index 0977235f..5d1e722c 100644 --- a/crates/types/src/pre_asap/resolve.rs +++ b/crates/types/src/pre_asap/resolve.rs @@ -518,8 +518,7 @@ fn resolve_agg_intent( AggIntent::Count { accuracy } => AggIntent::Count { accuracy: accuracy.clone(), }, - AggIntent::Bivariate { op, left, right } => AggIntent::Bivariate { - op: *op, + AggIntent::PearsonCorr { left, right } => AggIntent::PearsonCorr { left: resolve_column_ref(left, schema)?, right: resolve_column_ref(right, schema)?, }, @@ -616,14 +615,13 @@ mod tests { // Both sides resolve with qualifiers; an unknown right input is an error. #[test] - fn resolve_bivariate_aggregate_inputs() { - use crate::pre_asap::{BivariateAggOp, Column, DataType}; + fn resolve_pearson_corr_inputs() { + use crate::pre_asap::{Column, DataType}; let schema = Schema::new(vec![ Column::new("x", DataType::Float64, true).with_table("a"), Column::new("x", DataType::Float64, true).with_table("b"), ]); - let intent = AggIntent::Bivariate { - op: BivariateAggOp::Correlation, + let intent = AggIntent::PearsonCorr { left: ColumnRef::Qualified { table: "a".into(), name: "x".into(), @@ -635,14 +633,9 @@ mod tests { }; assert_eq!( resolve_agg_intent(&intent, &schema).unwrap(), - AggIntent::Bivariate { - op: BivariateAggOp::Correlation, - left: 0, - right: 1, - } + AggIntent::PearsonCorr { left: 0, right: 1 } ); - let missing = AggIntent::Bivariate { - op: BivariateAggOp::Correlation, + let missing = AggIntent::PearsonCorr { left: ColumnRef::Qualified { table: "a".into(), name: "x".into(), diff --git a/docs/design_docs/pre-asap-ir.md b/docs/design_docs/pre-asap-ir.md index 9f00e874..4bb52922 100644 --- a/docs/design_docs/pre-asap-ir.md +++ b/docs/design_docs/pre-asap-ir.md @@ -104,7 +104,7 @@ native-histogram accessors — this list is representative, not exhaustive: ```text Count, Sum(col), Min(col), Max(col), Avg(col), StdDev(col), Variance(col), -Quantile(col, q), TopK(k), Cardinality(col), Bivariate(op, left, right) // data-model-agnostic +Quantile(col, q), TopK(k), Cardinality(col), PearsonCorr(left, right) // data-model-agnostic Rate, Increase // counter derivatives Changes, Delta, IDelta, Deriv, Resets, PredictLinear(seconds), DoubleExpSmoothing(sf, tf) // range-vector functions @@ -113,20 +113,21 @@ HistogramStdVar, HistogramFraction(lo, hi), HistogramQuantile(q) // native-hist Math(func) // element-wise transform ``` -`Bivariate { op, left, right }` represents an aggregate over two value columns; -`BivariateAggOp::Correlation` is its first operation. Both references resolve to -positional column IDs, and `input_cols()` exposes both dependencies. SQL lowering -projects both arguments, preserving expressions, casts, and qualified join columns. -The result of correlation is nullable `Float64`, with pairwise null handling owned -by the executing engine. It remains exact: finalized correlation coefficients -cannot be combined as scalar rollups, and no sketch or maintained correlation -accumulator is selected. Physical costing accepts it as a hash aggregate with -provider-supplied accumulator size. - -Further two-input operations can extend `BivariateAggOp` without adding another -argument-carrying aggregate variant. Each needs explicit output and realization -semantics. SQL `corr` currently rejects `DISTINCT`, aggregate `FILTER`, aggregate -`ORDER BY`, explicit null treatment, and window usage (`OVER`). +`PearsonCorr { left, right }` is the only measure with two value inputs. Both +references resolve to positional column IDs, and `input_cols()` exposes both +dependencies — `input_col()` returns `None`, so single-column consumers cannot +pick up half of the pair. SQL lowering projects both arguments, preserving +expressions, casts, and qualified join columns. The result is nullable `Float64`, +with pairwise null handling owned by the executing engine. It remains exact: +finalized correlation coefficients cannot be combined as scalar rollups, and no +sketch or maintained correlation accumulator is selected. Physical costing accepts +it as a hash aggregate with provider-supplied accumulator size. + +SQL `corr` currently rejects `DISTINCT`, aggregate `FILTER`, aggregate `ORDER BY`, +explicit null treatment, and window usage (`OVER`). Further two-input statistics +(`covar`, the `regr_*` family) would each add their own variant, following the +`StdDev` / `Variance` precedent, once they have explicit output and realization +semantics. Additional measures can be added when there is a stable semantic distinction and a meaningful summary implementation.