Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 7 additions & 7 deletions datafusion/expr/src/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -367,18 +367,18 @@ fn get_excluded_columns(
};
idents.extend(exclude_owned.iter());
}
// Excluded columns should be unique
// Lookup uses the resolved name, so quoting differences cannot distinguish duplicates.
let n_elem = idents.len();
let unique_idents = idents.into_iter().collect::<HashSet<_>>();
// If HashSet size, and vector length are different, this means that some of the excluded columns
// are not unique. In this case return error.
if n_elem != unique_idents.len() {
let unique_names = idents
.into_iter()
.map(|ident| ident.value.as_str())
.collect::<HashSet<_>>();
if n_elem != unique_names.len() {
return plan_err!("EXCLUDE or EXCEPT contains duplicate column names");
}

let mut result = vec![];
for ident in unique_idents.into_iter() {
let col_name = ident.value.as_str();
for col_name in unique_names {
let (qualifier, field) = schema.qualified_field_with_name(qualifier, col_name)?;
result.push(Column::from((qualifier, field)));
}
Expand Down
45 changes: 24 additions & 21 deletions datafusion/sql/src/expr/function.rs
Original file line number Diff line number Diff line change
Expand Up @@ -40,25 +40,6 @@ use sqlparser::ast::{
WindowType,
};

fn function_wildcard_options(
options: WildcardAdditionalOptions,
) -> Result<Box<WildcardOptions>> {
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");
}

Ok(Box::new(WildcardOptions {
ilike: options.opt_ilike,
exclude: options.opt_exclude,
except: options.opt_except,
replace: None,
rename: options.opt_rename,
}))
}

/// Suggest a valid function based on an invalid input function name
///
/// Returns `None` if no valid matches are found. This happens when there are no
Expand Down Expand Up @@ -247,6 +228,28 @@ impl FunctionArgs {
type WithinGroupExtraction = (Vec<SortExpr>, Vec<Expr>, Vec<Option<ArgumentName>>);

impl<S: ContextProvider> SqlToRel<'_, S> {
fn function_wildcard_options(
&self,
options: WildcardAdditionalOptions,
) -> Result<Box<WildcardOptions>> {
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");
}

Ok(Box::new(WildcardOptions {
ilike: options.opt_ilike,
exclude: options
.opt_exclude
.map(|exclude| self.normalize_wildcard_exclude(exclude)),
except: options.opt_except,
replace: None,
rename: options.opt_rename,
}))
}

pub(super) fn sql_function_to_expr(
&self,
function: SQLFunction,
Expand Down Expand Up @@ -1109,7 +1112,7 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
#[expect(deprecated)]
let expr = Expr::Wildcard {
qualifier: None,
options: function_wildcard_options(options)?,
options: self.function_wildcard_options(options)?,
};
Ok((expr, None))
}
Expand Down Expand Up @@ -1140,7 +1143,7 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
#[expect(deprecated)]
let expr = Expr::Wildcard {
qualifier: qualifier.into(),
options: function_wildcard_options(options)?,
options: self.function_wildcard_options(options)?,
};
Ok((expr, None))
}
Expand Down
27 changes: 23 additions & 4 deletions datafusion/sql/src/select.rs
Original file line number Diff line number Diff line change
Expand Up @@ -54,9 +54,9 @@ use datafusion_expr::{

use indexmap::IndexMap;
use sqlparser::ast::{
Distinct, Expr as SQLExpr, GroupByExpr, NamedWindowExpr, OrderBy, RenameSelectItem,
SelectItemQualifiedWildcardKind, WildcardAdditionalOptions, WindowType,
visit_expressions_mut,
Distinct, ExcludeSelectItem, Expr as SQLExpr, GroupByExpr, NamedWindowExpr,
ObjectNamePart, OrderBy, RenameSelectItem, SelectItemQualifiedWildcardKind,
WildcardAdditionalOptions, WindowType, visit_expressions_mut,
};
use sqlparser::ast::{NamedWindowDefinition, Select, SelectItem, TableWithJoins};

Expand Down Expand Up @@ -1246,6 +1246,22 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
}
}

pub(crate) fn normalize_wildcard_exclude(
&self,
mut exclude: ExcludeSelectItem,
) -> ExcludeSelectItem {
let names = match &mut exclude {
ExcludeSelectItem::Single(name) => std::slice::from_mut(name),
ExcludeSelectItem::Multiple(names) => names,
};
for name in names {
if let [ObjectNamePart::Identifier(ident)] = name.0.as_mut_slice() {
ident.value = self.ident_normalizer.normalize(ident.clone());
}
}
exclude
}

/// If there is a REPLACE statement in the projected expression in the form of
/// "REPLACE (some_column_within_an_expr AS some_column)", we should plan the
/// replace expressions first.
Expand All @@ -1256,6 +1272,9 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
planner_context: &mut PlannerContext,
options: WildcardAdditionalOptions,
) -> Result<WildcardOptions> {
let exclude = options
.opt_exclude
.map(|exclude| self.normalize_wildcard_exclude(exclude));
let rename = options.opt_rename.map(|mut rename| {
let items = match &mut rename {
RenameSelectItem::Single(item) => std::slice::from_mut(item),
Expand All @@ -1269,7 +1288,7 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
});
let planned_option = WildcardOptions {
ilike: options.opt_ilike,
exclude: options.opt_exclude,
exclude,
except: options.opt_except,
replace: None,
rename,
Expand Down
21 changes: 21 additions & 0 deletions datafusion/sql/tests/sql_integration.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2477,6 +2477,27 @@ fn scalar_expr_planner_receives_wildcard_options() {
);
}

#[test]
fn scalar_expr_planner_normalizes_wildcard_exclude() {
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 scalar_expr_planner_applies_wildcard_ilike() {
let state = mock_session_state().with_expr_planner(Arc::new(ScalarWildcardPlanner));
Expand Down
18 changes: 18 additions & 0 deletions datafusion/sqllogictest/test_files/select.slt
Original file line number Diff line number Diff line change
Expand Up @@ -1289,6 +1289,24 @@ LIMIT 5
100 1000
200 2000

# Unquoted EXCLUDE targets follow the same normalization as input columns.
query III
SELECT * EXCLUDE B FROM table1 ORDER BY a
----
1 100 1000
2 200 2000

# Quoted targets preserve case and match quoted aliases.
query III
SELECT * EXCLUDE "B" FROM (SELECT a, b AS "B", c, d FROM table1) ORDER BY a
----
1 100 1000
2 200 2000

# Quoted and unquoted spellings that resolve to one field are duplicates.
statement error DataFusion error: Error during planning: EXCLUDE or EXCEPT contains duplicate column names
SELECT * EXCLUDE (B, "b") FROM table1

# when wildcard is prepended with table name, exclude should still work
# below query should emit all the columns except a and b
query II
Expand Down
Loading