use crate::aggregate_function;
use crate::built_in_function;
use crate::expr_fn::binary_expr;
use crate::logical_plan::Subquery;
use crate::window_frame;
use crate::window_function;
use crate::AggregateUDF;
use crate::Operator;
use crate::ScalarUDF;
use arrow::datatypes::DataType;
use datafusion_common::Column;
use datafusion_common::{DFSchema, Result};
use datafusion_common::{DataFusionError, ScalarValue};
use std::fmt;
use std::fmt::Write;
use std::hash::{BuildHasher, Hash, Hasher};
use std::ops::Not;
use std::sync::Arc;
#[derive(Clone, PartialEq, Hash)]
pub enum Expr {
Alias(Box<Expr>, String),
Column(Column),
ScalarVariable(DataType, Vec<String>),
Literal(ScalarValue),
BinaryExpr {
left: Box<Expr>,
op: Operator,
right: Box<Expr>,
},
Not(Box<Expr>),
IsNotNull(Box<Expr>),
IsNull(Box<Expr>),
Negative(Box<Expr>),
GetIndexedField {
expr: Box<Expr>,
key: ScalarValue,
},
Between {
expr: Box<Expr>,
negated: bool,
low: Box<Expr>,
high: Box<Expr>,
},
Case {
expr: Option<Box<Expr>>,
when_then_expr: Vec<(Box<Expr>, Box<Expr>)>,
else_expr: Option<Box<Expr>>,
},
Cast {
expr: Box<Expr>,
data_type: DataType,
},
TryCast {
expr: Box<Expr>,
data_type: DataType,
},
Sort {
expr: Box<Expr>,
asc: bool,
nulls_first: bool,
},
ScalarFunction {
fun: built_in_function::BuiltinScalarFunction,
args: Vec<Expr>,
},
ScalarUDF {
fun: Arc<ScalarUDF>,
args: Vec<Expr>,
},
AggregateFunction {
fun: aggregate_function::AggregateFunction,
args: Vec<Expr>,
distinct: bool,
},
WindowFunction {
fun: window_function::WindowFunction,
args: Vec<Expr>,
partition_by: Vec<Expr>,
order_by: Vec<Expr>,
window_frame: Option<window_frame::WindowFrame>,
},
AggregateUDF {
fun: Arc<AggregateUDF>,
args: Vec<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),
}
#[derive(Clone, PartialEq, 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 name(&self, input_schema: &DFSchema) -> Result<String> {
create_name(self, input_schema)
}
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::Literal(..) => "Literal",
Expr::Negative(..) => "Negative",
Expr::Not(..) => "Not",
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 {
binary_expr(self, Operator::Like, other)
}
pub fn not_like(self, other: Expr) -> Expr {
binary_expr(self, Operator::NotLike, other)
}
pub fn alias(self, name: &str) -> Expr {
Expr::Alias(Box::new(self), name.to_owned())
}
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 {
expr: Box::new(self),
asc,
nulls_first,
}
}
}
impl Not for Expr {
type Output = Self;
fn not(self) -> Self::Output {
Expr::Not(Box::new(self))
}
}
impl std::fmt::Display for Expr {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
match self {
Expr::BinaryExpr {
ref left,
ref right,
ref op,
} => write!(f, "{} {} {}", left, op, right),
Expr::AggregateFunction {
ref fun,
ref args,
ref distinct,
} => fmt_function(f, &fun.to_string(), *distinct, args, true),
Expr::ScalarFunction {
ref fun,
ref args,
} => fmt_function(f, &fun.to_string(), false, args, true),
_ => 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, "{:?} AS {}", expr, alias),
Expr::Column(c) => write!(f, "{}", c),
Expr::ScalarVariable(_, var_names) => write!(f, "{}", var_names.join(".")),
Expr::Literal(v) => write!(f, "{:?}", v),
Expr::Case {
expr,
when_then_expr,
else_expr,
..
} => {
write!(f, "CASE ")?;
if let Some(e) = expr {
write!(f, "{:?} ", e)?;
}
for (w, t) in when_then_expr {
write!(f, "WHEN {:?} THEN {:?} ", w, t)?;
}
if let Some(e) = else_expr {
write!(f, "ELSE {:?} ", e)?;
}
write!(f, "END")
}
Expr::Cast { expr, data_type } => {
write!(f, "CAST({:?} AS {:?})", expr, data_type)
}
Expr::TryCast { expr, data_type } => {
write!(f, "TRY_CAST({:?} AS {:?})", expr, data_type)
}
Expr::Not(expr) => write!(f, "NOT {:?}", expr),
Expr::Negative(expr) => write!(f, "(- {:?})", expr),
Expr::IsNull(expr) => write!(f, "{:?} IS NULL", expr),
Expr::IsNotNull(expr) => write!(f, "{:?} IS NOT NULL", expr),
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, "{:?} NOT IN ({:?})", expr, subquery),
Expr::InSubquery {
expr,
subquery,
negated: false,
} => write!(f, "{:?} IN ({:?})", expr, subquery),
Expr::ScalarSubquery(subquery) => write!(f, "({:?})", subquery),
Expr::BinaryExpr { left, op, right } => {
write!(f, "{:?} {} {:?}", left, op, right)
}
Expr::Sort {
expr,
asc,
nulls_first,
} => {
if *asc {
write!(f, "{:?} ASC", expr)?;
} else {
write!(f, "{:?} DESC", expr)?;
}
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 {
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)?;
}
if let Some(window_frame) = window_frame {
write!(
f,
" {} BETWEEN {} AND {}",
window_frame.units,
window_frame.start_bound,
window_frame.end_bound
)?;
}
Ok(())
}
Expr::AggregateFunction {
fun,
distinct,
ref args,
..
} => fmt_function(f, &fun.to_string(), *distinct, args, true),
Expr::AggregateUDF { fun, ref args, .. } => {
fmt_function(f, &fun.name, false, args, false)
}
Expr::Between {
expr,
negated,
low,
high,
} => {
if *negated {
write!(f, "{:?} NOT BETWEEN {:?} AND {:?}", expr, low, high)
} else {
write!(f, "{:?} BETWEEN {:?} AND {:?}", expr, low, high)
}
}
Expr::InList {
expr,
list,
negated,
} => {
if *negated {
write!(f, "{:?} NOT IN ({:?})", expr, list)
} else {
write!(f, "{:?} IN ({:?})", expr, list)
}
}
Expr::Wildcard => write!(f, "*"),
Expr::QualifiedWildcard { qualifier } => write!(f, "{}.*", qualifier),
Expr::GetIndexedField { ref expr, key } => {
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(", ")
)
}
},
}
}
}
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],
input_schema: &DFSchema,
) -> Result<String> {
let names: Vec<String> = args
.iter()
.map(|e| create_name(e, input_schema))
.collect::<Result<_>>()?;
let distinct_str = match distinct {
true => "DISTINCT ",
false => "",
};
Ok(format!("{}({}{})", fun, distinct_str, names.join(",")))
}
fn create_name(e: &Expr, input_schema: &DFSchema) -> 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 { left, op, right } => {
let left = create_name(left, input_schema)?;
let right = create_name(right, input_schema)?;
Ok(format!("{} {} {}", left, op, right))
}
Expr::Case {
expr,
when_then_expr,
else_expr,
} => {
let mut name = "CASE ".to_string();
if let Some(e) = expr {
let e = create_name(e, input_schema)?;
let _ = write!(name, "{} ", e);
}
for (w, t) in when_then_expr {
let when = create_name(w, input_schema)?;
let then = create_name(t, input_schema)?;
let _ = write!(name, "WHEN {} THEN {} ", when, then);
}
if let Some(e) = else_expr {
let e = create_name(e, input_schema)?;
let _ = write!(name, "ELSE {} ", e);
}
name += "END";
Ok(name)
}
Expr::Cast { expr, data_type } => {
let expr = create_name(expr, input_schema)?;
Ok(format!("CAST({} AS {:?})", expr, data_type))
}
Expr::TryCast { expr, data_type } => {
let expr = create_name(expr, input_schema)?;
Ok(format!("TRY_CAST({} AS {:?})", expr, data_type))
}
Expr::Not(expr) => {
let expr = create_name(expr, input_schema)?;
Ok(format!("NOT {}", expr))
}
Expr::Negative(expr) => {
let expr = create_name(expr, input_schema)?;
Ok(format!("(- {})", expr))
}
Expr::IsNull(expr) => {
let expr = create_name(expr, input_schema)?;
Ok(format!("{} IS NULL", expr))
}
Expr::IsNotNull(expr) => {
let expr = create_name(expr, input_schema)?;
Ok(format!("{} IS NOT NULL", expr))
}
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 { expr, key } => {
let expr = create_name(expr, input_schema)?;
Ok(format!("{}[{}]", expr, key))
}
Expr::ScalarFunction { fun, args, .. } => {
create_function_name(&fun.to_string(), false, args, input_schema)
}
Expr::ScalarUDF { fun, args, .. } => {
create_function_name(&fun.name, false, args, input_schema)
}
Expr::WindowFunction {
fun,
args,
window_frame,
partition_by,
order_by,
} => {
let mut parts: Vec<String> = vec![create_function_name(
&fun.to_string(),
false,
args,
input_schema,
)?];
if !partition_by.is_empty() {
parts.push(format!("PARTITION BY {:?}", partition_by));
}
if !order_by.is_empty() {
parts.push(format!("ORDER BY {:?}", order_by));
}
if let Some(window_frame) = window_frame {
parts.push(format!("{}", window_frame));
}
Ok(parts.join(" "))
}
Expr::AggregateFunction {
fun,
distinct,
args,
..
} => create_function_name(&fun.to_string(), *distinct, args, input_schema),
Expr::AggregateUDF { fun, args } => {
let mut names = Vec::with_capacity(args.len());
for e in args {
names.push(create_name(e, input_schema)?);
}
Ok(format!("{}({})", fun.name, names.join(",")))
}
Expr::GroupingSet(grouping_set) => match grouping_set {
GroupingSet::Rollup(exprs) => Ok(format!(
"ROLLUP ({})",
create_names(exprs.as_slice(), input_schema)?
)),
GroupingSet::Cube(exprs) => Ok(format!(
"CUBE ({})",
create_names(exprs.as_slice(), input_schema)?
)),
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(), input_schema)?
));
}
Ok(format!("GROUPING SETS ({})", list_of_names.join(", ")))
}
},
Expr::InList {
expr,
list,
negated,
} => {
let expr = create_name(expr, input_schema)?;
let list = list.iter().map(|expr| create_name(expr, input_schema));
if *negated {
Ok(format!("{} NOT IN ({:?})", expr, list))
} else {
Ok(format!("{} IN ({:?})", expr, list))
}
}
Expr::Between {
expr,
negated,
low,
high,
} => {
let expr = create_name(expr, input_schema)?;
let low = create_name(low, input_schema)?;
let high = create_name(high, input_schema)?;
if *negated {
Ok(format!("{} NOT BETWEEN {} AND {}", expr, low, high))
} else {
Ok(format!("{} BETWEEN {} AND {}", expr, low, high))
}
}
Expr::Sort { .. } => Err(DataFusionError::Internal(
"Create name does not support sort expression".to_string(),
)),
Expr::Wildcard => Err(DataFusionError::Internal(
"Create name does not support wildcard".to_string(),
)),
Expr::QualifiedWildcard { .. } => Err(DataFusionError::Internal(
"Create name does not support qualified wildcard".to_string(),
)),
}
}
fn create_names(exprs: &[Expr], input_schema: &DFSchema) -> Result<String> {
Ok(exprs
.iter()
.map(|e| create_name(e, input_schema))
.collect::<Result<Vec<String>>>()?
.join(", "))
}
#[cfg(test)]
mod test {
use crate::expr_fn::col;
use crate::lit;
#[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);
}
}