如何在Polars中展开列表列并按比例拆分整数列？

问题背景 #

现有如下Polars数据集，其中cat列为字符串列表，qty为对应数值：

import polars as pl

df = pl.DataFrame({
    'cat': [['green', 'blue'], ['green'], ['red'], ['blue'], ['red', 'blue']],
    'qty': [23, 23, 4, 5, 15]
})

数据集预览：

┌───────────────────┬─────┐
│ cat               ┆ qty │
│ ---               ┆ --- │
│ list[str]         ┆ i64 │
╞═══════════════════╪═════╡
│ ["green", "blue"] ┆ 23  │
│ ["green"]         ┆ 23  │
│ ["red"]           ┆ 4   │
│ ["blue"]          ┆ 5   │
│ ["red", "blue"]   ┆ 15  │
└───────────────────┴─────┘

需求是将cat列的列表展开为单独行，并按列表元素个数比例拆分qty值，得到如下结果：

┌───────┬──────┐
│ cat   ┆ qty  │
│ ---   ┆ ---  │
│ str   ┆ f64  │
╞═══════╪══════╡
│ green ┆ 11.5 │
│ blue  ┆ 11.5 │
│ green ┆ 23.0 │
│ red   ┆ 4.0  │
│ blue  ┆ 5.0  │
│ red   ┆ 7.5  │
│ blue  ┆ 7.5  │
└───────┴──────┘

纯Polars最优方案 #

立即执行版本 #

result = (
    df
    # 计算每个cat列表的元素个数
    .with_columns(pl.col('cat').list.lengths().alias('cat_len'))
    # 展开cat列的列表元素
    .explode('cat')
    # 按比例拆分qty，最后删除临时列
    .with_columns(pl.col('qty') / pl.col('cat_len'))
    .drop('cat_len')
)

print(result)

延迟模式（Lazy Mode）版本 #

lazy_result = (
    df.lazy()
    .with_columns(pl.col('cat').list.lengths().alias('cat_len'))
    .explode('cat')
    .with_columns(pl.col('qty') / pl.col('cat_len'))
    .drop('cat_len')
    # 触发执行
    .collect()
)

print(lazy_result)

方案优势 #

• 原生Polars操作：无需转换为Pandas，避免跨库数据拷贝开销
• 支持延迟模式：结合Polars的查询优化器，在处理大数据集时能显著提升性能、降低内存占用
• 无伪索引依赖：逻辑简洁直接，通过计算列表长度实现比例拆分，不需要额外创建索引列

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