Polars处理1700万条数据：List转Set分组聚合性能优化

原始数据与需求 #

给定如下Polars DataFrame：

df = pl.DataFrame({
    'col1': [["aaa", "aaa"], ["bbb", "ccc"], ["ccc", "ddd", "ddd"], ["ddd", "ddd", "ddd"]],
    'col2': ["a", "a", "a", "a"],
    'col3': ["x", "x", "y", "y"]
})

需求：按col2、col3分组，将col1中的所有元素聚合为无重复的字符串集合（对应Set[String]）。

低效实现及问题 #

最初采用的实现方式在处理1700万条记录时速度极慢，运行10分钟仍未完成：

(df
.group_by("col2", "col3")
.agg(pl.col("col1").flatten().map_elements(set).alias("result"))
)

核心问题是map_elements(set)属于Python层级的逐元素操作，无法利用Polars的矢量化优化能力，在大数据量场景下性能瓶颈显著。

高效优化方案 #

改用矢量化的字符串拼接、拆分与列表去重操作，替代Python层级的集合转换，性能得到大幅提升，代码如下：

df = (
   df
   .with_columns(
       pl.col("col1").list.join(",")
   )
   .group_by("col2", "col3")
   .agg(
       pl.col("col1").alias("col1")
   )
   .with_columns(
       pl.col("col1").list.join(",")
   )
   .with_columns(
       pl.col("col1").str.split(",").list.unique().alias("col1")
   )
)

执行结果 #

┌──────┬──────┬───────────────────────┐
│ col2 ┆ col3 ┆ col1                  │
│ ---  ┆ ---  ┆ ---                   │
│ str  ┆ str  ┆ list[str]             │
╞══════╪══════╪═══════════════════════╡
│ a    ┆ x    ┆ ["aaa", "bbb", "ccc"] │
│ a    ┆ y    ┆ ["ccc", "ddd"]        │
└──────┴──────┴───────────────────────┘

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