fix: keep a correlated filter below an aggregate with a grouping set

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<Column>,
+    ) -> bool {
+        let grouping_sets = group_expr
+            .iter()
+            .filter_map(|expr| match expr {
+                Expr::GroupingSet(grouping_set) => Some(grouping_set),
+                _ => None,
+            })
+            .collect::<Vec<_>>();
+        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::<BTreeSet<_>>();
+        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],

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 <name> WHERE
+    /// c = test.c GROUP BY <group_expr>`.
+    fn correlated_grouping_set_subquery(
+        name: &str,
+        group_expr: Expr,
+    ) -> Result<Arc<LogicalPlan>> {
+        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::<Expr>::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.
+    /// <https://github.com/apache/datafusion/issues/25519>
+    #[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 (<subquery>) [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.
+    /// <https://github.com/apache/datafusion/issues/25519>
+    #[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::<Expr>::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 (<subquery>) [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.
+    /// <https://github.com/apache/datafusion/issues/25519>
+    #[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 (<subquery>) [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.
+    /// <https://github.com/apache/datafusion/issues/25519>
+    #[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<()> {

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.
+    /// <https://github.com/apache/datafusion/issues/25519>
+    #[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 = (<subquery>) [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<()> {