Polars 中列表列分组求交集的策略与实践

在数据处理中，我们经常会遇到需要对 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 级别而不是列表元素级别进行聚合，导致结果不符合预期。

ChatX翻译

最实用、可靠的社交类实时翻译工具。 支持全球主流的20+款社交软件的聊天应用，全球200+语言随意切换。 让您彻底告别复制粘贴的翻译模式，与世界各地高效连接！

下载

例如，以下尝试：

# 尝试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，但这会将所有元素扁平化后再进行集合操作，失去了原始列表的边界信息，导致结果是所有元素的并集而非交集：