Polars中基于分组的列间逐元素列表操作技术问询

报错原因 #

你之前的代码报错是因为list.eval上下文里，pl.col("value")会获取整个分组的value列（长度等于分组行数），而pl.element()是列表中的单个元素，两者长度不匹配，导致计算时出现长度不一致的错误。

正确实现方式 #

Polars支持列表与标量的广播运算，直接对分组聚合后的列表列和当前行的value列做减法即可，无需使用list.eval：

import polars as pl

# 示例DataFrame
df = pl.DataFrame({
    "id": [1, 1, 1, 2, 2],
    "value": [10, 20, 30, 5, 15]
})

# 生成目标列
result = df.with_columns(
    (pl.col("value").implode().over("id") - pl.col("value")).alias("output")
)

print(result)

运行后输出：

shape: (5, 3)
┌─────┬───────┬─────────────────┐
│ id  ┆ value ┆ output          │
│ --- ┆ ---   ┆ ---             │
│ i64 ┆ i64   ┆ list[i64]       │
╞═════╪═══════╪═════════════════╡
│ 1   ┆ 10    ┆ [0, 10, 20]     │
│ 1   ┆ 20    ┆ [-10, 0, 10]    │
│ 1   ┆ 30    ┆ [-20, -10, 0]   │
│ 2   ┆ 5     ┆ [0, 10]         │
│ 2   ┆ 15    ┆ [-10, 0]        │
└─────┴───────┴─────────────────┘

扩展其他聚合操作 #

如果需要尝试其他跨列聚合操作，比如分组内列表元素与当前行value的乘积、除法等，都可以用同样的广播逻辑：

• 乘积：(pl.col("value").implode().over("id") * pl.col("value")).alias("product_output")
• 除法（注意处理除数为0）：(pl.col("value").implode().over("id") / pl.col("value")).alias("div_output")

如果需要更复杂的自定义操作，也可以用list.map结合pl.element()和当前行的列，比如：

df.with_columns(
    pl.col("value")
    .implode()
    .over("id")
    .list.map(lambda x: x - pl.col("value"))
    .alias("output")
)

这种写法和广播运算的效果一致，适合需要自定义逻辑的场景。

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