This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-25010-5b7cc8703138bad8b0e9696f428e35a87300d45e in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 7f3e657343ce4d8a0694c8a7449183912a8c5a1e Author: Goutam Adwant <[email protected]> AuthorDate: Sun Sep 20 10:32:09 2026 +0000 perf: reuse column normalization context across expressions (#25010) ## Which issue does this PR close? - Addresses fix (1) in #24777. Constructor-based logical-plan decoding remains separate. ## Rationale for this change Column normalization repeatedly collects fallback schemas and traverses the input plan for USING columns. Wide expression lists and projections repeat that work for the same immutable plan. ## What changes are included in this PR? - Introduce a private, lazy normalization context reused across columns and expression lists. - Share the context in sort normalization and validated projection construction, including wildcard expansion. - Keep already-qualified columns and expressions that do not need normalization on the existing fast path. - Add benchmarks for qualified and unqualified expressions and projection construction at several schema widths. ## What is the testing strategy for this PR? - Add `normalize_batch_schema_precedence`, `normalize_batch_using_join`, and `normalize_batch_skips_unused_plan_context` to cover schema precedence, USING joins, ambiguity/error order, sort options, and lazy handling of qualified columns and literals. - In balanced local `release-nonlto` runs, constructing a 2,000-column unqualified projection falls from about 100 ms to 35 ms. This measures projection construction, not full protobuf decoding; small controls remain noisy. - Reproduce with `cargo bench -p datafusion-expr --bench normalize_columns --profile release-nonlto`. - Focused expression and SQL tests pass. The required extended workspace test command also passes, including all 511 SQL logic-test files. - Expression-crate Clippy passes with all targets and features enabled. The complete documented `dev/rust_lint.sh` also passes, including strict workspace documentation checks. - Full-workspace Clippy with all features enabled hits the existing PostgreSQL decimal-formatting lint in #24974; the affected source is unchanged here. ## Are there any user-facing changes? No public API or name-resolution behavior changes are intended. Normalization reuses plan context instead of collecting it for each column; existing schema lookup costs remain. --- Cargo.lock | 1 + datafusion/expr/Cargo.toml | 5 + datafusion/expr/benches/normalize_columns.rs | 97 +++++++++++++++++ datafusion/expr/src/expr_rewriter/mod.rs | 156 ++++++++++++++++++++++++--- datafusion/expr/src/logical_plan/builder.rs | 22 ++-- 5 files changed, 249 insertions(+), 32 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 221bbf2583..bc8f497624 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2202,6 +2202,7 @@ dependencies = [ "arrow-schema", "async-trait", "chrono", + "criterion", "ctor", "datafusion-common", "datafusion-doc", diff --git a/datafusion/expr/Cargo.toml b/datafusion/expr/Cargo.toml index 4fe7b65f6d..74b85501dd 100644 --- a/datafusion/expr/Cargo.toml +++ b/datafusion/expr/Cargo.toml @@ -69,9 +69,14 @@ serde_json = { workspace = true } sqlparser = { workspace = true, optional = true } [dev-dependencies] +criterion = { workspace = true } ctor = { workspace = true } env_logger = { workspace = true } insta = { workspace = true } # Makes sure `test_display_pg_json` behaves in a consistent way regardless of # feature unification with dependencies serde_json = { workspace = true, features = ["preserve_order"] } + +[[bench]] +name = "normalize_columns" +harness = false diff --git a/datafusion/expr/benches/normalize_columns.rs b/datafusion/expr/benches/normalize_columns.rs new file mode 100644 index 0000000000..988c2310c3 --- /dev/null +++ b/datafusion/expr/benches/normalize_columns.rs @@ -0,0 +1,97 @@ +// 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 std::hint::black_box; + +use arrow::datatypes::{DataType, Field, Schema}; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_expr::expr_rewriter::normalize_cols; +use datafusion_expr::logical_plan::table_scan; +use datafusion_expr::{Expr, LogicalPlanBuilder, col, lit}; + +fn normalize_columns(c: &mut Criterion) { + let mut group = c.benchmark_group("normalize_columns"); + for width in [10, 100, 500, 2000] { + let schema = Schema::new( + (0..width) + .map(|i| Field::new(format!("c{i}"), DataType::Int32, false)) + .collect::<Vec<_>>(), + ); + let input = table_scan(Some("t"), &schema, None) + .unwrap() + .project( + (0..width) + .map(|i| (col(format!("t.c{i}")) + lit(1)).alias(format!("a{i}"))), + ) + .unwrap() + .build() + .unwrap(); + + for qualified in [false, true] { + let (input, qualifier) = if qualified { + ( + LogicalPlanBuilder::from(input.clone()) + .alias("s") + .unwrap() + .build() + .unwrap(), + "s.", + ) + } else { + (input.clone(), "") + }; + let exprs: Vec<Expr> = (0..width) + .map(|i| { + (col(format!("{qualifier}a{i}")) + lit(1)).alias(format!("b{i}")) + }) + .collect(); + let kind = if qualified { + "qualified" + } else { + "unqualified" + }; + + group.bench_with_input( + BenchmarkId::new(format!("expressions/{kind}"), width), + &width, + |b, _| { + b.iter(|| { + normalize_cols(black_box(exprs.clone()), black_box(&input)) + .unwrap() + }) + }, + ); + group.bench_with_input( + BenchmarkId::new(format!("projection/{kind}"), width), + &width, + |b, _| { + b.iter(|| { + LogicalPlanBuilder::from(black_box(input.clone())) + .project(black_box(exprs.clone())) + .unwrap() + .build() + .unwrap() + }) + }, + ); + } + } + group.finish(); +} + +criterion_group!(benches, normalize_columns); +criterion_main!(benches); diff --git a/datafusion/expr/src/expr_rewriter/mod.rs b/datafusion/expr/src/expr_rewriter/mod.rs index 4e9839e2f7..2552b61062 100644 --- a/datafusion/expr/src/expr_rewriter/mod.rs +++ b/datafusion/expr/src/expr_rewriter/mod.rs @@ -24,7 +24,7 @@ use std::sync::Arc; use crate::expr::{Alias, Sort, Unnest}; use crate::logical_plan::Projection; -use crate::{Expr, ExprSchemable, LogicalPlan, LogicalPlanBuilder}; +use crate::{Expr, ExprSchemable, LogicalPlan}; use datafusion_common::TableReference; use datafusion_common::config::ConfigOptions; @@ -64,20 +64,54 @@ pub trait FunctionRewrite: Debug { ) -> Result<Transformed<Expr>>; } -/// Recursively call `LogicalPlanBuilder::normalize` on all [`Column`] expressions -/// in the `expr` expression tree. +/// Recursively normalize all [`Column`] expressions in the `expr` expression tree. pub fn normalize_col(expr: Expr, plan: &LogicalPlan) -> Result<Expr> { - expr.transform(|expr| { - Ok({ - if let Expr::Column(c) = expr { - let col = LogicalPlanBuilder::normalize(plan, c)?; - Transformed::yes(Expr::Column(col)) - } else { - Transformed::no(expr) - } + ColumnNormalizer::new(plan).normalize(expr) +} + +/// Reuses the plan's normalization context across expressions. Initialize it +/// lazily so literals and already-qualified columns need no plan traversal. +pub(crate) struct ColumnNormalizer<'a> { + plan: &'a LogicalPlan, + context: Option<(Vec<&'a DFSchema>, Vec<HashSet<Column>>)>, +} + +impl<'a> ColumnNormalizer<'a> { + pub(crate) fn new(plan: &'a LogicalPlan) -> Self { + Self { + plan, + context: None, + } + } + + #[inline] + pub(crate) fn normalize_column(&mut self, column: Column) -> Result<Column> { + if column.relation.is_some() { + return Ok(column); + } + + let (fallback_schemas, using_columns) = match &mut self.context { + Some(context) => context, + context @ None => context.insert(( + self.plan.fallback_normalize_schemas(), + self.plan.using_columns()?, + )), + }; + column.normalize_with_schemas_and_ambiguity_check( + &[&[self.plan.schema()], fallback_schemas], + using_columns, + ) + } + + pub(crate) fn normalize(&mut self, expr: Expr) -> Result<Expr> { + expr.transform(|expr| match expr { + Expr::Column(column) => self + .normalize_column(column) + .map(|column| Transformed::yes(Expr::Column(column))), + _ => Ok(Transformed::no(expr)), }) - }) - .data() + .data() + } } /// See [`Column::normalize_with_schemas_and_ambiguity_check`] for usage @@ -118,9 +152,10 @@ pub fn normalize_cols( exprs: impl IntoIterator<Item = impl Into<Expr>>, plan: &LogicalPlan, ) -> Result<Vec<Expr>> { + let mut normalizer = ColumnNormalizer::new(plan); exprs .into_iter() - .map(|e| normalize_col(e.into(), plan)) + .map(|e| normalizer.normalize(e.into())) .collect() } @@ -128,11 +163,13 @@ pub fn normalize_sorts( sorts: impl IntoIterator<Item = impl Into<Sort>>, plan: &LogicalPlan, ) -> Result<Vec<Sort>> { + let mut normalizer = ColumnNormalizer::new(plan); sorts .into_iter() .map(|e| { let sort = e.into(); - normalize_col(sort.expr, plan) + normalizer + .normalize(sort.expr) .map(|expr| Sort::new(expr, sort.asc, sort.nulls_first)) }) .collect() @@ -382,7 +419,7 @@ mod test { use super::*; use crate::literal::lit_with_metadata; - use crate::{Cast, col, lit}; + use crate::{Cast, LogicalPlanBuilder, col, lit}; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::ScalarValue; use datafusion_common::tree_node::TreeNodeRewriter; @@ -491,6 +528,93 @@ mod test { assert_eq!(error, expected); } + #[test] + fn normalize_batch_schema_precedence() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ]); + let plan = crate::logical_plan::table_scan(Some("t"), &schema, None)? + .project([col("t.a").alias("b")])? + .build()?; + let exprs = vec![col("b") + col("a"), col("other.missing"), lit(1)]; + let expected = vec![col("b") + col("t.a"), col("other.missing"), lit(1)]; + assert_eq!(super::normalize_cols(exprs.clone(), &plan)?, expected); + assert_eq!(normalize_col(exprs[0].clone(), &plan)?, expected[0]); + let sorts = exprs.into_iter().map(|e| e.sort(false, true)); + assert_eq!( + normalize_sorts(sorts, &plan)?, + expected + .into_iter() + .map(|e| e.sort(false, true)) + .collect::<Vec<_>>() + ); + Ok(()) + } + + #[test] + fn normalize_batch_using_join() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ]); + let right = + crate::logical_plan::table_scan(Some("right"), &schema, None)?.build()?; + let plan = crate::logical_plan::table_scan(Some("left"), &schema, None)? + .join_using(right, crate::JoinType::Inner, vec![Column::from_name("a")])? + .build()?; + assert_eq!( + super::normalize_cols([col("a"), col("a") + col("right.a")], &plan)?, + vec![col("left.a"), col("left.a") + col("right.a")] + ); + let projected = LogicalPlanBuilder::from(plan.clone()) + .project([col("a").alias("key"), col("left.b").alias("value")])? + .build()?; + assert_eq!( + projected.expressions(), + vec![col("left.a").alias("key"), col("left.b").alias("value")] + ); + let err = super::normalize_cols([col("a"), col("b"), col("missing")], &plan) + .unwrap_err(); + assert!( + err.strip_backtrace() + .contains("Ambiguous reference to unqualified field b") + ); + Ok(()) + } + + #[test] + fn normalize_batch_skips_unused_plan_context() -> Result<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let right = + crate::logical_plan::table_scan(Some("right"), &schema, None)?.build()?; + let mut plan = crate::logical_plan::table_scan(Some("left"), &schema, None)? + .join_using(right, crate::JoinType::Inner, vec![Column::from_name("a")])? + .build()?; + // Invalid USING keys must only be inspected when an unqualified column + // needs normalization, just as with a single expression. + let LogicalPlan::Join(join) = &mut plan else { + unreachable!() + }; + join.on[0].0 = lit(1); + assert!(plan.using_columns().is_err()); + assert!(super::normalize_cols(Vec::<Expr>::new(), &plan)?.is_empty()); + let exprs = vec![lit(1), col("left.a"), col("other.missing")]; + assert_eq!(super::normalize_cols(exprs.clone(), &plan)?, exprs); + assert_eq!(normalize_col(col("left.a"), &plan)?, col("left.a")); + assert_eq!( + normalize_sorts([col("left.a").sort(true, false)], &plan)?, + vec![col("left.a").sort(true, false)] + ); + assert!(super::normalize_cols([col("left.a"), col("a")], &plan).is_err()); + let projected = LogicalPlanBuilder::from(plan.clone()) + .project([col("left.a")])? + .build()?; + assert_eq!(projected.expressions(), vec![col("left.a")]); + assert!(LogicalPlanBuilder::from(plan).project([col("a")]).is_err()); + Ok(()) + } + #[test] fn unnormalize_cols() { let expr = col("tableA.a") + col("tableB.b"); diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index ba16116136..6b55cabfa6 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -26,7 +26,7 @@ use std::sync::Arc; use crate::dml::CopyTo; use crate::expr::{Alias, PlannedReplaceSelectItem, Sort as SortExpr}; use crate::expr_rewriter::{ - coerce_plan_expr_for_schema, normalize_col, + ColumnNormalizer, coerce_plan_expr_for_schema, normalize_col, normalize_col_with_schemas_and_ambiguity_check, normalize_cols, normalize_sorts, rewrite_sort_cols_by_aggs, }; @@ -1132,18 +1132,7 @@ impl LogicalPlanBuilder { } pub(crate) fn normalize(plan: &LogicalPlan, column: Column) -> Result<Column> { - if column.relation.is_some() { - // column is already normalized - return Ok(column); - } - - let schema = plan.schema(); - let fallback_schemas = plan.fallback_normalize_schemas(); - let using_columns = plan.using_columns()?; - column.normalize_with_schemas_and_ambiguity_check( - &[&[schema], &fallback_schemas], - &using_columns, - ) + ColumnNormalizer::new(plan).normalize_column(column) } /// Apply a join with on constraint and specified null equality. @@ -2082,6 +2071,7 @@ fn project_with_validation( ) -> Result<LogicalPlan> { let mut projected_expr = vec![]; let mut has_wildcard = false; + let mut normalizer = ColumnNormalizer::new(&plan); for (e, validate) in expr { let e = e.into(); match e { @@ -2100,7 +2090,7 @@ fn project_with_validation( for e in expanded { if validate { projected_expr - .push(columnize_expr(normalize_col(e, &plan)?, &plan)?) + .push(columnize_expr(normalizer.normalize(e)?, &plan)?) } else { projected_expr.push(e) } @@ -2122,7 +2112,7 @@ fn project_with_validation( for e in expanded { if validate { projected_expr - .push(columnize_expr(normalize_col(e, &plan)?, &plan)?) + .push(columnize_expr(normalizer.normalize(e)?, &plan)?) } else { projected_expr.push(e) } @@ -2130,7 +2120,7 @@ fn project_with_validation( } SelectExpr::Expression(e) => { if validate { - projected_expr.push(columnize_expr(normalize_col(e, &plan)?, &plan)?) + projected_expr.push(columnize_expr(normalizer.normalize(e)?, &plan)?) } else { projected_expr.push(e) } --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
