Skip to content

Commit

Permalink
Minor: Remove unecessary `clone_wih_replacement
Browse files Browse the repository at this point in the history
  • Loading branch information
@alamb
alamb committed Aug 8, 2023
1 parent 627abd7 commit 408aefc
Showing 1 changed file with 15 additions and 305 deletions.
320 changes: 15 additions & 305 deletions datafusion/sql/src/utils.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,32 +18,28 @@
//! SQL Utility Functions

use arrow_schema::{DataType, DECIMAL128_MAX_PRECISION, DECIMAL_DEFAULT_SCALE};
use datafusion_common::tree_node::{Transformed, TreeNode};
use sqlparser::ast::Ident;

use datafusion_common::plan_err;
use datafusion_common::{DataFusionError, Result, ScalarValue};
use datafusion_expr::expr::{
AggregateFunction, AggregateUDF, Alias, Between, BinaryExpr, Case, GetIndexedField,
GroupingSet, InList, InSubquery, Like, Placeholder, ScalarFunction, ScalarUDF,
WindowFunction,
};
use datafusion_expr::expr::{Cast, Sort};
use datafusion_expr::expr::{Alias, GroupingSet, WindowFunction};
use datafusion_expr::expr_vec_fmt;
use datafusion_expr::utils::{expr_as_column_expr, find_column_exprs};
use datafusion_expr::{Expr, LogicalPlan, TryCast};
use datafusion_expr::{Expr, LogicalPlan};
use std::collections::HashMap;

/// Make a best-effort attempt at resolving all columns in the expression tree
pub(crate) fn resolve_columns(expr: &Expr, plan: &LogicalPlan) -> Result<Expr> {
clone_with_replacement(expr, &|nested_expr| {
expr.clone().transform_up(&|nested_expr| {
match nested_expr {
Expr::Column(col) => {
let field = plan.schema().field_from_column(col)?;
Ok(Some(Expr::Column(field.qualified_column())))
let field = plan.schema().field_from_column(&col)?;
Ok(Transformed::Yes(Expr::Column(field.qualified_column())))
}
_ => {
// keep recursing
Ok(None)
Ok(Transformed::No(nested_expr))
}
}
})
Expand All @@ -68,11 +64,11 @@ pub(crate) fn rebase_expr(
base_exprs: &[Expr],
plan: &LogicalPlan,
) -> Result<Expr> {
clone_with_replacement(expr, &|nested_expr| {
if base_exprs.contains(nested_expr) {
Ok(Some(expr_as_column_expr(nested_expr, plan)?))
expr.clone().transform_up(&|nested_expr| {
if base_exprs.contains(&nested_expr) {
Ok(Transformed::Yes(expr_as_column_expr(&nested_expr, plan)?))
} else {
Ok(None)
Ok(Transformed::No(nested_expr))
}
})
}
Expand Down Expand Up @@ -132,292 +128,6 @@ fn check_column_satisfies_expr(
Ok(())
}

/// Returns a cloned `Expr`, but any of the `Expr`'s in the tree may be
/// replaced/customized by the replacement function.
///
/// The replacement function is called repeatedly with `Expr`, starting with
/// the argument `expr`, then descending depth-first through its
/// descendants. The function chooses to replace or keep (clone) each `Expr`.
///
/// The function's return type is `Result<Option<Expr>>>`, where:
///
/// * `Ok(Some(replacement_expr))`: A replacement `Expr` is provided; it is
/// swapped in at the particular node in the tree. Any nested `Expr` are
/// not subject to cloning/replacement.
/// * `Ok(None)`: A replacement `Expr` is not provided. The `Expr` is
/// recreated, with all of its nested `Expr`'s subject to
/// cloning/replacement.
/// * `Err(err)`: Any error returned by the function is returned as-is by
/// `clone_with_replacement()`.
fn clone_with_replacement<F>(expr: &Expr, replacement_fn: &F) -> Result<Expr>
where
F: Fn(&Expr) -> Result<Option<Expr>>,
{
let replacement_opt = replacement_fn(expr)?;

match replacement_opt {
// If we were provided a replacement, use the replacement. Do not
// descend further.
Some(replacement) => Ok(replacement),
// No replacement was provided, clone the node and recursively call
// clone_with_replacement() on any nested expressions.
None => match expr {
Expr::AggregateFunction(AggregateFunction {
fun,
args,
distinct,
filter,
order_by,
}) => Ok(Expr::AggregateFunction(AggregateFunction::new(
fun.clone(),
args.iter()
.map(|e| clone_with_replacement(e, replacement_fn))
.collect::<Result<Vec<Expr>>>()?,
*distinct,
filter.clone(),
order_by.clone(),
))),
Expr::WindowFunction(WindowFunction {
fun,
args,
partition_by,
order_by,
window_frame,
}) => Ok(Expr::WindowFunction(WindowFunction::new(
fun.clone(),
args.iter()
.map(|e| clone_with_replacement(e, replacement_fn))
.collect::<Result<Vec<_>>>()?,
partition_by
.iter()
.map(|e| clone_with_replacement(e, replacement_fn))
.collect::<Result<Vec<_>>>()?,
order_by
.iter()
.map(|e| clone_with_replacement(e, replacement_fn))
.collect::<Result<Vec<_>>>()?,
window_frame.clone(),
))),
Expr::AggregateUDF(AggregateUDF {
fun,
args,
filter,
order_by,
}) => Ok(Expr::AggregateUDF(AggregateUDF::new(
fun.clone(),
args.iter()
.map(|e| clone_with_replacement(e, replacement_fn))
.collect::<Result<Vec<Expr>>>()?,
filter.clone(),
order_by.clone(),
))),
Expr::Alias(Alias { expr, name, .. }) => Ok(Expr::Alias(Alias::new(
clone_with_replacement(expr, replacement_fn)?,
name.clone(),
))),
Expr::Between(Between {
expr,
negated,
low,
high,
}) => Ok(Expr::Between(Between::new(
Box::new(clone_with_replacement(expr, replacement_fn)?),
*negated,
Box::new(clone_with_replacement(low, replacement_fn)?),
Box::new(clone_with_replacement(high, replacement_fn)?),
))),
Expr::InList(InList {
expr: nested_expr,
list,
negated,
}) => Ok(Expr::InList(InList::new(
Box::new(clone_with_replacement(nested_expr, replacement_fn)?),
list.iter()
.map(|e| clone_with_replacement(e, replacement_fn))
.collect::<Result<Vec<Expr>>>()?,
*negated,
))),
Expr::BinaryExpr(BinaryExpr { left, right, op }) => {
Ok(Expr::BinaryExpr(BinaryExpr::new(
Box::new(clone_with_replacement(left, replacement_fn)?),
*op,
Box::new(clone_with_replacement(right, replacement_fn)?),
)))
}
Expr::Like(Like {
negated,
expr,
pattern,
escape_char,
case_insensitive,
}) => Ok(Expr::Like(Like::new(
*negated,
Box::new(clone_with_replacement(expr, replacement_fn)?),
Box::new(clone_with_replacement(pattern, replacement_fn)?),
*escape_char,
*case_insensitive,
))),
Expr::SimilarTo(Like {
negated,
expr,
pattern,
escape_char,
case_insensitive,
}) => Ok(Expr::SimilarTo(Like::new(
*negated,
Box::new(clone_with_replacement(expr, replacement_fn)?),
Box::new(clone_with_replacement(pattern, replacement_fn)?),
*escape_char,
*case_insensitive,
))),
Expr::Case(case) => Ok(Expr::Case(Case::new(
match &case.expr {
Some(case_expr) => {
Some(Box::new(clone_with_replacement(case_expr, replacement_fn)?))
}
None => None,
},
case.when_then_expr
.iter()
.map(|(a, b)| {
Ok((
Box::new(clone_with_replacement(a, replacement_fn)?),
Box::new(clone_with_replacement(b, replacement_fn)?),
))
})
.collect::<Result<Vec<(_, _)>>>()?,
match &case.else_expr {
Some(else_expr) => {
Some(Box::new(clone_with_replacement(else_expr, replacement_fn)?))
}
None => None,
},
))),
Expr::ScalarFunction(ScalarFunction { fun, args }) => {
Ok(Expr::ScalarFunction(ScalarFunction::new(
*fun,
args.iter()
.map(|e| clone_with_replacement(e, replacement_fn))
.collect::<Result<Vec<Expr>>>()?,
)))
}
Expr::ScalarUDF(ScalarUDF { fun, args }) => {
Ok(Expr::ScalarUDF(ScalarUDF::new(
fun.clone(),
args.iter()
.map(|arg| clone_with_replacement(arg, replacement_fn))
.collect::<Result<Vec<Expr>>>()?,
)))
}
Expr::Negative(nested_expr) => Ok(Expr::Negative(Box::new(
clone_with_replacement(nested_expr, replacement_fn)?,
))),
Expr::Not(nested_expr) => Ok(Expr::Not(Box::new(clone_with_replacement(
nested_expr,
replacement_fn,
)?))),
Expr::IsNotNull(nested_expr) => Ok(Expr::IsNotNull(Box::new(
clone_with_replacement(nested_expr, replacement_fn)?,
))),
Expr::IsNull(nested_expr) => Ok(Expr::IsNull(Box::new(
clone_with_replacement(nested_expr, replacement_fn)?,
))),
Expr::IsTrue(nested_expr) => Ok(Expr::IsTrue(Box::new(
clone_with_replacement(nested_expr, replacement_fn)?,
))),
Expr::IsFalse(nested_expr) => Ok(Expr::IsFalse(Box::new(
clone_with_replacement(nested_expr, replacement_fn)?,
))),
Expr::IsUnknown(nested_expr) => Ok(Expr::IsUnknown(Box::new(
clone_with_replacement(nested_expr, replacement_fn)?,
))),
Expr::IsNotTrue(nested_expr) => Ok(Expr::IsNotTrue(Box::new(
clone_with_replacement(nested_expr, replacement_fn)?,
))),
Expr::IsNotFalse(nested_expr) => Ok(Expr::IsNotFalse(Box::new(
clone_with_replacement(nested_expr, replacement_fn)?,
))),
Expr::IsNotUnknown(nested_expr) => Ok(Expr::IsNotUnknown(Box::new(
clone_with_replacement(nested_expr, replacement_fn)?,
))),
Expr::Cast(Cast { expr, data_type }) => Ok(Expr::Cast(Cast::new(
Box::new(clone_with_replacement(expr, replacement_fn)?),
data_type.clone(),
))),
Expr::TryCast(TryCast {
expr: nested_expr,
data_type,
}) => Ok(Expr::TryCast(TryCast::new(
Box::new(clone_with_replacement(nested_expr, replacement_fn)?),
data_type.clone(),
))),
Expr::Sort(Sort {
expr: nested_expr,
asc,
nulls_first,
}) => Ok(Expr::Sort(Sort::new(
Box::new(clone_with_replacement(nested_expr, replacement_fn)?),
*asc,
*nulls_first,
))),
Expr::Column { .. }
| Expr::OuterReferenceColumn(_, _)
| Expr::Literal(_)
| Expr::ScalarVariable(_, _)
| Expr::Exists { .. }
| Expr::ScalarSubquery(_) => Ok(expr.clone()),
Expr::InSubquery(InSubquery {
expr: nested_expr,
subquery,
negated,
}) => Ok(Expr::InSubquery(InSubquery::new(
Box::new(clone_with_replacement(nested_expr, replacement_fn)?),
subquery.clone(),
*negated,
))),
Expr::Wildcard => Ok(Expr::Wildcard),
Expr::QualifiedWildcard { .. } => Ok(expr.clone()),
Expr::GetIndexedField(GetIndexedField { expr, field }) => {
Ok(Expr::GetIndexedField(GetIndexedField::new(
Box::new(clone_with_replacement(expr.as_ref(), replacement_fn)?),
field.clone(),
)))
}
Expr::GroupingSet(set) => match set {
GroupingSet::Rollup(exprs) => Ok(Expr::GroupingSet(GroupingSet::Rollup(
exprs
.iter()
.map(|e| clone_with_replacement(e, replacement_fn))
.collect::<Result<Vec<Expr>>>()?,
))),
GroupingSet::Cube(exprs) => Ok(Expr::GroupingSet(GroupingSet::Cube(
exprs
.iter()
.map(|e| clone_with_replacement(e, replacement_fn))
.collect::<Result<Vec<Expr>>>()?,
))),
GroupingSet::GroupingSets(lists_of_exprs) => {
let mut new_lists_of_exprs = vec![];
for exprs in lists_of_exprs {
new_lists_of_exprs.push(
exprs
.iter()
.map(|e| clone_with_replacement(e, replacement_fn))
.collect::<Result<Vec<Expr>>>()?,
);
}
Ok(Expr::GroupingSet(GroupingSet::GroupingSets(
new_lists_of_exprs,
)))
}
},
Expr::Placeholder(Placeholder { id, data_type }) => Ok(Expr::Placeholder(
Placeholder::new(id.clone(), data_type.clone()),
)),
},
}
}

/// Returns mapping of each alias (`String`) to the expression (`Expr`) it is
/// aliasing.
pub(crate) fn extract_aliases(exprs: &[Expr]) -> HashMap<String, Expr> {
Expand Down Expand Up @@ -460,15 +170,15 @@ pub(crate) fn resolve_aliases_to_exprs(
expr: &Expr,
aliases: &HashMap<String, Expr>,
) -> Result<Expr> {
clone_with_replacement(expr, &|nested_expr| match nested_expr {
expr.clone().transform_up(&|nested_expr| match nested_expr {
Expr::Column(c) if c.relation.is_none() => {
if let Some(aliased_expr) = aliases.get(&c.name) {
Ok(Some(aliased_expr.clone()))
Ok(Transformed::Yes(aliased_expr.clone()))
} else {
Ok(None)
Ok(Transformed::No(Expr::Column(c)))
}
}
_ => Ok(None),
_ => Ok(Transformed::No(nested_expr)),
})
}

Expand Down

0 comments on commit 408aefc

Please sign in to comment.