如何实现Polars 0.20.7版本之前的pivot()旧功能？

版本差异说明 #

在Polars 0.20.7之前，当pivot()方法的columns参数传入多个列时，会基于index列对columns中的每一列单独执行聚合逻辑，而非将多列作为一组处理。

旧版（0.20.7前）示例 #

df = pl.DataFrame(
    {
        "foo": ["one", "one", "two", "two", "one", "two"],
        "bar": ["y", "y", "y", "x", "x", "x"],
        "biz": ['m', 'f', 'm', 'f', 'm', 'f'],
        "baz": [1, 2, 3, 4, 5, 6],
    }
)

df.pivot(index='foo', values='baz', columns=('bar', 'biz'), aggregate_function='sum')

返回结果：

shape: (2, 5)
┌─────┬─────┬─────┬─────┬─────┐
│ foo ┆ y   ┆ x   ┆ m   ┆ f   │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╪═════╪═════╡
│ one ┆ 3   ┆ 5   ┆ 6   ┆ 2   │
│ two ┆ 3   ┆ 10  ┆ 3   ┆ 10  │
└─────┴─────┴─────┴─────┴─────┘

新版（0.20.7及之后）行为 #

官方将此变更归类为Bug修复，但新版会把columns参数中的多列视为一组，生成带组合列名的透视表，示例结果如下：

shape: (2, 5)
┌─────┬───────────┬───────────┬───────────┬───────────┐
│ foo ┆ {"y","m"} ┆ {"y","f"} ┆ {"x","f"} ┆ {"x","m"} │
│ --- ┆ ---       ┆ ---       ┆ ---       ┆ ---       │
│ str ┆ i64       ┆ i64       ┆ i64       ┆ i64       │
╞═════╪═══════════╪═══════════╪═══════════╪═══════════╡
│ one ┆ 1         ┆ 2         ┆ null      ┆ 5         │
│ two ┆ 3         ┆ null      ┆ 10        ┆ null      │
└─────┴───────────┴───────────┴───────────┴───────────┘

还原旧版pivot行为的方法 #

如果需要保留旧版对多列单独聚合的逻辑，可以对每一列分别执行pivot，再通过join合并结果：

# 对bar列执行pivot
pivot_bar = df.pivot(index='foo', values='baz', columns='bar', aggregate_function='sum')
# 对biz列执行pivot
pivot_biz = df.pivot(index='foo', values='baz', columns='biz', aggregate_function='sum')

# 合并两个透视表
result = pivot_bar.join(pivot_biz, on='foo')

执行后即可得到与旧版一致的结果。

也可以封装成通用函数，方便重复调用：

def old_style_pivot(df, index, values, columns, aggregate_function):
    pivot_dfs = []
    for col in columns:
        pivot_df = df.pivot(index=index, values=values, columns=col, aggregate_function=aggregate_function)
        pivot_dfs.append(pivot_df)
    # 依次合并所有透视表
    result = pivot_dfs[0]
    for df_to_join in pivot_dfs[1:]:
        result = result.join(df_to_join, on=index)
    return result

# 使用示例
old_style_pivot(df, index='foo', values='baz', columns=('bar', 'biz'), aggregate_function='sum')

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