use crate::aggregate_function;
use crate::built_in_function;
use crate::expr_fn::binary_expr;
use crate::logical_plan::Subquery;
use crate::udaf;
use crate::utils::{expr_to_columns, find_out_reference_exprs};
use crate::window_frame;
use crate::window_function;
use crate::Operator;
use arrow::datatypes::DataType;
use datafusion_common::internal_err;
use datafusion_common::{plan_err, Column, DataFusionError, Result, ScalarValue};
use std::collections::HashSet;
use std::fmt;
use std::fmt::{Display, Formatter, Write};
use std::hash::{BuildHasher, Hash, Hasher};
use std::sync::Arc;
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub enum Expr {
Alias(Alias),
Column(Column),
ScalarVariable(DataType, Vec<String>),
Literal(ScalarValue),
BinaryExpr(BinaryExpr),
Like(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(ScalarFunction),
ScalarUDF(ScalarUDF),
AggregateFunction(AggregateFunction),
WindowFunction(WindowFunction),
AggregateUDF(AggregateUDF),
InList(InList),
Exists(Exists),
InSubquery(InSubquery),
ScalarSubquery(Subquery),
Wildcard,
QualifiedWildcard { qualifier: String },
GroupingSet(GroupingSet),
Placeholder(Placeholder),
OuterReferenceColumn(DataType, Column),
}
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct Alias {
pub expr: Box<Expr>,
pub name: String,
}
impl Alias {
pub fn new(expr: Expr, name: impl Into<String>) -> Self {
Self {
expr: Box::new(expr),
name: name.into(),
}
}
}
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
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 }
}
}
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.op.precedence();
if p == 0 || p < precedence {
write!(f, "({child})")?;
} else {
write!(f, "{child}")?;
}
}
_ => write!(f, "{expr}")?,
}
Ok(())
}
let precedence = self.op.precedence();
write_child(f, self.left.as_ref(), precedence)?;
write!(f, " {} ", self.op)?;
write_child(f, self.right.as_ref(), precedence)
}
}
#[derive(Clone, Debug, 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, Debug)]
pub struct Like {
pub negated: bool,
pub expr: Box<Expr>,
pub pattern: Box<Expr>,
pub escape_char: Option<char>,
pub case_insensitive: bool,
}
impl Like {
pub fn new(
negated: bool,
expr: Box<Expr>,
pattern: Box<Expr>,
escape_char: Option<char>,
case_insensitive: bool,
) -> Self {
Self {
negated,
expr,
pattern,
escape_char,
case_insensitive,
}
}
}
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
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, Debug)]
pub struct ScalarFunction {
pub fun: built_in_function::BuiltinScalarFunction,
pub args: Vec<Expr>,
}
impl ScalarFunction {
pub fn new(fun: built_in_function::BuiltinScalarFunction, args: Vec<Expr>) -> Self {
Self { fun, args }
}
}
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct ScalarUDF {
pub fun: Arc<crate::ScalarUDF>,
pub args: Vec<Expr>,
}
impl ScalarUDF {
pub fn new(fun: Arc<crate::ScalarUDF>, args: Vec<Expr>) -> Self {
Self { fun, args }
}
}
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub enum GetFieldAccess {
NamedStructField { name: ScalarValue },
ListIndex { key: Box<Expr> },
ListRange { start: Box<Expr>, stop: Box<Expr> },
}
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct GetIndexedField {
pub expr: Box<Expr>,
pub field: GetFieldAccess,
}
impl GetIndexedField {
pub fn new(expr: Box<Expr>, field: GetFieldAccess) -> Self {
Self { expr, field }
}
}
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
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, Debug)]
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, Debug)]
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, Debug)]
pub struct AggregateFunction {
pub fun: aggregate_function::AggregateFunction,
pub args: Vec<Expr>,
pub distinct: bool,
pub filter: Option<Box<Expr>>,
pub order_by: Option<Vec<Expr>>,
}
impl AggregateFunction {
pub fn new(
fun: aggregate_function::AggregateFunction,
args: Vec<Expr>,
distinct: bool,
filter: Option<Box<Expr>>,
order_by: Option<Vec<Expr>>,
) -> Self {
Self {
fun,
args,
distinct,
filter,
order_by,
}
}
}
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
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, Debug)]
pub struct Exists {
pub subquery: Subquery,
pub negated: bool,
}
impl Exists {
pub fn new(subquery: Subquery, negated: bool) -> Self {
Self { subquery, negated }
}
}
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct AggregateUDF {
pub fun: Arc<udaf::AggregateUDF>,
pub args: Vec<Expr>,
pub filter: Option<Box<Expr>>,
pub order_by: Option<Vec<Expr>>,
}
impl AggregateUDF {
pub fn new(
fun: Arc<udaf::AggregateUDF>,
args: Vec<Expr>,
filter: Option<Box<Expr>>,
order_by: Option<Vec<Expr>>,
) -> Self {
Self {
fun,
args,
filter,
order_by,
}
}
}
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct InList {
pub expr: Box<Expr>,
pub list: Vec<Expr>,
pub negated: bool,
}
impl InList {
pub fn new(expr: Box<Expr>, list: Vec<Expr>, negated: bool) -> Self {
Self {
expr,
list,
negated,
}
}
}
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct InSubquery {
pub expr: Box<Expr>,
pub subquery: Subquery,
pub negated: bool,
}
impl InSubquery {
pub fn new(expr: Box<Expr>, subquery: Subquery, negated: bool) -> Self {
Self {
expr,
subquery,
negated,
}
}
}
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
pub struct Placeholder {
pub id: String,
pub data_type: Option<DataType>,
}
impl Placeholder {
pub fn new(id: String, data_type: Option<DataType>) -> Self {
Self { id, data_type }
}
}
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
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::OuterReferenceColumn(_, _) => "Outer",
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::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)
}
pub fn like(self, other: Expr) -> Expr {
Expr::Like(Like::new(
false,
Box::new(self),
Box::new(other),
None,
false,
))
}
pub fn not_like(self, other: Expr) -> Expr {
Expr::Like(Like::new(
true,
Box::new(self),
Box::new(other),
None,
false,
))
}
pub fn ilike(self, other: Expr) -> Expr {
Expr::Like(Like::new(
false,
Box::new(self),
Box::new(other),
None,
true,
))
}
pub fn not_ilike(self, other: Expr) -> Expr {
Expr::Like(Like::new(true, Box::new(self), Box::new(other), None, true))
}
pub fn name_for_alias(&self) -> Result<String> {
match self {
Expr::Sort(Sort { expr, .. }) => expr.name_for_alias(),
expr => expr.display_name(),
}
}
pub fn alias_if_changed(self, original_name: String) -> Result<Expr> {
let new_name = self.name_for_alias()?;
if new_name == original_name {
return Ok(self);
}
Ok(self.alias(original_name))
}
pub fn alias(self, name: impl Into<String>) -> Expr {
match self {
Expr::Sort(Sort {
expr,
asc,
nulls_first,
}) => Expr::Sort(Sort::new(Box::new(expr.alias(name)), asc, nulls_first)),
_ => Expr::Alias(Alias::new(self, name.into())),
}
}
pub fn unalias(self) -> Expr {
match self {
Expr::Alias(alias) => alias.expr.as_ref().clone(),
_ => self,
}
}
pub fn in_list(self, list: Vec<Expr>, negated: bool) -> Expr {
Expr::InList(InList::new(Box::new(self), list, negated))
}
pub fn is_null(self) -> Expr {
Expr::IsNull(Box::new(self))
}
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 between(self, low: Expr, high: Expr) -> Expr {
Expr::Between(Between::new(
Box::new(self),
false,
Box::new(low),
Box::new(high),
))
}
pub fn not_between(self, low: Expr, high: Expr) -> Expr {
Expr::Between(Between::new(
Box::new(self),
true,
Box::new(low),
Box::new(high),
))
}
pub fn field(self, name: impl Into<String>) -> Self {
Expr::GetIndexedField(GetIndexedField {
expr: Box::new(self),
field: GetFieldAccess::NamedStructField {
name: ScalarValue::Utf8(Some(name.into())),
},
})
}
pub fn index(self, key: Expr) -> Self {
Expr::GetIndexedField(GetIndexedField {
expr: Box::new(self),
field: GetFieldAccess::ListIndex { key: Box::new(key) },
})
}
pub fn range(self, start: Expr, stop: Expr) -> Self {
Expr::GetIndexedField(GetIndexedField {
expr: Box::new(self),
field: GetFieldAccess::ListRange {
start: Box::new(start),
stop: Box::new(stop),
},
})
}
pub fn try_into_col(&self) -> Result<Column> {
match self {
Expr::Column(it) => Ok(it.clone()),
_ => plan_err!("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)
}
pub fn contains_outer(&self) -> bool {
!find_out_reference_exprs(self).is_empty()
}
}
#[macro_export]
macro_rules! expr_vec_fmt {
( $ARRAY:expr ) => {{
$ARRAY
.iter()
.map(|e| format!("{e}"))
.collect::<Vec<String>>()
.join(", ")
}};
}
impl fmt::Display for Expr {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
Expr::Alias(Alias { expr, name, .. }) => write!(f, "{expr} AS {name}"),
Expr::Column(c) => write!(f, "{c}"),
Expr::OuterReferenceColumn(_, c) => write!(f, "outer_ref({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(Exists {
subquery,
negated: true,
}) => write!(f, "NOT EXISTS ({subquery:?})"),
Expr::Exists(Exists {
subquery,
negated: false,
}) => write!(f, "EXISTS ({subquery:?})"),
Expr::InSubquery(InSubquery {
expr,
subquery,
negated: true,
}) => write!(f, "{expr} NOT IN ({subquery:?})"),
Expr::InSubquery(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(func) => {
fmt_function(f, &func.fun.to_string(), false, &func.args, true)
}
Expr::ScalarUDF(ScalarUDF { fun, args }) => {
fmt_function(f, &fun.name, false, args, true)
}
Expr::WindowFunction(WindowFunction {
fun,
args,
partition_by,
order_by,
window_frame,
}) => {
fmt_function(f, &fun.to_string(), false, args, true)?;
if !partition_by.is_empty() {
write!(f, " PARTITION BY [{}]", expr_vec_fmt!(partition_by))?;
}
if !order_by.is_empty() {
write!(f, " ORDER BY [{}]", expr_vec_fmt!(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,
order_by,
..
}) => {
fmt_function(f, &fun.to_string(), *distinct, args, true)?;
if let Some(fe) = filter {
write!(f, " FILTER (WHERE {fe})")?;
}
if let Some(ob) = order_by {
write!(f, " ORDER BY [{}]", expr_vec_fmt!(ob))?;
}
Ok(())
}
Expr::AggregateUDF(AggregateUDF {
fun,
ref args,
filter,
order_by,
..
}) => {
fmt_function(f, &fun.name, false, args, true)?;
if let Some(fe) = filter {
write!(f, " FILTER (WHERE {fe})")?;
}
if let Some(ob) = order_by {
write!(f, " ORDER BY [{}]", expr_vec_fmt!(ob))?;
}
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,
case_insensitive,
}) => {
write!(f, "{expr}")?;
let op_name = if *case_insensitive { "ILIKE" } else { "LIKE" };
if *negated {
write!(f, " NOT")?;
}
if let Some(char) = escape_char {
write!(f, " {op_name} {pattern} ESCAPE '{char}'")
} else {
write!(f, " {op_name} {pattern}")
}
}
Expr::SimilarTo(Like {
negated,
expr,
pattern,
escape_char,
case_insensitive: _,
}) => {
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(InList {
expr,
list,
negated,
}) => {
if *negated {
write!(f, "{expr} NOT IN ([{}])", expr_vec_fmt!(list))
} else {
write!(f, "{expr} IN ([{}])", expr_vec_fmt!(list))
}
}
Expr::Wildcard => write!(f, "*"),
Expr::QualifiedWildcard { qualifier } => write!(f, "{qualifier}.*"),
Expr::GetIndexedField(GetIndexedField { field, expr }) => match field {
GetFieldAccess::NamedStructField { name } => {
write!(f, "({expr})[{name}]")
}
GetFieldAccess::ListIndex { key } => write!(f, "({expr})[{key}]"),
GetFieldAccess::ListRange { start, stop } => {
write!(f, "({expr})[{start}:{stop}]")
}
},
Expr::GroupingSet(grouping_sets) => match grouping_sets {
GroupingSet::Rollup(exprs) => {
write!(f, "ROLLUP ({})", expr_vec_fmt!(exprs))
}
GroupingSet::Cube(exprs) => {
write!(f, "CUBE ({})", expr_vec_fmt!(exprs))
}
GroupingSet::GroupingSets(lists_of_exprs) => {
write!(
f,
"GROUPING SETS ({})",
lists_of_exprs
.iter()
.map(|exprs| format!("({})", expr_vec_fmt!(exprs)))
.collect::<Vec<String>>()
.join(", ")
)
}
},
Expr::Placeholder(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(Alias { name, .. }) => Ok(name.clone()),
Expr::Column(c) => Ok(c.flat_name()),
Expr::OuterReferenceColumn(_, c) => Ok(format!("outer_ref({})", 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,
case_insensitive,
}) => {
let s = format!(
"{} {}{} {} {}",
expr,
if *negated { "NOT " } else { "" },
if *case_insensitive { "ILIKE" } else { "LIKE" },
pattern,
if let Some(char) = escape_char {
format!("CHAR '{char}'")
} else {
"".to_string()
}
);
Ok(s)
}
Expr::SimilarTo(Like {
negated,
expr,
pattern,
escape_char,
case_insensitive: _,
}) => {
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(Exists { negated: true, .. }) => Ok("NOT EXISTS".to_string()),
Expr::Exists(Exists { negated: false, .. }) => Ok("EXISTS".to_string()),
Expr::InSubquery(InSubquery { negated: true, .. }) => Ok("NOT IN".to_string()),
Expr::InSubquery(InSubquery { negated: false, .. }) => Ok("IN".to_string()),
Expr::ScalarSubquery(subquery) => {
Ok(subquery.subquery.schema().field(0).name().clone())
}
Expr::GetIndexedField(GetIndexedField { expr, field }) => {
let expr = create_name(expr)?;
match field {
GetFieldAccess::NamedStructField { name } => {
Ok(format!("{expr}[{name}]"))
}
GetFieldAccess::ListIndex { key } => {
let key = create_name(key)?;
Ok(format!("{expr}[{key}]"))
}
GetFieldAccess::ListRange { start, stop } => {
let start = create_name(start)?;
let stop = create_name(stop)?;
Ok(format!("{expr}[{start}:{stop}]"))
}
}
}
Expr::ScalarFunction(func) => {
create_function_name(&func.fun.to_string(), false, &func.args)
}
Expr::ScalarUDF(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 [{}]", expr_vec_fmt!(partition_by)));
}
if !order_by.is_empty() {
parts.push(format!("ORDER BY [{}]", expr_vec_fmt!(order_by)));
}
parts.push(format!("{window_frame}"));
Ok(parts.join(" "))
}
Expr::AggregateFunction(AggregateFunction {
fun,
distinct,
args,
filter,
order_by,
}) => {
let mut name = create_function_name(&fun.to_string(), *distinct, args)?;
if let Some(fe) = filter {
name = format!("{name} FILTER (WHERE {fe})");
};
if let Some(order_by) = order_by {
name = format!("{name} ORDER BY [{}]", expr_vec_fmt!(order_by));
};
Ok(name)
}
Expr::AggregateUDF(AggregateUDF {
fun,
args,
filter,
order_by,
}) => {
let mut names = Vec::with_capacity(args.len());
for e in args {
names.push(create_name(e)?);
}
let mut info = String::new();
if let Some(fe) = filter {
info += &format!(" FILTER (WHERE {fe})");
}
if let Some(ob) = order_by {
info += &format!(" ORDER BY ([{}])", expr_vec_fmt!(ob));
}
Ok(format!("{}({}){}", fun.name, names.join(","), info))
}
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(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 { .. } => {
internal_err!("Create name does not support sort expression")
}
Expr::Wildcard => Ok("*".to_string()),
Expr::QualifiedWildcard { .. } => {
internal_err!("Create name does not support qualified wildcard")
}
Expr::Placeholder(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, 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!("Float32(1.23)", expr.display_name()?);
Ok(())
}
#[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(())
}
}