本文探讨了在 Polars 中对包含字符串列表的列进行分组求交集的有效方法。针对直接使用列表操作可能遇到的挑战，教程提出了一种通过展开（explode）、行索引计数和过滤的巧妙策略。该方法将列表交集问题转化为元素在组内所有原始行中出现的计数问题，最终实现高效且准确的分组列表交集聚合。
在数据处理中，我们经常会遇到需要对 DataFrame 中的列表（list）类型列进行分组聚合的场景。一个常见的需求是，对于每个分组，找出其所有列表中元素的交集。例如，给定一个包含 id 和 values（字符串列表）的 DataFrame，我们希望得到每个 id 组内所有 values 列表的共同元素。
考虑以下示例 DataFrame：
import polars as pl df = pl.DataFrame( {"id": [1,1,2,2,3,3], "values": [["A", "B"], ["B", "C"], ["A", "B"], ["B", "C"], ["A", "B"], ["B", "C"]] } ) print(df)
其输出为：
shape: (6, 2) ┌─────┬───────────┐ │ id ┆ values │ │ --- ┆ --- │ │ i64 ┆ list[str] │ ╞═════╪═══════════╡ │ 1 ┆ ["A", "B"]│ │ 1 ┆ ["B", "C"]│ │ 2 ┆ ["A", "B"]│ │ 2 ┆ ["B", "C"]│ │ 3 ┆ ["A", "B"]│ │ 3 ┆ ["B", "C"]│ └─────┴───────────┘
我们期望的输出是每个 id 组内 values 列表的交集，即：
shape: (3, 2) ┌─────┬───────────┐ │ id ┆ values │ │ --- ┆ --- │ │ i64 ┆ list[str] │ ╞═════╪═══════════╡ │ 1 ┆ ["B"] │ │ 2 ┆ ["B"] │ │ 3 ┆ ["B"] │ └─────┴───────────┘
直接列表操作的局限性
初看起来，我们可能会尝试使用 pl.reduce 结合 Polars 提供的列表集合操作 list.set_intersection。然而，这种方法往往无法直接达到预期效果，因为它可能在 Series 级别而不是列表元素级别进行聚合，导致结果不符合预期。
例如，以下尝试：
# 尝试1：直接在聚合中使用pl.reduce result_reduce_1 = df.group_by("id").agg( pl.reduce(function=lambda acc, x: acc.list.set_intersection(x), exprs=pl.col("values")) ) print(result_reduce_1)
输出显示 values 列变成了 list[list[str]]，这并非我们想要的单个列表的交集：
shape: (3, 2) ┌─────┬──────────────────────────┐ │ id ┆ values │ │ --- ┆ --- │ │ i64 ┆ list[list[str]] │ ╞═════╪══════════════════════════╡ │ 1 ┆ [["A", "B"], ["B", "C"]] │ │ 3 ┆ [["A", "B"], ["B", "C"]] │ │ 2 ┆ [["A", "B"], ["B", "C"]] │ └─────┴──────────────────────────┘
另一个尝试是先 explode 再 reduce，但这会将所有元素扁平化后再进行集合操作，失去了原始列表的边界信息，导致结果是所有元素的并集而非交集：
© 版权声明
文章版权归作者所有，未经允许请勿转载。
THE END
