如何在Polars中无需添加索引将逆透视后的DataFrame转回原结构

问题背景 #

创建Polars DataFrame并执行unpivot得到长格式数据后，发现需要手动添加索引才能通过pivot转回原始宽格式结构，想知道是否存在无需额外索引的实现方式，同时提及R语言中可以完成该操作。

初始Polars DataFrame #

import polars as pl

df = pl.DataFrame({
    'A': range(1,4),
    'B': range(1,4),
    'C': range(1,4),
    'D': range(1,4)
})

print(df)

输出：

shape: (3, 4)
┌─────┬─────┬─────┬─────┐
│ A   ┆ B   ┆ C   ┆ D   │
│ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╪═════╡
│ 1   ┆ 1   ┆ 1   ┆ 1   │
│ 2   ┆ 2   ┆ 2   ┆ 2   │
│ 3   ┆ 3   ┆ 3   ┆ 3   │
└─────┴─────┴─────┴─────┘

执行unpivot得到长格式 #

df_long = df.unpivot(
    variable_name="recipe",
    value_name="revenue")

print(df_long)

输出：

shape: (12, 2)
┌────────┬─────────┐
│ recipe ┆ revenue │
│ ---    ┆ ---     │
│ str    ┆ i64     │
╞════════╪═════════╡
│ A      ┆ 1       │
│ A      ┆ 2       │
│ A      ┆ 3       │
│ B      ┆ 1       │
│ B      ┆ 2       │
│ …      ┆ …       │
│ C      ┆ 2       │
│ C      ┆ 3       │
│ D      ┆ 1       │
│ D      ┆ 2       │
│ D      ┆ 3       │
└────────┴─────────┘

当前添加索引转回宽格式的方法 #

目前需要手动添加索引列才能完成转回：

df_long = df_long.with_columns(index=pl.col("revenue").cum_count().over("recipe"))

df_long.pivot(
    on='recipe',
    index='index',
    values='revenue',
    aggregate_function='first'
)

输出：

shape: (3, 5)
┌───────┬─────┬─────┬─────┬─────┐
│ index ┆ A   ┆ B   ┆ C   ┆ D   │
│ ---   ┆ --- ┆ --- ┆ --- ┆ --- │
│ u32   ┆ i64 ┆ i64 ┆ i64 ┆ i64 │
╞═══════╪═════╪═════╪═════╪═════╡
│ 1     ┆ 1   ┆ 1   ┆ 1   ┆ 1   │
│ 2     ┆ 2   ┆ 2   ┆ 2   ┆ 2   │
│ 3     ┆ 3   ┆ 3   ┆ 3   ┆ 3   │
└───────┴─────┴─────┴─────┴─────┘

R语言中的实现方式 #

R中无需额外索引即可完成转换，步骤如下：
首先将Polars DataFrame转为Pandas格式：

df_pandas = df.to_pandas()

然后使用R的tidyverse工具链：

library(tidyverse)
library(reticulate)

df_long <-
    py$df_pandas |>
    pivot_longer(
        everything(),
        names_to = 'recipe',
        values_to = 'value'
    )

df_long |>
    pivot_wider(
        names_from='recipe',
        values_from='value'
    ) |>
    unnest(cols = c(A,B,C,D))

无需额外索引的Polars实现 #

可以通过聚合为列表再展开的方式，无需手动添加索引即可转回原始宽格式：

df_long.pivot(
    on='recipe',
    values='revenue',
    aggregate_function=list
).explode(pl.all())

输出结果：

shape: (3, 4)
┌─────┬─────┬─────┬─────┐
│ A   ┆ B   ┆ C   ┆ D   │
│ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╪═════╡
│ 1   ┆ 1   ┆ 1   ┆ 1   │
│ 2   ┆ 2   ┆ 2   ┆ 2   │
│ 3   ┆ 3   ┆ 3   ┆ 3   │
└─────┴─────┴─────┴─────┘

原理说明：aggregate_function=list会将每个recipe对应的所有值收集为列表，随后explode将列表展开为独立行，完美还原原始结构且无额外索引列。

内容的提问来源于stack exchange，提问作者Joe