Add pl.Expr.struct.unnest()

[mkleinbort]

Description

Structs are a fantastic data type, and I like working with them. However, I don't know a way to easily expand a struct column into multiple columns without using the dataframe level .unnest() (or looping over the field names)

I propose adding pl.Expr.struct.unnest

Example:

df = pl.DataFrame({'x': [{'x1':10, 'x2':20}]}) # Sample data

df.unnest('x') # Current solution to expand the dataframe

df.with_columns(pl.col('x').stuct.unnest()) # Proposed addition

Note that I envision pl.col('x').stuct.unnest() as the equivalent of pl.col('x1', 'x2')

To give a usecase, suppose I want to write an expression to increase the value of x1 in the above example.

Today one can do

df.unnest('x').select(x=pl.struct(pl.col('x1')+1, 'x2')) 
# This works but requires the df-level unnest, meaning it can't be saved as an expression

[avimallu]

I don't think this is possible the way Expr is implemented right now in Polars. From the user guide section that I wrote with Ritchie's input:

Polars expressions always have a Fn(Series) -> Series signature and Struct is thus the data type that allows us to provide multiple columns as input/output of an expression. In other words, all expressions have to return a Series object, and Struct allows us to stay consistent with that requirement.

[mkleinbort]

Interesting, but not strictly true I think. Or at least there is some behaviour that appears as an exception to that rule.

df = pl.DataFrame({'x':[1,2], 'y':[3,4]})
df.with_columns(pl.col('x','y').add(1).name.suffix('_incremented'))

shape: (2, 4)
┌─────┬─────┬───────────────┬───────────────┐
│ x   ┆ y   ┆ x_incremented ┆ y_incremented │
│ --- ┆ --- ┆ ---           ┆ ---           │
│ i64 ┆ i64 ┆ i64           ┆ i64           │
╞═════╪═════╪═══════════════╪═══════════════╡
│ 1   ┆ 3   ┆ 2             ┆ 4             │
│ 2   ┆ 4   ┆ 3             ┆ 5             │
└─────┴─────┴───────────────┴───────────────┘

[avimallu]

I believe that internally it's still translated to series operations with the same function signature. From the user guide that mentions expression expansion:

a single expression that specifies multiple columns expands into a list of expressions (depending on the DataFrame schema), resulting in being able to select multiple columns + run computation on them!

[cmdlineluser]

#10919 (comment)

I think what's really lacking here is Expr.unnest(). While this may not be feasible in the near-future until we have expressions that can return dataframes [...]

[mcrumiller]

Could we not internally translate pl.col('a').struct.unnest() to:

df.select(
    pl.col('a').struct.field("field1"),
    pl.col('a').struct.field("field2"),
)

[deanm0000]

@mcrumiller does df.schema "know" that 'a' has those fields or, at least, how many fields are in the struct?

[mkleinbort]

@deanm0000 - does this answer your question?

import polars as pl 
df = pl.DataFrame({'x': [{'x1':'A', 'x2':20}]})
df.schema
>>> OrderedDict([('x', Struct({'x1': String, 'x2': Int64}))])

[ritchie46]

This isn't possible naively. An expression can only return a single result series.

@mcrumiller might work, but then we must rely on CSE to not duplicate the work done before the struct.

[cmdlineluser]

I think the latest struct commits will close this:

• feat: Add struct.with_fields #16305
• feat: Add struct.field expansion (regex, wildcard, columns) #16320

df.with_columns(pl.col("x").struct.field("*"))
# shape: (1, 3)
# ┌───────────┬─────┬─────┐
# │ x         ┆ x1  ┆ x2  │
# │ ---       ┆ --- ┆ --- │
# │ struct[2] ┆ i64 ┆ i64 │
# ╞═══════════╪═════╪═════╡
# │ {10,20}   ┆ 10  ┆ 20  │
# └───────────┴─────┴─────┘

df.with_columns(
    pl.col("x").struct.with_fields(
        pl.col("x").struct.field("x1") + 1
    )
)
# shape: (1, 1)
# ┌───────────┐
# │ x         │
# │ ---       │
# │ struct[2] │
# ╞═══════════╡
# │ {11,20}   │
# └───────────┘

[ritchie46]

Yes, I think we have a more general form now with the two last commits. The unnest can be done with a wildcard as shown by @cmdlineluser. There is also regex support and multiple column names.