Resolve patch apply conflicts for CaseBuilder concurrency improvements

- Added CaseBuilderHandle guard that keeps the underlying CaseBuilder alive while holding the mutex and restores it on drop
- Updated when, otherwise, and end methods to operate through the guard and consume the builder explicitly
- This prevents transient None states during concurrent access and improves thread safety

Author: kosiew

python/tests/test_expr.py

@@ -16,6 +16,7 @@
 # under the License.
 
 import re
+from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime, timezone
 
 import pyarrow as pa
@@ -281,6 +282,26 @@ def test_case_builder_when_handles_are_independent():
     ]
 
 
+def test_case_builder_when_thread_safe():
+    case_builder = functions.when(lit(_TRUE), lit(1))
+
+    def build_expr(value: int) -> bool:
+        builder = case_builder.when(lit(_TRUE), lit(value))
+        builder.otherwise(lit(value))
+        return True
+
+    with ThreadPoolExecutor(max_workers=8) as executor:
+        futures = [executor.submit(build_expr, idx) for idx in range(16)]
+        results = [future.result() for future in futures]
+
+    assert all(results)
+
+    # Ensure the shared builder remains usable after concurrent `when` calls.
+    follow_up_builder = case_builder.when(lit(_TRUE), lit(42))
+    assert isinstance(follow_up_builder, type(case_builder))
+    follow_up_builder.otherwise(lit(7))
+
+
 def test_expr_getitem() -> None:
     ctx = SessionContext()
     data = {

src/expr/conditional_expr.rs

@@ -22,9 +22,46 @@ use crate::{
     expr::PyExpr,
 };
 use datafusion::logical_expr::conditional_expressions::CaseBuilder;
+use parking_lot::{Mutex, MutexGuard};
 use pyo3::prelude::*;
 
-use parking_lot::{Mutex, MutexGuard};
+struct CaseBuilderHandle<'a> {
+    guard: MutexGuard<'a, Option<CaseBuilder>>,
+    builder: Option<CaseBuilder>,
+}
+
+impl<'a> CaseBuilderHandle<'a> {
+    fn new(mut guard: MutexGuard<'a, Option<CaseBuilder>>) -> PyDataFusionResult<Self> {
+        let builder = guard.take().ok_or_else(|| {
+            PyDataFusionError::Common("CaseBuilder has already been consumed".to_string())
+        })?;
+
+        Ok(Self {
+            guard,
+            builder: Some(builder),
+        })
+    }
+
+    fn builder_mut(&mut self) -> &mut CaseBuilder {
+        self.builder
+            .as_mut()
+            .expect("builder should be present while handle is alive")
+    }
+
+    fn into_inner(mut self) -> CaseBuilder {
+        self.builder
+            .take()
+            .expect("builder should be present when consuming handle")
+    }
+}
+
+impl Drop for CaseBuilderHandle<'_> {
+    fn drop(&mut self) {
+        if let Some(builder) = self.builder.take() {
+            *self.guard = Some(builder);
+        }
+    }
+}
 
 #[pyclass(name = "CaseBuilder", module = "datafusion.expr", subclass, frozen)]
 #[derive(Clone)]
@@ -41,45 +78,38 @@ impl From<CaseBuilder> for PyCaseBuilder {
 }
 
 impl PyCaseBuilder {
-    fn lock_case_builder(&self) -> MutexGuard<'_, Option<CaseBuilder>> {
-        self.case_builder.lock()
+    fn case_builder_handle(&self) -> PyDataFusionResult<CaseBuilderHandle<'_>> {
+        let guard = self.case_builder.lock();
+        CaseBuilderHandle::new(guard)
     }
 
     pub fn into_case_builder(self) -> PyDataFusionResult<CaseBuilder> {
-        let mut guard = self.case_builder.lock();
-        guard.take().ok_or_else(|| {
-            PyDataFusionError::Common("CaseBuilder has already been consumed".to_string())
-        })
+        let guard = self.case_builder.lock();
+        CaseBuilderHandle::new(guard).map(CaseBuilderHandle::into_inner)
     }
 }
 
 #[pymethods]
 impl PyCaseBuilder {
     fn when(&self, when: PyExpr, then: PyExpr) -> PyDataFusionResult<PyCaseBuilder> {
-        let mut guard = self.lock_case_builder();
-        let builder = guard.as_mut().ok_or_else(|| {
-            PyDataFusionError::Common("CaseBuilder has already been consumed".to_string())
-        })?;
-        let next_builder = builder.when(when.expr, then.expr);
+        let mut handle = self.case_builder_handle()?;
+        let next_builder = handle.builder_mut().when(when.expr, then.expr);
         Ok(next_builder.into())
     }
 
     fn otherwise(&self, else_expr: PyExpr) -> PyDataFusionResult<PyExpr> {
-        let mut guard = self.lock_case_builder();
-        let builder = guard.as_mut().ok_or_else(|| {
-            PyDataFusionError::Common("CaseBuilder has already been consumed".to_string())
-        })?;
-        builder
-            .otherwise(else_expr.expr)
-            .map(|expr| expr.into())
-            .map_err(Into::into)
+        let mut handle = self.case_builder_handle()?;
+        match handle.builder_mut().otherwise(else_expr.expr) {
+            Ok(expr) => Ok(expr.clone().into()),
+            Err(err) => Err(err.into()),
+        }
     }
 
     fn end(&self) -> PyDataFusionResult<PyExpr> {
-        let mut guard = self.lock_case_builder();
-        let builder = guard.as_mut().ok_or_else(|| {
-            PyDataFusionError::Common("CaseBuilder has already been consumed".to_string())
-        })?;
-        builder.end().map(|expr| expr.into()).map_err(Into::into)
+        let mut handle = self.case_builder_handle()?;
+        match handle.builder_mut().end() {
+            Ok(expr) => Ok(expr.clone().into()),
+            Err(err) => Err(err.into()),
+        }
     }
 }