Skip to content

Commit

Permalink
Wildcard fix
Browse files Browse the repository at this point in the history
  • Loading branch information
peter-toth committed Nov 25, 2024
1 parent 2e05648 commit 951704c
Show file tree
Hide file tree
Showing 12 changed files with 72 additions and 66 deletions.
20 changes: 12 additions & 8 deletions datafusion/expr/src/expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -311,10 +311,7 @@ pub enum Expr {
///
/// This expr has to be resolved to a list of columns before translating logical
/// plan into physical plan.
Wildcard {
qualifier: Option<TableReference>,
options: WildcardOptions,
},
Wildcard(Wildcard),
/// List of grouping set expressions. Only valid in the context of an aggregate
/// GROUP BY expression list
GroupingSet(GroupingSet),
Expand Down Expand Up @@ -367,6 +364,13 @@ impl<'a> TreeNodeContainer<'a, Self> for Expr {
}
}

/// Wildcard expression.
#[derive(Clone, PartialEq, Eq, PartialOrd, Hash, Debug)]
pub struct Wildcard {
pub qualifier: Option<TableReference>,
pub options: WildcardOptions,
}

/// UNNEST expression.
#[derive(Clone, PartialEq, Eq, PartialOrd, Hash, Debug)]
pub struct Unnest {
Expand Down Expand Up @@ -1797,9 +1801,9 @@ impl HashNode for Expr {
Expr::ScalarSubquery(subquery) => {
subquery.hash(state);
}
Expr::Wildcard { qualifier, options } => {
qualifier.hash(state);
options.hash(state);
Expr::Wildcard(wildcard) => {
wildcard.hash(state);
wildcard.hash(state);
}
Expr::GroupingSet(grouping_set) => {
mem::discriminant(grouping_set).hash(state);
Expand Down Expand Up @@ -2321,7 +2325,7 @@ impl Display for Expr {
write!(f, "{expr} IN ([{}])", expr_vec_fmt!(list))
}
}
Expr::Wildcard { qualifier, options } => match qualifier {
Expr::Wildcard(Wildcard { qualifier, options }) => match qualifier {
Some(qualifier) => write!(f, "{qualifier}.*{options}"),
None => write!(f, "*{options}"),
},
Expand Down
18 changes: 9 additions & 9 deletions datafusion/expr/src/expr_fn.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
use crate::expr::{
AggregateFunction, BinaryExpr, Cast, Exists, GroupingSet, InList, InSubquery,
Placeholder, TryCast, Unnest, WildcardOptions, WindowFunction,
Placeholder, TryCast, Unnest, Wildcard, WildcardOptions, WindowFunction,
};
use crate::function::{
AccumulatorArgs, AccumulatorFactoryFunction, PartitionEvaluatorFactory,
Expand Down Expand Up @@ -121,18 +121,18 @@ pub fn placeholder(id: impl Into<String>) -> Expr {
/// assert_eq!(p.to_string(), "*")
/// ```
pub fn wildcard() -> Expr {
Expr::Wildcard {
Expr::Wildcard(Wildcard {
qualifier: None,
options: WildcardOptions::default(),
}
})
}

/// Create an '*' [`Expr::Wildcard`] expression with the wildcard options
pub fn wildcard_with_options(options: WildcardOptions) -> Expr {
Expr::Wildcard {
Expr::Wildcard(Wildcard {
qualifier: None,
options,
}
})
}

/// Create an 't.*' [`Expr::Wildcard`] expression that matches all columns from a specific table
Expand All @@ -146,21 +146,21 @@ pub fn wildcard_with_options(options: WildcardOptions) -> Expr {
/// assert_eq!(p.to_string(), "t.*")
/// ```
pub fn qualified_wildcard(qualifier: impl Into<TableReference>) -> Expr {
Expr::Wildcard {
Expr::Wildcard(Wildcard {
qualifier: Some(qualifier.into()),
options: WildcardOptions::default(),
}
})
}

/// Create an 't.*' [`Expr::Wildcard`] expression with the wildcard options
pub fn qualified_wildcard_with_options(
qualifier: impl Into<TableReference>,
options: WildcardOptions,
) -> Expr {
Expr::Wildcard {
Expr::Wildcard(Wildcard {
qualifier: Some(qualifier.into()),
options,
}
})
}

/// Return a new expression `left <op> right`
Expand Down
8 changes: 4 additions & 4 deletions datafusion/expr/src/logical_plan/plan.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ use std::sync::{Arc, OnceLock};
use super::dml::CopyTo;
use super::DdlStatement;
use crate::builder::{change_redundant_column, unnest_with_options};
use crate::expr::{Placeholder, Sort as SortExpr, WindowFunction};
use crate::expr::{Placeholder, Sort as SortExpr, Wildcard, WindowFunction};
use crate::expr_rewriter::{
create_col_from_scalar_expr, normalize_cols, normalize_sorts, NamePreserver,
};
Expand Down Expand Up @@ -3187,12 +3187,12 @@ fn calc_func_dependencies_for_project(
let proj_indices = exprs
.iter()
.map(|expr| match expr {
Expr::Wildcard { qualifier, options } => {
Expr::Wildcard(Wildcard { qualifier, options }) => {
let wildcard_fields = exprlist_to_fields(
vec![&Expr::Wildcard {
vec![&Expr::Wildcard(Wildcard {
qualifier: qualifier.clone(),
options: options.clone(),
}],
})],
input,
)?;
Ok::<_, DataFusionError>(
Expand Down
12 changes: 6 additions & 6 deletions datafusion/expr/src/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ use std::collections::HashSet;
use std::ops::Deref;
use std::sync::Arc;

use crate::expr::{Alias, Sort, WildcardOptions, WindowFunction};
use crate::expr::{Alias, Sort, Wildcard, WildcardOptions, WindowFunction};
use crate::expr_rewriter::strip_outer_reference;
use crate::{
and, BinaryExpr, Expr, ExprSchemable, Filter, GroupingSet, LogicalPlan, Operator,
Expand Down Expand Up @@ -703,7 +703,7 @@ pub fn exprlist_to_fields<'a>(
let result = exprs
.into_iter()
.map(|e| match e {
Expr::Wildcard { qualifier, options } => match qualifier {
Expr::Wildcard(Wildcard { qualifier, options }) => match qualifier {
None => {
let excluded: Vec<String> = get_excluded_columns(
options.exclude.as_ref(),
Expand Down Expand Up @@ -802,10 +802,10 @@ pub fn exprlist_len(
exprs
.iter()
.map(|e| match e {
Expr::Wildcard {
Expr::Wildcard(Wildcard {
qualifier: None,
options,
} => {
}) => {
let excluded = get_excluded_columns(
options.exclude.as_ref(),
options.except.as_ref(),
Expand All @@ -819,10 +819,10 @@ pub fn exprlist_len(
.len(),
)
}
Expr::Wildcard {
Expr::Wildcard(Wildcard {
qualifier: Some(qualifier),
options,
} => {
}) => {
let related_wildcard_schema = wildcard_schema.as_ref().map_or_else(
|| Ok(Arc::clone(schema)),
|schema| {
Expand Down
4 changes: 2 additions & 2 deletions datafusion/optimizer/src/analyzer/expand_wildcard_rule.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::PlannedReplaceSelectItem;
use datafusion_expr::expr::{PlannedReplaceSelectItem, Wildcard};
use datafusion_expr::utils::{
expand_qualified_wildcard, expand_wildcard, find_base_plan,
};
Expand Down Expand Up @@ -89,7 +89,7 @@ fn expand_exprlist(input: &LogicalPlan, expr: Vec<Expr>) -> Result<Vec<Expr>> {
let input = find_base_plan(input);
for e in expr {
match e {
Expr::Wildcard { qualifier, options } => {
Expr::Wildcard(Wildcard { qualifier, options }) => {
if let Some(qualifier) = qualifier {
let expanded = expand_qualified_wildcard(
&qualifier,
Expand Down
6 changes: 3 additions & 3 deletions datafusion/optimizer/src/analyzer/inline_table_scan.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ use crate::analyzer::AnalyzerRule;
use datafusion_common::config::ConfigOptions;
use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
use datafusion_common::{Column, Result};
use datafusion_expr::expr::WildcardOptions;
use datafusion_expr::expr::{Wildcard, WildcardOptions};
use datafusion_expr::{logical_plan::LogicalPlan, Expr, LogicalPlanBuilder};

/// Analyzed rule that inlines TableScan that provide a [`LogicalPlan`]
Expand Down Expand Up @@ -93,10 +93,10 @@ fn generate_projection_expr(
)));
}
} else {
exprs.push(Expr::Wildcard {
exprs.push(Expr::Wildcard(Wildcard {
qualifier: None,
options: WildcardOptions::default(),
});
}));
}
Ok(exprs)
}
Expand Down
6 changes: 3 additions & 3 deletions datafusion/proto/src/logical_plan/from_proto.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ use datafusion_common::{
exec_datafusion_err, internal_err, plan_datafusion_err, RecursionUnnestOption,
Result, ScalarValue, TableReference, UnnestOptions,
};
use datafusion_expr::expr::{Alias, Placeholder, Sort};
use datafusion_expr::expr::{Alias, Placeholder, Sort, Wildcard};
use datafusion_expr::expr::{Unnest, WildcardOptions};
use datafusion_expr::ExprFunctionExt;
use datafusion_expr::{
Expand Down Expand Up @@ -511,10 +511,10 @@ pub fn parse_expr(
))),
ExprType::Wildcard(protobuf::Wildcard { qualifier }) => {
let qualifier = qualifier.to_owned().map(|x| x.try_into()).transpose()?;
Ok(Expr::Wildcard {
Ok(Expr::Wildcard(Wildcard {
qualifier,
options: WildcardOptions::default(),
})
}))
}
ExprType::ScalarUdfExpr(protobuf::ScalarUdfExprNode {
fun_name,
Expand Down
4 changes: 2 additions & 2 deletions datafusion/proto/src/logical_plan/to_proto.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
use datafusion_common::{TableReference, UnnestOptions};
use datafusion_expr::expr::{
self, Alias, Between, BinaryExpr, Cast, GroupingSet, InList, Like, Placeholder,
ScalarFunction, Unnest,
ScalarFunction, Unnest, Wildcard,
};
use datafusion_expr::{
logical_plan::PlanType, logical_plan::StringifiedPlan, Expr, JoinConstraint,
Expand Down Expand Up @@ -552,7 +552,7 @@ pub fn serialize_expr(
expr_type: Some(ExprType::InList(expr)),
}
}
Expr::Wildcard { qualifier, .. } => protobuf::LogicalExprNode {
Expr::Wildcard(Wildcard { qualifier, .. }) => protobuf::LogicalExprNode {
expr_type: Some(ExprType::Wildcard(protobuf::Wildcard {
qualifier: qualifier.to_owned().map(|x| x.into()),
})),
Expand Down
10 changes: 5 additions & 5 deletions datafusion/proto/tests/cases/roundtrip_logical_plan.rs
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ use datafusion_common::{
use datafusion_expr::dml::CopyTo;
use datafusion_expr::expr::{
self, Between, BinaryExpr, Case, Cast, GroupingSet, InList, Like, ScalarFunction,
Unnest, WildcardOptions,
Unnest, Wildcard, WildcardOptions,
};
use datafusion_expr::logical_plan::{Extension, UserDefinedLogicalNodeCore};
use datafusion_expr::{
Expand Down Expand Up @@ -2059,21 +2059,21 @@ fn roundtrip_unnest() {

#[test]
fn roundtrip_wildcard() {
let test_expr = Expr::Wildcard {
let test_expr = Expr::Wildcard(Wildcard {
qualifier: None,
options: WildcardOptions::default(),
};
});

let ctx = SessionContext::new();
roundtrip_expr_test(test_expr, ctx);
}

#[test]
fn roundtrip_qualified_wildcard() {
let test_expr = Expr::Wildcard {
let test_expr = Expr::Wildcard(Wildcard {
qualifier: Some("foo".into()),
options: WildcardOptions::default(),
};
});

let ctx = SessionContext::new();
roundtrip_expr_test(test_expr, ctx);
Expand Down
20 changes: 11 additions & 9 deletions datafusion/sql/src/expr/function.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ use datafusion_common::{
internal_datafusion_err, internal_err, not_impl_err, plan_datafusion_err, plan_err,
DFSchema, Dependency, Result,
};
use datafusion_expr::expr::WildcardOptions;
use datafusion_expr::expr::{Wildcard, WildcardOptions};
use datafusion_expr::expr::{ScalarFunction, Unnest};
use datafusion_expr::planner::PlannerResult;
use datafusion_expr::{
Expand Down Expand Up @@ -413,28 +413,30 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
name: _,
arg: FunctionArgExpr::Wildcard,
operator: _,
} => Ok(Expr::Wildcard {
} => Ok(Expr::Wildcard(Wildcard {
qualifier: None,
options: WildcardOptions::default(),
}),
})),
FunctionArg::Unnamed(FunctionArgExpr::Expr(arg)) => {
self.sql_expr_to_logical_expr(arg, schema, planner_context)
}
FunctionArg::Unnamed(FunctionArgExpr::Wildcard) => Ok(Expr::Wildcard {
qualifier: None,
options: WildcardOptions::default(),
}),
FunctionArg::Unnamed(FunctionArgExpr::Wildcard) => {
Ok(Expr::Wildcard(Wildcard {
qualifier: None,
options: WildcardOptions::default(),
}))
}
FunctionArg::Unnamed(FunctionArgExpr::QualifiedWildcard(object_name)) => {
let qualifier = self.object_name_to_table_reference(object_name)?;
// Sanity check on qualifier with schema
let qualified_indices = schema.fields_indices_with_qualified(&qualifier);
if qualified_indices.is_empty() {
return plan_err!("Invalid qualifier {qualifier}");
}
Ok(Expr::Wildcard {
Ok(Expr::Wildcard(Wildcard {
qualifier: Some(qualifier),
options: WildcardOptions::default(),
})
}))
}
_ => not_impl_err!("Unsupported qualified wildcard argument: {sql:?}"),
}
Expand Down
10 changes: 5 additions & 5 deletions datafusion/sql/src/expr/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,8 +30,8 @@ use datafusion_common::{
internal_datafusion_err, internal_err, not_impl_err, plan_err, DFSchema, Result,
ScalarValue,
};
use datafusion_expr::expr::ScalarFunction;
use datafusion_expr::expr::{InList, WildcardOptions};
use datafusion_expr::expr::{ScalarFunction, Wildcard};
use datafusion_expr::{
lit, Between, BinaryExpr, Cast, Expr, ExprSchemable, GetFieldAccess, Like, Literal,
Operator, TryCast,
Expand Down Expand Up @@ -565,14 +565,14 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
}
not_impl_err!("AnyOp not supported by ExprPlanner: {binary_expr:?}")
}
SQLExpr::Wildcard => Ok(Expr::Wildcard {
SQLExpr::Wildcard => Ok(Expr::Wildcard(Wildcard {
qualifier: None,
options: WildcardOptions::default(),
}),
SQLExpr::QualifiedWildcard(object_name) => Ok(Expr::Wildcard {
})),
SQLExpr::QualifiedWildcard(object_name) => Ok(Expr::Wildcard(Wildcard {
qualifier: Some(self.object_name_to_table_reference(object_name)?),
options: WildcardOptions::default(),
}),
})),
SQLExpr::Tuple(values) => self.parse_tuple(schema, planner_context, values),
_ => not_impl_err!("Unsupported ast node in sqltorel: {sql:?}"),
}
Expand Down
Loading

0 comments on commit 951704c

Please sign in to comment.