allow branching null with struct dtype (#3856)

pola-rs · Jun 29, 2022 · 3b4ea61 · 3b4ea61
1 parent 6817e2c
commit 3b4ea61
Show file tree

Hide file tree

Showing 4 changed files with 59 additions and 8 deletions.
diff --git a/polars/polars-lazy/src/physical_plan/expressions/cast.rs b/polars/polars-lazy/src/physical_plan/expressions/cast.rs
@@ -44,6 +44,18 @@ impl CastExpr {
                         .into_duration(*tu)
                         .into_series())
                 }
+                #[cfg(feature = "dtype-struct")]
+                DataType::Struct(fields) => {
+                    let fields = fields
+                        .iter()
+                        .map(|field| {
+                            Series::full_null(field.name(), input.len(), field.data_type())
+                        })
+                        .collect::<Vec<_>>();
+                    return Ok(StructChunked::new(input.name(), &fields)
+                        .unwrap()
+                        .into_series());
+                }
                 _ => {}
             }
         }

diff --git a/py-polars/polars/internals/whenthen.py b/py-polars/polars/internals/whenthen.py
@@ -18,10 +18,11 @@ class WhenThenThen:
     def __init__(self, pywhenthenthen: Any):
         self.pywhenthenthen = pywhenthenthen
 
-    def when(self, predicate: pli.Expr) -> "WhenThenThen":
+    def when(self, predicate: Union[pli.Expr, bool]) -> "WhenThenThen":
         """
         Start another when, then, otherwise layer.
         """
+        predicate = pli.expr_to_lit_or_expr(predicate)
         return WhenThenThen(self.pywhenthenthen.when(predicate._pyexpr))
 
     def then(
@@ -32,6 +33,7 @@ def then(
             float,
             str,
             None,
+            pli.Series,
             Sequence[
                 Union[
                     int,
@@ -85,13 +87,14 @@ class WhenThen:
     def __init__(self, pywhenthen: Any):
         self._pywhenthen = pywhenthen
 
-    def when(self, predicate: pli.Expr) -> WhenThenThen:
+    def when(self, predicate: Union[pli.Expr, bool]) -> WhenThenThen:
         """
         Start another when, then, otherwise layer.
         """
+        predicate = pli.expr_to_lit_or_expr(predicate)
         return WhenThenThen(self._pywhenthen.when(predicate._pyexpr))
 
-    def otherwise(self, expr: Union[pli.Expr, int, float, str]) -> pli.Expr:
+    def otherwise(self, expr: Union[pli.Expr, int, float, str, None]) -> pli.Expr:
         """
         Values to return in case of the predicate being `False`.
 
@@ -112,7 +115,13 @@ def __init__(self, pywhen: "pywhen"):
     def then(
         self,
         expr: Union[
-            pli.Expr, int, float, str, None, Sequence[Union[None, int, float, str]]
+            pli.Expr,
+            pli.Series,
+            int,
+            float,
+            str,
+            None,
+            Sequence[Union[None, int, float, str]],
         ],
     ) -> WhenThen:
         """
@@ -125,7 +134,7 @@ def then(
         return WhenThen(pywhenthen)
 
 
-def when(expr: pli.Expr) -> When:
+def when(expr: Union[pli.Expr, bool]) -> When:
     """
     Start a when, then, otherwise expression.
 

diff --git a/py-polars/src/arrow_interop/to_py.rs b/py-polars/src/arrow_interop/to_py.rs
@@ -17,9 +17,7 @@ pub(crate) fn to_py_array(array: ArrayRef, py: Python, pyarrow: &PyModule) -> Py
     let schema_ptr: *const ffi::ArrowSchema = &*schema;
     let array_ptr: *const ffi::ArrowArray = &*array;
 
-    let pa = py.import("pyarrow")?;
-
-    let array = pa.getattr("Array")?.call_method1(
+    let array = pyarrow.getattr("Array")?.call_method1(
         "_import_from_c",
         (array_ptr as Py_uintptr_t, schema_ptr as Py_uintptr_t),
     )?;

diff --git a/py-polars/tests/test_queries.py b/py-polars/tests/test_queries.py
@@ -257,3 +257,35 @@ def test_opaque_filter_on_lists_3784() -> None:
             )
         )
     ).collect().to_dict(False) == {"group": [1], "str_list": [["A", "B", "B"]]}
+
+
+def test_ternary_none_struct() -> None:
+    ignore_nulls = False
+
+    def map_expr(name: str) -> pl.Expr:
+        return (
+            pl.when(ignore_nulls or pl.col(name).null_count() == 0)
+            .then(
+                pl.struct(
+                    [
+                        pl.sum(name).alias("sum"),
+                        (pl.count() - pl.col(name).null_count()).alias("count"),
+                    ]
+                ),
+            )
+            .otherwise(None)
+        ).alias("out")
+
+    assert (
+        pl.DataFrame({"groups": [1, 2, 3, 4], "values": [None, None, 1, 2]})
+        .groupby("groups", maintain_order=True)
+        .agg([map_expr("values")])
+    ).to_dict(False) == {
+        "groups": [1, 2, 3, 4],
+        "out": [
+            {"sum": None, "count": None},
+            {"sum": None, "count": None},
+            {"sum": 1, "count": 1},
+            {"sum": 2, "count": 1},
+        ],
+    }