use crate::aggregate_function;
use crate::built_in_function;
use crate::expr_fn::binary_expr;
use crate::logical_plan::Subquery;
use crate::utils::expr_to_columns;
use crate::window_frame;
use crate::window_function;
use crate::AggregateUDF;
use crate::Operator;
use crate::ScalarUDF;
use arrow::datatypes::DataType;
use datafusion_common::Result;
use datafusion_common::{plan_err, Column};
use datafusion_common::{DataFusionError, ScalarValue};
use std::collections::HashSet;
use std::fmt;
use std::fmt::{Display, Formatter, Write};
use std::hash::{BuildHasher, Hash, Hasher};
use std::ops::Not;
use std::sync::Arc;
#[derive(Clone, PartialEq, Eq, Hash)]
pub enum Expr {
Alias(Box<Expr>, String),
Column(Column),
ScalarVariable(DataType, Vec<String>),
Literal(ScalarValue),
BinaryExpr(BinaryExpr),
Like(Like),
ILike(Like),
SimilarTo(Like),
Not(Box<Expr>),
IsNotNull(Box<Expr>),
IsNull(Box<Expr>),
IsTrue(Box<Expr>),
IsFalse(Box<Expr>),
IsUnknown(Box<Expr>),
IsNotTrue(Box<Expr>),
IsNotFalse(Box<Expr>),
IsNotUnknown(Box<Expr>),
Negative(Box<Expr>),
GetIndexedField(GetIndexedField),
Between(Between),
Case(Case),
Cast(Cast),
TryCast(TryCast),
Sort(Sort),
ScalarFunction {
fun: built_in_function::BuiltinScalarFunction,
args: Vec<Expr>,
},
ScalarUDF {
fun: Arc<ScalarUDF>,
args: Vec<Expr>,
},
AggregateFunction(AggregateFunction),
WindowFunction(WindowFunction),
AggregateUDF {
fun: Arc<AggregateUDF>,
args: Vec<Expr>,
filter: Option<Box<Expr>>,
},
InList {
expr: Box<Expr>,
list: Vec<Expr>,
negated: bool,
},
Exists {
subquery: Subquery,
negated: bool,
},
InSubquery {
expr: Box<Expr>,
subquery: Subquery,
negated: bool,
},
ScalarSubquery(Subquery),
Wildcard,
QualifiedWildcard { qualifier: String },
GroupingSet(GroupingSet),
Placeholder {
id: String,
data_type: Option<DataType>,
},
}
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct BinaryExpr {
pub left: Box<Expr>,
pub op: Operator,
pub right: Box<Expr>,
}
impl BinaryExpr {
pub fn new(left: Box<Expr>, op: Operator, right: Box<Expr>) -> Self {
Self { left, op, right }
}
pub fn precedence(&self) -> u8 {
match self.op {
Operator::Or => 5,
Operator::And => 10,
Operator::NotEq
| Operator::Eq
| Operator::Lt
| Operator::LtEq
| Operator::Gt
| Operator::GtEq => 20,
Operator::Plus | Operator::Minus => 30,
Operator::Multiply | Operator::Divide | Operator::Modulo => 40,
Operator::IsDistinctFrom
| Operator::IsNotDistinctFrom
| Operator::RegexMatch
| Operator::RegexNotMatch
| Operator::RegexIMatch
| Operator::RegexNotIMatch
| Operator::BitwiseAnd
| Operator::BitwiseOr
| Operator::BitwiseShiftLeft
| Operator::BitwiseShiftRight
| Operator::BitwiseXor
| Operator::StringConcat => 0,
}
}
}
impl Display for BinaryExpr {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
fn write_child(
f: &mut Formatter<'_>,
expr: &Expr,
precedence: u8,
) -> fmt::Result {
match expr {
Expr::BinaryExpr(child) => {
let p = child.precedence();
if p == 0 || p < precedence {
write!(f, "({child})")?;
} else {
write!(f, "{child}")?;
}
}
_ => write!(f, "{expr}")?,
}
Ok(())
}
let precedence = self.precedence();
write_child(f, self.left.as_ref(), precedence)?;
write!(f, " {} ", self.op)?;
write_child(f, self.right.as_ref(), precedence)
}
}
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct Case {
pub expr: Option<Box<Expr>>,
pub when_then_expr: Vec<(Box<Expr>, Box<Expr>)>,
pub else_expr: Option<Box<Expr>>,
}
impl Case {
pub fn new(
expr: Option<Box<Expr>>,
when_then_expr: Vec<(Box<Expr>, Box<Expr>)>,
else_expr: Option<Box<Expr>>,
) -> Self {
Self {
expr,
when_then_expr,
else_expr,
}
}
}
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct Like {
pub negated: bool,
pub expr: Box<Expr>,
pub pattern: Box<Expr>,
pub escape_char: Option<char>,
}
impl Like {
pub fn new(
negated: bool,
expr: Box<Expr>,
pattern: Box<Expr>,
escape_char: Option<char>,
) -> Self {
Self {
negated,
expr,
pattern,
escape_char,
}
}
}
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct Between {
pub expr: Box<Expr>,
pub negated: bool,
pub low: Box<Expr>,
pub high: Box<Expr>,
}
impl Between {
pub fn new(expr: Box<Expr>, negated: bool, low: Box<Expr>, high: Box<Expr>) -> Self {
Self {
expr,
negated,
low,
high,
}
}
}
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct GetIndexedField {
pub expr: Box<Expr>,
pub key: ScalarValue,
}
impl GetIndexedField {
pub fn new(expr: Box<Expr>, key: ScalarValue) -> Self {
Self { expr, key }
}
}
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct Cast {
pub expr: Box<Expr>,
pub data_type: DataType,
}
impl Cast {
pub fn new(expr: Box<Expr>, data_type: DataType) -> Self {
Self { expr, data_type }
}
}
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct TryCast {
pub expr: Box<Expr>,
pub data_type: DataType,
}
impl TryCast {
pub fn new(expr: Box<Expr>, data_type: DataType) -> Self {
Self { expr, data_type }
}
}
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct Sort {
pub expr: Box<Expr>,
pub asc: bool,
pub nulls_first: bool,
}
impl Sort {
pub fn new(expr: Box<Expr>, asc: bool, nulls_first: bool) -> Self {
Self {
expr,
asc,
nulls_first,
}
}
}
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct AggregateFunction {
pub fun: aggregate_function::AggregateFunction,
pub args: Vec<Expr>,
pub distinct: bool,
pub filter: Option<Box<Expr>>,
}
impl AggregateFunction {
pub fn new(
fun: aggregate_function::AggregateFunction,
args: Vec<Expr>,
distinct: bool,
filter: Option<Box<Expr>>,
) -> Self {
Self {
fun,
args,
distinct,
filter,
}
}
}
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct WindowFunction {
pub fun: window_function::WindowFunction,
pub args: Vec<Expr>,
pub partition_by: Vec<Expr>,
pub order_by: Vec<Expr>,
pub window_frame: window_frame::WindowFrame,
}
impl WindowFunction {
pub fn new(
fun: window_function::WindowFunction,
args: Vec<Expr>,
partition_by: Vec<Expr>,
order_by: Vec<Expr>,
window_frame: window_frame::WindowFrame,
) -> Self {
Self {
fun,
args,
partition_by,
order_by,
window_frame,
}
}
}
#[derive(Clone, PartialEq, Eq, Hash)]
pub enum GroupingSet {
Rollup(Vec<Expr>),
Cube(Vec<Expr>),
GroupingSets(Vec<Vec<Expr>>),
}
impl GroupingSet {
pub fn distinct_expr(&self) -> Vec<Expr> {
match self {
GroupingSet::Rollup(exprs) => exprs.clone(),
GroupingSet::Cube(exprs) => exprs.clone(),
GroupingSet::GroupingSets(groups) => {
let mut exprs: Vec<Expr> = vec![];
for exp in groups.iter().flatten() {
if !exprs.contains(exp) {
exprs.push(exp.clone());
}
}
exprs
}
}
}
}
const SEED: ahash::RandomState = ahash::RandomState::with_seeds(0, 0, 0, 0);
impl PartialOrd for Expr {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
let mut hasher = SEED.build_hasher();
self.hash(&mut hasher);
let s = hasher.finish();
let mut hasher = SEED.build_hasher();
other.hash(&mut hasher);
let o = hasher.finish();
Some(s.cmp(&o))
}
}
impl Expr {
pub fn display_name(&self) -> Result<String> {
create_name(self)
}
#[deprecated(since = "14.0.0", note = "please use `display_name` instead")]
pub fn name(&self) -> Result<String> {
self.display_name()
}
pub fn canonical_name(&self) -> String {
format!("{self}")
}
pub fn variant_name(&self) -> &str {
match self {
Expr::AggregateFunction { .. } => "AggregateFunction",
Expr::AggregateUDF { .. } => "AggregateUDF",
Expr::Alias(..) => "Alias",
Expr::Between { .. } => "Between",
Expr::BinaryExpr { .. } => "BinaryExpr",
Expr::Case { .. } => "Case",
Expr::Cast { .. } => "Cast",
Expr::Column(..) => "Column",
Expr::Exists { .. } => "Exists",
Expr::GetIndexedField { .. } => "GetIndexedField",
Expr::GroupingSet(..) => "GroupingSet",
Expr::InList { .. } => "InList",
Expr::InSubquery { .. } => "InSubquery",
Expr::IsNotNull(..) => "IsNotNull",
Expr::IsNull(..) => "IsNull",
Expr::Like { .. } => "Like",
Expr::ILike { .. } => "ILike",
Expr::SimilarTo { .. } => "RLike",
Expr::IsTrue(..) => "IsTrue",
Expr::IsFalse(..) => "IsFalse",
Expr::IsUnknown(..) => "IsUnknown",
Expr::IsNotTrue(..) => "IsNotTrue",
Expr::IsNotFalse(..) => "IsNotFalse",
Expr::IsNotUnknown(..) => "IsNotUnknown",
Expr::Literal(..) => "Literal",
Expr::Negative(..) => "Negative",
Expr::Not(..) => "Not",
Expr::Placeholder { .. } => "Placeholder",
Expr::QualifiedWildcard { .. } => "QualifiedWildcard",
Expr::ScalarFunction { .. } => "ScalarFunction",
Expr::ScalarSubquery { .. } => "ScalarSubquery",
Expr::ScalarUDF { .. } => "ScalarUDF",
Expr::ScalarVariable(..) => "ScalarVariable",
Expr::Sort { .. } => "Sort",
Expr::TryCast { .. } => "TryCast",
Expr::WindowFunction { .. } => "WindowFunction",
Expr::Wildcard => "Wildcard",
}
}
pub fn eq(self, other: Expr) -> Expr {
binary_expr(self, Operator::Eq, other)
}
pub fn not_eq(self, other: Expr) -> Expr {
binary_expr(self, Operator::NotEq, other)
}
pub fn gt(self, other: Expr) -> Expr {
binary_expr(self, Operator::Gt, other)
}
pub fn gt_eq(self, other: Expr) -> Expr {
binary_expr(self, Operator::GtEq, other)
}
pub fn lt(self, other: Expr) -> Expr {
binary_expr(self, Operator::Lt, other)
}
pub fn lt_eq(self, other: Expr) -> Expr {
binary_expr(self, Operator::LtEq, other)
}
pub fn and(self, other: Expr) -> Expr {
binary_expr(self, Operator::And, other)
}
pub fn or(self, other: Expr) -> Expr {
binary_expr(self, Operator::Or, other)
}
#[allow(clippy::should_implement_trait)]
pub fn not(self) -> Expr {
!self
}
pub fn modulus(self, other: Expr) -> Expr {
binary_expr(self, Operator::Modulo, other)
}
pub fn like(self, other: Expr) -> Expr {
Expr::Like(Like::new(false, Box::new(self), Box::new(other), None))
}
pub fn not_like(self, other: Expr) -> Expr {
Expr::Like(Like::new(true, Box::new(self), Box::new(other), None))
}
pub fn ilike(self, other: Expr) -> Expr {
Expr::ILike(Like::new(false, Box::new(self), Box::new(other), None))
}
pub fn not_ilike(self, other: Expr) -> Expr {
Expr::ILike(Like::new(true, Box::new(self), Box::new(other), None))
}
pub fn alias(self, name: impl Into<String>) -> Expr {
Expr::Alias(Box::new(self), name.into())
}
pub fn unalias(self) -> Expr {
match self {
Expr::Alias(expr, _) => expr.as_ref().clone(),
_ => self,
}
}
pub fn in_list(self, list: Vec<Expr>, negated: bool) -> Expr {
Expr::InList {
expr: Box::new(self),
list,
negated,
}
}
#[allow(clippy::wrong_self_convention)]
pub fn is_null(self) -> Expr {
Expr::IsNull(Box::new(self))
}
#[allow(clippy::wrong_self_convention)]
pub fn is_not_null(self) -> Expr {
Expr::IsNotNull(Box::new(self))
}
pub fn sort(self, asc: bool, nulls_first: bool) -> Expr {
Expr::Sort(Sort::new(Box::new(self), asc, nulls_first))
}
pub fn is_true(self) -> Expr {
Expr::IsTrue(Box::new(self))
}
pub fn is_not_true(self) -> Expr {
Expr::IsNotTrue(Box::new(self))
}
pub fn is_false(self) -> Expr {
Expr::IsFalse(Box::new(self))
}
pub fn is_not_false(self) -> Expr {
Expr::IsNotFalse(Box::new(self))
}
pub fn is_unknown(self) -> Expr {
Expr::IsUnknown(Box::new(self))
}
pub fn is_not_unknown(self) -> Expr {
Expr::IsNotUnknown(Box::new(self))
}
pub fn try_into_col(&self) -> Result<Column> {
match self {
Expr::Column(it) => Ok(it.clone()),
_ => plan_err!(format!("Could not coerce '{self}' into Column!")),
}
}
pub fn to_columns(&self) -> Result<HashSet<Column>> {
let mut using_columns = HashSet::new();
expr_to_columns(self, &mut using_columns)?;
Ok(using_columns)
}
}
impl Not for Expr {
type Output = Self;
fn not(self) -> Self::Output {
match self {
Expr::Like(Like {
negated,
expr,
pattern,
escape_char,
}) => Expr::Like(Like::new(!negated, expr, pattern, escape_char)),
Expr::ILike(Like {
negated,
expr,
pattern,
escape_char,
}) => Expr::ILike(Like::new(!negated, expr, pattern, escape_char)),
Expr::SimilarTo(Like {
negated,
expr,
pattern,
escape_char,
}) => Expr::SimilarTo(Like::new(!negated, expr, pattern, escape_char)),
_ => Expr::Not(Box::new(self)),
}
}
}
impl fmt::Display for Expr {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
write!(f, "{self:?}")
}
}
impl fmt::Debug for Expr {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
Expr::Alias(expr, alias) => write!(f, "{expr:?} AS {alias}"),
Expr::Column(c) => write!(f, "{c}"),
Expr::ScalarVariable(_, var_names) => write!(f, "{}", var_names.join(".")),
Expr::Literal(v) => write!(f, "{v:?}"),
Expr::Case(case) => {
write!(f, "CASE ")?;
if let Some(e) = &case.expr {
write!(f, "{e:?} ")?;
}
for (w, t) in &case.when_then_expr {
write!(f, "WHEN {w:?} THEN {t:?} ")?;
}
if let Some(e) = &case.else_expr {
write!(f, "ELSE {e:?} ")?;
}
write!(f, "END")
}
Expr::Cast(Cast { expr, data_type }) => {
write!(f, "CAST({expr:?} AS {data_type:?})")
}
Expr::TryCast(TryCast { expr, data_type }) => {
write!(f, "TRY_CAST({expr:?} AS {data_type:?})")
}
Expr::Not(expr) => write!(f, "NOT {expr:?}"),
Expr::Negative(expr) => write!(f, "(- {expr:?})"),
Expr::IsNull(expr) => write!(f, "{expr:?} IS NULL"),
Expr::IsNotNull(expr) => write!(f, "{expr:?} IS NOT NULL"),
Expr::IsTrue(expr) => write!(f, "{expr:?} IS TRUE"),
Expr::IsFalse(expr) => write!(f, "{expr:?} IS FALSE"),
Expr::IsUnknown(expr) => write!(f, "{expr:?} IS UNKNOWN"),
Expr::IsNotTrue(expr) => write!(f, "{expr:?} IS NOT TRUE"),
Expr::IsNotFalse(expr) => write!(f, "{expr:?} IS NOT FALSE"),
Expr::IsNotUnknown(expr) => write!(f, "{expr:?} IS NOT UNKNOWN"),
Expr::Exists {
subquery,
negated: true,
} => write!(f, "NOT EXISTS ({subquery:?})"),
Expr::Exists {
subquery,
negated: false,
} => write!(f, "EXISTS ({subquery:?})"),
Expr::InSubquery {
expr,
subquery,
negated: true,
} => write!(f, "{expr:?} NOT IN ({subquery:?})"),
Expr::InSubquery {
expr,
subquery,
negated: false,
} => write!(f, "{expr:?} IN ({subquery:?})"),
Expr::ScalarSubquery(subquery) => write!(f, "({subquery:?})"),
Expr::BinaryExpr(expr) => write!(f, "{expr}"),
Expr::Sort(Sort {
expr,
asc,
nulls_first,
}) => {
if *asc {
write!(f, "{expr:?} ASC")?;
} else {
write!(f, "{expr:?} DESC")?;
}
if *nulls_first {
write!(f, " NULLS FIRST")
} else {
write!(f, " NULLS LAST")
}
}
Expr::ScalarFunction { fun, args, .. } => {
fmt_function(f, &fun.to_string(), false, args, false)
}
Expr::ScalarUDF { fun, ref args, .. } => {
fmt_function(f, &fun.name, false, args, false)
}
Expr::WindowFunction(WindowFunction {
fun,
args,
partition_by,
order_by,
window_frame,
}) => {
fmt_function(f, &fun.to_string(), false, args, false)?;
if !partition_by.is_empty() {
write!(f, " PARTITION BY {partition_by:?}")?;
}
if !order_by.is_empty() {
write!(f, " ORDER BY {order_by:?}")?;
}
write!(
f,
" {} BETWEEN {} AND {}",
window_frame.units, window_frame.start_bound, window_frame.end_bound
)?;
Ok(())
}
Expr::AggregateFunction(AggregateFunction {
fun,
distinct,
ref args,
filter,
..
}) => {
fmt_function(f, &fun.to_string(), *distinct, args, true)?;
if let Some(fe) = filter {
write!(f, " FILTER (WHERE {fe})")?;
}
Ok(())
}
Expr::AggregateUDF {
fun,
ref args,
filter,
..
} => {
fmt_function(f, &fun.name, false, args, false)?;
if let Some(fe) = filter {
write!(f, " FILTER (WHERE {fe})")?;
}
Ok(())
}
Expr::Between(Between {
expr,
negated,
low,
high,
}) => {
if *negated {
write!(f, "{expr:?} NOT BETWEEN {low:?} AND {high:?}")
} else {
write!(f, "{expr:?} BETWEEN {low:?} AND {high:?}")
}
}
Expr::Like(Like {
negated,
expr,
pattern,
escape_char,
}) => {
write!(f, "{expr:?}")?;
if *negated {
write!(f, " NOT")?;
}
if let Some(char) = escape_char {
write!(f, " LIKE {pattern:?} ESCAPE '{char}'")
} else {
write!(f, " LIKE {pattern:?}")
}
}
Expr::ILike(Like {
negated,
expr,
pattern,
escape_char,
}) => {
write!(f, "{expr:?}")?;
if *negated {
write!(f, " NOT")?;
}
if let Some(char) = escape_char {
write!(f, " ILIKE {pattern:?} ESCAPE '{char}'")
} else {
write!(f, " ILIKE {pattern:?}")
}
}
Expr::SimilarTo(Like {
negated,
expr,
pattern,
escape_char,
}) => {
write!(f, "{expr:?}")?;
if *negated {
write!(f, " NOT")?;
}
if let Some(char) = escape_char {
write!(f, " SIMILAR TO {pattern:?} ESCAPE '{char}'")
} else {
write!(f, " SIMILAR TO {pattern:?}")
}
}
Expr::InList {
expr,
list,
negated,
} => {
if *negated {
write!(f, "{expr:?} NOT IN ({list:?})")
} else {
write!(f, "{expr:?} IN ({list:?})")
}
}
Expr::Wildcard => write!(f, "*"),
Expr::QualifiedWildcard { qualifier } => write!(f, "{qualifier}.*"),
Expr::GetIndexedField(GetIndexedField { key, expr }) => {
write!(f, "({expr:?})[{key}]")
}
Expr::GroupingSet(grouping_sets) => match grouping_sets {
GroupingSet::Rollup(exprs) => {
write!(
f,
"ROLLUP ({})",
exprs
.iter()
.map(|e| format!("{e}"))
.collect::<Vec<String>>()
.join(", ")
)
}
GroupingSet::Cube(exprs) => {
write!(
f,
"CUBE ({})",
exprs
.iter()
.map(|e| format!("{e}"))
.collect::<Vec<String>>()
.join(", ")
)
}
GroupingSet::GroupingSets(lists_of_exprs) => {
write!(
f,
"GROUPING SETS ({})",
lists_of_exprs
.iter()
.map(|exprs| format!(
"({})",
exprs
.iter()
.map(|e| format!("{e}"))
.collect::<Vec<String>>()
.join(", ")
))
.collect::<Vec<String>>()
.join(", ")
)
}
},
Expr::Placeholder { id, .. } => write!(f, "{id}"),
}
}
}
fn fmt_function(
f: &mut fmt::Formatter,
fun: &str,
distinct: bool,
args: &[Expr],
display: bool,
) -> fmt::Result {
let args: Vec<String> = match display {
true => args.iter().map(|arg| format!("{arg}")).collect(),
false => args.iter().map(|arg| format!("{arg:?}")).collect(),
};
let distinct_str = match distinct {
true => "DISTINCT ",
false => "",
};
write!(f, "{}({}{})", fun, distinct_str, args.join(", "))
}
fn create_function_name(fun: &str, distinct: bool, args: &[Expr]) -> Result<String> {
let names: Vec<String> = args.iter().map(create_name).collect::<Result<_>>()?;
let distinct_str = match distinct {
true => "DISTINCT ",
false => "",
};
Ok(format!("{}({}{})", fun, distinct_str, names.join(",")))
}
fn create_name(e: &Expr) -> Result<String> {
match e {
Expr::Alias(_, name) => Ok(name.clone()),
Expr::Column(c) => Ok(c.flat_name()),
Expr::ScalarVariable(_, variable_names) => Ok(variable_names.join(".")),
Expr::Literal(value) => Ok(format!("{value:?}")),
Expr::BinaryExpr(binary_expr) => {
let left = create_name(binary_expr.left.as_ref())?;
let right = create_name(binary_expr.right.as_ref())?;
Ok(format!("{} {} {}", left, binary_expr.op, right))
}
Expr::Like(Like {
negated,
expr,
pattern,
escape_char,
}) => {
let s = format!(
"{} {} {} {}",
expr,
if *negated { "NOT LIKE" } else { "LIKE" },
pattern,
if let Some(char) = escape_char {
format!("CHAR '{char}'")
} else {
"".to_string()
}
);
Ok(s)
}
Expr::ILike(Like {
negated,
expr,
pattern,
escape_char,
}) => {
let s = format!(
"{} {} {} {}",
expr,
if *negated { "NOT ILIKE" } else { "ILIKE" },
pattern,
if let Some(char) = escape_char {
format!("CHAR '{char}'")
} else {
"".to_string()
}
);
Ok(s)
}
Expr::SimilarTo(Like {
negated,
expr,
pattern,
escape_char,
}) => {
let s = format!(
"{} {} {} {}",
expr,
if *negated {
"NOT SIMILAR TO"
} else {
"SIMILAR TO"
},
pattern,
if let Some(char) = escape_char {
format!("CHAR '{char}'")
} else {
"".to_string()
}
);
Ok(s)
}
Expr::Case(case) => {
let mut name = "CASE ".to_string();
if let Some(e) = &case.expr {
let e = create_name(e)?;
let _ = write!(name, "{e} ");
}
for (w, t) in &case.when_then_expr {
let when = create_name(w)?;
let then = create_name(t)?;
let _ = write!(name, "WHEN {when} THEN {then} ");
}
if let Some(e) = &case.else_expr {
let e = create_name(e)?;
let _ = write!(name, "ELSE {e} ");
}
name += "END";
Ok(name)
}
Expr::Cast(Cast { expr, .. }) => {
create_name(expr)
}
Expr::TryCast(TryCast { expr, .. }) => {
create_name(expr)
}
Expr::Not(expr) => {
let expr = create_name(expr)?;
Ok(format!("NOT {expr}"))
}
Expr::Negative(expr) => {
let expr = create_name(expr)?;
Ok(format!("(- {expr})"))
}
Expr::IsNull(expr) => {
let expr = create_name(expr)?;
Ok(format!("{expr} IS NULL"))
}
Expr::IsNotNull(expr) => {
let expr = create_name(expr)?;
Ok(format!("{expr} IS NOT NULL"))
}
Expr::IsTrue(expr) => {
let expr = create_name(expr)?;
Ok(format!("{expr} IS TRUE"))
}
Expr::IsFalse(expr) => {
let expr = create_name(expr)?;
Ok(format!("{expr} IS FALSE"))
}
Expr::IsUnknown(expr) => {
let expr = create_name(expr)?;
Ok(format!("{expr} IS UNKNOWN"))
}
Expr::IsNotTrue(expr) => {
let expr = create_name(expr)?;
Ok(format!("{expr} IS NOT TRUE"))
}
Expr::IsNotFalse(expr) => {
let expr = create_name(expr)?;
Ok(format!("{expr} IS NOT FALSE"))
}
Expr::IsNotUnknown(expr) => {
let expr = create_name(expr)?;
Ok(format!("{expr} IS NOT UNKNOWN"))
}
Expr::Exists { negated: true, .. } => Ok("NOT EXISTS".to_string()),
Expr::Exists { negated: false, .. } => Ok("EXISTS".to_string()),
Expr::InSubquery { negated: true, .. } => Ok("NOT IN".to_string()),
Expr::InSubquery { negated: false, .. } => Ok("IN".to_string()),
Expr::ScalarSubquery(subquery) => {
Ok(subquery.subquery.schema().field(0).name().clone())
}
Expr::GetIndexedField(GetIndexedField { key, expr }) => {
let expr = create_name(expr)?;
Ok(format!("{expr}[{key}]"))
}
Expr::ScalarFunction { fun, args, .. } => {
create_function_name(&fun.to_string(), false, args)
}
Expr::ScalarUDF { fun, args, .. } => create_function_name(&fun.name, false, args),
Expr::WindowFunction(WindowFunction {
fun,
args,
window_frame,
partition_by,
order_by,
}) => {
let mut parts: Vec<String> =
vec![create_function_name(&fun.to_string(), false, args)?];
if !partition_by.is_empty() {
parts.push(format!("PARTITION BY {partition_by:?}"));
}
if !order_by.is_empty() {
parts.push(format!("ORDER BY {order_by:?}"));
}
parts.push(format!("{window_frame}"));
Ok(parts.join(" "))
}
Expr::AggregateFunction(AggregateFunction {
fun,
distinct,
args,
filter,
}) => {
let name = create_function_name(&fun.to_string(), *distinct, args)?;
if let Some(fe) = filter {
Ok(format!("{name} FILTER (WHERE {fe})"))
} else {
Ok(name)
}
}
Expr::AggregateUDF { fun, args, filter } => {
let mut names = Vec::with_capacity(args.len());
for e in args {
names.push(create_name(e)?);
}
let filter = if let Some(fe) = filter {
format!(" FILTER (WHERE {fe})")
} else {
"".to_string()
};
Ok(format!("{}({}){}", fun.name, names.join(","), filter))
}
Expr::GroupingSet(grouping_set) => match grouping_set {
GroupingSet::Rollup(exprs) => {
Ok(format!("ROLLUP ({})", create_names(exprs.as_slice())?))
}
GroupingSet::Cube(exprs) => {
Ok(format!("CUBE ({})", create_names(exprs.as_slice())?))
}
GroupingSet::GroupingSets(lists_of_exprs) => {
let mut list_of_names = vec![];
for exprs in lists_of_exprs {
list_of_names.push(format!("({})", create_names(exprs.as_slice())?));
}
Ok(format!("GROUPING SETS ({})", list_of_names.join(", ")))
}
},
Expr::InList {
expr,
list,
negated,
} => {
let expr = create_name(expr)?;
let list = list.iter().map(create_name);
if *negated {
Ok(format!("{expr} NOT IN ({list:?})"))
} else {
Ok(format!("{expr} IN ({list:?})"))
}
}
Expr::Between(Between {
expr,
negated,
low,
high,
}) => {
let expr = create_name(expr)?;
let low = create_name(low)?;
let high = create_name(high)?;
if *negated {
Ok(format!("{expr} NOT BETWEEN {low} AND {high}"))
} else {
Ok(format!("{expr} BETWEEN {low} AND {high}"))
}
}
Expr::Sort { .. } => Err(DataFusionError::Internal(
"Create name does not support sort expression".to_string(),
)),
Expr::Wildcard => Ok("*".to_string()),
Expr::QualifiedWildcard { .. } => Err(DataFusionError::Internal(
"Create name does not support qualified wildcard".to_string(),
)),
Expr::Placeholder { id, .. } => Ok((*id).to_string()),
}
}
fn create_names(exprs: &[Expr]) -> Result<String> {
Ok(exprs
.iter()
.map(create_name)
.collect::<Result<Vec<String>>>()?
.join(", "))
}
#[cfg(test)]
mod test {
use crate::expr::Cast;
use crate::expr_fn::col;
use crate::{case, lit, Expr};
use arrow::datatypes::DataType;
use datafusion_common::Column;
use datafusion_common::{Result, ScalarValue};
#[test]
fn format_case_when() -> Result<()> {
let expr = case(col("a"))
.when(lit(1), lit(true))
.when(lit(0), lit(false))
.otherwise(lit(ScalarValue::Null))?;
let expected = "CASE a WHEN Int32(1) THEN Boolean(true) WHEN Int32(0) THEN Boolean(false) ELSE NULL END";
assert_eq!(expected, expr.canonical_name());
assert_eq!(expected, format!("{expr}"));
assert_eq!(expected, format!("{expr:?}"));
assert_eq!(expected, expr.display_name()?);
Ok(())
}
#[test]
fn format_cast() -> Result<()> {
let expr = Expr::Cast(Cast {
expr: Box::new(Expr::Literal(ScalarValue::Float32(Some(1.23)))),
data_type: DataType::Utf8,
});
let expected_canonical = "CAST(Float32(1.23) AS Utf8)";
assert_eq!(expected_canonical, expr.canonical_name());
assert_eq!(expected_canonical, format!("{expr}"));
assert_eq!(expected_canonical, format!("{expr:?}"));
assert_eq!("Float32(1.23)", expr.display_name()?);
Ok(())
}
#[test]
fn test_not() {
assert_eq!(lit(1).not(), !lit(1));
}
#[test]
fn test_partial_ord() {
let exp1 = col("a") + lit(1);
let exp2 = col("a") + lit(2);
let exp3 = !(col("a") + lit(2));
assert!(exp1 < exp2);
assert!(exp2 > exp1);
assert!(exp2 > exp3);
assert!(exp3 < exp2);
}
#[test]
fn test_collect_expr() -> Result<()> {
{
let expr = &Expr::Cast(Cast::new(Box::new(col("a")), DataType::Float64));
let columns = expr.to_columns()?;
assert_eq!(1, columns.len());
assert!(columns.contains(&Column::from_name("a")));
}
{
let expr = col("a") + col("b") + lit(1);
let columns = expr.to_columns()?;
assert_eq!(2, columns.len());
assert!(columns.contains(&Column::from_name("a")));
assert!(columns.contains(&Column::from_name("b")));
}
Ok(())
}
}