fix: SQL bind a scalar subquery's aggregate to its own relation

crates/polars-sql/src/sql_visitors.rs

@@ -340,13 +340,18 @@ pub(crate) fn expr_references_any_column(expr: &SQLExpr) -> bool {
     expr.visit(&mut ColumnRefFinder).is_break()
 }
 
+/// Whether this expression node is itself a subquery, in any of its forms.
+pub(crate) fn is_subquery_expr(expr: &SQLExpr) -> bool {
+    matches!(
+        expr,
+        SQLExpr::Subquery(_) | SQLExpr::Exists { .. } | SQLExpr::InSubquery { .. }
+    )
+}
+
 /// Check if a SQL expression contains a subquery, in any of its forms.
 pub(crate) fn expr_contains_subquery(expr: &SQLExpr) -> bool {
     visit_expressions(expr, |e| {
-        if matches!(
-            e,
-            SQLExpr::Subquery(_) | SQLExpr::Exists { .. } | SQLExpr::InSubquery { .. }
-        ) {
+        if is_subquery_expr(e) {
             ControlFlow::Break(())
         } else {
             ControlFlow::Continue(())

crates/polars-sql/src/subquery.rs

@@ -21,7 +21,7 @@ use sqlparser::ast::{
 use crate::SQLContext;
 use crate::context::{CORRELATED_COL_PREFIX, FilterMode, get_table_name};
 use crate::sql_expr::{parse_sql_expr, sql_in_membership};
-use crate::sql_visitors::expr_contains_subquery;
+use crate::sql_visitors::{expr_contains_subquery, is_subquery_expr};
 
 impl SQLContext {
     // Entry point: offer each WHERE conjunct to the rewrite, returning the
@@ -274,15 +274,9 @@ impl SQLContext {
             return Ok(None);
         };
 
-        let left_key = parse_sql_expr(lhs, self, Some(outer_schema))?
-            .meta()
-            .undo_aliases();
-        if has_expr(&left_key, |e| matches!(e, Expr::SubPlan(_, _)))
-            || !expr_to_leaf_column_names_iter(&left_key)
-                .all(|name| outer_schema.contains(name.as_str()))
-        {
+        let Some(left_key) = self.try_parse_outer_only_expr(lhs, outer_schema)? else {
             return Ok(None);
-        }
+        };
 
         let mut ctx = self.isolated();
         let Some((inner_names, inner_lf, inner_schema)) =
@@ -419,6 +413,26 @@ impl SQLContext {
         }))
     }
 
+    // Parse an outer-query expression, or `None` if it doesn't stand on the outer
+    // relation alone: an unlowered subquery, or a column the outer frame lacks.
+    // Any alias is cosmetic here and stripped.
+    fn try_parse_outer_only_expr(
+        &mut self,
+        sql_expr: &SQLExpr,
+        outer_schema: &Schema,
+    ) -> PolarsResult<Option<Expr>> {
+        let expr = parse_sql_expr(sql_expr, self, Some(outer_schema))?
+            .meta()
+            .undo_aliases();
+        if has_expr(&expr, |e| matches!(e, Expr::SubPlan(_, _)))
+            || !expr_to_leaf_column_names_iter(&expr)
+                .all(|name| outer_schema.contains(name.as_str()))
+        {
+            return Ok(None);
+        }
+        Ok(Some(expr))
+    }
+
     // Parse a subquery expression as one over the inner relation only, or `None`
     // if it references any outer column (a correlation shape we don't handle) or
     // contains a nested subquery.
@@ -509,10 +523,11 @@ impl SQLContext {
         };
 
         // The projection must be a scalar aggregate over the inner relation.
-        let agg_expr = parse_sql_expr(proj, &mut ctx, Some(&inner_schema))?;
-        if has_expr(&agg_expr, |e| matches!(e, Expr::SubPlan(_, _)))
-            || !has_expr(&agg_expr, |e| matches!(e, Expr::Agg(_) | Expr::Len))
-        {
+        let Some(agg_expr) = ctx.try_parse_inner_only_expr(proj, &inner_names, &inner_schema)?
+        else {
+            return Ok(None);
+        };
+        if !has_expr(&agg_expr, |e| matches!(e, Expr::Agg(_) | Expr::Len)) {
             return Ok(None);
         }
         let count_like = matches!(
@@ -632,15 +647,9 @@ impl SQLContext {
             return Ok(None);
         };
 
-        let needle = parse_sql_expr(lhs, self, Some(outer_schema))?
-            .meta()
-            .undo_aliases();
-        if has_expr(&needle, |e| matches!(e, Expr::SubPlan(_, _)))
-            || !expr_to_leaf_column_names_iter(&needle)
-                .all(|name| outer_schema.contains(name.as_str()))
-        {
+        let Some(needle) = self.try_parse_outer_only_expr(lhs, outer_schema)? else {
             return Ok(None);
-        }
+        };
 
         let mut ctx = self.isolated();
         let Some((inner_names, inner_lf, inner_schema)) =
@@ -1049,12 +1058,10 @@ fn binds_to_inner_relation(
     inner_names: &PlHashSet<String>,
     inner_schema: &Schema,
 ) -> bool {
-    // A nested subquery is its own scope, which this cannot resolve against.
-    if expr_contains_subquery(expr) {
-        return false;
-    }
     visit_expressions(expr, |e| {
         let resolves = match e {
+            // A nested subquery is its own scope, which this cannot resolve against.
+            _ if is_subquery_expr(e) => false,
             SQLExpr::Identifier(_) | SQLExpr::CompoundIdentifier(_) => qualifier_and_name(e)
                 .is_some_and(|(qualifier, name)| {
                     qualifier.is_none_or(|q| inner_names.contains(q)) && inner_schema.contains(name)

py-polars/tests/unit/sql/test_subqueries.py

@@ -770,3 +770,33 @@ def test_qualifier_naming_a_declared_relation(query: str) -> None:
         query=query,
         compare_with="duckdb",
     )
+
+
+def _scalar_subquery_frames() -> dict[str, pl.DataFrame]:
+    return {
+        "t1": pl.DataFrame({"k": [1, 2], "v": [10, 20]}),
+        "t2": pl.DataFrame({"b": [1, 2], "a": [100, 200]}),
+        "foo": pl.DataFrame({"a": [7, 7, 7]}),
+    }
+
+
+@pytest.mark.parametrize("aggregate", ["SUM(x.a)", "SUM(a)", "COUNT(*)", "MAX(x.a)"])
+def test_correlated_scalar_subquery_aggregate(aggregate: str) -> None:
+    assert_sql_matches(
+        frames=_scalar_subquery_frames(),
+        query=(
+            f"SELECT k, (SELECT {aggregate} FROM t2 x WHERE x.b = t1.k) AS s"
+            f" FROM t1 ORDER BY k"
+        ),
+        compare_with="duckdb",
+    )
+
+
+def test_correlated_scalar_subquery_rejects_foreign_qualifier() -> None:
+    # `foo` is registered but is not a relation of the subquery, so its column
+    # must not be read as one of the subquery's own.
+    ctx = pl.SQLContext(frames=_scalar_subquery_frames())
+    with pytest.raises(SQLInterfaceError, match="no table or struct column named"):
+        ctx.execute(
+            "SELECT k, (SELECT SUM(foo.a) FROM t2 x WHERE x.b = t1.k) AS s FROM t1"
+        ).collect()