fix: Don't panic in string_addition_to_linear_concat

pola-rs · Mar 12, 2024 · 8d70175 · 8d70175
1 parent f2a18cd
commit 8d70175
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/crates/polars-plan/src/logical_plan/optimizer/simplify_expr.rs b/crates/polars-plan/src/logical_plan/optimizer/simplify_expr.rs
@@ -300,12 +300,8 @@ fn string_addition_to_linear_concat(
         let schema = lp_arena.get(input).schema(lp_arena);
 
         let get_type = |ae: &AExpr| ae.get_type(&schema, Context::Default, expr_arena).ok();
-        let type_a = get_type(left_aexpr)
-            .or_else(|| get_type(right_aexpr))
-            .unwrap();
-        let type_b = get_type(right_aexpr)
-            .or_else(|| get_type(right_aexpr))
-            .unwrap();
+        let type_a = get_type(left_aexpr).or_else(|| get_type(right_aexpr))?;
+        let type_b = get_type(right_aexpr).or_else(|| get_type(right_aexpr))?;
 
         if type_a != type_b {
             return None;

diff --git a/py-polars/tests/unit/test_errors.py b/py-polars/tests/unit/test_errors.py
@@ -699,3 +699,9 @@ def test_error_lazyframe_not_repeating() -> None:
 
     match = "Error originated just after this operation:"
     assert str(exc_info).count(match) == 1
+
+
+def test_raise_not_found_in_simplify_14974() -> None:
+    df = pl.DataFrame()
+    with pytest.raises(pl.ColumnNotFoundError):
+        df.select(1 / (1 + pl.col("a")))