From 70c9ec7387b09966f7779be2f32260acbf90e76c Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 6 Oct 2026 10:01:08 -0400 Subject: [PATCH 1/2] fix: `PushDownLeafProjections` yields when the extraction cannot move A projection of a `MoveTowardsLeafNodes` expression directly over a node that the extraction cannot go through (for example a `TableScan` or an `Aggregate`) was split into a recovery projection over an in-place extraction projection. `OptimizeProjections` then merged the two back into the original projection. The two rules undid each other in every optimizer pass, and the loop stopped only because the plan signature repeated. Now the rule pushes the in-place extraction projection immediately. If it cannot move below the input, the rule leaves the original projection unchanged. The final plans do not change, and the source still absorbs the leaf expressions (for example `DataSourceExec projection=[get_field(...)]`). `filter_and_extraction_projection_reach_fixed_point` now checks that no rule changes the plan in the last pass, and covers a projection over a scan. Co-Authored-By: Claude Opus 5.5 --- .../optimizer/src/extract_leaf_expressions.rs | 81 ++++++++++--------- .../optimizer/src/optimize_projections/mod.rs | 11 +++ datafusion/optimizer/src/push_down_filter.rs | 20 +++-- datafusion/sqllogictest/test_files/unnest.slt | 2 +- .../library-user-guide/query-optimizer.md | 3 +- 5 files changed, 71 insertions(+), 46 deletions(-) diff --git a/datafusion/optimizer/src/extract_leaf_expressions.rs b/datafusion/optimizer/src/extract_leaf_expressions.rs index ca12fb0bbb52..d8fc3888b69d 100644 --- a/datafusion/optimizer/src/extract_leaf_expressions.rs +++ b/datafusion/optimizer/src/extract_leaf_expressions.rs @@ -53,7 +53,19 @@ //! that runs before the extraction projection exists, so row group pruning and //! source level filtering still happen. //! +//! # Yields to [`OptimizeProjections`] +//! +//! When the extraction projection cannot move below the input of a projection +//! (for example a `TableScan` or an `Aggregate`), [`PushDownLeafProjections`] +//! leaves the projection unchanged. Splitting it in place would give a +//! recovery projection over an extraction projection on the same input, and +//! `OptimizeProjections` merges those two back into the original projection. +//! The two rules would then undo each other on every optimizer pass. The split +//! has no use there: the source absorbs the leaf expressions of the original +//! projection just as well. +//! //! [`PushDownFilter`]: crate::push_down_filter::PushDownFilter +//! [`OptimizeProjections`]: crate::optimize_projections::OptimizeProjections use indexmap::{IndexMap, IndexSet}; use std::collections::{BTreeSet, HashMap}; @@ -1150,18 +1162,33 @@ fn split_and_push_projection( // Only pre-existing __datafusion_extracted aliases and columns, no new // extractions from routing_extract. The original projection is // already an extraction projection that couldn't be pushed - // further. Return None. + // further. Return None. This return also ends the recursion + // of the `try_push_input` call below, which comes back here + // with a pure extraction projection. return Ok(None); } - // Build extraction projection in-place (couldn't push down) + // Build the extraction projection in place, then push it. It is a + // pure extraction projection, which can go through nodes that the + // original projection cannot. let input_arc = Arc::clone(input); - let extraction = build_extraction_projection_impl( + let extraction = LogicalPlan::Projection(build_extraction_projection_impl( &extraction_pairs, columns_needed, &input_arc, input_schema.as_ref(), - )?; - LogicalPlan::Projection(extraction) + )?); + match try_push_input(&extraction, alias_generator)? { + Some(pushed) => pushed, + // The extraction projection cannot move below `input`. Leave + // the original projection unchanged: the split would put the + // extraction projection on the same input, and + // `OptimizeProjections` merges the recovery and extraction + // projections back into the original projection. The rules + // would then undo each other in every optimizer pass. The + // source absorbs the leaf expressions of the original + // projection just as well. + None => return Ok(None), + } } }; @@ -1730,13 +1757,10 @@ mod tests { (same as original) ## After Pushdown - Projection: __datafusion_extracted_1 AS leaf_udf(test.user,Utf8("name")) - Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1, test.user - TableScan: test projection=[user] + (same as after extraction) ## Optimized - Projection: leaf_udf(test.user, Utf8("name")) - TableScan: test projection=[user] + (same as after pushdown) "#) } @@ -1760,13 +1784,10 @@ mod tests { (same as original) ## After Pushdown - Projection: __datafusion_extracted_1 IS NOT NULL AS has_name - Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1, test.user - TableScan: test projection=[user] + (same as after extraction) ## Optimized - Projection: leaf_udf(test.user, Utf8("name")) IS NOT NULL AS has_name - TableScan: test projection=[user] + (same as after pushdown) "#) } @@ -1965,13 +1986,10 @@ mod tests { (same as original) ## After Pushdown - Projection: __datafusion_extracted_1 AS username - Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1, test.user - TableScan: test projection=[user] + (same as after extraction) ## Optimized - Projection: leaf_udf(test.user, Utf8("name")) AS username - TableScan: test projection=[user] + (same as after pushdown) "#) } @@ -2031,13 +2049,10 @@ mod tests { (same as original) ## After Pushdown - Projection: __datafusion_extracted_1 AS leaf_udf(test.user,Utf8("name")), __datafusion_extracted_1 AS name2 - Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1, test.user - TableScan: test projection=[user] + (same as after extraction) ## Optimized - Projection: leaf_udf(test.user, Utf8("name")), leaf_udf(test.user, Utf8("name")) AS name2 - TableScan: test projection=[user] + (same as after pushdown) "#) } @@ -2245,13 +2260,10 @@ mod tests { (same as original) ## After Pushdown - Projection: __datafusion_extracted_1 AS leaf_udf(test.user,Utf8("name")) - Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1, test.user - TableScan: test projection=[user] + (same as after extraction) ## Optimized - Projection: leaf_udf(test.user, Utf8("name")) - TableScan: test projection=[user] + (same as after pushdown) "#) } @@ -2432,15 +2444,10 @@ mod tests { (same as original) ## After Pushdown - Projection: __datafusion_extracted_1 IS NOT NULL AS has_name, COUNT(Int32(1)) - Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1, test.user, COUNT(Int32(1)) - Aggregate: groupBy=[[test.user]], aggr=[[COUNT(Int32(1))]] - TableScan: test projection=[user] + (same as after extraction) ## Optimized - Projection: leaf_udf(test.user, Utf8("name")) IS NOT NULL AS has_name, COUNT(Int32(1)) - Aggregate: groupBy=[[test.user]], aggr=[[COUNT(Int32(1))]] - TableScan: test projection=[user] + (same as after pushdown) "#) } diff --git a/datafusion/optimizer/src/optimize_projections/mod.rs b/datafusion/optimizer/src/optimize_projections/mod.rs index c913a4ce961c..8fdb5775c971 100644 --- a/datafusion/optimizer/src/optimize_projections/mod.rs +++ b/datafusion/optimizer/src/optimize_projections/mod.rs @@ -16,6 +16,17 @@ // under the License. //! [`OptimizeProjections`] identifies and eliminates unused columns +//! +//! # Precedence over `PushDownLeafProjections` +//! +//! `OptimizeProjections` merges adjacent projections, including a recovery +//! projection over an extraction projection that +//! [`PushDownLeafProjections`] creates. To keep the two rules from undoing each +//! other, `PushDownLeafProjections` only creates that pair when the extraction +//! projection moves below the input. See the module documentation of +//! [`extract_leaf_expressions`](crate::extract_leaf_expressions). +//! +//! [`PushDownLeafProjections`]: crate::extract_leaf_expressions::PushDownLeafProjections mod required_indices; diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs index 54f220a1fe9d..dac58e67ea7f 100644 --- a/datafusion/optimizer/src/push_down_filter.rs +++ b/datafusion/optimizer/src/push_down_filter.rs @@ -4864,9 +4864,11 @@ mod tests { /// `PushDownFilter` and `PushDownLeafProjections` must not undo each other /// for a filter next to a pure extraction projection - /// (). Neither rule may - /// change the plan in the last optimizer pass. The source does not absorb - /// filters, so the `Filter` node stays in the plan. + /// (). The source does + /// not absorb filters, so the `Filter` node stays in the plan. Also, + /// `PushDownLeafProjections` must not split a projection that + /// `OptimizeProjections` merges back. No rule may change the plan in the + /// last optimizer pass. #[test] fn filter_and_extraction_projection_reach_fixed_point() -> Result<()> { let scan = || { @@ -4887,13 +4889,17 @@ mod tests { .project(vec![leaf_udf_expr(col("a")), col("b")])? .build()?; - for plan in [simple, two_filters] { + // No filter: the extraction projection cannot move below the scan. + let no_filter = LogicalPlanBuilder::from(test_table_scan()?) + .project(vec![leaf_udf_expr(col("a"))])? + .build()?; + + for plan in [simple, two_filters, no_filter] { let passes = rules_that_changed_plan_per_pass(plan)?; let last = passes.last().unwrap(); assert!( - !last.iter().any(|rule| rule == "push_down_filter" - || rule == "push_down_leaf_projections"), - "pushdown rules changed the plan in the last pass: {passes:?}" + last.is_empty(), + "rules changed the plan in the last pass: {passes:?}" ); } Ok(()) diff --git a/datafusion/sqllogictest/test_files/unnest.slt b/datafusion/sqllogictest/test_files/unnest.slt index 39ed64a52675..cf190949aad0 100644 --- a/datafusion/sqllogictest/test_files/unnest.slt +++ b/datafusion/sqllogictest/test_files/unnest.slt @@ -668,7 +668,7 @@ explain select unnest(unnest(unnest(column3)['c1'])), column3 from recursive_unn logical_plan 01)Projection: __unnest_placeholder(UNNEST(recursive_unnest_table.column3)[c1],depth=2) AS UNNEST(UNNEST(UNNEST(recursive_unnest_table.column3)[c1])), recursive_unnest_table.column3 02)--Unnest: lists[__unnest_placeholder(UNNEST(recursive_unnest_table.column3)[c1])|depth=2] structs[] -03)----Projection: get_field(__unnest_placeholder(recursive_unnest_table.column3,depth=1), Utf8("c1")) AS __unnest_placeholder(UNNEST(recursive_unnest_table.column3)[c1]), recursive_unnest_table.column3 +03)----Projection: get_field(__unnest_placeholder(recursive_unnest_table.column3,depth=1) AS UNNEST(recursive_unnest_table.column3), Utf8("c1")) AS __unnest_placeholder(UNNEST(recursive_unnest_table.column3)[c1]), recursive_unnest_table.column3 04)------Unnest: lists[__unnest_placeholder(recursive_unnest_table.column3)|depth=1] structs[] 05)--------Projection: recursive_unnest_table.column3 AS __unnest_placeholder(recursive_unnest_table.column3), recursive_unnest_table.column3 06)----------TableScan: recursive_unnest_table projection=[column3] diff --git a/docs/source/library-user-guide/query-optimizer.md b/docs/source/library-user-guide/query-optimizer.md index 73c3ed6ed11d..03fb6f2e33fe 100644 --- a/docs/source/library-user-guide/query-optimizer.md +++ b/docs/source/library-user-guide/query-optimizer.md @@ -168,9 +168,10 @@ Two rules can want the opposite order for the same pair of adjacent plan nodes. Do not let two rules compete. Give one rule precedence, make the competing rule yield, and record the decision in the module documentation of both rules. -There is one such decision today: +There are two such decisions today: - `PushDownFilter` yields to a _pure extraction projection_. A pure extraction projection is a projection whose expressions are only `__datafusion_extracted_N` aliases and pass-through columns. `ExtractLeafExpressions` creates it, and `PushDownLeafProjections` moves it towards the leaves. `PushDownFilter` does not move a filter below such a projection, so the projection stays next to the scan. A Parquet scan then merges the projection into the file projection and reads only the struct leaf. The filter loses nothing, because `PushDownFilter` records the predicate in `TableScan::filters` in the pass that runs before the extraction projection exists. See [issue #14540](https://github.com/apache/datafusion/issues/14540). +- `PushDownLeafProjections` yields to `OptimizeProjections`. When the extraction projection cannot move below the input of a projection, for example a `TableScan` or an `Aggregate`, `PushDownLeafProjections` leaves the projection unchanged. A split in place would give a recovery projection over an extraction projection on the same input, which `OptimizeProjections` merges back. The source absorbs the leaf expressions of the original projection just as well. See [PR #26085](https://github.com/apache/datafusion/pull/26085). ### Expression Naming From f4ac83c70f7c2b630d1c515454c42daf7cf243bd Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 22 Oct 2026 16:52:22 -0700 Subject: [PATCH 2/2] test: cover aggregate barrier in fixed-point test, fix stale comment Co-Authored-By: Claude Opus 5.5 --- .../optimizer/src/extract_leaf_expressions.rs | 4 ++-- datafusion/optimizer/src/push_down_filter.rs | 13 +++++++++++-- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/datafusion/optimizer/src/extract_leaf_expressions.rs b/datafusion/optimizer/src/extract_leaf_expressions.rs index d8fc3888b69d..dad3cc03042c 100644 --- a/datafusion/optimizer/src/extract_leaf_expressions.rs +++ b/datafusion/optimizer/src/extract_leaf_expressions.rs @@ -2417,8 +2417,8 @@ mod tests { } /// Projection containing a MoveTowardsLeafNodes sub-expression above an - /// Aggregate. Aggregate blocks pushdown, so the (None, true) recovery - /// fallback path fires: in-place extraction + recovery projection. + /// Aggregate. The Aggregate blocks pushdown, so the extraction cannot move + /// and the original projection is kept unchanged. #[test] fn test_projection_with_leaf_expr_above_aggregate() -> Result<()> { use datafusion_expr::test::function_stub::count; diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs index dac58e67ea7f..ccde8aa07f64 100644 --- a/datafusion/optimizer/src/push_down_filter.rs +++ b/datafusion/optimizer/src/push_down_filter.rs @@ -1751,7 +1751,7 @@ mod tests { use crate::simplify_expressions::SimplifyExpressions; use crate::test::udfs::leaf_udf_expr; use crate::test::*; - use datafusion_expr::test::function_stub::sum; + use datafusion_expr::test::function_stub::{count, sum}; use insta::assert_snapshot; use super::*; @@ -4894,7 +4894,16 @@ mod tests { .project(vec![leaf_udf_expr(col("a"))])? .build()?; - for plan in [simple, two_filters, no_filter] { + // Leaf expression above an Aggregate: the Aggregate blocks movement. + let above_aggregate = LogicalPlanBuilder::from(test_table_scan_with_struct()?) + .aggregate(vec![col("user")], vec![count(lit(1))])? + .project(vec![ + leaf_udf_expr(col("user")).is_not_null().alias("has_name"), + col("COUNT(Int32(1))"), + ])? + .build()?; + + for plan in [simple, two_filters, no_filter, above_aggregate] { let passes = rules_that_changed_plan_per_pass(plan)?; let last = passes.last().unwrap(); assert!(