diff --git a/crates/polars-python/src/dataframe/general.rs b/crates/polars-python/src/dataframe/general.rs index 543a32eb98d9..af0ff8928bdb 100644 --- a/crates/polars-python/src/dataframe/general.rs +++ b/crates/polars-python/src/dataframe/general.rs @@ -651,13 +651,6 @@ impl PyDataFrame { }) } - pub fn unnest(&self, py: Python, columns: Vec) -> PyResult { - let df = py - .allow_threads(|| self.df.unnest(columns)) - .map_err(PyPolarsErr::from)?; - Ok(df.into()) - } - pub fn clear(&self, py: Python) -> Self { py.allow_threads(|| self.df.clear()).into() } diff --git a/py-polars/polars/dataframe/frame.py b/py-polars/polars/dataframe/frame.py index fc6f892a25ca..3c04447a9a6b 100644 --- a/py-polars/polars/dataframe/frame.py +++ b/py-polars/polars/dataframe/frame.py @@ -11261,8 +11261,7 @@ def unnest( │ bar ┆ 2 ┆ b ┆ null ┆ [3] ┆ womp │ └────────┴─────┴─────┴──────┴───────────┴───────┘ """ - columns = _expand_selectors(self, columns, *more_columns) - return self._from_pydf(self._df.unnest(columns)) + return self.lazy().unnest(columns, *more_columns).collect(_eager=True) def corr(self, **kwargs: Any) -> DataFrame: """ diff --git a/py-polars/tests/unit/constructors/test_constructors.py b/py-polars/tests/unit/constructors/test_constructors.py index dbeb26484bd0..af801301b16a 100644 --- a/py-polars/tests/unit/constructors/test_constructors.py +++ b/py-polars/tests/unit/constructors/test_constructors.py @@ -15,6 +15,7 @@ from pydantic import BaseModel, Field, TypeAdapter import polars as pl +import polars.selectors as cs from polars._utils.construction.utils import try_get_type_hints from polars.datatypes import numpy_char_code_to_dtype from polars.dependencies import dataclasses, pydantic @@ -1397,24 +1398,23 @@ def test_from_records_nullable_structs() -> None: assert series.to_list() == [] -def test_from_categorical_in_struct_defined_by_schema() -> None: +@pytest.mark.parametrize("unnest_column", ["a", pl.col("a"), cs.by_name("a")]) +def test_from_categorical_in_struct_defined_by_schema(unnest_column: Any) -> None: df = pl.DataFrame( - { - "a": [ - {"value": "foo", "counts": 1}, - {"value": "bar", "counts": 2}, - ] - }, + {"a": [{"value": "foo", "counts": 1}, {"value": "bar", "counts": 2}]}, schema={"a": pl.Struct({"value": pl.Categorical, "counts": pl.UInt32})}, ) - result = df.unnest("a") - expected = pl.DataFrame( {"value": ["foo", "bar"], "counts": [1, 2]}, schema={"value": pl.Categorical, "counts": pl.UInt32}, ) - assert_frame_equal(result, expected, categorical_as_str=True) + + res_eager = df.unnest(unnest_column) + assert_frame_equal(res_eager, expected, categorical_as_str=True) + + res_lazy = df.lazy().unnest(unnest_column) + assert_frame_equal(res_lazy.collect(), expected, categorical_as_str=True) def test_nested_schema_construction() -> None: