use std::fmt::{Debug, Display, Formatter};
use std::hash::{Hash, Hasher};
use polars_core::prelude::*;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
pub use super::expr_dyn_fn::*;
use crate::prelude::*;
#[derive(PartialEq, Clone)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum AggExpr {
Min {
input: Box<Expr>,
propagate_nans: bool,
},
Max {
input: Box<Expr>,
propagate_nans: bool,
},
Median(Box<Expr>),
NUnique(Box<Expr>),
First(Box<Expr>),
Last(Box<Expr>),
Mean(Box<Expr>),
Implode(Box<Expr>),
Count(Box<Expr>, bool),
Quantile {
expr: Box<Expr>,
quantile: Box<Expr>,
interpol: QuantileInterpolOptions,
},
Sum(Box<Expr>),
AggGroups(Box<Expr>),
Std(Box<Expr>, u8),
Var(Box<Expr>, u8),
}
impl AsRef<Expr> for AggExpr {
fn as_ref(&self) -> &Expr {
use AggExpr::*;
match self {
Min { input, .. } => input,
Max { input, .. } => input,
Median(e) => e,
NUnique(e) => e,
First(e) => e,
Last(e) => e,
Mean(e) => e,
Implode(e) => e,
Count(e, _) => e,
Quantile { expr, .. } => expr,
Sum(e) => e,
AggGroups(e) => e,
Std(e, _) => e,
Var(e, _) => e,
}
}
}
#[derive(Clone, PartialEq)]
#[must_use]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum Expr {
Alias(Box<Expr>, Arc<str>),
Column(Arc<str>),
Columns(Vec<String>),
DtypeColumn(Vec<DataType>),
Literal(LiteralValue),
BinaryExpr {
left: Box<Expr>,
op: Operator,
right: Box<Expr>,
},
Cast {
expr: Box<Expr>,
data_type: DataType,
strict: bool,
},
Sort {
expr: Box<Expr>,
options: SortOptions,
},
Gather {
expr: Box<Expr>,
idx: Box<Expr>,
returns_scalar: bool,
},
SortBy {
expr: Box<Expr>,
by: Vec<Expr>,
descending: Vec<bool>,
},
Agg(AggExpr),
Ternary {
predicate: Box<Expr>,
truthy: Box<Expr>,
falsy: Box<Expr>,
},
Function {
input: Vec<Expr>,
function: FunctionExpr,
options: FunctionOptions,
},
Explode(Box<Expr>),
Filter {
input: Box<Expr>,
by: Box<Expr>,
},
Window {
function: Box<Expr>,
partition_by: Vec<Expr>,
options: WindowType,
},
Wildcard,
Slice {
input: Box<Expr>,
offset: Box<Expr>,
length: Box<Expr>,
},
Exclude(Box<Expr>, Vec<Excluded>),
KeepName(Box<Expr>),
Len,
Nth(i64),
#[cfg_attr(feature = "serde", serde(skip))]
RenameAlias {
function: SpecialEq<Arc<dyn RenameAliasFn>>,
expr: Box<Expr>,
},
AnonymousFunction {
input: Vec<Expr>,
function: SpecialEq<Arc<dyn SeriesUdf>>,
#[cfg_attr(feature = "serde", serde(skip))]
output_type: GetOutput,
options: FunctionOptions,
},
SubPlan(SpecialEq<Arc<LogicalPlan>>, Vec<String>),
Selector(super::selector::Selector),
}
#[allow(clippy::derived_hash_with_manual_eq)]
impl Hash for Expr {
fn hash<H: Hasher>(&self, state: &mut H) {
let d = std::mem::discriminant(self);
d.hash(state);
match self {
Expr::Column(name) => name.hash(state),
Expr::Columns(names) => names.hash(state),
Expr::DtypeColumn(dtypes) => dtypes.hash(state),
Expr::Literal(lv) => std::mem::discriminant(lv).hash(state),
Expr::Selector(s) => s.hash(state),
Expr::Nth(v) => v.hash(state),
Expr::Filter { input, by } => {
input.hash(state);
by.hash(state);
},
Expr::BinaryExpr { left, op, right } => {
left.hash(state);
right.hash(state);
std::mem::discriminant(op).hash(state)
},
Expr::Cast {
expr,
data_type,
strict,
} => {
expr.hash(state);
data_type.hash(state);
strict.hash(state)
},
Expr::Sort { expr, options } => {
expr.hash(state);
options.hash(state);
},
Expr::Alias(input, name) => {
input.hash(state);
name.hash(state)
},
Expr::KeepName(input) => input.hash(state),
Expr::Ternary {
predicate,
truthy,
falsy,
} => {
predicate.hash(state);
truthy.hash(state);
falsy.hash(state);
},
Expr::Function {
input,
function,
options,
} => {
input.hash(state);
std::mem::discriminant(function).hash(state);
options.hash(state);
},
Expr::Wildcard | Expr::Len => {},
#[allow(unreachable_code)]
_ => {
#[cfg(debug_assertions)]
{
todo!("IMPLEMENT")
}
let s = format!("{self:?}");
s.hash(state)
},
}
}
}
impl Eq for Expr {}
impl Default for Expr {
fn default() -> Self {
Expr::Literal(LiteralValue::Null)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum Excluded {
Name(Arc<str>),
Dtype(DataType),
}
impl Expr {
pub fn to_field(&self, schema: &Schema, ctxt: Context) -> PolarsResult<Field> {
let mut arena = Arena::with_capacity(5);
self.to_field_amortized(schema, ctxt, &mut arena)
}
pub(crate) fn to_field_amortized(
&self,
schema: &Schema,
ctxt: Context,
expr_arena: &mut Arena<AExpr>,
) -> PolarsResult<Field> {
let root = to_aexpr(self.clone(), expr_arena);
expr_arena.get(root).to_field(schema, ctxt, expr_arena)
}
}
#[derive(Copy, Clone, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum Operator {
Eq,
EqValidity,
NotEq,
NotEqValidity,
Lt,
LtEq,
Gt,
GtEq,
Plus,
Minus,
Multiply,
Divide,
TrueDivide,
FloorDivide,
Modulus,
And,
Or,
Xor,
LogicalAnd,
LogicalOr,
}
impl Display for Operator {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
use Operator::*;
let tkn = match self {
Eq => "==",
EqValidity => "==v",
NotEq => "!=",
NotEqValidity => "!=v",
Lt => "<",
LtEq => "<=",
Gt => ">",
GtEq => ">=",
Plus => "+",
Minus => "-",
Multiply => "*",
Divide => "//",
TrueDivide => "/",
FloorDivide => "floor_div",
Modulus => "%",
And | LogicalAnd => "&",
Or | LogicalOr => "|",
Xor => "^",
};
write!(f, "{tkn}")
}
}
impl Operator {
pub(crate) fn is_comparison(&self) -> bool {
matches!(
self,
Self::Eq
| Self::NotEq
| Self::Lt
| Self::LtEq
| Self::Gt
| Self::GtEq
| Self::And
| Self::Or
| Self::Xor
| Self::EqValidity
| Self::NotEqValidity
)
}
pub(crate) fn is_arithmetic(&self) -> bool {
!(self.is_comparison())
}
}