如何在Polars流模式下对聚合操作应用多过滤器

问题背景 #

我有一个大型数据集，希望在不将全量数据加载到内存的前提下执行多组聚合操作。具体来说，我有近百个形如pl.col('x1') == 'y1'的简单过滤器，需要分别获取每个过滤器下按a分组后的b列求和结果。

目前单个过滤器的实现方式如下（能正常以流模式运行）：

dataset.filter(
  pl.col(x) == y
).groupby(pl.col('a')).agg(
  pl.col('b').sum()
).collect(streaming=True)

但近百个过滤器需要多次遍历数据集，耗时过长。于是尝试单遍分组聚合的写法：

dataset.groupby(pl.col('a')).agg(
  *[pl.col('b').filter(pl.col(x) == y).sum().alias(f'{x}={y}') for x, y in filters]
).collect(streaming=True)

这段代码在小数据集上可行，但在完整数据集上即使单个过滤器也会耗尽内存，推测Polars未以流模式执行该查询。想知道是否有办法通过单遍流模式实现需求，这是否是Polars暂不支持的功能？

解决方案：替换聚合内的filter()为pl.when() #

Polars的流模式对聚合阶段内的filter()操作支持有限，这类操作会强制加载全量数据到内存。改用pl.when().then().otherwise(0).sum()的写法可以让Polars在流模式下高效处理，实现单遍遍历且不占过多内存。

修改后的代码：

dataset.groupby(pl.col('a')).agg(
  *[pl.when(pl.col(x) == y).then(pl.col('b')).otherwise(0).sum().alias(f'{x}={y}') for x, y in filters]
).collect(streaming=True)

原理说明 #

• 原逻辑中pl.col('b').filter(cond).sum()等价于pl.when(cond).then(b).otherwise(0).sum()：满足条件时取b的值，否则取0，最终求和结果完全一致。
• 后者的写法能被Polars的流模式优化，每处理一块数据就完成对应分组的累加计算，无需缓存全量数据，完美适配大型数据集的单遍遍历需求。

验证建议 #

可以先用小数据集对比两种写法的输出结果，确认聚合逻辑一致后，再应用到全量数据集上。

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