diff --git a/datafusion/optimizer/src/decorrelate.rs b/datafusion/optimizer/src/decorrelate.rs index 0c37f00b6435..a288ded88b10 100644 --- a/datafusion/optimizer/src/decorrelate.rs +++ b/datafusion/optimizer/src/decorrelate.rs @@ -28,7 +28,7 @@ use datafusion_common::tree_node::{ use datafusion_common::{ Column, DFSchemaRef, HashMap, Result, ScalarValue, assert_or_internal_err, plan_err, }; -use datafusion_expr::expr::Alias; +use datafusion_expr::expr::{Alias, GroupingSet}; use datafusion_expr::simplify::SimplifyContext; use datafusion_expr::utils::{ collect_subquery_cols, conjunction, find_join_exprs, split_conjunction, @@ -299,11 +299,51 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr { &self.correlated_subquery_cols_map, &mut local_correlated_cols, ); - // add missing columns to Aggregation's group expressions - let mut missing_exprs = self.collect_missing_exprs( - &aggregate.group_expr, - &local_correlated_cols, - )?; + + // A grouping set cannot take the columns the pull up adds. + // `LogicalPlanBuilder::aggregate` cross joins a plain group + // expression with the sets that are already there, so `ROLLUP(i.k)`, + // which is `GROUPING SETS ((i.k), ())`, becomes + // `GROUPING SETS ((i.k), (i.k, i.k))`. The empty set is gone, and + // with it the grand total row the subquery returns for every outer + // row, including the rows whose correlated filter matches nothing. + // The join that replaces the filter cannot bring those rows back, + // so the subquery stays correlated unless every set already groups + // by each column the pull up would add. + let mut missing_exprs = if aggregate + .group_expr + .iter() + .any(|expr| matches!(expr, Expr::GroupingSet(_))) + { + if self.grouping_sets_cover_pull_up_cols( + &aggregate.group_expr, + &local_correlated_cols, + ) { + // Every set already groups by them, so the sets stay as + // they are. Adding the columns again would repeat them + // inside every set. + aggregate.group_expr.to_vec() + } else { + self.can_pull_up = false; + // The rewrite still runs, the same way the + // `can_pull_over_aggregation` case above does. The callers + // read `can_pull_up` only after the whole rewrite has + // finished, and the nodes above this one still expect the + // pulled up columns in its output, so leaving them out here + // would fail the rewrite with a schema error instead. They + // drop this plan and keep the correlated subquery. + self.collect_missing_exprs( + &aggregate.group_expr, + &local_correlated_cols, + )? + } + } else { + // add missing columns to Aggregation's group expressions + self.collect_missing_exprs( + &aggregate.group_expr, + &local_correlated_cols, + )? + }; // if the original group expressions are empty, need to handle the Count bug let mut expr_result_map_for_count_bug = HashMap::new(); @@ -404,6 +444,64 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr { } impl PullUpCorrelatedExpr { + /// Whether the pull up can add its columns to `group_expr` without changing + /// what the aggregate returns. + /// + /// `true` when `group_expr` holds no grouping set, and when every set of every + /// grouping set it holds already groups by each column + /// [`Self::collect_missing_exprs`] would add. In the second case the pull up + /// adds nothing and the aggregate keeps the sets it has. + /// + /// `ROLLUP` and `CUBE` always contain the empty set, which yields a row for + /// outer rows the correlated filter matches nothing for, so they are only safe + /// when there is nothing to add. + /// + /// A non-empty set that leaves a column out fills it with NULL. Adding the + /// column would give it a value instead, which a `HAVING` or a projection + /// above the aggregate can read, so such a set is rejected as well. + fn grouping_sets_cover_pull_up_cols( + &self, + group_expr: &[Expr], + correlated_subquery_cols: &BTreeSet, + ) -> bool { + let grouping_sets = group_expr + .iter() + .filter_map(|expr| match expr { + Expr::GroupingSet(grouping_set) => Some(grouping_set), + _ => None, + }) + .collect::>(); + if grouping_sets.is_empty() { + return true; + } + + // The same columns `collect_missing_exprs` appends: the correlated columns + // and the columns of a pulled up HAVING, minus the ones `group_expr` + // already lists on their own, which it leaves alone. + let mut required_cols = correlated_subquery_cols.iter().collect::>(); + if let Some(pull_up_having) = &self.pull_up_having_expr { + required_cols.extend(pull_up_having.column_refs()); + } + required_cols.retain(|col| { + !group_expr + .iter() + .any(|expr| matches!(expr, Expr::Column(c) if c == *col)) + }); + if required_cols.is_empty() { + return true; + } + + grouping_sets.iter().all(|grouping_set| match grouping_set { + GroupingSet::Rollup(_) | GroupingSet::Cube(_) => false, + GroupingSet::GroupingSets(sets) => sets.iter().all(|set| { + required_cols.iter().all(|col| { + set.iter() + .any(|expr| matches!(expr, Expr::Column(c) if c == *col)) + }) + }), + }) + } + fn collect_missing_exprs( &self, exprs: &[Expr], diff --git a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs index 0ad8b44c40de..d4eafaf1fa4b 100644 --- a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs +++ b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs @@ -713,7 +713,9 @@ mod tests { use crate::assert_optimized_plan_eq_display_indent_snapshot; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_expr::builder::table_source; - use datafusion_expr::{and, binary_expr, col, out_ref_col, table_scan}; + use datafusion_expr::{ + and, binary_expr, col, cube, grouping_set, out_ref_col, rollup, table_scan, + }; macro_rules! assert_optimized_plan_equal { ( @@ -775,6 +777,138 @@ mod tests { optimizer.optimize(plan, &crate::OptimizerContext::new(), |_, _| {}) } + /// A grouping set subquery for the tests below: `SELECT c FROM WHERE + /// c = test.c GROUP BY `. + fn correlated_grouping_set_subquery( + name: &str, + group_expr: Expr, + ) -> Result> { + Ok(Arc::new( + LogicalPlanBuilder::from(test_table_scan_with_name(name)?) + .filter( + col(format!("{name}.c")).eq(out_ref_col(DataType::UInt32, "test.c")), + )? + .aggregate(vec![group_expr], Vec::::new())? + .project(vec![col(format!("{name}.c"))])? + .build()?, + )) + } + + /// `ROLLUP(c)` is `GROUPING SETS ((c), ())`. Adding the correlated column to + /// every set drops the empty one, so the subquery is left correlated. + /// + #[test] + fn exists_subquery_with_rollup_is_not_decorrelated() -> Result<()> { + let subquery = correlated_grouping_set_subquery("sq", rollup(vec![col("sq.c")]))?; + let plan = LogicalPlanBuilder::from(test_table_scan()?) + .filter(exists(subquery))? + .project(vec![col("test.b")])? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Projection: test.b [b:UInt32] + Filter: EXISTS () [a:UInt32, b:UInt32, c:UInt32] + Subquery: [c:UInt32;N] + Projection: sq.c [c:UInt32;N] + Aggregate: groupBy=[[ROLLUP (sq.c)]], aggr=[[]] [c:UInt32;N, __grouping_id:UInt8] + Filter: sq.c = outer_ref(test.c) [a:UInt32, b:UInt32, c:UInt32] + TableScan: sq [a:UInt32, b:UInt32, c:UInt32] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + /// `CUBE(c)` holds the empty set for the same reason. The correlation is on + /// `a` rather than on the `IN` key, so it stays a filter of its own instead + /// of being folded into the `IN` predicate. + /// + #[test] + fn in_subquery_with_cube_is_not_decorrelated() -> Result<()> { + let subquery = Arc::new( + LogicalPlanBuilder::from(test_table_scan_with_name("sq")?) + .filter(col("sq.a").eq(out_ref_col(DataType::UInt32, "test.a")))? + .aggregate(vec![cube(vec![col("sq.c")])], Vec::::new())? + .project(vec![col("sq.c")])? + .build()?, + ); + let plan = LogicalPlanBuilder::from(test_table_scan()?) + .filter(in_subquery(col("test.c"), subquery))? + .project(vec![col("test.b")])? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Projection: test.b [b:UInt32] + Filter: test.c IN () [a:UInt32, b:UInt32, c:UInt32] + Subquery: [c:UInt32;N] + Projection: sq.c [c:UInt32;N] + Aggregate: groupBy=[[CUBE (sq.c)]], aggr=[[]] [c:UInt32;N, __grouping_id:UInt8] + Filter: sq.a = outer_ref(test.a) [a:UInt32, b:UInt32, c:UInt32] + TableScan: sq [a:UInt32, b:UInt32, c:UInt32] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + /// A set that groups by another column does not carry the correlated one. + /// + #[test] + fn exists_subquery_with_partial_grouping_set_is_not_decorrelated() -> Result<()> { + let subquery = correlated_grouping_set_subquery( + "sq", + grouping_set(vec![vec![col("sq.c")], vec![col("sq.b")]]), + )?; + let plan = LogicalPlanBuilder::from(test_table_scan()?) + .filter(exists(subquery))? + .project(vec![col("test.b")])? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Projection: test.b [b:UInt32] + Filter: EXISTS () [a:UInt32, b:UInt32, c:UInt32] + Subquery: [c:UInt32;N] + Projection: sq.c [c:UInt32;N] + Aggregate: groupBy=[[GROUPING SETS ((sq.c), (sq.b))]], aggr=[[]] [c:UInt32;N, b:UInt32;N, __grouping_id:UInt8] + Filter: sq.c = outer_ref(test.c) [a:UInt32, b:UInt32, c:UInt32] + TableScan: sq [a:UInt32, b:UInt32, c:UInt32] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + /// Every set already groups by the correlated column, so the pull up adds + /// nothing and the subquery decorrelates as it did before. + /// + #[test] + fn exists_subquery_with_covering_grouping_set_is_decorrelated() -> Result<()> { + let subquery = correlated_grouping_set_subquery( + "sq", + grouping_set(vec![vec![col("sq.c")], vec![col("sq.c"), col("sq.b")]]), + )?; + let plan = LogicalPlanBuilder::from(test_table_scan()?) + .filter(exists(subquery))? + .project(vec![col("test.b")])? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Projection: test.b [b:UInt32] + LeftSemi Join: Filter: __correlated_sq_1.c = test.c [a:UInt32, b:UInt32, c:UInt32] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + SubqueryAlias: __correlated_sq_1 [c:UInt32;N] + Projection: sq.c [c:UInt32;N] + Aggregate: groupBy=[[GROUPING SETS ((sq.c), (sq.c, sq.b))]], aggr=[[]] [c:UInt32;N, b:UInt32;N, __grouping_id:UInt8] + TableScan: sq [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + /// Test for several IN subquery expressions #[test] fn in_subquery_multiple() -> Result<()> { diff --git a/datafusion/optimizer/src/scalar_subquery_to_join.rs b/datafusion/optimizer/src/scalar_subquery_to_join.rs index 44011a125ba9..ed5e0a6def1c 100644 --- a/datafusion/optimizer/src/scalar_subquery_to_join.rs +++ b/datafusion/optimizer/src/scalar_subquery_to_join.rs @@ -445,7 +445,7 @@ mod tests { use datafusion_expr::test::function_stub::sum; use crate::assert_optimized_plan_eq_display_indent_snapshot; - use datafusion_expr::{Between, col, expr, out_ref_col, scalar_subquery}; + use datafusion_expr::{Between, col, expr, out_ref_col, rollup, scalar_subquery}; use datafusion_functions_aggregate::min_max::{max, min}; macro_rules! assert_optimized_plan_equal { @@ -462,6 +462,46 @@ mod tests { }}; } + /// A correlated scalar subquery whose aggregate uses `ROLLUP` keeps its + /// correlation: the empty set yields a row for outer rows the filter matches + /// nothing for, and the join that would replace the filter cannot produce it. + /// + #[test] + fn scalar_subquery_with_rollup_is_not_decorrelated() -> Result<()> { + let sq = Arc::new( + LogicalPlanBuilder::from(scan_tpch_table("orders")) + .filter( + col("orders.o_custkey") + .eq(out_ref_col(DataType::Int64, "customer.c_custkey")), + )? + .aggregate( + vec![rollup(vec![col("orders.o_custkey")])], + vec![max(col("orders.o_custkey"))], + )? + .project(vec![max(col("orders.o_custkey"))])? + .build()?, + ); + + let plan = LogicalPlanBuilder::from(scan_tpch_table("customer")) + .filter(col("customer.c_custkey").eq(scalar_subquery(sq)))? + .project(vec![col("customer.c_custkey")])? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Projection: customer.c_custkey [c_custkey:Int64] + Filter: customer.c_custkey = () [c_custkey:Int64, c_name:Utf8] + Subquery: [max(orders.o_custkey):Int64;N] + Projection: max(orders.o_custkey) [max(orders.o_custkey):Int64;N] + Aggregate: groupBy=[[ROLLUP (orders.o_custkey)]], aggr=[[max(orders.o_custkey)]] [o_custkey:Int64;N, __grouping_id:UInt8, max(orders.o_custkey):Int64;N] + Filter: orders.o_custkey = outer_ref(customer.c_custkey) [o_orderkey:Int64, o_custkey:Int64, o_orderstatus:Utf8, o_totalprice:Float64;N] + TableScan: orders [o_orderkey:Int64, o_custkey:Int64, o_orderstatus:Utf8, o_totalprice:Float64;N] + TableScan: customer [c_custkey:Int64, c_name:Utf8] + " + ) + } + /// Test multiple correlated subqueries #[test] fn multiple_subqueries() -> Result<()> { diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 626ef60762b9..3c052a4afffb 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -2744,3 +2744,107 @@ b 400 statement ok DROP TABLE metrics; + +# Regression test for #25519: a correlated filter that sits below an aggregate +# with a grouping set must not be pulled above it. The pull up adds the +# correlated column to every set, so `ROLLUP(k)`, which is +# `GROUPING SETS ((k), ())`, turns into `GROUPING SETS ((k), (k, k))`. The empty +# set is gone, and with it the grand total row the subquery returns for every +# outer row, including the rows whose filter matches nothing. +statement ok +CREATE TABLE gs_outer(k INT) AS VALUES (1), (2), (NULL), (4), (5); + +statement ok +CREATE TABLE gs_inner(k INT, j INT) AS VALUES (1, 10), (NULL, 20), (5, 30), (2, 40); + +# ROLLUP holds the empty set, so the subquery stays correlated. +statement error DataFusion error: This feature is not implemented: Physical plan does not support logical expression Exists +SELECT gs_outer.k, EXISTS (SELECT 1 FROM gs_inner WHERE gs_inner.k = gs_outer.k GROUP BY ROLLUP(gs_inner.k)) FROM gs_outer; + +# So does CUBE. +statement error DataFusion error: This feature is not implemented: Physical plan does not support logical expression Exists +SELECT gs_outer.k, EXISTS (SELECT 1 FROM gs_inner WHERE gs_inner.k = gs_outer.k GROUP BY CUBE(gs_inner.k)) FROM gs_outer; + +# And an explicit grouping set that lists the empty set. +statement error DataFusion error: This feature is not implemented: Physical plan does not support logical expression Exists +SELECT gs_outer.k, EXISTS (SELECT 1 FROM gs_inner WHERE gs_inner.k = gs_outer.k GROUP BY GROUPING SETS ((gs_inner.k), ())) FROM gs_outer; + +# A set that leaves out the correlated column fills it with NULL. The pull up +# would turn `(j)` into `(j, k)`, and `k` would then carry a value in the rows +# where the set fills it with NULL. Anything above the aggregate that reads `k` +# sees the difference, so the subquery stays correlated. +# +# Known limitation: when nothing reads `k`, as here, the pull up was correct +# before this guard and the query now fails to plan. Telling the two cases apart +# needs the correlated column added to each set under an alias. +statement error DataFusion error: This feature is not implemented: Physical plan does not support logical expression Exists +SELECT gs_outer.k, EXISTS (SELECT 1 FROM gs_inner WHERE gs_inner.k = gs_outer.k GROUP BY GROUPING SETS ((gs_inner.k), (gs_inner.j))) FROM gs_outer; + +# The same sets with a HAVING that reads the NULL filled column. For k = 1 the +# `(j)` set yields the row `(NULL, 10)`, which passes the HAVING, so EXISTS is +# true; with `(j, k)` that row holds k = 1 and is filtered out. On main this +# query answers false for every row. +statement error DataFusion error: This feature is not implemented: Physical plan does not support logical expression Exists +SELECT gs_outer.k, EXISTS (SELECT 1 FROM gs_inner WHERE gs_inner.k = gs_outer.k GROUP BY GROUPING SETS ((gs_inner.k), (gs_inner.j)) HAVING gs_inner.k IS NULL) FROM gs_outer; + +# The correlated column does not have to appear in the grouping set at all. +statement error DataFusion error: This feature is not implemented: Physical plan does not support logical expression Exists +SELECT gs_outer.k, EXISTS (SELECT 1 FROM gs_inner WHERE gs_inner.k = gs_outer.k GROUP BY ROLLUP(gs_inner.j)) FROM gs_outer; + +# The `IN` form of the same subquery. +statement error DataFusion error: This feature is not implemented: Physical plan does not support logical expression InSubquery +SELECT gs_outer.k, gs_outer.k IN (SELECT gs_inner.k FROM gs_inner WHERE gs_inner.k = gs_outer.k GROUP BY ROLLUP(gs_inner.k)) FROM gs_outer; + +# `NOT EXISTS` too. +statement error DataFusion error: This feature is not implemented: Physical plan does not support logical expression Exists +SELECT gs_outer.k, NOT EXISTS (SELECT 1 FROM gs_inner WHERE gs_inner.k = gs_outer.k GROUP BY ROLLUP(gs_inner.k)) FROM gs_outer; + +# And a correlated scalar subquery. +statement error DataFusion error: This feature is not implemented: Physical plan does not support logical expression ScalarSubquery +SELECT gs_outer.k, (SELECT count(*) FROM gs_inner WHERE gs_inner.k = gs_outer.k GROUP BY ROLLUP(gs_inner.k) LIMIT 1) FROM gs_outer; + +# When every set already groups by the correlated column the pull up adds +# nothing, so the subquery still decorrelates and keeps its results. +query IB +SELECT gs_outer.k, EXISTS (SELECT 1 FROM gs_inner WHERE gs_inner.k = gs_outer.k GROUP BY GROUPING SETS ((gs_inner.k))) FROM gs_outer ORDER BY gs_outer.k; +---- +1 true +2 true +4 false +5 true +NULL false + +query IB +SELECT gs_outer.k, EXISTS (SELECT 1 FROM gs_inner WHERE gs_inner.k = gs_outer.k GROUP BY GROUPING SETS ((gs_inner.k), (gs_inner.k, gs_inner.j))) FROM gs_outer ORDER BY gs_outer.k; +---- +1 true +2 true +4 false +5 true +NULL false + +# A plain GROUP BY is unaffected. +query IB +SELECT gs_outer.k, EXISTS (SELECT 1 FROM gs_inner WHERE gs_inner.k = gs_outer.k GROUP BY gs_inner.k) FROM gs_outer ORDER BY gs_outer.k; +---- +1 true +2 true +4 false +5 true +NULL false + +# An uncorrelated subquery keeps its grouping set. +query IB +SELECT gs_outer.k, EXISTS (SELECT 1 FROM gs_inner GROUP BY ROLLUP(gs_inner.k)) FROM gs_outer ORDER BY gs_outer.k; +---- +1 true +2 true +4 true +5 true +NULL true + +statement ok +DROP TABLE gs_outer; + +statement ok +DROP TABLE gs_inner;