Support SELECT wildcard RENAME modifiers

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<Expr>,
     replace: &PlannedReplaceSelectItem,
 ) -> Result<Vec<Expr>> {
-    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<Expr>, rename: &RenameSelectItem) -> Result<Vec<Expr>> {
+    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<String, Option<usize>> {
+    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<String, Option<usize>>,
+    name: &str,
+    modifier: &str,
+) -> Result<usize> {
+    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,

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<String>),
+    Single(IdentWithAlias),
+    Multiple(Vec<IdentWithAlias>),
 }
 
 impl Display for RenameSelectItem {

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<S: ContextProvider> 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<S: ContextProvider> SqlToRel<'_, S> {
         planner_context: &mut PlannerContext,
         options: WildcardAdditionalOptions,
     ) -> Result<WildcardOptions> {
+        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<S: ContextProvider> SqlToRel<'_, S> {
                 })
                 .collect::<Vec<_>>();
 
+            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))

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<_>>(),
+        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<_>>(),
+        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<_>>(),
+        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";