refactor[rust]: not/isnull/isnotnull to FunctionExpr (#4696)

pola-rs · Sep 2, 2022 · d90f197 · d90f197
1 parent e81a904
commit d90f197
Show file tree

Hide file tree

Showing 25 changed files with 217 additions and 375 deletions.
diff --git a/polars/polars-lazy/src/dsl/expr.rs b/polars/polars-lazy/src/dsl/expr.rs
@@ -268,9 +268,6 @@ pub enum Expr {
         op: Operator,
         right: Box<Expr>,
     },
-    Not(Box<Expr>),
-    IsNotNull(Box<Expr>),
-    IsNull(Box<Expr>),
     Cast {
         expr: Box<Expr>,
         data_type: DataType,

diff --git a/polars/polars-lazy/src/dsl/function_expr/dispatch.rs b/polars/polars-lazy/src/dsl/function_expr/dispatch.rs
@@ -1,3 +1,5 @@
+use std::ops::Not;
+
 use super::*;
 
 pub(super) fn shift(s: &Series, periods: i64) -> Result<Series> {
@@ -7,3 +9,15 @@ pub(super) fn shift(s: &Series, periods: i64) -> Result<Series> {
 pub(super) fn reverse(s: &Series) -> Result<Series> {
     Ok(s.reverse())
 }
+
+pub(super) fn is_null(s: &Series) -> Result<Series> {
+    Ok(s.is_null().into_series())
+}
+
+pub(super) fn is_not_null(s: &Series) -> Result<Series> {
+    Ok(s.is_not_null().into_series())
+}
+
+pub(super) fn is_not(s: &Series) -> Result<Series> {
+    Ok(s.bool()?.not().into_series())
+}
diff --git a/polars/polars-lazy/src/dsl/function_expr/mod.rs b/polars/polars-lazy/src/dsl/function_expr/mod.rs
@@ -99,53 +99,59 @@ pub enum FunctionExpr {
     },
     Shift(i64),
     Reverse,
+    IsNull,
+    IsNotNull,
+    Not,
 }
 
 impl Display for FunctionExpr {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        use self::*;
+        use FunctionExpr::*;
 
         match self {
-            FunctionExpr::NullCount => write!(f, "null_count"),
-            FunctionExpr::Pow => write!(f, "pow"),
+            NullCount => write!(f, "null_count"),
+            Pow => write!(f, "pow"),
             #[cfg(feature = "row_hash")]
-            FunctionExpr::Hash(_, _, _, _) => write!(f, "hash"),
+            Hash(_, _, _, _) => write!(f, "hash"),
             #[cfg(feature = "is_in")]
-            FunctionExpr::IsIn => write!(f, "is_in"),
+            IsIn => write!(f, "is_in"),
             #[cfg(feature = "arg_where")]
-            FunctionExpr::ArgWhere => write!(f, "arg_where"),
+            ArgWhere => write!(f, "arg_where"),
             #[cfg(feature = "search_sorted")]
-            FunctionExpr::SearchSorted => write!(f, "search_sorted"),
+            SearchSorted => write!(f, "search_sorted"),
             #[cfg(feature = "strings")]
-            FunctionExpr::StringExpr(s) => write!(f, "{}", s),
+            StringExpr(s) => write!(f, "{}", s),
             #[cfg(feature = "date_offset")]
-            FunctionExpr::DateOffset(_) => write!(f, "dt.offset_by"),
+            DateOffset(_) => write!(f, "dt.offset_by"),
             #[cfg(feature = "trigonometry")]
-            FunctionExpr::Trigonometry(func) => write!(f, "{}", func),
+            Trigonometry(func) => write!(f, "{}", func),
             #[cfg(feature = "sign")]
-            FunctionExpr::Sign => write!(f, "sign"),
-            FunctionExpr::FillNull { .. } => write!(f, "fill_null"),
+            Sign => write!(f, "sign"),
+            FillNull { .. } => write!(f, "fill_null"),
             #[cfg(feature = "is_in")]
-            FunctionExpr::ListContains => write!(f, "arr.contains"),
+            ListContains => write!(f, "arr.contains"),
             #[cfg(all(feature = "rolling_window", feature = "moment"))]
-            FunctionExpr::RollingSkew { .. } => write!(f, "rolling_skew"),
-            FunctionExpr::ShiftAndFill { .. } => write!(f, "shift_and_fill"),
-            FunctionExpr::Nan(_) => write!(f, "nan"),
+            RollingSkew { .. } => write!(f, "rolling_skew"),
+            ShiftAndFill { .. } => write!(f, "shift_and_fill"),
+            Nan(_) => write!(f, "nan"),
             #[cfg(feature = "round_series")]
-            FunctionExpr::Clip { min, max } => match (min, max) {
+            Clip { min, max } => match (min, max) {
                 (Some(_), Some(_)) => write!(f, "clip"),
                 (None, Some(_)) => write!(f, "clip_max"),
                 (Some(_), None) => write!(f, "clip_min"),
                 _ => unreachable!(),
             },
             #[cfg(feature = "list")]
-            FunctionExpr::ListExpr(func) => write!(f, "{}", func),
+            ListExpr(func) => write!(f, "{}", func),
             #[cfg(feature = "dtype-struct")]
-            FunctionExpr::StructExpr(func) => write!(f, "{}", func),
+            StructExpr(func) => write!(f, "{}", func),
             #[cfg(feature = "top_k")]
-            FunctionExpr::TopK { .. } => write!(f, "top_k"),
-            FunctionExpr::Shift(_) => write!(f, "shift"),
-            FunctionExpr::Reverse => write!(f, "reverse"),
+            TopK { .. } => write!(f, "top_k"),
+            Shift(_) => write!(f, "shift"),
+            Reverse => write!(f, "reverse"),
+            Not => write!(f, "is_not"),
+            IsNull => write!(f, "is_null"),
+            IsNotNull => write!(f, "is_not_null"),
         }
     }
 }
@@ -300,6 +306,9 @@ impl From<FunctionExpr> for SpecialEq<Arc<dyn SeriesUdf>> {
             }
             Shift(periods) => map!(dispatch::shift, periods),
             Reverse => map!(dispatch::reverse),
+            IsNull => map!(dispatch::is_null),
+            IsNotNull => map!(dispatch::is_not_null),
+            Not => map!(dispatch::is_not),
         }
     }
 }

