使用Polars将多行数据重映射为列格式的实现方法

需求场景 #

现有稀疏数据集的长格式Polars DataFrame：

df = pl.from_repr("""
┌─────┬─────┬───────┐
│ id  ┆ key ┆ value │
│ --- ┆ --- ┆ ---   │
│ str ┆ str ┆ i64   │
╞═════╪═════╪═══════╡
│ a   ┆ m1  ┆ 1     │
│ a   ┆ m2  ┆ 2     │
│ a   ┆ m3  ┆ 1     │
│ b   ┆ m2  ┆ 4     │
│ c   ┆ m1  ┆ 2     │
│ c   ┆ m3  ┆ 6     │
│ d   ┆ m4  ┆ 4     │
│ e   ┆ m2  ┆ 1     │
└─────┴─────┴───────┘
""")

需要转换为如下列格式的稀疏表示：

shape: (5, 5)
┌─────┬──────┬──────┬──────┬──────┐
│ id  ┆ m1   ┆ m2   ┆ m3   ┆ m4   │
│ --- ┆ ---  ┆ ---  ┆ ---  ┆ ---  │
│ str ┆ i64  ┆ i64  ┆ i64  ┆ i64  │
╞═════╪══════╪══════╪══════╪══════╡
│ a   ┆ 1    ┆ 2    ┆ 1    ┆ null │
│ b   ┆ null ┆ 4    ┆ null ┆ null │
│ c   ┆ 2    ┆ null ┆ 6    ┆ null │
│ d   ┆ null ┆ null ┆ null ┆ 4    │
│ e   ┆ null ┆ 1    ┆ null ┆ null │
└─────┴──────┴──────┴──────┴──────┘

高效实现方法 #

直接使用Polars内置的pivot方法即可完成转换，该方法针对大数据量做了优化，天然适配稀疏场景：

wide_df = df.pivot(index='id', columns='key', values='value')

参数说明 #

• index：指定作为行标识的列（这里是id）
• columns：指定要转成列的字段（这里是key）
• values：指定填充到新列中的值（这里是value）

转换后得到的wide_df完全符合目标格式，缺失值自动以null填充，完美匹配稀疏数据的存储需求。

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