From 6d24d9dec7079634de5be0a4b4d6c91f35fb66c9 Mon Sep 17 00:00:00 2001 From: ritchie Date: Mon, 9 Sep 2024 09:01:47 +0200 Subject: [PATCH] refactor(rust): One simplify expression module and keep utility local --- crates/polars-plan/src/plans/aexpr/mod.rs | 13 ---------- crates/polars-plan/src/plans/optimizer/mod.rs | 1 - .../mod.rs} | 24 +++++++++++++++---- .../{ => simplify_expr}/simplify_functions.rs | 8 +++---- 4 files changed, 24 insertions(+), 22 deletions(-) rename crates/polars-plan/src/plans/optimizer/{simplify_expr.rs => simplify_expr/mod.rs} (97%) rename crates/polars-plan/src/plans/optimizer/{ => simplify_expr}/simplify_functions.rs (97%) diff --git a/crates/polars-plan/src/plans/aexpr/mod.rs b/crates/polars-plan/src/plans/aexpr/mod.rs index 35035ffee02e..70c6335bcd1e 100644 --- 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 index 4215347f2e7d..34dc6dca9a29 100644 --- 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/crates/polars-plan/src/plans/optimizer/simplify_expr.rs b/crates/polars-plan/src/plans/optimizer/simplify_expr/mod.rs similarity index 97% rename from crates/polars-plan/src/plans/optimizer/simplify_expr.rs rename to crates/polars-plan/src/plans/optimizer/simplify_expr/mod.rs index aec74543a5dc..1df68a0adcfa 100644 --- a/crates/polars-plan/src/plans/optimizer/simplify_expr.rs +++ b/crates/polars-plan/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/crates/polars-plan/src/plans/optimizer/simplify_functions.rs b/crates/polars-plan/src/plans/optimizer/simplify_expr/simplify_functions.rs similarity index 97% rename from crates/polars-plan/src/plans/optimizer/simplify_functions.rs rename to crates/polars-plan/src/plans/optimizer/simplify_expr/simplify_functions.rs index 3d9518193276..2b5493c62e6b 100644 --- a/crates/polars-plan/src/plans/optimizer/simplify_functions.rs +++ b/crates/polars-plan/src/plans/optimizer/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))), }),