feat: support corr with a general bivariate aggregate type

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

@@ -1187,6 +1187,7 @@ fn supports_hash_aggregate(
                     | AggIntent::Avg { .. }
                     | AggIntent::StdDev { .. }
                     | AggIntent::Variance { .. }
+                    | AggIntent::PearsonCorr { .. }
                     | AggIntent::Group
                     | AggIntent::CountValues { .. }
             )
@@ -1443,6 +1444,56 @@ 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, 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::PearsonCorr { 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};

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::PearsonCorr { .. } => {
             vec![Implementation::PassThrough]
         }
 
@@ -6062,6 +6065,17 @@ mod tests {
         }
     }
 
+    // Correlation must never acquire a single-input sketch or scalar accumulator.
+    #[test]
+    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]
+        ));
+        assert!(summary_candidates(&intent).is_empty());
+    }
+
     #[test]
     fn accuracy_target_drives_the_boundary() {
         // Same intent, three targets → three different decisions.

crates/frontend-sql/src/sql/mod.rs

@@ -1653,6 +1653,25 @@ fn lower_agg_intent(expr: &Expr) -> Result<AggIntent<ColumnRef>, 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::PearsonCorr {
+                        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(),
@@ -2051,6 +2070,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;

crates/frontend-sql/tests/data_quality_check/tpch_deequ.rs

@@ -91,13 +91,11 @@ async fn lowers_the_warehouse_ingestion_check_set() {
             Err(error) => panic!("unexpected failure for U-{id}: {error}"),
         }
     }
+    // P4d and P4l are Pearson correlation checks; they lower since `corr`
+    // became `AggIntent::PearsonCorr`.
     assert_eq!(
         rejected,
-        vec![
-            ("P2b", "multi-column COUNT(DISTINCT)".into()),
-            ("P4d", "corr".into()),
-            ("P4l", "corr".into()),
-        ]
+        vec![("P2b", "multi-column COUNT(DISTINCT)".into())]
     );
-    assert_eq!(lowered, 47);
+    assert_eq!(lowered, 49);
 }

crates/frontend-sql/tests/pearson_corr.rs

@@ -0,0 +1,136 @@
+//! 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, 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::PearsonCorr { 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::PearsonCorr { .. }))
+        .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();
+}

crates/frontend-sql/tests/sql_lowering.rs

@@ -2455,6 +2455,16 @@ 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);
+}
+
 // A multi-column DISTINCT must not silently count only the first column.
 #[tokio::test]
 async fn composite_distinct_is_rejected() {

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 with no `AggIntent`
+    // variant; `AggIntent::PearsonCorr` covers `corr` only.
     "covar",
     "covar_pop",
     "covar_samp",