diff --git a/polars/polars-lazy/src/dsl/function_expr/schema.rs b/polars/polars-lazy/src/dsl/function_expr/schema.rs
@@ -155,6 +155,7 @@ impl FunctionExpr {
             #[cfg(feature = "top_k")]
             TopK { .. } => same_type(),
             Shift(..) | Reverse => same_type(),
+            IsNotNull | IsNull | Not => with_dtype(DataType::Boolean),
         }
     }
 }
diff --git a/polars/polars-lazy/src/dsl/functions.rs b/polars/polars-lazy/src/dsl/functions.rs
@@ -801,17 +801,17 @@ pub fn all_exprs<E: AsRef<[Expr]>>(exprs: E) -> Expr {
 
 /// [Not](Expr::Not) expression.
 pub fn not(expr: Expr) -> Expr {
-    Expr::Not(Box::new(expr))
+    expr.not()
 }
 
 /// [IsNull](Expr::IsNotNull) expression
 pub fn is_null(expr: Expr) -> Expr {
-    Expr::IsNull(Box::new(expr))
+    expr.is_null()
 }
 
 /// [IsNotNull](Expr::IsNotNull) expression.
 pub fn is_not_null(expr: Expr) -> Expr {
-    Expr::IsNotNull(Box::new(expr))
+    expr.is_not_null()
 }
 
 /// [Cast](Expr::Cast) expression.

diff --git a/polars/polars-lazy/src/dsl/mod.rs b/polars/polars-lazy/src/dsl/mod.rs
@@ -274,7 +274,7 @@ impl Expr {
     /// Negate `Expr`
     #[allow(clippy::should_implement_trait)]
     pub fn not(self) -> Expr {
-        Expr::Not(Box::new(self))
+        self.map_private(FunctionExpr::Not)
     }
 
     /// Rename Column.
@@ -285,13 +285,13 @@ impl Expr {
     /// Run is_null operation on `Expr`.
     #[allow(clippy::wrong_self_convention)]
     pub fn is_null(self) -> Self {
-        Expr::IsNull(Box::new(self))
+        self.map_private(FunctionExpr::IsNull)
     }
 
     /// Run is_not_null operation on `Expr`.
     #[allow(clippy::wrong_self_convention)]
     pub fn is_not_null(self) -> Self {
-        Expr::IsNotNull(Box::new(self))
+        self.map_private(FunctionExpr::IsNotNull)
     }
 
     /// Drop null values

diff --git a/polars/polars-lazy/src/logical_plan/aexpr.rs b/polars/polars-lazy/src/logical_plan/aexpr.rs
@@ -46,9 +46,6 @@ pub enum AExpr {
         op: Operator,
         right: Node,
     },
-    Not(Node),
-    IsNotNull(Node),
-    IsNull(Node),
     Cast {
         expr: Node,
         data_type: DataType,
@@ -142,9 +139,6 @@ impl AExpr {
             // to determine if the whole expr. is group sensitive
             | BinaryExpr { .. }
             | Ternary { .. }
-            | Not(_)
-            | IsNotNull(_)
-            | IsNull(_)
             | Wildcard
             | Cast { .. }
             | Filter { .. } => false,
@@ -166,8 +160,6 @@ impl AExpr {
         use AExpr::*;
         match self {
             Alias(_, name) => Alias(input, name),
-            IsNotNull(_) => IsNotNull(input),
-            IsNull(_) => IsNull(input),
             Cast {
                 expr: _,
                 data_type,
@@ -185,8 +177,6 @@ impl AExpr {
         use AExpr::*;
         match self {
             Alias(input, _) => *input,
-            IsNotNull(input) => *input,
-            IsNull(input) => *input,
             Cast { expr, .. } => *expr,
             _ => todo!(),
         }
@@ -282,9 +272,6 @@ impl AExpr {
 
                 Ok(Field::new(out_name, expr_type))
             }
-            Not(_) => Ok(Field::new("not", DataType::Boolean)),
-            IsNull(_) => Ok(Field::new("is_null", DataType::Boolean)),
-            IsNotNull(_) => Ok(Field::new("is_not_null", DataType::Boolean)),
             Sort { expr, .. } => arena.get(*expr).to_field(schema, ctxt, arena),
             Take { expr, .. } => arena.get(*expr).to_field(schema, ctxt, arena),
             SortBy { expr, .. } => arena.get(*expr).to_field(schema, ctxt, arena),

diff --git a/polars/polars-lazy/src/logical_plan/conversion.rs b/polars/polars-lazy/src/logical_plan/conversion.rs
@@ -24,10 +24,6 @@ pub(crate) fn to_aexpr(expr: Expr, arena: &mut Arena<AExpr>) -> Node {
                 right: r,
             }
         }
-        Expr::Not(e) => AExpr::Not(to_aexpr(*e, arena)),
-        Expr::IsNotNull(e) => AExpr::IsNotNull(to_aexpr(*e, arena)),
-        Expr::IsNull(e) => AExpr::IsNull(to_aexpr(*e, arena)),
-
         Expr::Cast {
             expr,
             data_type,
@@ -465,18 +461,6 @@ pub(crate) fn node_to_expr(node: Node, expr_arena: &Arena<AExpr>) -> Expr {
                 right: Box::new(r),
             }
         }
-        AExpr::Not(expr) => {
-            let exp = node_to_expr(expr, expr_arena);
-            Expr::Not(Box::new(exp))
-        }
-        AExpr::IsNotNull(expr) => {
-            let exp = node_to_expr(expr, expr_arena);
-            Expr::IsNotNull(Box::new(exp))
-        }
-        AExpr::IsNull(expr) => {
-            let exp = node_to_expr(expr, expr_arena);
-            Expr::IsNull(Box::new(exp))
-        }
         AExpr::Cast {
             expr,
             data_type,

diff --git a/polars/polars-lazy/src/logical_plan/format.rs b/polars/polars-lazy/src/logical_plan/format.rs
@@ -217,9 +217,6 @@ impl fmt::Debug for Expr {
                 }
             }
             BinaryExpr { left, op, right } => write!(f, "[({:?}) {:?} ({:?})]", left, op, right),
-            Not(expr) => write!(f, "not({:?})", expr),
-            IsNull(expr) => write!(f, "{:?}.is_null()", expr),
-            IsNotNull(expr) => write!(f, "{:?}.is_not_null()", expr),
             Sort { expr, options } => match options.descending {
                 true => write!(f, "{:?} DESC", expr),
                 false => write!(f, "{:?} ASC", expr),

diff --git a/polars/polars-lazy/src/logical_plan/iterator.rs b/polars/polars-lazy/src/logical_plan/iterator.rs
@@ -5,14 +5,11 @@ macro_rules! push_expr {
         match $current_expr {
             Nth(_) | Column(_) | Literal(_) | Wildcard | Columns(_) | DtypeColumn(_) | Count => {}
             Alias(e, _) => $push(e),
-            Not(e) => $push(e),
             BinaryExpr { left, op: _, right } => {
                 // reverse order so that left is popped first
                 $push(right);
                 $push(left);
             }
-            IsNull(e) => $push(e),
-            IsNotNull(e) => $push(e),
             Cast { expr, .. } => $push(expr),
             Sort { expr, .. } => $push(expr),
             Take { expr, idx } => {
@@ -171,14 +168,11 @@ impl AExpr {
         match self {
             Nth(_) | Column(_) | Literal(_) | Wildcard | Count => {}
             Alias(e, _) => push(e),
-            Not(e) => push(e),
             BinaryExpr { left, op: _, right } => {
                 // reverse order so that left is popped first
                 push(right);
                 push(left);
             }
-            IsNull(e) => push(e),
-            IsNotNull(e) => push(e),
             Cast { expr, .. } => push(expr),
             Sort { expr, .. } => push(expr),
             Take { expr, idx } => {

diff --git a/polars/polars-lazy/src/logical_plan/mod.rs b/polars/polars-lazy/src/logical_plan/mod.rs
@@ -258,9 +258,13 @@ impl LogicalPlan {
             }
             Error { err, .. } => {
                 // We just take the error. The LogicalPlan should not be used anymore once this
-                // is taken.
                 let mut err = err.lock();
-                Err(err.take().unwrap())
+                match err.take() {
+                    Some(err) => Err(err),
+                    None => Err(PolarsError::ComputeError(
+                        "LogicalPlan already failed".into(),
+                    )),
+                }
             }
             ExtContext { schema, .. } => Ok(Cow::Borrowed(schema)),
         }

diff --git a/polars/polars-lazy/src/logical_plan/optimizer/drop_nulls.rs b/polars/polars-lazy/src/logical_plan/optimizer/drop_nulls.rs
@@ -2,6 +2,7 @@ use std::sync::Arc;
 
 use polars_core::prelude::*;
 
+use crate::dsl::function_expr::FunctionExpr;
 use crate::logical_plan::iterator::*;
 use crate::prelude::stack_opt::OptimizationRule;
 use crate::prelude::*;
@@ -37,7 +38,15 @@ impl OptimizationRule for ReplaceDropNulls {
                         }
                     )
                 };
-                let is_not_null = |e: &AExpr| matches!(e, &AExpr::IsNotNull(_));
+                let is_not_null = |e: &AExpr| {
+                    matches!(
+                        e,
+                        &AExpr::Function {
+                            function: FunctionExpr::IsNotNull,
+                            ..
+                        }
+                    )
+                };
                 let is_column = |e: &AExpr| matches!(e, &AExpr::Column(_));
                 let is_lit_true =
                     |e: &AExpr| matches!(e, &AExpr::Literal(LiteralValue::Boolean(true)));

diff --git a/polars/polars-lazy/src/logical_plan/optimizer/predicate_pushdown/mod.rs b/polars/polars-lazy/src/logical_plan/optimizer/predicate_pushdown/mod.rs
@@ -4,6 +4,7 @@ use polars_core::datatypes::PlHashMap;
 use polars_core::prelude::*;
 use utils::*;
 
+use crate::dsl::function_expr::FunctionExpr;
 use crate::logical_plan::{optimizer, Context};
 use crate::prelude::*;
 use crate::utils::{aexpr_to_root_names, aexprs_to_schema, check_input_node, has_aexpr};
@@ -428,7 +429,7 @@ impl PredicatePushDown {
                         |e: &AExpr| matches!(e, AExpr::IsUnique(_) | AExpr::Duplicated(_));
 
                     let checks_nulls =
-                        |e: &AExpr| matches!(e, AExpr::IsNull(_) | AExpr::IsNotNull(_)) ||
+                        |e: &AExpr| matches!(e, AExpr::Function{function: FunctionExpr::IsNotNull | FunctionExpr::IsNull, ..} ) ||
                             // any operation that checks for equality or ordering can be wrong because
                             // the join can produce null values
                             matches!(e, AExpr::BinaryExpr {op, ..} if !matches!(op, Operator::NotEq));

diff --git a/polars/polars-lazy/src/logical_plan/optimizer/predicate_pushdown/utils.rs b/polars/polars-lazy/src/logical_plan/optimizer/predicate_pushdown/utils.rs
@@ -165,7 +165,12 @@ pub(super) fn project_other_column_is_predicate_pushdown_boundary(
     has_aexpr(node, expr_arena, matches)
 }
 
-/// This checks the boundary of same columns. So that means columns that are referred in the predicate
+/// This checks the boundary of same columns.
+/// So that means columns that are referred in the predicate
+/// for instance `predicate = col(A) == col(B).`
+/// and `col().some_func().alias(B)` is projected.
+/// then the projection can not pass, as column `B` maybe
+/// changed by `some_func`
 pub(super) fn projection_column_is_predicate_pushdown_boundary(
     node: Node,
     expr_arena: &Arena<AExpr>,