如何高效对Polars DataFrame执行多互斥条件筛选？

问题背景 #

使用Polars处理百万行级别的DataFrame，需要按三个互斥条件将行分为三类：

1. b和c列均为0的行
2. b和c列均不为0的行
3. b和c列一个为0、一个不为0的行

示例数据：

import polars as pl

df = pl.from_repr("""
┌─────┬───────┬────────┐
│ a   ┆   b   ┆   c    │
│ --- ┆ ---   ┆ ---    │
│ i64 ┆ i64   ┆ i64    │
╞═════╪═══════╪════════╡
│ 1   ┆   0   ┆   0    │
│ 2   ┆   1   ┆   1    │
│ 3   ┆   0   ┆   1    │
└─────┴───────┴────────┘
""")

现有方法是多次调用filter，但每次筛选都要遍历全量数据，效率较低。

高效实现方案 #

1. 先添加分类标签（仅遍历一次DataFrame） #

通过pl.when().then().otherwise()给每行打上唯一的类别标签，这一步只需要遍历数据一次，后续所有操作都基于标签完成：

df = df.with_columns(
    category=pl.when(pl.all_horizontal(pl.col('b','c') == 0))
              .then("all_zeros")
              .when(pl.all_horizontal(pl.col('b','c') != 0))
              .then("no_zeros")
              .otherwise("mixed")
)

2. 拆分出三类子集 #

方法一：基于标签筛选

标签已提前计算完成，后续筛选无需重复计算条件，速度极快：

all_zeros = df.filter(pl.col("category") == "all_zeros")
no_zeros = df.filter(pl.col("category") == "no_zeros")
mixed = df.filter(pl.col("category") == "mixed")

方法二：用partition_by直接拆分（更高效）

Polars的partition_by是专门的分组拆分优化API，内部处理更高效，适合大规模数据：

all_zeros, no_zeros, mixed = df.partition_by("category", maintain_order=True)

3. 用group_by做聚合分析 #

如果不需要拆分数据，而是要对每个类别做统计聚合，直接基于标签分组即可：

agg_result = df.group_by("category").agg(
    pl.col("a").count().alias("row_count"),
    pl.col("a").sum().alias("total_a")
)

方案优势 #

• 仅需一次遍历计算条件，避免多次filter重复扫描全量数据，大幅提升百万行级数据的处理效率
• 复杂条件只需编写一次，后续操作无需重复复制条件，降低维护成本和出错概率

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