[Python] A faster way to unstack columns by groups with Polars?

I need to unstack some columns in polars, but I want to group the columns by column 'a'

import polars as pl

df = pl.from_repr("""
┌─────┬─────┬─────┐
│ a ┆ b ┆ c │
│ --- ┆ --- ┆ --- │
│ i64 ┆ str ┆ str │
╞═════╪═════╪═════╡
│ 1 ┆ a ┆ f │
│ 2 ┆ b ┆ g │
│ 1 ┆ c ┆ h │
│ 1 ┆ d ┆ i │
│ 2 ┆ e ┆ j │
└─────┴─────┴─────┘
""")


My current approach:

seq_ = df["a"].unique()
l = len(seq_)
dfs = []

for idx, a in enumerate(seq_):
df_tmp = df.filter(pl.col("a") == a)
steps = df_tmp.shape[0]
df_tmp = df_tmp.unstack(step=steps, how="horizontal", columns=["b", "c"])
df_tmp = df_tmp.with_columns(pl.lit(a).alias("a"))
dfs.append(df_tmp)

df_tmp = pl.concat(dfs, how="diagonal")


This produces the expected output but is too slow.

┌─────┬─────┬─────┬──────┬─────┬─────┬──────┐
│ a ┆ b_0 ┆ b_1 ┆ b_2 ┆ c_0 ┆ c_1 ┆ c_2 │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ str ┆ str ┆ str ┆ str ┆ str ┆ str │
╞═════╪═════╪═════╪══════╪═════╪═════╪══════╡
│ 1 ┆ a ┆ c ┆ d ┆ f ┆ h ┆ i │
│ 2 ┆ b ┆ e ┆ null ┆ g ┆ j ┆ null │
└─────┴─────┴─────┴──────┴─────┴─────┴──────┘


Is there a faster way to do this in polars?

Continue reading...