fixes

Author: JanKaul

datafusion/optimizer/src/decorrelate_predicate_subquery.rs

@@ -37,8 +37,8 @@ use datafusion_expr::logical_plan::{
 };
 use datafusion_expr::utils::{conjunction, expr_to_columns, split_conjunction_owned};
 use datafusion_expr::{
-    BinaryExpr, Expr, Filter, LogicalPlan, LogicalPlanBuilder, Operator, exists,
-    in_subquery, lit, not, not_exists, not_in_subquery,
+    Aggregate, BinaryExpr, Expr, Filter, LogicalPlan, LogicalPlanBuilder, Operator,
+    exists, in_subquery, lit, not, not_exists, not_in_subquery,
 };
 
 use log::debug;
@@ -198,10 +198,9 @@ impl OptimizerRule for DecorrelatePredicateSubquery {
                 }
             }
 
-            let new_plan = LogicalPlanBuilder::from(cur_input)
-                .aggregate(new_group_exprs, new_aggr_exprs)?
-                .build()?;
-            return Ok(Transformed::yes(new_plan));
+            let new_agg =
+                Aggregate::try_new(Arc::new(cur_input), new_group_exprs, new_aggr_exprs)?;
+            return Ok(Transformed::yes(LogicalPlan::Aggregate(new_agg)));
         }
 
         // Handle Projection nodes with subqueries in expressions
@@ -678,45 +677,6 @@ mod tests {
         ))
     }
 
-    /// Aggregation with CASE WHEN ... IN (subquery) should be decorrelated under the Aggregate
-    #[test]
-    fn aggregate_case_in_subquery() -> Result<()> {
-        let table_scan = test_table_scan_with_name("distinct_source")?;
-        use datafusion_expr::expr_fn::when;
-        use datafusion_functions_aggregate::expr_fn::max as agg_max;
-
-        let agg_b: Expr = agg_max(col("distinct_source.b"));
-        let subq = LogicalPlanBuilder::from(table_scan.clone())
-            .aggregate(Vec::<Expr>::new(), vec![agg_b])?
-            .project(vec![col("max(distinct_source.b)")])?
-            .build()?;
-
-        let case_expr = when(
-            in_subquery(col("distinct_source.b"), Arc::new(subq)),
-            lit(1),
-        )
-        .otherwise(lit(0))?;
-
-        let plan = LogicalPlanBuilder::from(table_scan)
-            .aggregate(
-                vec![col("distinct_source.a").alias("primary_key")],
-                vec![
-                    agg_max(case_expr).alias("is_in_most_recent_task"),
-                    agg_max(col("distinct_source.c")).alias("max_timestamp"),
-                ],
-            )?
-            .build()?;
-
-        use crate::{OptimizerContext, OptimizerRule};
-        let optimized = DecorrelatePredicateSubquery::new()
-            .rewrite(plan, &OptimizerContext::new())?
-            .data;
-        let lp = optimized.display_indent().to_string();
-        assert!(lp.contains("Aggregate:"));
-        assert!(lp.contains("Left"));
-        Ok(())
-    }
-
     /// Test for several IN subquery expressions
     #[test]
     fn in_subquery_multiple() -> Result<()> {
@@ -834,9 +794,10 @@ mod tests {
           LeftMark Join:  Filter: a = __correlated_sq_1.ua [a:Int32;N, mark:Boolean]
             Projection: column1 AS a [a:Int32;N]
               Values: (Int32(1)), (Int32(2)) [column1:Int32;N]
-            SubqueryAlias: __correlated_sq_1 [ua:Int32;N]
-              Projection: column1 AS ua [ua:Int32;N]
-                Values: (Int32(2)) [column1:Int32;N]
+            Projection: __correlated_sq_1.ua [ua:Int32;N]
+              SubqueryAlias: __correlated_sq_1 [ua:Int32;N]
+                Projection: column1 AS ua [ua:Int32;N]
+                  Values: (Int32(2)) [column1:Int32;N]
         "
         )
     }
