如何用Polars实现Pandas的groupby+sum+unstack等效操作？

需求背景 #

在Pandas中可通过以下代码实现分组聚合后展开（unstack）的操作：

import pandas as pd
pandasdf = pd.DataFrame(
    {
        "A": [1, 2, 3, 4, 5],
        "fruits": ["banana", "banana", "apple", "apple", "banana"],
        "B": [5, 4, 3, 2, 1],
        "cars": ["beetle", "audi", "beetle", "beetle", "beetle"],
        "optional": [28, 300, None, 2, -30],
    }
)
pandasdf.groupby(["fruits","cars"])['B'].sum().unstack()

执行后得到结果：

cars    audi  beetle
fruits              
apple    NaN     5.0
banana   4.0     6.0

尝试的Polars代码及问题 #

我尝试了以下Polars代码：

import polars as pl
df = pl.DataFrame(
    {
        "A": [1, 2, 3, 4, 5],
        "fruits": ["banana", "banana", "apple", "apple", "banana"],
        "B": [5, 4, 3, 2, 1],
        "cars": ["beetle", "audi", "beetle", "beetle", "beetle"],
        "optional": [28, 300, None, 2, -30],
    }
)
df.group_by("fruits", "cars").agg(pl.col("B").sum())

得到的结果是：

shape: (3, 3)
┌────────┬────────┬─────┐
│ fruits ┆ cars   ┆ B   │
│ ---    ┆ ---    ┆ --- │
│ str    ┆ str    ┆ i64 │
╞════════╪════════╪═════╡
│ banana ┆ audi   ┆ 4   │
│ banana ┆ beetle ┆ 6   │
│ apple  ┆ beetle ┆ 5   │
└────────┴────────┴─────┘

但Polars的unstack函数行为与Pandas不一致，请问如何在Polars中创建等效的结果表？同时需要Polars中替代pd.crosstab的方案，由于数据集较大（用于Apriori算法），代码效率十分重要。

解决方案 #

方法1：用pivot实现等效unstack效果 #

Polars可直接通过pivot一步完成分组聚合与展开，无需先分组再unstack，性能更优：

result = df.pivot(
    index="fruits",
    columns="cars",
    values="B",
    aggregate_function=pl.sum
)
print(result)

执行结果：

shape: (2, 3)
┌────────┬──────┬────────┐
│ fruits ┆ audi ┆ beetle │
│ ---    ┆ ---  ┆ ---    │
│ str    ┆ i64  ┆ i64    │
╞════════╪══════╪════════╡
│ apple  ┆ null ┆ 5      │
│ banana ┆ 4    ┆ 6      │
└────────┴──────┴────────┘

若要将null转为NaN（与Pandas结果完全一致），可追加类型转换：

result = df.pivot(
    index="fruits",
    columns="cars",
    values="B",
    aggregate_function=pl.sum
).with_columns(pl.all().cast(pl.Float64))
print(result)

得到：

shape: (2, 3)
┌────────┬──────┬────────┐
│ fruits ┆ audi ┆ beetle │
│ ---    ┆ ---  ┆ ---    │
│ str    ┆ f64  ┆ f64    │
╞════════╪══════╪════════╡
│ apple  ┆ NaN  ┆ 5.0    │
│ banana ┆ 4.0  ┆ 6.0    │
└────────┴──────┴────────┘

方法2：替代pd.crosstab的方案 #

Polars无直接的crosstab函数，但可通过pivot结合对应聚合函数实现等效功能。例如实现pd.crosstab(df['fruits'], df['cars'])的计数效果：

crosstab_result = df.pivot(
    index="fruits",
    columns="cars",
    values="B",  # 任意数值列均可，此处用B列
    aggregate_function=pl.count
).fill_null(0)
print(crosstab_result)

执行结果：

shape: (2, 3)
┌────────┬──────┬────────┐
│ fruits ┆ audi ┆ beetle │
│ ---    ┆ ---  ┆ ---    │
│ str    ┆ u32  ┆ u32    │
╞════════╪══════╪════════╡
│ apple  ┆ 0    ┆ 2      │
│ banana ┆ 1    ┆ 2      │
└────────┴──────┴────────┘

效率说明 #

Polars的pivot基于向量化实现，处理大规模数据集时性能优于Pandas的groupby+unstack或crosstab，适配Apriori算法这类对数据处理效率要求较高的场景。

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