[pull] main from apache:main

datafusion/core/tests/physical_optimizer/enforce_distribution.rs

@@ -52,6 +52,7 @@ use datafusion_physical_plan::aggregates::{
     AggregateExec, AggregateMode, PhysicalGroupBy,
 };
 use datafusion_physical_plan::coalesce_batches::CoalesceBatchesExec;
+use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use datafusion_physical_plan::execution_plan::ExecutionPlan;
 use datafusion_physical_plan::expressions::col;
 use datafusion_physical_plan::filter::FilterExec;
@@ -3471,3 +3472,47 @@ fn optimize_away_unnecessary_repartition2() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn test_replace_order_preserving_variants_with_fetch() -> Result<()> {
+    // Create a base plan
+    let parquet_exec = parquet_exec();
+
+    let sort_expr = PhysicalSortExpr {
+        expr: Arc::new(Column::new("id", 0)),
+        options: SortOptions::default(),
+    };
+
+    let ordering = LexOrdering::new(vec![sort_expr]);
+
+    // Create a SortPreservingMergeExec with fetch=5
+    let spm_exec = Arc::new(
+        SortPreservingMergeExec::new(ordering, parquet_exec.clone()).with_fetch(Some(5)),
+    );
+
+    // Create distribution context
+    let dist_context = DistributionContext::new(
+        spm_exec,
+        true,
+        vec![DistributionContext::new(parquet_exec, false, vec![])],
+    );
+
+    // Apply the function
+    let result = replace_order_preserving_variants(dist_context)?;
+
+    // Verify the plan was transformed to CoalescePartitionsExec
+    result
+        .plan
+        .as_any()
+        .downcast_ref::<CoalescePartitionsExec>()
+        .expect("Expected CoalescePartitionsExec");
+
+    // Verify fetch was preserved
+    assert_eq!(
+        result.plan.fetch(),
+        Some(5),
+        "Fetch value was not preserved after transformation"
+    );
+
+    Ok(())
+}

datafusion/optimizer/Cargo.toml

@@ -55,10 +55,15 @@ regex-syntax = "0.8.0"
 
 [dev-dependencies]
 async-trait = { workspace = true }
+criterion = { workspace = true }
 ctor = { workspace = true }
 datafusion-functions-aggregate = { workspace = true }
 datafusion-functions-window = { workspace = true }
 datafusion-functions-window-common = { workspace = true }
 datafusion-sql = { workspace = true }
 env_logger = { workspace = true }
 insta = { workspace = true }
+
+[[bench]]
+name = "projection_unnecessary"
+harness = false

datafusion/optimizer/benches/projection_unnecessary.rs

@@ -0,0 +1,79 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::datatypes::{DataType, Field, Schema};
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
+use datafusion_common::ToDFSchema;
+use datafusion_common::{Column, TableReference};
+use datafusion_expr::{logical_plan::LogicalPlan, projection_schema, Expr};
+use datafusion_optimizer::optimize_projections::is_projection_unnecessary;
+use std::sync::Arc;
+
+fn is_projection_unnecessary_old(
+    input: &LogicalPlan,
+    proj_exprs: &[Expr],
+) -> datafusion_common::Result<bool> {
+    // First check if all expressions are trivial (cheaper operation than `projection_schema`)
+    if !proj_exprs
+        .iter()
+        .all(|expr| matches!(expr, Expr::Column(_) | Expr::Literal(_)))
+    {
+        return Ok(false);
+    }
+    let proj_schema = projection_schema(input, proj_exprs)?;
+    Ok(&proj_schema == input.schema())
+}
+
+fn create_plan_with_many_exprs(num_exprs: usize) -> (LogicalPlan, Vec<Expr>) {
+    // Create schema with many fields
+    let fields = (0..num_exprs)
+        .map(|i| Field::new(format!("col{}", i), DataType::Int32, false))
+        .collect::<Vec<_>>();
+    let schema = Schema::new(fields);
+
+    // Create table scan
+    let table_scan = LogicalPlan::EmptyRelation(datafusion_expr::EmptyRelation {
+        produce_one_row: true,
+        schema: Arc::new(schema.clone().to_dfschema().unwrap()),
+    });
+
+    // Create projection expressions (just column references)
+    let exprs = (0..num_exprs)
+        .map(|i| Expr::Column(Column::new(None::<TableReference>, format!("col{}", i))))
+        .collect();
+
+    (table_scan, exprs)
+}
+
+fn benchmark_is_projection_unnecessary(c: &mut Criterion) {
+    let (plan, exprs) = create_plan_with_many_exprs(1000);
+
+    let mut group = c.benchmark_group("projection_unnecessary_comparison");
+
+    group.bench_function("is_projection_unnecessary_new", |b| {
+        b.iter(|| black_box(is_projection_unnecessary(&plan, &exprs).unwrap()))
+    });
+
+    group.bench_function("is_projection_unnecessary_old", |b| {
+        b.iter(|| black_box(is_projection_unnecessary_old(&plan, &exprs).unwrap()))
+    });
+
+    group.finish();
+}
+
+criterion_group!(benches, benchmark_is_projection_unnecessary);
+criterion_main!(benches);

