From b1b322f94ed4b4d18bc1b6de6a3d1c41c22df488 Mon Sep 17 00:00:00 2001
From: deanm0000 <37878412+deanm0000@users.noreply.github.com>
Date: Wed, 2 Oct 2024 03:35:27 -0400
Subject: [PATCH] fix(python): Don't ignore multiple columns in
 LazyFrame.unnest (#19035)

---
 py-polars/polars/lazyframe/frame.py           |  2 +-
 .../tests/unit/lazyframe/test_lazyframe.py    | 27 +++++++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/py-polars/polars/lazyframe/frame.py b/py-polars/polars/lazyframe/frame.py
index 7aa9a870b9c6..d5c56c93ca0a 100644
--- a/py-polars/polars/lazyframe/frame.py
+++ b/py-polars/polars/lazyframe/frame.py
@@ -6528,7 +6528,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
index 026d4f157e2f..fbc9b5e1ae30 100644
--- 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)