Skip to content

Commit

Permalink
feat(python): Add Expr.struct.unnest() as alias for Expr.struct.field…
Browse files Browse the repository at this point in the history
…("*") (#19212)
  • Loading branch information
alonme authored Oct 13, 2024
1 parent b481597 commit 0628e0e
Show file tree
Hide file tree
Showing 3 changed files with 52 additions and 0 deletions.
1 change: 1 addition & 0 deletions py-polars/docs/source/reference/expressions/struct.rst
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ The following methods are available under the `expr.struct` attribute.
:template: autosummary/accessor_method.rst

Expr.struct.field
Expr.struct.unnest
Expr.struct.json_encode
Expr.struct.rename_fields
Expr.struct.with_fields
37 changes: 37 additions & 0 deletions py-polars/polars/expr/struct.py
Original file line number Diff line number Diff line change
Expand Up @@ -152,6 +152,43 @@ def field(self, name: str | list[str], *more_names: str) -> Expr:

return wrap_expr(self._pyexpr.struct_field_by_name(name))

def unnest(self) -> Expr:
"""
Expand the struct into its individual fields.
Alias for `Expr.struct.field("*")`.
>>> df = pl.DataFrame(
... {
... "aaa": [1, 2],
... "bbb": ["ab", "cd"],
... "ccc": [True, None],
... "ddd": [[1, 2], [3]],
... }
... ).select(pl.struct("aaa", "bbb", "ccc", "ddd").alias("struct_col"))
>>> df
shape: (2, 1)
┌──────────────────────┐
│ struct_col │
│ --- │
│ struct[4] │
╞══════════════════════╡
│ {1,"ab",true,[1, 2]} │
│ {2,"cd",null,[3]} │
└──────────────────────┘
>>> df.select(pl.col("struct_col").struct.unnest())
shape: (2, 4)
┌─────┬─────┬──────┬───────────┐
│ aaa ┆ bbb ┆ ccc ┆ ddd │
│ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ str ┆ bool ┆ list[i64] │
╞═════╪═════╪══════╪═══════════╡
│ 1 ┆ ab ┆ true ┆ [1, 2] │
│ 2 ┆ cd ┆ null ┆ [3] │
└─────┴─────┴──────┴───────────┘
"""
return self.field("*")

def rename_fields(self, names: Sequence[str]) -> Expr:
"""
Rename the fields of the struct.
Expand Down
14 changes: 14 additions & 0 deletions py-polars/tests/unit/test_expansion.py
Original file line number Diff line number Diff line change
Expand Up @@ -138,6 +138,20 @@ def test_struct_field_expand_star() -> None:
assert_frame_equal(struct_df.select(pl.col("struct_col").struct.field("*")), df)


def test_struct_unnest() -> None:
"""Same as test_struct_field_expand_star but using the unnest alias."""
df = pl.DataFrame(
{
"aaa": [1, 2],
"bbb": ["ab", "cd"],
"ccc": [True, None],
"ddd": [[1, 2], [3]],
}
)
struct_df = df.select(pl.struct(["aaa", "bbb", "ccc", "ddd"]).alias("struct_col"))
assert_frame_equal(struct_df.select(pl.col("struct_col").struct.unnest()), df)


def test_struct_field_expand_rewrite() -> None:
df = pl.DataFrame({"A": [1], "B": [2]})
assert df.select(
Expand Down

0 comments on commit 0628e0e

Please sign in to comment.