如何高效将Polars DataFrame拆分为符合与不符合条件的两个数据集？

你原来的两次filter操作会遍历DataFrame两次，在处理大型数据集时会浪费资源。以下两种方法都只需遍历一次DataFrame就能完成拆分：

方法一：使用partition_by直接按谓词分组 #

partition_by可以直接根据布尔谓词的结果将DataFrame拆分成多个子DataFrame，全程仅遍历一次：

import polars as pl

predicate = pl.col("a").gt(pl.col("b"))
# 按谓词结果分组，返回包含两个DataFrame的列表
dfx, dfy = df.partition_by(predicate, maintain_order=True)

• maintain_order=True会保留原DataFrame的行顺序，若不需要顺序可以省略该参数以获得轻微性能提升
• 返回的列表中，第一个元素是满足predicate（即a > b）的行，第二个是不满足的行

方法二：添加临时标记列后拆分 #

先通过with_columns添加一个布尔标记列（仅遍历一次DataFrame），后续基于标记列过滤时无需重复计算谓词：

import polars as pl

predicate = pl.col("a").gt(pl.col("b"))
# 添加临时标记列
df_with_flag = df.with_columns(is_gt=predicate)
# 基于标记列拆分并删除临时列
dfx = df_with_flag.filter(pl.col("is_gt")).drop("is_gt")
dfy = df_with_flag.filter(~pl.col("is_gt")).drop("is_gt")

两种方法都能避免重复遍历，大幅提升大型数据集的处理效率。

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