perf: Remove cast to boolean after comparison in optimizer (#21022)

Author: nameexhaustion

Diff for: crates/polars-plan/src/plans/conversion/dsl_to_ir.rs

@@ -117,7 +117,7 @@ pub(super) fn run_conversion(
 ) -> PolarsResult<Node> {
     let lp_node = ctxt.lp_arena.add(lp);
     ctxt.conversion_optimizer
-        .coerce_types(ctxt.expr_arena, ctxt.lp_arena, lp_node)
+        .optimize_exprs(ctxt.expr_arena, ctxt.lp_arena, lp_node)
         .map_err(|e| e.context(format!("'{name}' failed").into()))?;
 
     Ok(lp_node)

Diff for: crates/polars-plan/src/plans/conversion/join.rs

@@ -129,12 +129,12 @@ pub fn resolve_join(
     ctxt.conversion_optimizer
         .fill_scratch(&left_on, ctxt.expr_arena);
     ctxt.conversion_optimizer
-        .coerce_types(ctxt.expr_arena, ctxt.lp_arena, input_left)
+        .optimize_exprs(ctxt.expr_arena, ctxt.lp_arena, input_left)
         .map_err(|e| e.context("'join' failed".into()))?;
     ctxt.conversion_optimizer
         .fill_scratch(&right_on, ctxt.expr_arena);
     ctxt.conversion_optimizer
-        .coerce_types(ctxt.expr_arena, ctxt.lp_arena, input_right)
+        .optimize_exprs(ctxt.expr_arena, ctxt.lp_arena, input_right)
         .map_err(|e| e.context("'join' failed".into()))?;
 
     // Re-evaluate because of mutable borrows earlier.
@@ -447,11 +447,20 @@ fn resolve_join_where(
     for e in predicates {
         let predicate = to_expr_ir_ignore_alias(e, ctxt.expr_arena)?;
 
+        ctxt.conversion_optimizer
+            .push_scratch(predicate.node(), ctxt.expr_arena);
+
         let ir = IR::Filter {
             input: last_node,
             predicate,
         };
+
         last_node = ctxt.lp_arena.add(ir);
     }
+
+    ctxt.conversion_optimizer
+        .optimize_exprs(ctxt.expr_arena, ctxt.lp_arena, last_node)
+        .map_err(|e| e.context("'join_where' failed".into()))?;
+
     Ok((last_node, join_node))
 }

Diff for: crates/polars-plan/src/plans/conversion/stack_opt.rs

@@ -61,17 +61,19 @@ impl ConversionOptimizer {
         }
     }
 
-    pub(super) fn coerce_types(
+    /// Optimizes the expressions in the scratch space. This should be called after filling the
+    /// scratch space with the expressions that you want to optimize.
+    pub(super) fn optimize_exprs(
         &mut self,
         expr_arena: &mut Arena<AExpr>,
         ir_arena: &mut Arena<IR>,
-        current_node: Node,
+        current_ir_node: Node,
     ) -> PolarsResult<()> {
         // Different from the stack-opt in the optimizer phase, this does a single pass until fixed point per expression.
 
         if let Some(rule) = &mut self.check {
-            while let Some(x) = rule.optimize_plan(ir_arena, expr_arena, current_node)? {
-                ir_arena.replace(current_node, x);
+            while let Some(x) = rule.optimize_plan(ir_arena, expr_arena, current_ir_node)? {
+                ir_arena.replace(current_ir_node, x);
             }
         }
 
@@ -85,14 +87,14 @@ impl ConversionOptimizer {
 
             if let Some(rule) = &mut self.simplify {
                 while let Some(x) =
-                    rule.optimize_expr(expr_arena, current_expr_node, ir_arena, current_node)?
+                    rule.optimize_expr(expr_arena, current_expr_node, ir_arena, current_ir_node)?
                 {
                     expr_arena.replace(current_expr_node, x);
                 }
             }
             if let Some(rule) = &mut self.coerce {
                 while let Some(x) =
-                    rule.optimize_expr(expr_arena, current_expr_node, ir_arena, current_node)?
+                    rule.optimize_expr(expr_arena, current_expr_node, ir_arena, current_ir_node)?
                 {
                     expr_arena.replace(current_expr_node, x);
                 }

Diff for: crates/polars-plan/src/plans/conversion/type_coercion/mod.rs

@@ -103,14 +103,7 @@ impl OptimizationRule for TypeCoercionRule {
             } => {
                 let input = expr_arena.get(expr);
 
-                inline_or_prune_cast(
-                    input,
-                    dtype,
-                    options.strict(),
-                    lp_node,
-                    lp_arena,
-                    expr_arena,
-                )?
+                inline_or_prune_cast(input, dtype, options, lp_node, lp_arena, expr_arena)?
             },
             AExpr::Ternary {
                 truthy: truthy_node,
@@ -323,7 +316,13 @@ impl OptimizationRule for TypeCoercionRule {
                             DataType::Categorical(_, _) if dtype.is_string() => {
                                 // pass
                             },
-                            _ => cast_expr_ir(e, &dtype, &super_type, expr_arena, false)?,
+                            _ => cast_expr_ir(
+                                e,
+                                &dtype,
+                                &super_type,
+                                expr_arena,
+                                CastOptions::NonStrict,
+                            )?,
                         }
                     }
                 }
@@ -355,50 +354,57 @@ impl OptimizationRule for TypeCoercionRule {
 fn inline_or_prune_cast(
     aexpr: &AExpr,
     dtype: &DataType,
-    strict: bool,
+    options: CastOptions,
     lp_node: Node,
     lp_arena: &Arena<IR>,
     expr_arena: &Arena<AExpr>,
 ) -> PolarsResult<Option<AExpr>> {
     if !dtype.is_known() {
         return Ok(None);
     }
-    let lv = match aexpr {
+
+    let out = match aexpr {
         // PRUNE
-        AExpr::BinaryExpr {
-            op: Operator::LogicalOr | Operator::LogicalAnd,
-            ..
-        } => {
-            if let Some(schema) = lp_arena.get(lp_node).input_schema(lp_arena) {
-                let field = aexpr.to_field(&schema, Context::Default, expr_arena)?;
-                if field.dtype == *dtype {
-                    return Ok(Some(aexpr.clone()));
-                }
+        AExpr::BinaryExpr { op, .. } => {
+            use Operator::*;
+
+            match op {
+                LogicalOr | LogicalAnd => {
+                    if let Some(schema) = lp_arena.get(lp_node).input_schema(lp_arena) {
+                        let field = aexpr.to_field(&schema, Context::Default, expr_arena)?;
+                        if field.dtype == *dtype {
+                            return Ok(Some(aexpr.clone()));
+                        }
+                    }
+
+                    None
+                },
+                Eq | EqValidity | NotEq | NotEqValidity | Lt | LtEq | Gt | GtEq => {
+                    if dtype.is_bool() {
+                        Some(aexpr.clone())
+                    } else {
+                        None
+                    }
+                },
+                _ => None,
             }
-            return Ok(None);
         },
         // INLINE
-        AExpr::Literal(lv) => match try_inline_literal_cast(lv, dtype, strict)? {
-            None => return Ok(None),
-            Some(lv) => lv,
-        },
-        _ => return Ok(None),
+        AExpr::Literal(lv) => try_inline_literal_cast(lv, dtype, options)?.map(AExpr::Literal),
+        _ => None,
     };
-    Ok(Some(AExpr::Literal(lv)))
+
+    Ok(out)
 }
 
 fn try_inline_literal_cast(
     lv: &LiteralValue,
     dtype: &DataType,
-    strict: bool,
+    options: CastOptions,
 ) -> PolarsResult<Option<LiteralValue>> {
     let lv = match lv {
         LiteralValue::Series(s) => {
-            let s = if strict {
-                s.strict_cast(dtype)
-            } else {
-                s.cast(dtype)
-            }?;
+            let s = s.cast_with_options(dtype, options)?;
             LiteralValue::Series(SpecialEq::new(s))
         },
         LiteralValue::StrCat(s) => {
@@ -477,7 +483,7 @@ fn cast_expr_ir(
     from_dtype: &DataType,
     to_dtype: &DataType,
     expr_arena: &mut Arena<AExpr>,
-    strict: bool,
+    options: CastOptions,
 ) -> PolarsResult<()> {
     if from_dtype == to_dtype {
         return Ok(());
@@ -486,7 +492,7 @@ fn cast_expr_ir(
     check_cast(from_dtype, to_dtype)?;
 
     if let AExpr::Literal(lv) = expr_arena.get(e.node()) {
-        if let Some(literal) = try_inline_literal_cast(lv, to_dtype, strict)? {
+        if let Some(literal) = try_inline_literal_cast(lv, to_dtype, options)? {
             e.set_node(expr_arena.add(AExpr::Literal(literal)));
             e.set_dtype(to_dtype.clone());
             return Ok(());

Diff for: py-polars/tests/unit/lazyframe/test_lazyframe.py

@@ -1462,3 +1462,29 @@ def test_lf_unnest() -> None:
         ]
     )
     assert_frame_equal(lf.unnest("a", "b").collect(), expected)
+
+
+def test_type_coercion_cast_boolean_after_comparison() -> None:
+    import operator
+
+    lf = pl.LazyFrame({"a": 1, "b": 2})
+
+    for op in [
+        operator.eq,
+        operator.ne,
+        operator.lt,
+        operator.le,
+        operator.gt,
+        operator.ge,
+        pl.Expr.eq_missing,
+        pl.Expr.ne_missing,
+    ]:
+        e = op(pl.col("a"), pl.col("b")).cast(pl.Boolean).alias("o")
+        assert "cast" not in lf.with_columns(e).explain()
+
+        e = op(pl.col("a"), pl.col("b")).cast(pl.Boolean).cast(pl.Boolean).alias("o")
+        assert "cast" not in lf.with_columns(e).explain()
+
+    for op in [operator.and_, operator.or_, operator.xor]:
+        e = op(pl.col("a"), pl.col("b")).cast(pl.Boolean)
+        assert "cast" in lf.with_columns(e).explain()

Diff for: py-polars/tests/unit/operations/test_cast.py

@@ -698,3 +698,19 @@ def test_cast_python_dtypes() -> None:
     assert s.cast(float).dtype == pl.Float64
     assert s.cast(bool).dtype == pl.Boolean
     assert s.cast(str).dtype == pl.String
+
+
+def test_overflowing_cast_literals_21023() -> None:
+    for no_optimization in [True, False]:
+        assert_frame_equal(
+            (
+                pl.LazyFrame()
+                .select(
+                    pl.lit(pl.Series([128], dtype=pl.Int64)).cast(
+                        pl.Int8, wrap_numerical=True
+                    )
+                )
+                .collect(no_optimization=no_optimization)
+            ),
+            pl.Series([-128], dtype=pl.Int8).to_frame(),
+        )

Diff for: py-polars/tests/unit/operations/test_join.py

@@ -1440,3 +1440,43 @@ def test_no_collapse_join_when_maintain_order_20725() -> None:
     df_pl_eager = ldf.collect().filter(pl.col("Fraction_1") == 100)
 
     assert_frame_equal(df_pl_lazy, df_pl_eager)
+
+
+def test_join_where_predicate_type_coercion_21009() -> None:
+    left_frame = pl.LazyFrame(
+        {
+            "left_match": ["A", "B", "C", "D", "E", "F"],
+            "left_date_start": range(6),
+        }
+    )
+
+    right_frame = pl.LazyFrame(
+        {
+            "right_match": ["D", "E", "F", "G", "H", "I"],
+            "right_date": range(6),
+        }
+    )
+
+    # Note: Cannot eq the plans as the operand sides are non-deterministic
+
+    q1 = left_frame.join_where(
+        right_frame,
+        pl.col("left_match") == pl.col("right_match"),
+        pl.col("right_date") >= pl.col("left_date_start"),
+    )
+
+    plan = q1.explain().splitlines()
+    assert plan[0].strip().startswith("FILTER")
+    assert plan[1].strip().startswith("INNER JOIN")
+
+    q2 = left_frame.join_where(
+        right_frame,
+        pl.all_horizontal(pl.col("left_match") == pl.col("right_match")),
+        pl.col("right_date") >= pl.col("left_date_start"),
+    )
+
+    plan = q2.explain().splitlines()
+    assert plan[0].strip().startswith("FILTER")
+    assert plan[1].strip().startswith("INNER JOIN")
+
+    assert_frame_equal(q1.collect(), q2.collect())

Diff for: py-polars/tests/unit/test_cwc.py

@@ -146,8 +146,7 @@ def test_cwc_with_internal_aliases() -> None:
     explain = df.explain()
 
     assert (
-        """[[(col("a")) == (2)].cast(Boolean).alias("c"), [(col("b")) * (3)].alias("d")]"""
-        in explain
+        """[[(col("a")) == (2)].alias("c"), [(col("b")) * (3)].alias("d")]""" in explain
     )