Polars快速求交集方法详解

一分耕耘，一分收获！既然打开了这篇文章《Polars快速聚合列表交集方法》，就坚持看下去吧！文中内容包含等等知识点...希望你能在阅读本文后，能真真实实学到知识或者帮你解决心中的疑惑，也欢迎大佬或者新人朋友们多留言评论，多给建议！谢谢！

本文深入探讨了如何使用 Polars 数据框高效地对分组内的字符串列表进行交集操作。面对直接使用 reduce 和 list.set_intersection 的局限性，文章提出了一种基于元素计数和过滤的创新方法。通过计算每个元素在组内出现的唯一行数，并与组的总行数进行比较，我们能准确识别出所有列表的共同元素，最终实现预期的聚合交集结果，并提供详细的 Polars 代码示例和解释。

1. 问题背景与传统方法尝试

在数据处理中，我们经常会遇到需要对分组内的列表类型数据进行聚合操作的场景。一个常见的需求是找到某个分组（例如，按 id 分组）中所有列表的交集。例如，给定一个包含 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)

期望的输出是每个 id 组内 values 列表的交集：

shape: (3, 2)
┌─────┬───────────┐
│ id  ┆ values    │
│ --- ┆ ---       │
│ i64 ┆ list[str] │
╞═════╪═══════════╡
│ 1   ┆ ["B"]     │
│ 2   ┆ ["B"]     │
│ 3   ┆ ["B"]     │
└─────┴───────────┘

初次尝试时，可能会想到使用 pl.reduce 结合 list.set_intersection：

# 尝试1：直接在列表列上使用 reduce
result_attempt1 = df.group_by("id").agg(
    pl.reduce(function=lambda acc, x: acc.list.set_intersection(x),
              exprs=pl.col("values"))
)
print("尝试1结果:")
print(result_attempt1)
# 结果类型为 list[list[str]]，不符合预期

此尝试的结果是一个 list[list[str]] 类型，因为 pl.reduce 在这里接收的是整个 values 列，而不是单个列表元素，导致无法正确进行交集运算。

另一种尝试是先 explode 再 reduce：

# 尝试2：explode 后再 reduce
result_attempt2 = df.group_by("id").agg(
    pl.reduce(function=lambda acc, x: acc.list.set_intersection(x),
              exprs=pl.col("values").explode())
)
print("\n尝试2结果:")
print(result_attempt2)
# 结果为所有元素的并集，并非交集

这个尝试将列表扁平化，但 explode 操作改变了数据的结构，使得 reduce 无法在原始列表的上下文进行交集运算，反而得到了类似并集的结果。这表明直接在列表类型上进行聚合交集操作并非易事，需要一种更巧妙的方法。

2. 基于元素计数的交集实现

为了克服上述挑战，我们可以转换思路：将列表交集问题转化为识别那些在组内所有原始列表中都出现的元素。这可以通过以下步骤实现：

2.1 准备阶段：计算组长度并添加行索引

首先，我们需要知道每个 id 组中有多少行，并为每行添加一个唯一的索引，以便在 explode 后追踪元素来源。

# 计算每个id组的行数 (group_len)
# 添加行索引 (index)
df_processed = df.with_columns(
    pl.len().over("id").alias("group_len")
).with_row_index()

print("步骤1：添加组长度和行索引")
print(df_processed)

group_len 列记录了每个 id 组中原始行的数量。index 列为 DataFrame 中的每一行提供了一个唯一的标识符。

2.2 扁平化与元素溯源

接下来，我们将 values 列进行 explode 操作，将每个列表元素拆分成独立的行。然后，对于每个 (id, value) 对，我们需要计算它在多少个原始行中出现过。

# 扁平化 values 列
# 计算每个 (id, value) 对在多少个唯一的原始行中出现 (n_unique)
df_exploded = df_processed.explode("values").with_columns(
    pl.col("index").n_unique().over("id", "values").alias("n_unique")
)

print("\n步骤2：扁平化并计算唯一行数")
print(df_exploded)

在 df_exploded 中：

• index 列仍然保留了原始行的索引。
• values 列现在包含了单个字符串元素。
• n_unique 列表示当前 (id, values) 组合（即某个 id 下的特定字符串元素）在原始 DataFrame 中属于多少个不同的 index。

例如，对于 id=1，元素 "B" 出现了两次，其对应的 index 分别是 0 和 1。因此，n_unique 为 2。而元素 "A" 只出现在 index=0 的行中，所以 n_unique 为 1。

2.3 过滤与最终聚合

现在，关键的逻辑来了：如果一个元素在某个 id 组中所有原始行（即所有列表）中都出现过，那么它的 n_unique 值应该等于该组的 group_len。我们可以利用这个条件来过滤出共同元素。

# 过滤出 n_unique 等于 group_len 的行，这些是所有列表的共同元素
# 再次按 id 分组，并收集去重后的 values
final_result = df_exploded.filter(
    pl.col("n_unique") == pl.col("group_len")
).group_by("id", maintain_order=True).agg(
    pl.col("values").unique()
)

print("\n步骤3：过滤并聚合最终结果")
print(final_result)

filter 操作确保只保留那些在组内所有列表中都存在的元素。最后，group_by("id").agg(pl.col("values").unique()) 将这些共同元素重新聚合回列表形式，并使用 unique() 确保列表中不包含重复项。

3. 完整代码示例

将上述步骤整合在一起，得到完整的解决方案：

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"]]
   }
)

# 解决方案
result = (
    df.with_columns(
        pl.len().over("id").alias("group_len") # 1. 计算每个id组的行数
    )
    .with_row_index() # 2. 添加行索引
    .explode("values") # 3. 扁平化列表
    .filter(
        pl.col("index").n_unique().over("id", "values") # 4. 计算每个(id, value)对在多少个原始行中出现
        == pl.col("group_len") # 5. 过滤出在所有原始行中都出现的元素
    )
    .group_by("id", maintain_order=True) # 6. 按id分组
    .agg(pl.col("values").unique()) # 7. 收集去重后的共同元素
)

print("\n最终结果:")
print(result)

4. 注意事项与总结

• 性能考量： 这种方法通过 explode 和 over 操作将列表处理转换为更通用的 Polars 列操作，通常比自定义 reduce 函数更高效，尤其是在处理大型数据集时。
• 逻辑清晰： 将复杂的列表交集问题分解为“元素是否在所有原始行中出现”的判断，使得逻辑更易于理解和实现。
• 适用性： 这种技术不仅适用于字符串列表，理论上也可扩展到其他可哈希的数据类型列表的交集。
• maintain_order=True： 在最终的 group_by 中使用 maintain_order=True 可以确保输出的 id 顺序与原始 DataFrame 中的首次出现顺序一致，这在某些场景下非常有用。

通过这种基于元素计数和过滤的策略，我们能够有效地在 Polars 中实现分组内列表的交集操作，提供了一个强大且灵活的数据处理方案。

到这里，我们也就讲完了《Polars快速求交集方法详解》的内容了。个人认为，基础知识的学习和巩固，是为了更好的将其运用到项目中，欢迎关注golang学习网公众号，带你了解更多关于的知识点！