Replace case builter with keeping the expressions and then applying as required

Author: timsaucer

src/expr/conditional_expr.rs

@@ -16,56 +16,65 @@
 // under the License.
 
 use crate::{errors::PyDataFusionResult, expr::PyExpr};
-use datafusion::common::{exec_err, DataFusionError};
 use datafusion::logical_expr::conditional_expressions::CaseBuilder;
 use datafusion::prelude::Expr;
 use pyo3::prelude::*;
 
+// TODO(tsaucer) replace this all with CaseBuilder after it implements Clone
+#[derive(Clone, Debug)]
 #[pyclass(name = "CaseBuilder", module = "datafusion.expr", subclass, frozen)]
 pub struct PyCaseBuilder {
-    case_builder: CaseBuilder,
+    expr: Option<Expr>,
+    when: Vec<Expr>,
+    then: Vec<Expr>,
 }
 
-impl From<CaseBuilder> for PyCaseBuilder {
-    fn from(case_builder: CaseBuilder) -> PyCaseBuilder {
-        PyCaseBuilder { case_builder }
+#[pymethods]
+impl PyCaseBuilder {
+    #[new]
+    pub fn new(expr: Option<PyExpr>) -> Self {
+        Self {
+            expr: expr.map(Into::into),
+            when: vec![],
+            then: vec![],
+        }
     }
-}
-
-// TODO(tsaucer) upstream make CaseBuilder impl Clone
-fn builder_clone(case_builder: &CaseBuilder) -> Result<CaseBuilder, DataFusionError> {
-    let Expr::Case(case) = case_builder.end()? else {
-        return exec_err!("CaseBuilder returned an invalid expression");
-    };
 
-    let (when_expr, then_expr) = case
-        .when_then_expr
-        .iter()
-        .map(|(w, t)| (w.as_ref().to_owned(), t.as_ref().to_owned()))
-        .unzip();
+    pub fn when(&self, when: PyExpr, then: PyExpr) -> PyCaseBuilder {
+        println!("when called {self:?}");
+        let mut case_builder = self.clone();
+        case_builder.when.push(when.into());
+        case_builder.then.push(then.into());
 
-    Ok(CaseBuilder::new(
-        case.expr,
-        when_expr,
-        then_expr,
-        case.else_expr,
-    ))
-}
-
-#[pymethods]
-impl PyCaseBuilder {
-    fn when(&self, when: PyExpr, then: PyExpr) -> PyDataFusionResult<PyCaseBuilder> {
-        let case_builder = builder_clone(&self.case_builder)?.when(when.expr, then.expr);
-        Ok(PyCaseBuilder { case_builder })
+        case_builder
     }
 
     fn otherwise(&self, else_expr: PyExpr) -> PyDataFusionResult<PyExpr> {
-        Ok(builder_clone(&self.case_builder)?
-            .otherwise(else_expr.expr)?
-            .into())
+        println!("otherwise called {self:?}");
+        let case_builder = CaseBuilder::new(
+            self.expr.clone().map(Box::new),
+            self.when.clone(),
+            self.then.clone(),
+            Some(Box::new(else_expr.into())),
+        );
+
+        let expr = case_builder.end()?;
+
+        Ok(expr.into())
     }
 
     fn end(&self) -> PyDataFusionResult<PyExpr> {
-        Ok(builder_clone(&self.case_builder)?.end()?.into())
+        println!("end called {self:?}");
+
+        let case_builder = CaseBuilder::new(
+            self.expr.clone().map(Box::new),
+            self.when.clone(),
+            self.then.clone(),
+            None,
+        );
+
+        let expr = case_builder.end()?;
+
+        Ok(expr.into())
     }
 }

src/functions.rs

@@ -230,13 +230,13 @@ fn col(name: &str) -> PyResult<PyExpr> {
 /// Create a CASE WHEN statement with literal WHEN expressions for comparison to the base expression.
 #[pyfunction]
 fn case(expr: PyExpr) -> PyResult<PyCaseBuilder> {
-    Ok(datafusion::logical_expr::case(expr.expr).into())
+    Ok(PyCaseBuilder::new(Some(expr)))
 }
 
 /// Create a CASE WHEN statement with literal WHEN expressions for comparison to the base expression.
 #[pyfunction]
 fn when(when: PyExpr, then: PyExpr) -> PyResult<PyCaseBuilder> {
-    Ok(datafusion::logical_expr::when(when.expr, then.expr).into())
+    Ok(PyCaseBuilder::new(None).when(when, then))
 }
 
 /// Helper function to find the appropriate window function.