Polars 0.18.4稀疏DataFrame压缩为列表列的实现需求

针对稀疏结构的Polars DataFrame，在0.18.4版本下按以下规则压缩：

• 生成column_names（符合条件的列名列表）和column_values（对应值列表）两列
• 过滤掉值低于指定阈值的条目（示例中保留>1的值）
• column_values按降序排列，column_names同步对应排序
• 全程优化内存占用

代码实现 #

import polars as pl

# 示例数据（实际场景中column_names为给定变量）
df = pl.DataFrame({
    "column1": [2, 1, 3],
    "column2": [0, 2, 0],
    "column3": [0, 0, 4]
})
column_names = df.columns
threshold = 1

# 核心处理逻辑
result = (
    df
    .with_row_index("row_idx")  # 添加临时行索引，用于按原行分组
    .melt(id_vars="row_idx", value_vars=column_names)  # 宽表转长表，仅处理目标列
    .filter(pl.col("value") > threshold)  # 过滤不符合阈值条件的行，减少后续计算量
    .sort(["row_idx", "value"], descending=[False, True])  # 先按行索引排序，再按值降序，确保分组后顺序正确
    .group_by("row_idx", maintain_order=True)  # 按原行分组，保持输出顺序与原DataFrame一致
    .agg(
        column_names=pl.col("variable").list(),
        column_values=pl.col("value").list()
    )
    .drop("row_idx")  # 删除临时索引列，精简结果
)

print(result)

内存优化说明 #

1. 提前过滤：在转长表后立即过滤掉不符合阈值的行，减少后续分组和聚合的数据量
2. 精准列处理：melt仅指定column_names作为目标列，避免无关列参与计算
3. 减少中间数据：临时行索引仅用于分组，处理完成后立即删除，不保留冗余数据
4. 有序分组：group_by时使用maintain_order=True，避免额外排序开销，同时保证输出行顺序与原DataFrame一致

输出结果 #

┌────────────────────────┬───────────────┐
│ column_names           ┆ column_values │
│ ---                    ┆ ---           │
│ list[str]              ┆ list[i64]     │
╞════════════════════════╪═══════════════╡
│ ["column1"]            ┆ [2]           │
│ ["column2"]            ┆ [2]           │
│ ["column3", "column1"] ┆ [4, 3]        │
└────────────────────────┴───────────────┘

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