perf(substrait): avoid quadratic field lookup when consuming ReadRel

datafusion/substrait/Cargo.toml

@@ -49,12 +49,17 @@ url = { workspace = true }
 tokio = { workspace = true, features = ["fs"] }
 
 [dev-dependencies]
+criterion = { workspace = true, features = ["async_tokio"] }
 datafusion = { workspace = true, features = ["nested_expressions", "unicode_expressions"] }
 datafusion-functions-aggregate = { workspace = true }
 serde_json = { workspace = true }
 tokio = { workspace = true }
 insta = { workspace = true }
 
+[[bench]]
+harness = false
+name = "consume_wide_read"
+
 [features]
 default = ["physical"]
 physical = ["datafusion/parquet"]

datafusion/substrait/benches/consume_wide_read.rs

@@ -0,0 +1,75 @@
+// 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.
+
+//! Benchmarks consuming a Substrait `ReadRel` over a wide table, which exercises
+//! the schema compatibility check and projection in the consumer.
+
+use std::hint::black_box;
+use std::sync::Arc;
+
+use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
+use datafusion::arrow::datatypes::{DataType, Field, Schema};
+use datafusion::datasource::empty::EmptyTable;
+use datafusion::prelude::SessionContext;
+use datafusion_substrait::logical_plan::consumer::from_substrait_plan;
+use datafusion_substrait::logical_plan::producer::to_substrait_plan;
+
+fn wide_table_context(num_columns: usize) -> SessionContext {
+    let fields: Vec<Field> = (0..num_columns)
+        .map(|i| Field::new(format!("c{i}"), DataType::Int64, true))
+        .collect();
+    let ctx = SessionContext::new();
+    ctx.register_table(
+        "wide",
+        Arc::new(EmptyTable::new(Arc::new(Schema::new(fields)))),
+    )
+    .unwrap();
+    ctx
+}
+
+fn consume_wide_read(c: &mut Criterion) {
+    let rt = tokio::runtime::Builder::new_current_thread()
+        .build()
+        .unwrap();
+
+    let mut group = c.benchmark_group("consume_wide_read");
+    for num_columns in [100, 1_000, 10_000] {
+        let ctx = wide_table_context(num_columns);
+        // The producer emits the full table schema as `base_schema` and a mask
+        // selecting the single projected column
+        let plan = rt.block_on(async {
+            let df = ctx.sql("SELECT c0 FROM wide").await.unwrap();
+            let logical_plan = df.into_optimized_plan().unwrap();
+            to_substrait_plan(&logical_plan, &ctx.state()).unwrap()
+        });
+        let state = ctx.state();
+
+        group.bench_with_input(
+            BenchmarkId::from_parameter(num_columns),
+            &plan,
+            |b, plan| {
+                b.to_async(&rt).iter(|| async {
+                    black_box(from_substrait_plan(&state, plan).await.unwrap())
+                })
+            },
+        );
+    }
+    group.finish();
+}
+
+criterion_group!(benches, consume_wide_read);
+criterion_main!(benches);

datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs

@@ -18,7 +18,9 @@
 use crate::logical_plan::consumer::SubstraitConsumer;
 use crate::logical_plan::consumer::from_substrait_literal;
 use crate::logical_plan::consumer::from_substrait_named_struct;
-use crate::logical_plan::consumer::utils::ensure_schema_compatibility;
+use crate::logical_plan::consumer::utils::{
+    ensure_schema_compatibility, unqualified_name_index,
+};
 use datafusion::common::{
     DFSchema, DFSchemaRef, TableReference, not_impl_err, plan_err,
     substrait_datafusion_err, substrait_err,
@@ -286,14 +288,17 @@ fn apply_projection(
 
     match plan {
         LogicalPlan::TableScan(mut scan) => {
+            let name_index = unqualified_name_index(&df_schema);
             let column_indices: Vec<usize> = substrait_schema
                 .strip_qualifiers()
                 .fields()
                 .iter()
                 .map(|substrait_field| {
-                    Ok(df_schema
-                        .index_of_column_by_name(None, substrait_field.name().as_str())
-                        .unwrap())
+                    let name = substrait_field.name().as_str();
+                    Ok(match name_index.get(name) {
+                        Some(Some(idx)) => *idx,
+                        _ => df_schema.index_of_column_by_name(None, name).unwrap(),
+                    })
                 })
                 .collect::<datafusion::common::Result<_>>()?;
 

datafusion/substrait/src/logical_plan/consumer/utils.rs

@@ -23,7 +23,7 @@ use datafusion::common::{
 };
 use datafusion::logical_expr::expr::Sort;
 use datafusion::logical_expr::{Cast, Expr, ExprSchemable};
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 use substrait::proto::SortField;
 use substrait::proto::sort_field::SortDirection;
@@ -298,17 +298,38 @@ pub(super) fn ensure_schema_compatibility(
     table_schema: &DFSchema,
     substrait_schema: DFSchema,
 ) -> datafusion::common::Result<()> {
+    let name_index = unqualified_name_index(table_schema);
     substrait_schema
         .strip_qualifiers()
         .fields()
         .iter()
         .try_for_each(|substrait_field| {
-            let df_field =
-                table_schema.field_with_unqualified_name(substrait_field.name())?;
+            let df_field = match name_index.get(substrait_field.name().as_str()) {
+                Some(Some(idx)) => table_schema.field(*idx),
+                // Missing or duplicated names are resolved by DFSchema, which
+                // reports the appropriate error
+                _ => table_schema.field_with_unqualified_name(substrait_field.name())?,
+            };
             ensure_field_compatibility(df_field, substrait_field)
         })
 }
 
+/// Maps each unqualified field name of the schema to its index, or to `None` if
+/// the name occurs more than once.
+///
+/// Used to avoid a linear scan of the schema for every looked up field, which is
+/// quadratic for wide schemas.
+pub(super) fn unqualified_name_index(schema: &DFSchema) -> HashMap<&str, Option<usize>> {
+    let mut index = HashMap::with_capacity(schema.fields().len());
+    for (idx, field) in schema.fields().iter().enumerate() {
+        index
+            .entry(field.name().as_str())
+            .and_modify(|entry| *entry = None)
+            .or_insert(Some(idx));
+    }
+    index
+}
+
 /// Ensures that the given Substrait field is compatible with the given DataFusion field
 ///
 /// A field is compatible between Substrait and DataFusion if: