new correlation aggregation supported

Author: Selvomega

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

@@ -1602,6 +1602,13 @@ fn lower_agg_intent(expr: &Expr) -> Result<AggIntent<ColumnRef>, LoweringError>
             if let Some(intent) = lower_arg_selector(&name, &agg_fn.args)? {
                 return Ok(intent);
             }
+            // `corr` IS a native DataFusion aggregate — what it lacks is an
+            // `AggSemantic`, since no core `AggIntent` reads two columns.
+            // Handled here for the same reason as the arg selectors above:
+            // the `lookup_native` lookup below has nothing to return for it.
+            if let Some(intent) = lower_corr(&name, &agg_fn.args)? {
+                return Ok(intent);
+            }
             let semantic = asap_sql_function_catalog::lookup_native(&name)
                 .ok_or_else(|| LoweringError::UnsupportedAggregate(name.clone()))?;
             // The canonical intent algebra has no DISTINCT modifier for the
@@ -1776,6 +1783,38 @@ fn lower_arg_selector(
     }))
 }
 
+/// Pearson correlation `corr(x, y)` — a two-column statistic.
+///
+/// Every core `AggIntent` reducer folds *one* column (`col: Option<C>`);
+/// correlation reads two and would be the first binary core variant. Per
+/// `AggIntent::Extension`'s own "core only grows for intents ≥2 deployment
+/// models actually use" bar, and a search that found no second model wanting
+/// it — PromQL has no correlation — this lowers to `Extension`, the same
+/// judgement `lower_arg_selector` records for `argMax`/`argMin`.
+///
+/// Both columns are kept as validated bare-column `ColumnRef`s in `payload`
+/// (`reducer_col`'s "no expression arguments" rule, issue #115); core carries
+/// them without resolving them, since `Extension` has no typed column field.
+/// The order is the order written: a later rewrite into co-moments has to tell
+/// x from y.
+fn lower_corr(name: &str, args: &[Expr]) -> Result<Option<AggIntent<ColumnRef>>, LoweringError> {
+    if name != "corr" {
+        return Ok(None);
+    }
+    let [x, y] = args else {
+        unreachable!(
+            "corr's DataFusion signature fixes its arity at 2 -- the planner already rejected              any other argument count before lower_agg_intent runs"
+        );
+    };
+    Ok(Some(AggIntent::Extension {
+        ext_kind: "corr".to_string(),
+        payload: serde_json::json!({
+            "x_col": reducer_col(name, std::slice::from_ref(x))?,
+            "y_col": reducer_col(name, std::slice::from_ref(y))?,
+        }),
+    }))
+}
+
 /// The name an `IN (subquery)`'s key column is projected under, so the join
 /// predicate cannot bind it to a same-named column of the outer relation.
 const IN_SUBQUERY_KEY: &str = "__asap_in_key";

crates/frontend-sql/tests/sql_lowering.rs

@@ -2454,3 +2454,91 @@ async fn clickhouse_tuple_element_preserves_declared_field_metadata() {
         );
     }
 }
