Add error handling to CaseBuilder methods to preserve builder state

Author: kosiew

python/tests/test_expr.py

@@ -200,6 +200,24 @@ def traverse_logical_plan(plan):
     assert not variant.negated()
 
 
+def test_case_builder_error_preserves_builder_state():
+    case_builder = functions.when(lit(True), lit(1))
+
+    with pytest.raises(Exception) as exc_info:
+        case_builder.otherwise(lit("bad"))
+
+    err_msg = str(exc_info.value)
+    assert "multiple data types" in err_msg
+    assert "CaseBuilder has already been consumed" not in err_msg
+
+    with pytest.raises(Exception) as exc_info:
+        case_builder.end()
+
+    err_msg = str(exc_info.value)
+    assert "multiple data types" in err_msg
+    assert "CaseBuilder has already been consumed" not in err_msg
+
+
 def test_expr_getitem() -> None:
     ctx = SessionContext()
     data = {

src/expr/conditional_expr.rs

@@ -83,13 +83,23 @@ impl PyCaseBuilder {
 
     fn otherwise(&self, else_expr: PyExpr) -> PyDataFusionResult<PyExpr> {
         let mut builder = self.take_case_builder()?;
-        let expr = builder.otherwise(else_expr.expr)?;
-        Ok(expr.clone().into())
+        match builder.otherwise(else_expr.expr) {
+            Ok(expr) => Ok(expr.clone().into()),
+            Err(err) => {
+                self.store_case_builder(builder)?;
+                Err(err.into())
+            }
+        }
     }
 
     fn end(&self) -> PyDataFusionResult<PyExpr> {
         let builder = self.take_case_builder()?;
-        let expr = builder.end()?;
-        Ok(expr.clone().into())
+        match builder.end() {
+            Ok(expr) => Ok(expr.clone().into()),
+            Err(err) => {
+                self.store_case_builder(builder)?;
+                Err(err.into())
+            }
+        }
     }
 }