Speed up optimize_projection (apache#15787)

* save

* fmt

Author: xudong963

Cargo.lock

@@ -55,9 +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/Cargo.toml

@@ -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/benches/projection_unnecessary.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;
@@ -782,9 +781,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> {
-    let proj_schema = projection_schema(input, proj_exprs)?;
-    Ok(&proj_schema == input.schema() && 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);
+    }
+    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)]