From 8461571125c601702345b27a53d67d7120592579 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Tue, 1 Sep 2026 23:56:59 +0300 Subject: [PATCH 1/2] Expose input schema to scalar expression planners --- datafusion/expr/src/planner.rs | 18 ++++++++++ datafusion/expr/src/utils.rs | 37 +++++++++++++++++--- datafusion/sql/src/expr/function.rs | 37 ++++++++++++++++++-- datafusion/sql/tests/common/mod.rs | 30 +++++++++++++++-- datafusion/sql/tests/sql_integration.rs | 45 ++++++++++++++++++++++++- 5 files changed, 158 insertions(+), 9 deletions(-) diff --git a/datafusion/expr/src/planner.rs b/datafusion/expr/src/planner.rs index 99cf5bef19fbf..4658d22325d2f 100644 --- a/datafusion/expr/src/planner.rs +++ b/datafusion/expr/src/planner.rs @@ -154,6 +154,17 @@ pub trait ContextProvider { /// /// [Extending SQL in DataFusion: from ->> to TABLESAMPLE blog]: https://datafusion.apache.org/blog/2026/01/12/extending-sql pub trait ExprPlanner: Debug + Send + Sync { + /// Plans scalar functions, such as `CONCAT(, ...)`, with access to the input schema. + /// + /// Returns the original scalar function if planning is not possible. + fn plan_scalar_with_schema( + &self, + expr: RawScalarExpr, + _schema: &DFSchema, + ) -> Result> { + Ok(PlannerResult::Original(expr)) + } + /// Plan the binary operation between two expressions, returns original /// BinaryExpr if not possible fn plan_binary_op( @@ -343,6 +354,13 @@ pub struct RawAggregateExpr { pub null_treatment: Option, } +/// This structure is used by scalar function expression planners. +#[derive(Debug, Clone)] +pub struct RawScalarExpr { + pub func: Arc, + pub args: Vec, +} + /// This structure is used by `WindowFunctionPlanner` to plan operators with /// custom expressions. #[derive(Debug, Clone)] diff --git a/datafusion/expr/src/utils.rs b/datafusion/expr/src/utils.rs index 7f79c5cf18c4a..646beffbc173d 100644 --- a/datafusion/expr/src/utils.rs +++ b/datafusion/expr/src/utils.rs @@ -449,19 +449,48 @@ pub fn expand_wildcard( wildcard_options: Option<&WildcardOptions>, ) -> Result> { let mut columns_to_skip = exclude_using_columns(plan)?; + columns_to_skip.extend(excluded_columns_from_schema( + schema, + wildcard_options, + None, + )?); + Ok(get_exprs_except_skipped(schema, &columns_to_skip)) +} + +/// Resolves an unqualified wildcard using only the input schema. +/// +/// Unlike [`expand_wildcard`], this helper cannot account for duplicate columns introduced by a +/// join with a `USING` clause. It is intended for schema-aware expression planners that receive a +/// wildcard as a function argument. +pub fn expand_wildcard_from_schema( + schema: &DFSchema, + wildcard_options: Option<&WildcardOptions>, +) -> Result> { + let columns_to_skip = excluded_columns_from_schema(schema, wildcard_options, None)?; + Ok(get_exprs_except_skipped(schema, &columns_to_skip)) +} + +fn excluded_columns_from_schema( + schema: &DFSchema, + wildcard_options: Option<&WildcardOptions>, + qualifier: Option<&TableReference>, +) -> Result> { let excluded_columns = if let Some(WildcardOptions { exclude: opt_exclude, except: opt_except, .. }) = wildcard_options { - get_excluded_columns(opt_exclude.as_ref(), opt_except.as_ref(), schema, None)? + get_excluded_columns( + opt_exclude.as_ref(), + opt_except.as_ref(), + schema, + qualifier, + )? } else { vec![] }; - // Add each excluded `Column` to columns_to_skip - columns_to_skip.extend(excluded_columns); - Ok(get_exprs_except_skipped(schema, &columns_to_skip)) + Ok(excluded_columns.into_iter().collect()) } /// Resolves an `Expr::Wildcard` to a collection of qualified `Expr::Column`'s. diff --git a/datafusion/sql/src/expr/function.rs b/datafusion/sql/src/expr/function.rs index 4269263dea1d9..0b340ff9001dc 100644 --- a/datafusion/sql/src/expr/function.rs +++ b/datafusion/sql/src/expr/function.rs @@ -30,7 +30,7 @@ use datafusion_expr::{ self, HigherOrderFunction, Lambda, NullTreatment, ScalarFunction, Unnest, WildcardOptions, WindowFunction, }, - planner::{PlannerResult, RawAggregateExpr, RawWindowExpr}, + planner::{PlannerResult, RawAggregateExpr, RawScalarExpr, RawWindowExpr}, type_coercion::functions::value_fields_with_higher_order_udf, }; use sqlparser::ast::{ @@ -348,7 +348,19 @@ impl SqlToRel<'_, S> { }; // After resolution, all arguments are positional - let inner = ScalarFunction::new_udf(fm, resolved_args); + let mut scalar_expr = RawScalarExpr { + func: fm, + args: resolved_args, + }; + for planner in self.context_provider.get_expr_planners().iter() { + match planner.plan_scalar_with_schema(scalar_expr, schema)? { + PlannerResult::Planned(expr) => return Ok(expr), + PlannerResult::Original(expr) => scalar_expr = expr, + } + } + + let RawScalarExpr { func, args } = scalar_expr; + let inner = ScalarFunction::new_udf(func, args); if name.eq_ignore_ascii_case(inner.name()) { return Ok(Expr::ScalarFunction(inner)); @@ -1063,6 +1075,27 @@ impl SqlToRel<'_, S> { }; Ok((expr, None)) } + FunctionArg::Unnamed(FunctionArgExpr::WildcardWithOptions(options)) => { + if options.opt_alias.is_some() { + return not_impl_err!("wildcard function argument with AS alias"); + } + if options.opt_replace.is_some() { + return not_impl_err!("wildcard function argument with REPLACE"); + } + + #[expect(deprecated)] + let expr = Expr::Wildcard { + qualifier: None, + options: Box::new(WildcardOptions { + ilike: options.opt_ilike, + exclude: options.opt_exclude, + except: options.opt_except, + replace: None, + rename: options.opt_rename, + }), + }; + Ok((expr, None)) + } FunctionArg::Unnamed(FunctionArgExpr::QualifiedWildcard(object_name)) => { let qualifier = self.object_name_to_table_reference(object_name)?; // Sanity check on qualifier with schema diff --git a/datafusion/sql/tests/common/mod.rs b/datafusion/sql/tests/common/mod.rs index a9e08440c318c..dfe393c1952ac 100644 --- a/datafusion/sql/tests/common/mod.rs +++ b/datafusion/sql/tests/common/mod.rs @@ -27,10 +27,11 @@ use datafusion_common::datatype::DataTypeExt; use datafusion_common::file_options::file_type::FileType; use datafusion_common::{Column, DFSchema, GetExt, Result, TableReference, plan_err}; use datafusion_expr::planner::{ - ExprPlanner, PlannerResult, RawAggregateExpr, TypePlanner, + ExprPlanner, PlannerResult, RawAggregateExpr, RawScalarExpr, TypePlanner, }; +use datafusion_expr::utils::expand_wildcard_from_schema; use datafusion_expr::{ - AggregateUDF, Expr, HigherOrderUDF, ScalarUDF, TableSource, WindowUDF, + AggregateUDF, Expr, ExprSchemable, HigherOrderUDF, ScalarUDF, TableSource, WindowUDF, }; use datafusion_functions_nested::expr_fn::make_array; use datafusion_sql::planner::ContextProvider; @@ -448,3 +449,28 @@ impl ExprPlanner for QualifiedWildcardCountPlanner { Ok(PlannerResult::Original(expr)) } } + +#[derive(Debug)] +pub struct ScalarWildcardPlanner; + +impl ExprPlanner for ScalarWildcardPlanner { + fn plan_scalar_with_schema( + &self, + mut expr: RawScalarExpr, + schema: &DFSchema, + ) -> Result> { + #[expect(deprecated)] + let [Expr::Wildcard { options, .. }] = expr.args.as_slice() else { + return Ok(PlannerResult::Original(expr)); + }; + if expr.func.name() != "concat" { + return Ok(PlannerResult::Original(expr)); + } + + expr.args = expand_wildcard_from_schema(schema, Some(options))? + .into_iter() + .filter(|expr| expr.get_type(schema).is_ok_and(|ty| ty == DataType::Utf8)) + .collect(); + Ok(PlannerResult::Original(expr)) + } +} diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index 39a17024fb632..e5897eaa33fee 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -44,7 +44,8 @@ use datafusion_sql::{ }; use crate::common::{ - CustomExprPlanner, CustomTypePlanner, MockSessionState, QualifiedWildcardCountPlanner, + CustomExprPlanner, CustomTypePlanner, MockSessionState, + QualifiedWildcardCountPlanner, ScalarWildcardPlanner, }; use datafusion_functions::core::planner::CoreFunctionPlanner; use datafusion_functions_aggregate::{ @@ -2180,6 +2181,48 @@ fn aggregate_expr_planner_can_resolve_qualified_wildcard_from_schema() { ); } +#[test] +fn scalar_expr_planner_can_resolve_wildcard_from_schema() { + let state = mock_session_state().with_expr_planner(Arc::new(ScalarWildcardPlanner)); + let plan = logical_plan_from_state( + "SELECT concat(*) FROM person AS p", + &GenericDialect {}, + ParserOptions::default(), + state, + ) + .unwrap(); + + assert_snapshot!( + plan, + @r" + Projection: concat(p.first_name, p.last_name, p.state) + SubqueryAlias: p + TableScan: person + " + ); +} + +#[test] +fn scalar_expr_planner_receives_wildcard_options() { + let state = mock_session_state().with_expr_planner(Arc::new(ScalarWildcardPlanner)); + let plan = logical_plan_from_state( + "SELECT concat(* EXCLUDE first_name) FROM person AS p", + &GenericDialect {}, + ParserOptions::default(), + state, + ) + .unwrap(); + + assert_snapshot!( + plan, + @r" + Projection: concat(p.last_name, p.state) + SubqueryAlias: p + TableScan: person + " + ); +} + #[test] fn select_approx_median() { let sql = "SELECT approx_median(age) FROM person"; From 9c8e6192ac5615951a99d20353d55038c5719389 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Wed, 2 Sep 2026 00:44:34 +0300 Subject: [PATCH 2/2] Apply ILIKE when expanding wildcard arguments --- datafusion/expr/src/utils.rs | 34 ++++++++++++++++++++++++- datafusion/sql/tests/sql_integration.rs | 21 +++++++++++++++ 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/datafusion/expr/src/utils.rs b/datafusion/expr/src/utils.rs index 646beffbc173d..297fbce14c795 100644 --- a/datafusion/expr/src/utils.rs +++ b/datafusion/expr/src/utils.rs @@ -490,7 +490,39 @@ fn excluded_columns_from_schema( } else { vec![] }; - Ok(excluded_columns.into_iter().collect()) + let mut excluded_columns = excluded_columns.into_iter().collect::>(); + if let Some(ilike) = wildcard_options.and_then(|options| options.ilike.as_ref()) { + excluded_columns.extend( + schema + .columns() + .into_iter() + .filter(|column| !matches_ilike(&column.name, &ilike.pattern)), + ); + } + Ok(excluded_columns) +} + +fn matches_ilike(value: &str, pattern: &str) -> bool { + let value = value.to_lowercase().chars().collect::>(); + let pattern = pattern.to_lowercase().chars().collect::>(); + let mut previous = vec![false; value.len() + 1]; + previous[0] = true; + + for pattern_char in pattern { + let mut current = vec![false; value.len() + 1]; + if pattern_char == '%' { + current[0] = previous[0]; + } + for (index, value_char) in value.iter().enumerate() { + current[index + 1] = match pattern_char { + '%' => previous[index + 1] || current[index], + '_' => previous[index], + literal => previous[index] && literal == *value_char, + }; + } + previous = current; + } + previous[value.len()] } /// Resolves an `Expr::Wildcard` to a collection of qualified `Expr::Column`'s. diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index e5897eaa33fee..71e7ee7ad9bdf 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -2223,6 +2223,27 @@ fn scalar_expr_planner_receives_wildcard_options() { ); } +#[test] +fn scalar_expr_planner_applies_wildcard_ilike() { + let state = mock_session_state().with_expr_planner(Arc::new(ScalarWildcardPlanner)); + let plan = logical_plan_from_state( + "SELECT concat(* ILIKE '%name') FROM person AS p", + &GenericDialect {}, + ParserOptions::default(), + state, + ) + .unwrap(); + + assert_snapshot!( + plan, + @r" + Projection: concat(p.first_name, p.last_name) + SubqueryAlias: p + TableScan: person + " + ); +} + #[test] fn select_approx_median() { let sql = "SELECT approx_median(age) FROM person";