python polars: df partition with pivot and concat

Question

my goal was to groupby/partition by one column (a below), create a string identifier (b and c columns) then use this b_c identifier as a name for a column in a pivoted data frame.
Code below works OK as far as I can tell, but the path to get the result is a bit twisted. So my question is: can this be done in a simpler way?
BTW, at this tiny scale (max 1k of rows so far) I am not obsessed to make it faster.

data = {
    "a": [1, 1, 1, 2, 2, 3],
    "b": [11, 12, 13, 11, 12, 11],
    "c": ["x1", "x2", "x3", "x1", "x2", "x1"],
    "val": [101, 102, 102, 201, 202, 301],
}
df = pl.DataFrame(data)

print(df)

counter = 0
for tmp_df in df.partition_by("a"):
    grp_df = (
        tmp_df.with_columns((pl.col("b") + "_" + pl.col("c")).alias("col_id"))
        .drop(["b", "c"])
        .pivot(values="val", index="a", columns="col_id")
    )

    if counter == 0:
        result_df = grp_df.select(pl.all())
    else:
        result_df = pl.concat([result_df, grp_df], how="diagonal")
    counter += 1

print(result_df)

Asked By: darked89

||

Source

Answer 1

you can do this in 2 step: first a select step to create the new id column, then the pivot.

Example 1:

(
    df.select(
        'a','val',
        id = pl.col('b').cast(pl.Utf8) + '_' + pl.col('c'))
    .pivot(values='val',index='a', columns='id')
)

# Result
shape: (3, 4)
┌─────┬───────┬───────┬───────┐
│ a   ┆ 11_x1 ┆ 12_x2 ┆ 13_x3 │
│ --- ┆ ---   ┆ ---   ┆ ---   │
│ i64 ┆ i64   ┆ i64   ┆ i64   │
╞═════╪═══════╪═══════╪═══════╡
│ 1   ┆ 101   ┆ 102   ┆ 102   │
│ 2   ┆ 201   ┆ 202   ┆ null  │
│ 3   ┆ 301   ┆ null  ┆ null  │
└─────┴───────┴───────┴───────┘

Example 2:
(suggested by @jqurious), using pl.format

(
    df.select(
        'a','val',
        id = pl.format("{}_{}", "b", "c"))
    .pivot(values='val',index='a', columns='id')
)

Answered By: Luca

python polars: df partition with pivot and concat

Question:

Answers: