From 278b91d187400fec9037a3377412aa5c07ef8a8f Mon Sep 17 00:00:00 2001 From: osipovartem Date: Sat, 12 Sep 2026 21:43:57 +0300 Subject: [PATCH] Support wildcard RENAME projection modifier --- Cargo.lock | 4 +- datafusion/expr/src/logical_plan/builder.rs | 79 ++++++++++-- datafusion/expr/src/sql.rs | 16 ++- datafusion/sql/src/select.rs | 33 +++-- datafusion/sql/tests/sql_integration.rs | 132 ++++++++++++++++++++ 5 files changed, 243 insertions(+), 21 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 04e4b0f029697..086b05c79b73d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5876,7 +5876,7 @@ dependencies = [ [[package]] name = "sqlparser" version = "0.62.0" -source = "git+https://github.com/Embucket/datafusion-sqlparser-rs.git?rev=5c9be32322d55493797722f70f86dfd81baa1de4#5c9be32322d55493797722f70f86dfd81baa1de4" +source = "git+https://github.com/Embucket/datafusion-sqlparser-rs.git?rev=98560cdc40fb7c13bd9617cbdcc77e6ea795df7e#98560cdc40fb7c13bd9617cbdcc77e6ea795df7e" dependencies = [ "log", "recursive", @@ -5886,7 +5886,7 @@ dependencies = [ [[package]] name = "sqlparser_derive" version = "0.5.0" -source = "git+https://github.com/Embucket/datafusion-sqlparser-rs.git?rev=5c9be32322d55493797722f70f86dfd81baa1de4#5c9be32322d55493797722f70f86dfd81baa1de4" +source = "git+https://github.com/Embucket/datafusion-sqlparser-rs.git?rev=98560cdc40fb7c13bd9617cbdcc77e6ea795df7e#98560cdc40fb7c13bd9617cbdcc77e6ea795df7e" dependencies = [ "proc-macro2", "quote", diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index a3d2c6e17adf9..82734ebbcadd4 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -24,7 +24,7 @@ use std::iter::once; use std::sync::Arc; use crate::dml::CopyTo; -use crate::expr::{Alias, PlannedReplaceSelectItem, Sort as SortExpr}; +use crate::expr::{Alias, PlannedReplaceSelectItem, RenameSelectItem, Sort as SortExpr}; use crate::expr_rewriter::{ coerce_plan_expr_for_schema, normalize_col, normalize_col_with_schemas_and_ambiguity_check, normalize_cols, normalize_sorts, @@ -1975,6 +1975,11 @@ fn project_with_validation( } else { expanded }; + let expanded = if let Some(rename) = opt.rename { + rename_columns(expanded, &rename)? + } else { + expanded + }; for e in expanded { if validate { @@ -1997,6 +2002,11 @@ fn project_with_validation( } else { expanded }; + let expanded = if let Some(rename) = opt.rename { + rename_columns(expanded, &rename)? + } else { + expanded + }; for e in expanded { if validate { @@ -2047,20 +2057,69 @@ fn replace_columns( mut exprs: Vec, replace: &PlannedReplaceSelectItem, ) -> Result> { - for expr in exprs.iter_mut() { - if let Expr::Column(Column { name, .. }) = expr - && let Some((_, new_expr)) = replace - .items() - .iter() - .zip(replace.expressions().iter()) - .find(|(item, _)| item.column_name.value == *name) - { - *expr = new_expr.clone().alias(name.clone()) + let indices = wildcard_column_indices(&exprs); + let mut seen = HashSet::new(); + for (item, new_expr) in replace.items().iter().zip(replace.expressions()) { + let name = &item.column_name.value; + if !seen.insert(name) { + return plan_err!("Column '{name}' is specified more than once in REPLACE"); } + let index = wildcard_column_index(&indices, name, "REPLACE")?; + exprs[index] = new_expr.clone().alias(name.clone()); } Ok(exprs) } +/// Rename columns produced by a wildcard without changing their order or values. +fn rename_columns(mut exprs: Vec, rename: &RenameSelectItem) -> Result> { + let indices = wildcard_column_indices(&exprs); + let items = match rename { + RenameSelectItem::Single(item) => std::slice::from_ref(item), + RenameSelectItem::Multiple(items) => items, + }; + let mut seen = HashSet::new(); + for item in items { + let name = &item.ident.value; + if !seen.insert(name) { + return plan_err!("Column '{name}' is specified more than once in RENAME"); + } + let index = wildcard_column_index(&indices, name, "RENAME")?; + let expr = std::mem::take(&mut exprs[index]); + exprs[index] = match expr { + Expr::Alias(mut alias) => { + alias.name.clone_from(&item.alias.value); + Expr::Alias(alias) + } + expr => expr.alias(item.alias.value.clone()), + }; + } + Ok(exprs) +} + +fn wildcard_column_indices(exprs: &[Expr]) -> HashMap> { + let mut indices = HashMap::with_capacity(exprs.len()); + for (index, expr) in exprs.iter().enumerate() { + let (_, name) = expr.qualified_name(); + indices + .entry(name) + .and_modify(|existing| *existing = None) + .or_insert(Some(index)); + } + indices +} + +fn wildcard_column_index( + indices: &HashMap>, + name: &str, + modifier: &str, +) -> Result { + match indices.get(name) { + Some(Some(index)) => Ok(*index), + Some(None) => plan_err!("Column '{name}' specified in {modifier} is ambiguous"), + None => plan_err!("Column '{name}' specified in {modifier} does not exist"), + } +} + /// Create a SubqueryAlias to wrap a LogicalPlan. pub fn subquery_alias( plan: LogicalPlan, diff --git a/datafusion/expr/src/sql.rs b/datafusion/expr/src/sql.rs index 23e8d2f63d941..2a6394184a89e 100644 --- a/datafusion/expr/src/sql.rs +++ b/datafusion/expr/src/sql.rs @@ -118,10 +118,22 @@ impl Display for ExceptSelectItem { } } +#[derive(Clone, PartialEq, Eq, PartialOrd, Hash, Debug)] +pub struct IdentWithAlias { + pub ident: Ident, + pub alias: Ident, +} + +impl Display for IdentWithAlias { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + write!(f, "{} AS {}", self.ident, self.alias) + } +} + #[derive(Clone, PartialEq, Eq, PartialOrd, Hash, Debug)] pub enum RenameSelectItem { - Single(String), - Multiple(Vec), + Single(IdentWithAlias), + Multiple(Vec), } impl Display for RenameSelectItem { diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index 544b54bf59032..1d69e14225493 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -54,7 +54,7 @@ use datafusion_expr::{ use indexmap::IndexMap; use sqlparser::ast::{ - Distinct, Expr as SQLExpr, GroupByExpr, NamedWindowExpr, OrderBy, + Distinct, Expr as SQLExpr, GroupByExpr, NamedWindowExpr, OrderBy, RenameSelectItem, SelectItemQualifiedWildcardKind, WildcardAdditionalOptions, WindowType, visit_expressions_mut, }; @@ -1232,16 +1232,14 @@ impl SqlToRel<'_, S> { // opt_exclude is handled opt_exclude: _opt_exclude, opt_except: _opt_except, - opt_rename, + opt_rename: _opt_rename, opt_replace: _opt_replace, opt_ilike: _opt_ilike, opt_alias, wildcard_token: _wildcard_token, } = options; - if opt_rename.is_some() { - not_impl_err!("wildcard * with RENAME not supported ") - } else if opt_alias.is_some() { + if opt_alias.is_some() { not_impl_err!("wildcard * with AS alias not supported") } else { Ok(()) @@ -1258,12 +1256,23 @@ impl SqlToRel<'_, S> { planner_context: &mut PlannerContext, options: WildcardAdditionalOptions, ) -> Result { + let rename = options.opt_rename.map(|mut rename| { + let items = match &mut rename { + RenameSelectItem::Single(item) => std::slice::from_mut(item), + RenameSelectItem::Multiple(items) => items, + }; + for item in items { + item.ident.value = self.ident_normalizer.normalize(item.ident.clone()); + item.alias.value = self.ident_normalizer.normalize(item.alias.clone()); + } + rename + }); let planned_option = WildcardOptions { ilike: options.opt_ilike, exclude: options.opt_exclude, except: options.opt_except, replace: None, - rename: options.opt_rename, + rename, }; if let Some(replace) = options.opt_replace { let replace_expr = replace @@ -1285,8 +1294,18 @@ impl SqlToRel<'_, S> { }) .collect::>(); + let items = replace + .items + .into_iter() + .map(|item| { + let mut item = *item; + item.column_name.value = + self.ident_normalizer.normalize(item.column_name.clone()); + item + }) + .collect(); let planned_replace = PlannedReplaceSelectItem { - items: replace.items.into_iter().map(|i| *i).collect(), + items, planned_expressions: replace_expr, }; Ok(planned_option.with_replace(planned_replace)) diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index 55303e08f5ba6..899b642f03c27 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -2452,6 +2452,138 @@ fn scalar_expr_planner_applies_qualified_wildcard_options() { ); } +#[test] +fn select_wildcard_rename_columns() { + let plan = logical_plan_with_dialect( + "SELECT * RENAME (id AS person_id, state AS region) FROM person", + &SnowflakeDialect {}, + ) + .unwrap(); + + assert_eq!( + plan.schema() + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>(), + vec![ + "person_id", + "first_name", + "last_name", + "age", + "region", + "salary", + "birth_date", + "😀" + ] + ); + assert_snapshot!( + plan, + @r" + Projection: person.id AS person_id, person.first_name, person.last_name, person.age, person.state AS region, person.salary, person.birth_date, person.😀 + TableScan: person + " + ); +} + +#[test] +fn select_wildcard_combines_replace_and_rename() { + let plan = logical_plan_with_dialect( + "SELECT * EXCLUDE salary REPLACE (age + 1 AS age) \ + RENAME (id AS person_id, age AS years) FROM person", + &SnowflakeDialect {}, + ) + .unwrap(); + + assert_eq!( + plan.schema() + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>(), + vec![ + "person_id", + "first_name", + "last_name", + "years", + "state", + "birth_date", + "😀" + ] + ); + assert_snapshot!( + plan, + @r" + Projection: person.id AS person_id, person.first_name, person.last_name, person.age + Int64(1) AS years, person.state, person.birth_date, person.😀 + TableScan: person + " + ); +} + +#[test] +fn select_qualified_wildcard_rename() { + let plan = logical_plan_with_dialect( + "SELECT p.* EXCLUDE age RENAME id AS person_id FROM person AS p", + &SnowflakeDialect {}, + ) + .unwrap(); + + assert_eq!( + plan.schema() + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>(), + vec![ + "person_id", + "first_name", + "last_name", + "state", + "salary", + "birth_date", + "😀" + ] + ); +} + +#[test] +fn select_wildcard_rename_preserves_quoted_identifiers() { + let plan = logical_plan_with_dialect( + "SELECT * RENAME (\"First Name\" AS \"Given Name\") \ + FROM person_quoted_cols", + &SnowflakeDialect {}, + ) + .unwrap(); + + assert_eq!(plan.schema().field(1).name(), "Given Name"); +} + +#[rstest] +#[case( + "SELECT * REPLACE (id + 1 AS missing) FROM person", + "Column 'missing' specified in REPLACE does not exist" +)] +#[case( + "SELECT * RENAME missing AS renamed FROM person", + "Column 'missing' specified in RENAME does not exist" +)] +#[case( + "SELECT * RENAME (id AS one, id AS two) FROM person", + "Column 'id' is specified more than once in RENAME" +)] +#[case( + "SELECT * REPLACE (id + 1 AS id, id + 2 AS id) FROM person", + "Column 'id' is specified more than once in REPLACE" +)] +#[case( + "SELECT * RENAME id AS person_id FROM person p JOIN person q ON p.id = q.id", + "Column 'id' specified in RENAME is ambiguous" +)] +fn select_wildcard_modifiers_validate_targets(#[case] sql: &str, #[case] expected: &str) { + let err = logical_plan_with_dialect(sql, &SnowflakeDialect {}).unwrap_err(); + assert_contains!(err.strip_backtrace(), expected); +} + #[test] fn select_approx_median() { let sql = "SELECT approx_median(age) FROM person";