datafusion/optimizer/src/optimize_projections/mod.rs

@@ -31,8 +31,7 @@ use datafusion_common::{
 use datafusion_expr::expr::Alias;
 use datafusion_expr::Unnest;
 use datafusion_expr::{
-    logical_plan::LogicalPlan, projection_schema, Aggregate, Distinct, Expr, Projection,
-    TableScan, Window,
+    logical_plan::LogicalPlan, Aggregate, Distinct, Expr, Projection, TableScan, Window,
 };
 
 use crate::optimize_projections::required_indices::RequiredIndices;
@@ -785,13 +784,24 @@ fn rewrite_projection_given_requirements(
 /// Projection is unnecessary, when
 /// - input schema of the projection, output schema of the projection are same, and
 /// - all projection expressions are either Column or Literal
-fn is_projection_unnecessary(input: &LogicalPlan, proj_exprs: &[Expr]) -> Result<bool> {
-    // First check if all expressions are trivial (cheaper operation than `projection_schema`)
-    if !proj_exprs.iter().all(is_expr_trivial) {
+pub fn is_projection_unnecessary(
+    input: &LogicalPlan,
+    proj_exprs: &[Expr],
+) -> Result<bool> {
+    // First check if the number of expressions is equal to the number of fields in the input schema.
+    if proj_exprs.len() != input.schema().fields().len() {
         return Ok(false);
     }
-    let proj_schema = projection_schema(input, proj_exprs)?;
-    Ok(&proj_schema == input.schema())
+    Ok(input.schema().iter().zip(proj_exprs.iter()).all(
+        |((field_relation, field_name), expr)| {
+            // Check if the expression is a column and if it matches the field name
+            if let Expr::Column(col) = expr {
+                col.relation.as_ref() == field_relation && col.name.eq(field_name.name())
+            } else {
+                false
+            }
+        },
+    ))
 }
 
 #[cfg(test)]

datafusion/physical-optimizer/src/enforce_distribution.rs

@@ -1018,7 +1018,7 @@ fn remove_dist_changing_operators(
 /// "    RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2",
 /// "      DataSourceExec: file_groups={2 groups: \[\[x], \[y]]}, projection=\[a, b, c, d, e], output_ordering=\[a@0 ASC], file_type=parquet",
 /// ```
-fn replace_order_preserving_variants(
+pub fn replace_order_preserving_variants(
     mut context: DistributionContext,
 ) -> Result<DistributionContext> {
     context.children = context
@@ -1035,7 +1035,10 @@ fn replace_order_preserving_variants(
 
     if is_sort_preserving_merge(&context.plan) {
         let child_plan = Arc::clone(&context.children[0].plan);
-        context.plan = Arc::new(CoalescePartitionsExec::new(child_plan));
+        // It's safe to unwrap because `CoalescePartitionsExec` supports `fetch`.
+        context.plan = CoalescePartitionsExec::new(child_plan)
+            .with_fetch(context.plan.fetch())
+            .unwrap();
         return Ok(context);
     } else if let Some(repartition) =
         context.plan.as_any().downcast_ref::<RepartitionExec>()