fix_unnest_multiple_cols

pola-rs · Sep 30, 2024 · b3b1873 · b3b1873
1 parent 71a8b05
commit b3b1873
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 1 deletion.
diff --git a/py-polars/polars/lazyframe/frame.py b/py-polars/polars/lazyframe/frame.py
@@ -6524,7 +6524,7 @@ def unnest(
         │ bar    ┆ 2   ┆ b   ┆ null ┆ [3]       ┆ womp  │
         └────────┴─────┴─────┴──────┴───────────┴───────┘
         """
-        columns = parse_into_list_of_expressions(columns)
+        columns = parse_into_list_of_expressions(columns, *more_columns)
         return self._from_pyldf(self._ldf.unnest(columns))
 
     def merge_sorted(self, other: LazyFrame, key: str) -> LazyFrame:

diff --git a/py-polars/tests/unit/lazyframe/test_lazyframe.py b/py-polars/tests/unit/lazyframe/test_lazyframe.py
@@ -1397,3 +1397,30 @@ def test_lf_properties() -> None:
         assert lf.dtypes == [pl.Int64, pl.Float64, pl.String]
     with pytest.warns(PerformanceWarning):
         assert lf.width == 3
+
+
+def test_lf_unnest() -> None:
+    lf = pl.DataFrame(
+        [
+            pl.Series(
+                "a",
+                [{"ab": [1, 2, 3], "ac": [3, 4, 5]}],
+                dtype=pl.Struct({"ab": pl.List(pl.Int64), "ac": pl.List(pl.Int64)}),
+            ),
+            pl.Series(
+                "b",
+                [{"ba": [5, 6, 7], "bb": [7, 8, 9]}],
+                dtype=pl.Struct({"ba": pl.List(pl.Int64), "bb": pl.List(pl.Int64)}),
+            ),
+        ]
+    ).lazy()
+
+    expected = pl.DataFrame(
+        [
+            pl.Series("ab", [[1, 2, 3]], dtype=pl.List(pl.Int64)),
+            pl.Series("ac", [[3, 4, 5]], dtype=pl.List(pl.Int64)),
+            pl.Series("ba", [[5, 6, 7]], dtype=pl.List(pl.Int64)),
+            pl.Series("bb", [[7, 8, 9]], dtype=pl.List(pl.Int64)),
+        ]
+    )
+    assert_frame_equal(lf.unnest("a", "b").collect(), expected)