@@ -1924,14 +1885,13 @@ mod tests {
             plan,
             @r"
         Projection: customer.c_custkey [c_custkey:Int64]
-          Projection: customer.c_custkey, customer.c_name [c_custkey:Int64, c_name:Utf8]
-            Filter: __correlated_sq_1.mark OR customer.c_custkey = Int32(1) [c_custkey:Int64, c_name:Utf8, mark:Boolean]
-              LeftMark Join:  Filter: Boolean(true) [c_custkey:Int64, c_name:Utf8, mark:Boolean]
-                TableScan: customer [c_custkey:Int64, c_name:Utf8]
-                SubqueryAlias: __correlated_sq_1 [o_custkey:Int64]
-                  Projection: orders.o_custkey [o_custkey:Int64]
-                    Filter: customer.c_custkey = orders.o_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]
+          Filter: __correlated_sq_1.mark OR customer.c_custkey = Int32(1) [c_custkey:Int64, c_name:Utf8, mark:Boolean]
+            LeftMark Join:  Filter: Boolean(true) [c_custkey:Int64, c_name:Utf8, mark:Boolean]
+              TableScan: customer [c_custkey:Int64, c_name:Utf8]
+              SubqueryAlias: __correlated_sq_1 [o_custkey:Int64]
+                Projection: orders.o_custkey [o_custkey:Int64]
+                  Filter: customer.c_custkey = orders.o_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]
         "
         )
     }

datafusion/optimizer/src/optimize_projections/mod.rs

@@ -402,9 +402,8 @@ fn optimize_projections(
         }
         LogicalPlan::Join(join) => {
             let left_len = join.left.schema().fields().len();
-            let right_len = join.right.schema().fields().len();
             let (left_req_indices, right_req_indices) =
-                split_join_requirements(left_len, right_len, indices, &join.join_type);
+                split_join_requirements(left_len, indices, &join.join_type);
             let left_indices =
                 left_req_indices.with_plan_exprs(&plan, join.left.schema())?;
             let right_indices =
@@ -729,29 +728,21 @@ fn rewrite_expr(expr: Expr, input: &Projection) -> Result<Transformed<Expr>> {
 /// adjusted based on the join type.
 fn split_join_requirements(
     left_len: usize,
-    right_len: usize,
     indices: RequiredIndices,
     join_type: &JoinType,
 ) -> (RequiredIndices, RequiredIndices) {
     match join_type {
         // In these cases requirements are split between left/right children:
-        JoinType::Inner | JoinType::Left | JoinType::Right | JoinType::Full => {
+        JoinType::Inner
+        | JoinType::Left
+        | JoinType::Right
+        | JoinType::Full
+        | JoinType::LeftMark
+        | JoinType::RightMark => {
             // Decrease right side indices by `left_len` so that they point to valid
             // positions within the right child:
             indices.split_off(left_len)
         }
-        JoinType::LeftMark => {
-            // LeftMark output: [left_cols(0..left_len), mark]
-            // The mark column is synthetic (produced by the join itself),
-            // so discard it and route only to the left child.
-            let (left_indices, _mark) = indices.split_off(left_len);
-            (left_indices, RequiredIndices::new())
-        }
-        JoinType::RightMark => {
-            // Same as LeftMark, but for the right child.
-            let (right_indices, _mark) = indices.split_off(right_len);
-            (RequiredIndices::new(), right_indices)
-        }
         // All requirements can be re-routed to left child directly.
         JoinType::LeftAnti | JoinType::LeftSemi => (indices, RequiredIndices::new()),
         // All requirements can be re-routed to right side directly.
@@ -2366,42 +2357,6 @@ mod tests {
     }
 
     // Regression test for https://github.com/apache/datafusion/issues/20083
-    // Optimizer must not fail when LeftMark joins from EXISTS OR EXISTS
-    // feed into a Left join.
-    #[test]
-    fn optimize_projections_exists_or_exists_with_outer_join() -> Result<()> {
-        use datafusion_expr::utils::disjunction;
-        use datafusion_expr::{exists, out_ref_col};
-
-        let table_a = test_table_scan_with_name("a")?;
-        let table_b = test_table_scan_with_name("b")?;
-
-        let sq_a = Arc::new(
-            LogicalPlanBuilder::from(test_table_scan_with_name("sq_a")?)
-                .filter(col("sq_a.a").eq(out_ref_col(DataType::UInt32, "a.a")))?
-                .project(vec![lit(1)])?
-                .build()?,
-        );
-
-        let sq_b = Arc::new(
-            LogicalPlanBuilder::from(test_table_scan_with_name("sq_b")?)
-                .filter(col("sq_b.b").eq(out_ref_col(DataType::UInt32, "a.b")))?
-                .project(vec![lit(1)])?
-                .build()?,
-        );
-
-        let plan = LogicalPlanBuilder::from(table_a)
-            .filter(disjunction(vec![exists(sq_a), exists(sq_b)]).unwrap())?
-            .join(table_b, JoinType::Left, (vec!["a"], vec!["a"]), None)?
-            .build()?;
-
-        let optimizer = Optimizer::new();
-        let config = OptimizerContext::new();
-        optimizer.optimize(plan, &config, observe)?;
-
-        Ok(())
-    }
-
     #[test]
     fn optimize_projections_left_mark_join_with_projection() -> Result<()> {
         let table_a = test_table_scan_with_name("a")?;