diff --git a/crates/polars-sql/src/sql_visitors.rs b/crates/polars-sql/src/sql_visitors.rs index 469b372a847c..d689c239e67b 100644 --- a/crates/polars-sql/src/sql_visitors.rs +++ b/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(()) diff --git a/crates/polars-sql/src/subquery.rs b/crates/polars-sql/src/subquery.rs index c7867f65c6e6..462e5c044be3 100644 --- a/crates/polars-sql/src/subquery.rs +++ b/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> { + 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, 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) diff --git a/py-polars/tests/unit/sql/test_subqueries.py b/py-polars/tests/unit/sql/test_subqueries.py index dc7040c1348b..1deef3f1c82a 100644 --- a/py-polars/tests/unit/sql/test_subqueries.py +++ b/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()