refactor(rust): One simplify expression module and keep utility local

pola-rs · Sep 9, 2024 · 6d24d9d · 6d24d9d
1 parent aa3b2c3
commit 6d24d9d
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 22 deletions.
diff --git a/crates/polars-plan/src/plans/aexpr/mod.rs b/crates/polars-plan/src/plans/aexpr/mod.rs
@@ -426,19 +426,6 @@ impl AExpr {
     pub(crate) fn is_leaf(&self) -> bool {
         matches!(self, AExpr::Column(_) | AExpr::Literal(_) | AExpr::Len)
     }
-    pub(crate) fn new_null_count(input: &[ExprIR]) -> Self {
-        AExpr::Function {
-            input: input.to_vec(),
-            function: FunctionExpr::NullCount,
-            options: FunctionOptions {
-                collect_groups: ApplyOptions::GroupWise,
-                fmt_str: "",
-                cast_to_supertypes: None,
-                check_lengths: UnsafeBool::default(),
-                flags: FunctionFlags::ALLOW_GROUP_AWARE | FunctionFlags::RETURNS_SCALAR,
-            },
-        }
-    }
 }
 
 impl IRAggExpr {

diff --git a/crates/polars-plan/src/plans/optimizer/mod.rs b/crates/polars-plan/src/plans/optimizer/mod.rs
@@ -18,7 +18,6 @@ mod join_utils;
 mod predicate_pushdown;
 mod projection_pushdown;
 mod simplify_expr;
-mod simplify_functions;
 mod slice_pushdown_expr;
 mod slice_pushdown_lp;
 mod stack_opt;

diff --git a/...plan/src/plans/optimizer/simplify_expr.rs → .../src/plans/optimizer/simplify_expr/mod.rs b/...plan/src/plans/optimizer/simplify_expr.rs → .../src/plans/optimizer/simplify_expr/mod.rs
@@ -1,8 +1,24 @@
+mod simplify_functions;
+
 use polars_utils::floor_divmod::FloorDivMod;
 use polars_utils::total_ord::ToTotalOrd;
+use simplify_functions::optimize_functions;
 
 use crate::plans::*;
-use crate::prelude::optimizer::simplify_functions::optimize_functions;
+
+fn new_null_count(input: &[ExprIR]) -> AExpr {
+    AExpr::Function {
+        input: input.to_vec(),
+        function: FunctionExpr::NullCount,
+        options: FunctionOptions {
+            collect_groups: ApplyOptions::GroupWise,
+            fmt_str: "",
+            cast_to_supertypes: None,
+            check_lengths: UnsafeBool::default(),
+            flags: FunctionFlags::ALLOW_GROUP_AWARE | FunctionFlags::RETURNS_SCALAR,
+        },
+    }
+}
 
 macro_rules! eval_binary_same_type {
     ($lhs:expr, $rhs:expr, |$l: ident, $r: ident| $ret: expr) => {{
@@ -457,7 +473,7 @@ impl OptimizationRule for SimplifyExprRule {
                             match expr_arena.get(drop_nulls_input_node) {
                                 AExpr::Column(_) => Some(AExpr::BinaryExpr {
                                     op: Operator::Minus,
-                                    right: expr_arena.add(AExpr::new_null_count(input)),
+                                    right: expr_arena.add(new_null_count(input)),
                                     left: expr_arena.add(AExpr::Agg(IRAggExpr::Count(
                                         drop_nulls_input_node,
                                         true,
@@ -481,7 +497,7 @@ impl OptimizationRule for SimplifyExprRule {
                         input,
                         function: FunctionExpr::Boolean(BooleanFunction::IsNull),
                         options: _,
-                    } => Some(AExpr::new_null_count(input)),
+                    } => Some(new_null_count(input)),
                     AExpr::Function {
                         input,
                         function: FunctionExpr::Boolean(BooleanFunction::IsNotNull),
@@ -494,7 +510,7 @@ impl OptimizationRule for SimplifyExprRule {
                             match expr_arena.get(is_not_null_input_node) {
                                 AExpr::Column(_) => Some(AExpr::BinaryExpr {
                                     op: Operator::Minus,
-                                    right: expr_arena.add(AExpr::new_null_count(input)),
+                                    right: expr_arena.add(new_null_count(input)),
                                     left: expr_arena.add(AExpr::Agg(IRAggExpr::Count(
                                         is_not_null_input_node,
                                         true,

diff --git a/...src/plans/optimizer/simplify_functions.rs → ...mizer/simplify_expr/simplify_functions.rs b/...src/plans/optimizer/simplify_functions.rs → ...mizer/simplify_expr/simplify_functions.rs
@@ -18,7 +18,7 @@ pub(super) fn optimize_functions(
                     function: FunctionExpr::Boolean(BooleanFunction::IsNull),
                     options: _,
                 } => Some(AExpr::BinaryExpr {
-                    left: expr_arena.add(AExpr::new_null_count(input)),
+                    left: expr_arena.add(new_null_count(input)),
                     op: Operator::Gt,
                     right: expr_arena.add(AExpr::Literal(LiteralValue::new_idxsize(0))),
                 }),
@@ -34,7 +34,7 @@ pub(super) fn optimize_functions(
                         match expr_arena.get(is_not_null_input_node) {
                             AExpr::Column(_) => Some(AExpr::BinaryExpr {
                                 op: Operator::Lt,
-                                left: expr_arena.add(AExpr::new_null_count(input)),
+                                left: expr_arena.add(new_null_count(input)),
                                 right: expr_arena.add(AExpr::Agg(IRAggExpr::Count(
                                     is_not_null_input_node,
                                     true,
@@ -66,7 +66,7 @@ pub(super) fn optimize_functions(
                         match expr_arena.get(is_null_input_node) {
                             AExpr::Column(_) => Some(AExpr::BinaryExpr {
                                 op: Operator::Eq,
-                                right: expr_arena.add(AExpr::new_null_count(input)),
+                                right: expr_arena.add(new_null_count(input)),
                                 left: expr_arena
                                     .add(AExpr::Agg(IRAggExpr::Count(is_null_input_node, true))),
                             }),
@@ -81,7 +81,7 @@ pub(super) fn optimize_functions(
                     function: FunctionExpr::Boolean(BooleanFunction::IsNotNull),
                     options: _,
                 } => Some(AExpr::BinaryExpr {
-                    left: expr_arena.add(AExpr::new_null_count(input)),
+                    left: expr_arena.add(new_null_count(input)),
                     op: Operator::Eq,
                     right: expr_arena.add(AExpr::Literal(LiteralValue::new_idxsize(0))),
                 }),