如何在Polars中对DataFrame按唯一ID实现透视转换并补零？

可以通过以下步骤完成需求：

1. 过滤无效数据：移除id为'N/A'的行，仅保留有效id对应的记录
2. 执行透视操作：以id为行维度，type为列维度，area为对应值进行透视
3. 补全列与填充空值：确保所有出现过的type都作为列存在，无对应值的位置填充0

完整代码示例 #

import polars as pl

data = {
    'id': ['N/A', 'N/A', '1', '1', '2'], 
    'type': ['red', 'blue', 'yellow', 'green', 'yellow'], 
    'area': [0, 0, 3, 4, 5]
}

df = pl.DataFrame(data)

# 过滤id为N/A的行
filtered_df = df.filter(pl.col('id') != 'N/A')

# 获取所有出现过的type种类
all_types = df['type'].unique().to_list()

# 透视操作，用first聚合（每个id+type组合唯一，无需复杂聚合）
pivoted_df = filtered_df.pivot(
    index='id',
    columns='type',
    values='area',
    aggregate_function='first'
).fill_null(0)

# 补全缺失的type列，确保所有类型都被包含
pivoted_df = pivoted_df.with_columns(
    pl.lit(0).alias(type) for type in all_types if type not in pivoted_df.columns
)

# 调整列顺序，与原数据中type的出现顺序一致
result = pivoted_df.select(['id'] + all_types)

print(result)

执行结果 #

shape: (2, 5)
┌─────┬─────┬──────┬────────┬───────┐
│ id  ┆ red ┆ blue ┆ yellow ┆ green │
│ --- ┆ --- ┆ ---  ┆ ---    ┆ ---   │
│ str ┆ i64 ┆ i64  ┆ i64    ┆ i64   │
╞═════╪═════╪══════╪════════╪═══════╡
│ 1   ┆ 0   ┆ 0    ┆ 3      ┆ 4     │
│ 2   ┆ 0   ┆ 0    ┆ 5      ┆ 0     │
└─────┴─────┴──────┴────────┴───────┘

如果希望id作为行索引而非数据列，可添加set_index('id')：

result = result.set_index('id')
print(result)

输出结果：

shape: (2, 4)
┌─────┬─────┬──────┬────────┬───────┐
│ id  ┆ red ┆ blue ┆ yellow ┆ green │
│ --- ┆ --- ┆ ---  ┆ ---    ┆ ---   │
│ str ┆ i64 ┆ i64  ┆ i64    ┆ i64   │
╞═════╪═════╪══════╪════════╪═══════╡
│ 1   ┆ 0   ┆ 0    ┆ 3      ┆ 4     │
│ 2   ┆ 0   ┆ 0    ┆ 5      ┆ 0     │
└─────┴─────┴──────┴────────┴───────┘

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