diff --git a/Cargo.lock b/Cargo.lock index a94228d01d2a1..7ab4231d0793b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2163,6 +2163,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 4fe7b65f6d05f..74b85501dde91 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 0000000000000..988c2310c3e45 --- /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::>(), + ); + 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 = (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 4e9839e2f7479..2552b61062fec 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>; } -/// 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.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>)>, +} + +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 { + 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.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>, plan: &LogicalPlan, ) -> Result> { + 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>, plan: &LogicalPlan, ) -> Result> { + 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::>() + ); + 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::::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 a6a254ea25bac..04a43df44b5bb 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, }; @@ -1070,18 +1070,7 @@ impl LogicalPlanBuilder { } pub(crate) fn normalize(plan: &LogicalPlan, column: Column) -> Result { - 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. @@ -2033,6 +2022,7 @@ fn project_with_validation( ) -> Result { 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 { @@ -2051,7 +2041,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) } @@ -2073,7 +2063,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) } @@ -2081,7 +2071,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) }