use std::fmt::{Debug, Formatter};
use polars_core::prelude::{Field, Schema};
use polars_utils::unitvec;
use super::*;
use crate::prelude::*;
impl TreeWalker for Expr {
type Arena = ();
fn apply_children<F: FnMut(&Self, &Self::Arena) -> PolarsResult<VisitRecursion>>(
&self,
op: &mut F,
arena: &Self::Arena,
) -> PolarsResult<VisitRecursion> {
let mut scratch = unitvec![];
self.nodes(&mut scratch);
for &child in scratch.as_slice() {
match op(child, arena)? {
VisitRecursion::Continue | VisitRecursion::Skip => {},
VisitRecursion::Stop => return Ok(VisitRecursion::Stop),
}
}
Ok(VisitRecursion::Continue)
}
fn map_children<F: FnMut(Self, &mut Self::Arena) -> PolarsResult<Self>>(
self,
f: &mut F,
_arena: &mut Self::Arena,
) -> PolarsResult<Self> {
use polars_utils::arc::try_arc_map as am;
let mut f = |expr| f(expr, &mut ());
use AggExpr::*;
use Expr::*;
#[rustfmt::skip]
let ret = match self {
Alias(l, r) => Alias(am(l, f)?, r),
Column(_) => self,
Literal(_) => self,
DataTypeFunction(_) => self,
#[cfg(feature = "dtype-struct")]
Field(_) => self,
BinaryExpr { left, op, right } => {
BinaryExpr { left: am(left, &mut f)? , op, right: am(right, f)?}
},
Cast { expr, dtype, options: strict } => Cast { expr: am(expr, f)?, dtype, options: strict },
Sort { expr, options } => Sort { expr: am(expr, f)?, options },
Gather { expr, idx, returns_scalar, null_on_oob } => Gather {
expr: am(expr, &mut f)?,
idx: am(idx, f)?,
returns_scalar,
null_on_oob,
},
SortBy { expr, by, sort_options } => SortBy { expr: am(expr, &mut f)?, by: by.into_iter().map(f).collect::<Result<_, _>>()?, sort_options },
Agg(agg_expr) => Agg(match agg_expr {
Min { input, propagate_nans } => Min { input: am(input, f)?, propagate_nans },
Max { input, propagate_nans } => Max { input: am(input, f)?, propagate_nans },
Median(x) => Median(am(x, f)?),
NUnique(x) => NUnique(am(x, f)?),
First(x) => First(am(x, f)?),
FirstNonNull(x) => FirstNonNull(am(x, f)?),
Last(x) => Last(am(x, f)?),
LastNonNull(x) => LastNonNull(am(x, f)?),
Item { input, allow_empty } => Item { input: am(input, f)?, allow_empty },
Mean(x) => Mean(am(x, f)?),
Implode { input, maintain_order } => Implode { input: am(input, f)?, maintain_order },
Count { input, include_nulls } => Count { input: am(input, f)?, include_nulls },
Sum(x) => Sum(am(x, f)?),
AggGroups(x) => AggGroups(am(x, f)?),
Std(x, ddf) => Std(am(x, f)?, ddf),
Var(x, ddf) => Var(am(x, f)?, ddf),
}),
Ternary { predicate, truthy, falsy } => Ternary { predicate: am(predicate, &mut f)?, truthy: am(truthy, &mut f)?, falsy: am(falsy, f)? },
Function { input, function } => Function { input: input.into_iter().map(f).collect::<Result<_, _>>()?, function },
Explode { input, options } => Explode { input: am(input, f)?, options },
Filter { input, by } => Filter { input: am(input, &mut f)?, by: am(by, f)? },
#[cfg(feature = "dynamic_group_by")]
Rolling { function, index_column, period, offset, closed_window } => Rolling { function: am(function, &mut f)?, index_column: am(index_column, &mut f)?, period, offset, closed_window },
Over { function, partition_by, order_by, mapping } => {
let partition_by = partition_by.into_iter().map(&mut f).collect::<Result<_, _>>()?;
Over { function: am(function, f)?, partition_by, order_by, mapping }
},
Slice { input, offset, length } => Slice { input: am(input, &mut f)?, offset: am(offset, &mut f)?, length: am(length, f)? },
KeepName(expr) => KeepName(am(expr, f)?),
Element => Element,
Len => Len,
RenameAlias { function, expr } => RenameAlias { function, expr: am(expr, f)? },
Display { inputs, fmt_str } => {
Display { inputs: inputs.into_iter().map(f).collect::<Result<_, _>>()?, fmt_str }
},
AnonymousFunction { input, function, options, fmt_str } => {
AnonymousFunction { input: input.into_iter().map(f).collect::<Result<_, _>>()?, function, options, fmt_str }
},
Eval { expr: input, evaluation, variant } => Eval { expr: am(input, &mut f)?, evaluation: am(evaluation, f)?, variant },
#[cfg(feature = "dtype-struct")]
StructEval { expr: input, evaluation } => {
StructEval { expr: am(input, &mut f)?, evaluation: evaluation.into_iter().map(f).collect::<Result<_, _>>()? }
},
SubPlan(_, _) => self,
Selector(_) => self,
};
Ok(ret)
}
}
#[derive(Copy, Clone, Debug)]
pub struct AexprNode {
node: Node,
}
impl AexprNode {
pub fn new(node: Node) -> Self {
Self { node }
}
pub fn node(&self) -> Node {
self.node
}
pub fn to_aexpr<'a>(&self, arena: &'a Arena<AExpr>) -> &'a AExpr {
arena.get(self.node)
}
pub fn to_expr(&self, arena: &Arena<AExpr>) -> Expr {
node_to_expr(self.node, arena)
}
pub fn to_field(&self, schema: &Schema, arena: &Arena<AExpr>) -> PolarsResult<Field> {
let aexpr = arena.get(self.node);
aexpr.to_field(&ToFieldContext::new(arena, schema))
}
pub fn assign(&mut self, ae: AExpr, arena: &mut Arena<AExpr>) {
let node = arena.add(ae);
self.node = node;
}
pub(crate) fn is_leaf(&self, arena: &Arena<AExpr>) -> bool {
matches!(self.to_aexpr(arena), AExpr::Column(_) | AExpr::Literal(_))
}
pub(crate) fn hashable_and_cmp<'a>(&self, arena: &'a Arena<AExpr>) -> AExprArena<'a> {
AExprArena {
node: self.node,
arena,
}
}
}
pub struct AExprArena<'a> {
node: Node,
arena: &'a Arena<AExpr>,
}
impl Debug for AExprArena<'_> {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "AexprArena: {}", self.node.0)
}
}
impl<'a> AExprArena<'a> {
pub fn new(node: Node, arena: &'a Arena<AExpr>) -> Self {
Self { node, arena }
}
pub fn to_aexpr(&self) -> &'a AExpr {
self.arena.get(self.node)
}
}
impl PartialEq for AExprArena<'_> {
fn eq(&self, other: &Self) -> bool {
self.to_aexpr()
.is_expr_equal_to(other.to_aexpr(), self.arena)
}
}
impl TreeWalker for AexprNode {
type Arena = Arena<AExpr>;
fn apply_children<F: FnMut(&Self, &Self::Arena) -> PolarsResult<VisitRecursion>>(
&self,
op: &mut F,
arena: &Self::Arena,
) -> PolarsResult<VisitRecursion> {
let mut scratch = unitvec![];
self.to_aexpr(arena).inputs_rev(&mut scratch);
for node in scratch.as_slice() {
let aenode = AexprNode::new(*node);
match op(&aenode, arena)? {
VisitRecursion::Continue | VisitRecursion::Skip => {},
VisitRecursion::Stop => return Ok(VisitRecursion::Stop),
}
}
Ok(VisitRecursion::Continue)
}
fn map_children<F: FnMut(Self, &mut Self::Arena) -> PolarsResult<Self>>(
mut self,
op: &mut F,
arena: &mut Self::Arena,
) -> PolarsResult<Self> {
let mut scratch = unitvec![];
let ae = arena.get(self.node).clone();
ae.inputs_rev(&mut scratch);
for node in scratch.as_mut_slice() {
let aenode = AexprNode::new(*node);
*node = op(aenode, arena)?.node;
}
scratch.as_mut_slice().reverse();
let ae = ae.replace_inputs(&scratch);
self.node = arena.add(ae);
Ok(self)
}
}