fix invalid concat dtype (#3622)

pola-rs · Jun 8, 2022 · f7aa58c · f7aa58c
1 parent b27cb9c
commit f7aa58c
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 1 deletion.
diff --git a/polars/polars-lazy/src/dsl/functions.rs b/polars/polars-lazy/src/dsl/functions.rs
@@ -225,7 +225,7 @@ pub fn concat_lst<E: AsRef<[IE]>, IE: Into<Expr> + Clone>(s: E) -> Expr {
     Expr::AnonymousFunction {
         input: s,
         function,
-        output_type: GetOutput::map_dtype(|dt| DataType::List(Box::new(dt.clone()))),
+        output_type: GetOutput::same_type(),
         options: FunctionOptions {
             collect_groups: ApplyOptions::ApplyFlat,
             input_wildcard_expansion: true,

diff --git a/py-polars/tests/test_lists.py b/py-polars/tests/test_lists.py
@@ -270,3 +270,34 @@ def test_empty_list_construction() -> None:
     assert pl.DataFrame([{"array": [], "not_array": 1234}], orient="row").to_dict(
         False
     ) == {"array": [[]], "not_array": [1234]}
+
+
+def test_list_ternary_concat() -> None:
+    df = pl.DataFrame(
+        {
+            "list1": [["123", "456"], None],
+            "list2": [["789"], ["zzz"]],
+        }
+    )
+
+    assert df.with_column(
+        pl.when(pl.col("list1").is_null())
+        .then(pl.col("list1").arr.concat(pl.col("list2")))
+        .otherwise(pl.col("list2"))
+        .alias("result")
+    ).to_dict(False) == {
+        "list1": [["123", "456"], None],
+        "list2": [["789"], ["zzz"]],
+        "result": [["789"], None],
+    }
+
+    assert df.with_column(
+        pl.when(pl.col("list1").is_null())
+        .then(pl.col("list2"))
+        .otherwise(pl.col("list1").arr.concat(pl.col("list2")))
+        .alias("result")
+    ).to_dict(False) == {
+        "list1": [["123", "456"], None],
+        "list2": [["789"], ["zzz"]],
+        "result": [["123", "456", "789"], ["zzz"]],
+    }