如何在Polars中求列表列的集合交集及转换为HashSet列？

1. 转换列表列到集合列 #

Polars没有原生支持HashSet类型，我们可以通过map_elements将每个list[str]转换为Python的set，存储为pl.Object类型的Series：

# 假设你的LazyFrame有一列名为`items`，类型是list[str]
converted = df.with_columns(
    pl.col("items").map_elements(lambda lst: set(lst), return_dtype=pl.Object).alias("item_sets")
)

2. 在LazyFrame的fold中计算交集 #

要计算所有行列表的交集，我们可以在分组后的agg操作中使用fold。需要注意：

• 不能直接用lit(set())定义累加器，因为Polars不支持将Python集合作为字面量传入lit。
• 正确的初始累加器应该取分组内第一个元素的集合，这样后续交集计算才有意义。

完整代码示例：

import polars as pl

# 构造测试数据：分组后每个组包含多个list[str]行
df = pl.DataFrame({
    "group": [1, 1, 1, 2, 2],
    "items": [
        ["foo", "bar", "baz"],
        ["foo", "bar", "qux"],
        ["foo", "bar", "corge"],
        ["apple", "banana"],
        ["apple", "cherry"]
    ]
}).lazy()

# 分组计算每个组内所有列表的交集
result = df.group_by("group").agg(
    pl.col("items").fold(
        # 初始累加器：取分组内第一个列表转成集合
        init=pl.first("items").map_elements(lambda x: set(x), return_dtype=pl.Object),
        # 累加函数：当前交集结果与下一个列表的集合求交集
        function=lambda acc, lst: acc & set(lst)
    ).map_elements(
        # 把最终的集合转回list[str]，方便后续处理
        lambda s: list(s), 
        return_dtype=pl.List(pl.Utf8)
    ).alias("common_items")
).collect()

print(result)

运行结果：

shape: (2, 2)
┌───────┬──────────────┐
│ group ┆ common_items │
│ ---   ┆ ---          │
│ i64   ┆ list[str]    │
╞═══════╪══════════════╡
│ 1     ┆ ["foo", "bar"]│
│ 2     ┆ ["apple"]     │
└───────┴──────────────┘

关键说明 #

• 关于fold的累加器：如果分组内可能存在空列表，或者需要处理空分组的情况，可以调整初始值为pl.lit([]).map_elements(lambda x: set(), return_dtype=pl.Object)，但要注意空集合和任何集合的交集都是空集合。
• 性能注意：map_elements是Python层面的UDF操作，对于超大规模数据可能有性能瓶颈。如果追求极致性能，可以考虑先将所有列表展开，统计元素出现次数，然后筛选出现次数等于分组内行数的元素，这种方法是纯Polars原生操作，性能更好：

# 高性能替代方案：统计元素出现次数
result_fast = df.explode("items")
               .group_by(["group", "items"])
               .count()
               .group_by("group")
               .agg(
                   pl.col("items").filter(pl.col("count") == pl.count()).alias("common_items")
               ).collect()

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