diff --git a/Cargo.lock b/Cargo.lock index 174010138c3c..2baad63bb939 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2758,6 +2758,7 @@ dependencies = [ "async-recursion", "async-trait", "chrono", + "criterion", "datafusion", "datafusion-functions-aggregate", "half", diff --git a/datafusion/substrait/Cargo.toml b/datafusion/substrait/Cargo.toml index a0f203cec8db..374f8bef8384 100644 --- a/datafusion/substrait/Cargo.toml +++ b/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"] diff --git a/datafusion/substrait/benches/consume_wide_read.rs b/datafusion/substrait/benches/consume_wide_read.rs new file mode 100644 index 000000000000..b2a2ce3ab57b --- /dev/null +++ b/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 = (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); diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs index f871fd4f4a2e..411edaf5c921 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs +++ b/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 = 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::>()?; diff --git a/datafusion/substrait/src/logical_plan/consumer/utils.rs b/datafusion/substrait/src/logical_plan/consumer/utils.rs index 04fafdc691d2..3c6e4c67cc2e 100644 --- a/datafusion/substrait/src/logical_plan/consumer/utils.rs +++ b/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> { + 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: