fix(rust, python): don't set fast_explode if null values in list (#5838)

pola-rs · Dec 17, 2022 · 33d2a11 · 33d2a11
1 parent 51be576
commit 33d2a11
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/polars/polars-core/src/chunked_array/list/iterator.rs b/polars/polars-core/src/chunked_array/list/iterator.rs
@@ -129,7 +129,7 @@ impl ListChunked {
         if self.is_empty() {
             return self.clone();
         }
-        let mut fast_explode = true;
+        let mut fast_explode = self.null_count() == 0;
         let mut ca: ListChunked = self
             .amortized_iter()
             .map(|opt_v| {
@@ -157,7 +157,7 @@ impl ListChunked {
         if self.is_empty() {
             return Ok(self.clone());
         }
-        let mut fast_explode = true;
+        let mut fast_explode = self.null_count() == 0;
         let mut ca: ListChunked = self
             .amortized_iter()
             .map(|opt_v| {

diff --git a/py-polars/tests/unit/test_lists.py b/py-polars/tests/unit/test_lists.py
@@ -564,3 +564,12 @@ def test_empty_eval_dtype_5546() -> None:
 def test_fast_explode_flag() -> None:
     df1 = pl.DataFrame({"values": [[[1, 2]]]})
     assert df1.clone().vstack(df1)["values"].flags["FAST_EXPLODE"]
+
+
+def test_list_amortized_apply_explode_5812() -> None:
+    s = pl.Series([None, [1, 3], [0, -3], [1, 2, 2]])
+    assert s.arr.sum().to_list() == [None, 4, -3, 5]
+    assert s.arr.min().to_list() == [None, 1, -3, 1]
+    assert s.arr.max().to_list() == [None, 3, 0, 2]
+    assert s.arr.arg_min().to_list() == [None, 0, 1, 0]
+    assert s.arr.arg_max().to_list() == [None, 1, 0, 1]