From 9da864f31ba186846c90b9d52f96636d3ac7e3a0 Mon Sep 17 00:00:00 2001 From: jeadie Date: Mon, 20 Apr 2026 17:03:10 +1000 Subject: [PATCH 1/3] projection: forward fetch through ProjectionExec --- datafusion/physical-plan/src/projection.rs | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index e8608f17a1b20..20e301b45577e 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -324,6 +324,17 @@ impl ExecutionPlan for ProjectionExec { true } + fn with_fetch(&self, limit: Option) -> Option> { + let child_with_fetch = self.input().with_fetch(limit)?; + ProjectionExec::try_new(self.expr().to_vec(), child_with_fetch) + .ok() + .map(|projection| Arc::new(projection) as Arc) + } + + fn fetch(&self) -> Option { + self.input().fetch() + } + fn cardinality_effect(&self) -> CardinalityEffect { CardinalityEffect::Equal } From 21803a2d8ef0a78d1e7629cfbbc1151002880dd2 Mon Sep 17 00:00:00 2001 From: jeadie Date: Mon, 20 Apr 2026 17:58:16 +1000 Subject: [PATCH 2/3] test: cover exact fetch propagation through projection --- .../tests/physical_optimizer/pushdown_sort.rs | 31 +++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/datafusion/core/tests/physical_optimizer/pushdown_sort.rs b/datafusion/core/tests/physical_optimizer/pushdown_sort.rs index efe3c53487fcc..583047e1f07ef 100644 --- a/datafusion/core/tests/physical_optimizer/pushdown_sort.rs +++ b/datafusion/core/tests/physical_optimizer/pushdown_sort.rs @@ -1093,3 +1093,34 @@ fn test_sort_pushdown_exact_preserves_fetch() { " ); } + +#[test] +fn test_sort_pushdown_exact_preserves_fetch_through_projection() { + // Regression test: in the Exact pushdown path, SortExec is removed and + // fetch must be propagated through ProjectionExec to the source via with_fetch(). + let schema = schema(); + let source = exact_test_scan(schema.clone()); + + // Projection reorders columns: SELECT b, a + let projection = simple_projection_exec(source, vec![1, 0]); + + // Sort on projected column b with LIMIT. + let b_expr_at_0 = sort_expr_named("b", 0); + let ordering = LexOrdering::new(vec![b_expr_at_0]).unwrap(); + let plan = sort_exec_with_fetch(ordering, Some(10), projection); + + insta::assert_snapshot!( + OptimizationTest::new(plan, PushdownSort::new(), true), + @r" + OptimizationTest: + input: + - SortExec: TopK(fetch=10), expr=[b@0 ASC], preserve_partitioning=[false] + - ProjectionExec: expr=[b@1 as b, a@0 as a] + - ExactTestScan + output: + Ok: + - ProjectionExec: expr=[b@1 as b, a@0 as a] + - ExactTestScan: ordered=[b@1 ASC], fetch=10 + " + ); +} From 3bdc0769feca3191ae1aa66fbe7bc5392d809a88 Mon Sep 17 00:00:00 2001 From: jeadie Date: Tue, 21 Apr 2026 10:06:22 +1000 Subject: [PATCH 3/3] projection: validate sort column alias before pushdown --- .../tests/physical_optimizer/pushdown_sort.rs | 37 +++++++++++++++++++ datafusion/physical-plan/src/projection.rs | 10 ++++- 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/datafusion/core/tests/physical_optimizer/pushdown_sort.rs b/datafusion/core/tests/physical_optimizer/pushdown_sort.rs index 583047e1f07ef..07d777e1244bb 100644 --- a/datafusion/core/tests/physical_optimizer/pushdown_sort.rs +++ b/datafusion/core/tests/physical_optimizer/pushdown_sort.rs @@ -746,6 +746,43 @@ fn test_sort_pushdown_through_projection_with_alias() { ); } +#[test] +fn test_no_sort_pushdown_for_projection_name_index_mismatch() { + // Regression: if a sort column's index points into a projection output but + // its name does not match the projected alias, do not rewrite by index only. + let schema = schema(); + + // Source has [a ASC] ordering + let a = sort_expr("a", &schema); + let source_ordering = LexOrdering::new(vec![a.clone()]).unwrap(); + let source = parquet_exec_with_sort(schema.clone(), vec![source_ordering]); + + // Projection: SELECT a, b + let projection = simple_projection_exec(source, vec![0, 1]); + + // Mismatched column metadata: name "_score" at index 0. + // Even though index 0 maps to projected column `a`, this must not push down. + let mismatched = sort_expr_named("_score", 0); + let ordering = LexOrdering::new(vec![mismatched.reverse()]).unwrap(); + let plan = sort_exec(ordering, projection); + + insta::assert_snapshot!( + OptimizationTest::new(plan, PushdownSort::new(), true), + @r" + OptimizationTest: + input: + - SortExec: expr=[_score@0 DESC NULLS LAST], preserve_partitioning=[false] + - ProjectionExec: expr=[a@0 as a, b@1 as b] + - DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet + output: + Ok: + - SortExec: expr=[_score@0 DESC NULLS LAST], preserve_partitioning=[false] + - ProjectionExec: expr=[a@0 as a, b@1 as b] + - DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet + " + ); +} + #[test] fn test_no_sort_pushdown_through_computed_projection() { use datafusion_expr::Operator; diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index 20e301b45577e..62048c797877b 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -395,7 +395,15 @@ impl ExecutionPlan for ProjectionExec { let proj_expr = &self.expr()[col.index()]; - // Check if projection expression is a simple column + // Guard against stale/mismatched column metadata (e.g. same index + // but different field name). In that case we must not rewrite by + // index only, as it can push down an order on a different column. + if col.name() != proj_expr.alias { + can_pushdown = false; + return Ok(Transformed::no(expr)); + } + + // Check if projection expression is a simple column. // We cannot push down order by clauses that depend on // projected computations as they would have nothing to reference. if let Some(child_col) =