fix list concat (#3636)

pola-rs · Jun 9, 2022 · 32f7716 · 32f7716
1 parent 82d6134
commit 32f7716
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 8 deletions.
diff --git a/polars/polars-ops/src/chunked_array/list/namespace.rs b/polars/polars-ops/src/chunked_array/list/namespace.rs
@@ -300,22 +300,25 @@ pub trait ListNameSpaceImpl: AsList {
                         continue;
                     }
                 };
-                let mut already_null = false;
+
+                let mut has_nulls = false;
                 for it in &mut iters {
                     match it.next().unwrap() {
                         Some(s) => {
-                            acc.append(s.as_ref())?;
+                            if !has_nulls {
+                                acc.append(s.as_ref())?;
+                            }
                         }
                         None => {
-                            if !already_null {
-                                builder.append_null();
-                                already_null = true;
-                            }
-
-                            continue;
+                            has_nulls = true;
                         }
                     }
                 }
+                if has_nulls {
+                    builder.append_null();
+                    continue;
+                }
+
                 match inner_type {
                     // structs don't have chunks, so we must first rechunk the underlying series
                     #[cfg(feature = "dtype-struct")]

diff --git a/py-polars/tests/test_lists.py b/py-polars/tests/test_lists.py
@@ -340,3 +340,17 @@ def test_list_ternary_concat() -> None:
         "list2": [["789"], ["zzz"]],
         "result": [["123", "456", "789"], ["zzz"]],
     }
+
+
+def test_list_concat_nulls() -> None:
+    assert pl.DataFrame(
+        {
+            "a": [["a", "b"], None, ["c", "d", "e"], None],
+            "t": [["x"], ["y"], None, None],
+        }
+    ).with_column(pl.concat_list(["a", "t"]).alias("concat"))["concat"].to_list() == [
+        ["a", "b", "x"],
+        None,
+        None,
+        None,
+    ]