Add tests for CaseBuilder to ensure builder state is preserved on success

Author: kosiew

python/tests/test_expr.py

@@ -218,6 +218,29 @@ def test_case_builder_error_preserves_builder_state():
     assert "CaseBuilder has already been consumed" not in err_msg
 
 
+def test_case_builder_success_preserves_builder_state():
+    ctx = SessionContext()
+    df = ctx.from_pydict({"flag": [False]}, name="tbl")
+
+    case_builder = functions.when(col("flag"), lit("true"))
+
+    expr_default_one = case_builder.otherwise(lit("default-1")).alias("result")
+    result_one = df.select(expr_default_one).collect()
+    assert result_one[0].column(0).to_pylist() == ["default-1"]
+
+    expr_default_two = case_builder.otherwise(lit("default-2")).alias("result")
+    result_two = df.select(expr_default_two).collect()
+    assert result_two[0].column(0).to_pylist() == ["default-2"]
+
+    expr_end_one = case_builder.end().alias("result")
+    end_one = df.select(expr_end_one).collect()
+    assert end_one[0].column(0).to_pylist() == ["default-2"]
+
+    expr_end_two = case_builder.end().alias("result")
+    end_two = df.select(expr_end_two).collect()
+    assert end_two[0].column(0).to_pylist() == ["default-2"]
+
+
 def test_expr_getitem() -> None:
     ctx = SessionContext()
     data = {

src/expr/conditional_expr.rs

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