diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs index daff5edd07fe..a93dfa5903b1 100644 --- a/datafusion/expr/src/expr.rs +++ b/datafusion/expr/src/expr.rs @@ -978,28 +978,40 @@ impl GroupingSet { #[derive(Clone, PartialEq, Eq, Hash, Debug, Default)] pub struct WildcardOptions { - pub opt_ilike: Option, - pub opt_exclude: Option, - pub opt_except: Option, - pub opt_replace: Option, - pub opt_rename: Option, + pub ilike: Option, + pub exclude: Option, + pub except: Option, + pub replace: Option, + pub rename: Option, +} + +impl WildcardOptions { + pub fn with_replace(self, replace: PlannedReplaceSelectItem) -> Self { + WildcardOptions { + ilike: self.ilike, + exclude: self.exclude, + except: self.except, + replace: Some(replace), + rename: self.rename, + } + } } impl Display for WildcardOptions { fn fmt(&self, f: &mut Formatter) -> fmt::Result { - if let Some(ilike) = &self.opt_ilike { + if let Some(ilike) = &self.ilike { write!(f, " {ilike}")?; } - if let Some(exclude) = &self.opt_exclude { + if let Some(exclude) = &self.exclude { write!(f, " {exclude}")?; } - if let Some(except) = &self.opt_except { + if let Some(except) = &self.except { write!(f, " {except}")?; } - if let Some(replace) = &self.opt_replace { + if let Some(replace) = &self.replace { write!(f, " {replace}")?; } - if let Some(rename) = &self.opt_rename { + if let Some(rename) = &self.rename { write!(f, " {rename}")?; } Ok(()) @@ -1008,7 +1020,7 @@ impl Display for WildcardOptions { #[derive(Clone, PartialEq, Eq, Hash, Debug, Default)] pub struct PlannedReplaceSelectItem { - pub items: Vec>, + pub items: Vec, pub planned_expressions: Vec, } @@ -1021,7 +1033,7 @@ impl Display for PlannedReplaceSelectItem { } impl PlannedReplaceSelectItem { - pub fn items(&self) -> &[Box] { + pub fn items(&self) -> &[ReplaceSelectElement] { &self.items } @@ -2986,11 +2998,11 @@ mod test { None, None, Some(PlannedReplaceSelectItem { - items: vec![Box::new(ReplaceSelectElement { + items: vec![ReplaceSelectElement { expr: ast::Expr::Identifier(Ident::from("c1")), column_name: Ident::from("a1"), as_keyword: false - })], + }], planned_expressions: vec![] }), None @@ -3024,11 +3036,11 @@ mod test { opt_rename: Option, ) -> WildcardOptions { WildcardOptions { - opt_ilike, - opt_exclude, - opt_except, - opt_replace, - opt_rename, + ilike: opt_ilike, + exclude: opt_exclude, + except: opt_except, + replace: opt_replace, + rename: opt_rename, } } } diff --git a/datafusion/expr/src/utils.rs b/datafusion/expr/src/utils.rs index efca62910a35..4db5061e8fe7 100644 --- a/datafusion/expr/src/utils.rs +++ b/datafusion/expr/src/utils.rs @@ -402,8 +402,8 @@ pub fn expand_wildcard( }) .collect::>(); let excluded_columns = if let Some(WildcardOptions { - opt_exclude, - opt_except, + exclude: opt_exclude, + except: opt_except, .. }) = wildcard_options { @@ -436,8 +436,8 @@ pub fn expand_qualified_wildcard( DFSchema::try_from_qualified_schema(qualifier.clone(), &qualified_schema)? .with_functional_dependencies(projected_func_dependencies)?; let excluded_columns = if let Some(WildcardOptions { - opt_exclude, - opt_except, + exclude: opt_exclude, + except: opt_except, .. }) = wildcard_options { @@ -739,8 +739,8 @@ pub fn exprlist_to_fields<'a>( Expr::Wildcard { qualifier, options } => match qualifier { None => { let excluded: Vec = get_excluded_columns( - options.opt_exclude.as_ref(), - options.opt_except.as_ref(), + options.exclude.as_ref(), + options.except.as_ref(), wildcard_schema, None, )? @@ -748,26 +748,22 @@ pub fn exprlist_to_fields<'a>( .map(|c| c.flat_name()) .collect(); Ok::<_, DataFusionError>( - (0..wildcard_schema.fields().len()) - .filter_map(|i| { - let (qualifier, field) = - wildcard_schema.qualified_field(i); - let flat_name = qualifier - .map(|t| format!("{}.{}", t, field.name())) - .unwrap_or(field.name().clone()); - if excluded.contains(&flat_name) { - None - } else { - Some((qualifier.cloned(), Arc::new(field.to_owned()))) - } + wildcard_schema + .field_names() + .iter() + .enumerate() + .filter(|(_, s)| !excluded.contains(s)) + .map(|(i, _)| wildcard_schema.qualified_field(i)) + .map(|(qualifier, f)| { + (qualifier.cloned(), Arc::new(f.to_owned())) }) .collect::>(), ) } Some(qualifier) => { let excluded: Vec = get_excluded_columns( - options.opt_exclude.as_ref(), - options.opt_except.as_ref(), + options.exclude.as_ref(), + options.except.as_ref(), wildcard_schema, Some(qualifier), )? @@ -826,8 +822,8 @@ pub fn exprlist_len( options, } => { let excluded = get_excluded_columns( - options.opt_exclude.as_ref(), - options.opt_except.as_ref(), + options.exclude.as_ref(), + options.except.as_ref(), wildcard_schema.unwrap_or(schema), None, )? @@ -843,8 +839,8 @@ pub fn exprlist_len( options, } => { let excluded = get_excluded_columns( - options.opt_exclude.as_ref(), - options.opt_except.as_ref(), + options.exclude.as_ref(), + options.except.as_ref(), wildcard_schema.unwrap_or(schema), Some(qualifier), )? diff --git a/datafusion/optimizer/src/analyzer/expand_wildcard_rule.rs b/datafusion/optimizer/src/analyzer/expand_wildcard_rule.rs index c1eecab8ae3d..53ba3042f522 100644 --- a/datafusion/optimizer/src/analyzer/expand_wildcard_rule.rs +++ b/datafusion/optimizer/src/analyzer/expand_wildcard_rule.rs @@ -22,7 +22,7 @@ use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::{Transformed, TransformedResult}; use datafusion_common::{Column, Result}; use datafusion_expr::builder::validate_unique_names; -use datafusion_expr::expr::{Alias, PlannedReplaceSelectItem}; +use datafusion_expr::expr::PlannedReplaceSelectItem; use datafusion_expr::utils::{ expand_qualified_wildcard, expand_wildcard, find_base_plan, }; @@ -83,7 +83,7 @@ fn expand_exprlist(input: &LogicalPlan, expr: Vec) -> Result> { )?; // If there is a REPLACE statement, replace that column with the given // replace expression. Column name remains the same. - let replaced = if let Some(replace) = options.opt_replace { + let replaced = if let Some(replace) = options.replace { replace_columns(expanded, replace)? } else { expanded @@ -94,7 +94,7 @@ fn expand_exprlist(input: &LogicalPlan, expr: Vec) -> Result> { expand_wildcard(input.schema(), input, Some(&options))?; // If there is a REPLACE statement, replace that column with the given // replace expression. Column name remains the same. - let replaced = if let Some(replace) = options.opt_replace { + let replaced = if let Some(replace) = options.replace { replace_columns(expanded, replace)? } else { expanded @@ -149,11 +149,7 @@ fn replace_columns( .zip(replace.expressions().iter()) .find(|(item, _)| item.column_name.value == *name) { - *expr = Expr::Alias(Alias { - expr: Box::new(new_expr.clone()), - relation: None, - name: name.clone(), - }) + *expr = new_expr.clone().alias(name.clone()) } } } diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index 740d37a73be2..339234d9965c 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -625,6 +625,13 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { planner_context: &mut PlannerContext, options: WildcardAdditionalOptions, ) -> Result { + let planned_option = WildcardOptions { + ilike: options.opt_ilike, + exclude: options.opt_exclude, + except: options.opt_except, + replace: None, + rename: options.opt_rename, + }; if let Some(replace) = options.opt_replace { let replace_expr = replace .items @@ -639,25 +646,13 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { .clone()) }) .collect::>>()?; - let planned_option = PlannedReplaceSelectItem { - items: replace.items, + let planned_replace = PlannedReplaceSelectItem { + items: replace.items.into_iter().map(|i| *i).collect(), planned_expressions: replace_expr, }; - Ok(WildcardOptions { - opt_ilike: options.opt_ilike, - opt_exclude: options.opt_exclude, - opt_except: options.opt_except, - opt_replace: Some(planned_option), - opt_rename: options.opt_rename, - }) + Ok(planned_option.with_replace(planned_replace)) } else { - Ok(WildcardOptions { - opt_ilike: options.opt_ilike, - opt_exclude: options.opt_exclude, - opt_except: options.opt_except, - opt_replace: None, - opt_rename: options.opt_rename, - }) + Ok(planned_option) } } diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index 3ee98d063331..292e5bd040b6 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/datafusion/sqllogictest/test_files/window.slt @@ -3929,10 +3929,6 @@ b 1 3 a 1 4 b 5 5 -# expected failed message: -# DataFusion error: expand_wildcard_rule -# caused by -# Schema error: No field named aggregate_test_100.c1. Valid fields are rn. statement error SELECT * FROM (SELECT c1, c2, ROW_NUMBER() OVER(PARTITION BY c1) as rn