+
+// ── `corr(x, y)`: a two-column statistic with no core `AggIntent` shape.
+// Every core reducer folds ONE column (`col: Option<C>`); correlation reads
+// two and would be the first binary core variant. A repo-wide search turned up
+// no second deployment model wanting it — PromQL has no correlation — so per
+// `AggIntent::Extension`'s own "core only grows for intents ≥2 deployment
+// models actually use" bar it lowers to an `Extension`, exactly as `argMax`
+// does. Unlike `argMax` it IS a native DataFusion aggregate; what is missing is
+// an `AggSemantic` for it, which is why `lookup_native` rejected it. ─────────
+
+#[tokio::test]
+async fn corr_lowers_to_an_extension_intent() {
+    let qe = lower("SELECT corr(latency, bytes) AS r FROM metrics").await;
+    let (by, measures) = find_aggregate(&qe).expect("expected an Aggregate");
+    assert!(by.is_empty());
+    assert!(
+        matches!(
+            measures.as_slice(),
+            [AggIntent::Extension { ext_kind, .. }] if ext_kind == "corr"
+        ),
+        "expected Extension {{ ext_kind: \"corr\", .. }}, got {measures:?}"
+    );
+}
+
+#[tokio::test]
+async fn corr_payload_preserves_both_column_names() {
+    // Core never resolves an `Extension`'s payload, so both columns stay as
+    // validated bare-column `ColumnRef`s — and in the order written, since
+    // a later rewrite into co-moments needs to tell the two apart.
+    let qe = lower("SELECT corr(latency, bytes) AS r FROM metrics").await;
+    let (_, measures) = find_aggregate(&qe).expect("expected an Aggregate");
+    let AggIntent::Extension { payload, .. } = &measures[0] else {
+        panic!("expected an Extension intent, got {:?}", measures[0]);
+    };
+    let named = |key: &str| {
+        payload
+            .get(key)
+            .and_then(|c| c.get("Named"))
+            .and_then(|n| n.as_str())
+            .map(str::to_string)
+    };
+    assert_eq!(named("x_col"), Some("latency".to_string()));
+    assert_eq!(named("y_col"), Some("bytes".to_string()));
+}
+
+#[tokio::test]
+async fn corr_over_an_expression_binds_the_derived_column() {
+    // `reducer_col`'s "bare column only" rule (issue #115) exists so an
+    // expression argument is never silently DROPPED. Here nothing is dropped:
+    // `corr` is a native DataFusion aggregate, so the planner materializes
+    // `latency * 2` into the Project below the Aggregate and hands the
+    // aggregate a column reference to it. The payload names that derived
+    // column, which resolves in the child's own output schema — so the
+    // expression is carried, not lost.
+    //
+    // `argMax` behaves differently only because it reaches `lower_agg_intent`
+    // as a stub ClickHouse builtin, before that rewrite applies.
+    let qe = lower("SELECT corr(latency * 2, bytes) AS r FROM metrics").await;
+    let (_, measures) = find_aggregate(&qe).expect("expected an Aggregate");
+    let AggIntent::Extension { payload, .. } = &measures[0] else {
+        panic!("expected an Extension intent, got {:?}", measures[0]);
+    };
+    let x = payload
+        .get("x_col")
+        .and_then(|c| c.get("Named"))
+        .and_then(|n| n.as_str())
+        .expect("x_col is a Named ColumnRef");
+    assert!(
+        x.contains('*'),
+        "expected the derived column carrying `latency * 2`, got {x:?}"
+    );
+}
+
+#[tokio::test]
+async fn corr_output_column_is_a_nullable_float() {
+    // `Extension`'s generic guess is `Utf8` — correct only for a shape core
+    // knows nothing about. Correlation always yields a float, and NULL when
+    // a variance is zero or fewer than two rows contributed, so the schema
+    // says so rather than carrying the placeholder downstream.
+    let qe = lower("SELECT corr(latency, bytes) AS r FROM metrics").await;
+    let schema = qe.output_schema().expect("output schema");
+    let out = schema.columns.last().expect("at least one output column");
+    assert_eq!(out.dtype, DataType::Float64, "got {:?}", schema.columns);
+    assert!(
+        out.nullable,
+        "corr is NULL on zero variance / fewer than 2 rows"
+    );
+}

crates/types/src/pre_asap/agg_intent.rs

@@ -568,6 +568,14 @@ impl<C: Clone> AggIntent<C> {
             // the column after `kind` and leave the type unconstrained.
             // The owning deployment model is expected to re-derive the
             // real schema itself rather than rely on this generic guess.
+            // Correlation always yields a float, and NULL when a variance is
+            // zero or fewer than two rows contributed — unlike the arg
+            // selectors, whose output type follows the selected column and is
+            // patched in during aggregate schema derivation, this one needs no
+            // schema and so is settled here.
+            AggIntent::Extension { ext_kind, .. } if ext_kind == "corr" => {
+                col("corr", DataType::Float64, true)
+            }
             AggIntent::Extension { ext_kind, .. } => col(ext_kind, DataType::Utf8, true),
         }
     }