本文探讨了在 Polars 数据帧中，如何高效地对分组内的字符串列表进行交集操作。面对 reduce 函数在处理嵌套列表时的局限性，文章提供了一种创新的解决方案：通过扁平化列表、引入行索引、计算元素在各原始行中出现的唯一性，并结合过滤和重新聚合，实现精确的按组列表交集计算。
引言
在数据处理中，我们经常会遇到需要对分组数据进行复杂聚合的场景。当聚合目标是包含字符串列表的列，并且需要计算每个组内所有列表的交集时，直观地使用 reduce 结合列表的集合操作可能无法达到预期效果。本文将详细介绍如何在 Polars 中优雅地解决这一挑战，确保获得正确且高效的列表交集结果。
问题描述与初始尝试
假设我们有一个 Polars DataFrame，其中包含 id 列和 values 列（类型为 list[str]），目标是计算每个 id 组内所有 values 列表的交集。
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)
原始 DataFrame 如下：
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"] │ └─────┴────────────┘
我们期望的输出是：
shape: (3, 2) ┌─────┬───────────┐ │ id ┆ values │ │ --- ┆ --- │ │ i64 ┆ list[str] │ ╞═════╪═══════════╡ │ 1 ┆ ["B"] │ │ 2 ┆ ["B"] │ │ 3 ┆ ["B"] │ └─────┴───────────┘
初次尝试可能会使用 group_by 结合 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]] 类型，因为 reduce 接收的是列表的列表，而不是列表中的单个元素。
# 尝试 2: 结合 explode 试图扁平化 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 在这种情况下也无法正确地将所有扁平化后的元素进行集合交集操作，因为它仍然在处理一个单一的扁平化列表，而不是多个独立的列表。
解决方案：基于行索引的交集计算
Polars 提供了一种更灵活、更“Polars 化”的解决方案，它将列表操作转换为更通用的聚合计数和过滤逻辑。核心思想是：一个元素如果在组内所有原始列表中都出现，那么它就是这些列表交集的一部分。我们可以通过跟踪每个元素在组内出现的原始行数来判断这一点。
以下是实现步骤：
-
计算每组的行数 (group_len): 使用 pl.len().over(“id”) 为每个 id 组添加一个新列，表示该组包含的原始行数。这是判断一个元素是否出现在组内所有列表中的基准。
-
添加行索引 (row_index): 使用 with_row_index() 为 DataFrame 的每一行添加一个唯一的索引。这个索引将帮助我们追踪每个元素来源于哪一行原始数据。
-
扁平化 values 列 (explode): 使用 explode(“values”) 将 list[str] 列转换为 str 列，使得每个列表中的元素都成为单独的行。此时，每一行将包含原始的 id、扁平化的 value、原始行索引 row_index 和组长度 group_len。
-
计算元素在组内出现的唯一原始行数 (n_unique): 对于每个 (id, value) 对，计算其对应的 row_index 的唯一数量，即 pl.col.index.n_unique().over(“id”, “values”)。如果 n_unique 等于 group_len，则说明这个 value 在该 id 组的所有原始列表中都出现过。
-
过滤: 根据上一步的判断，过滤出那些 n_unique 等于 group_len 的行。这些行中的 value 就是组内所有列表的交集元素。
-
重新聚合: 最后，使用 group_by(“id”) 并 agg(pl.col.values.unique()) 将过滤后的 values 重新聚合回列表形式。使用 unique() 是为了确保每个交集元素只出现一次。maintain_order=True 可以保持 id 的原始顺序。
完整代码示例
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"]] } ) # 步骤分解展示中间结果 (可选) intermediate_df = ( df.with_columns(pl.len().over("id").alias("group_len")) # 1. 计算每组的行数 .with_row_index() # 2. 添加行索引 .explode("values") # 3. 扁平化 values 列 .with_columns( pl.col.index.n_unique().over("id", "values").alias("n_unique") # 4. 计算元素在组内出现的唯一原始行数 ) ) print("n中间结果 (扁平化并计算 n_unique):") print(intermediate_df) # 最终解决方案 final_result = ( df.with_columns(pl.len().over("id").alias("group_len")) .with_row_index() .explode("values") .Filter( pl.col.index.n_unique().over("id", "values") == pl.col.group_len ) # 5. 过滤 .group_by("id", maintain_order=True) .agg(pl.col.values.unique()) # 6. 重新聚合 ) print("n最终交集结果:") print(final_result)
中间结果 (扁平化并计算 n_unique) 示例：
shape: (12, 5) ┌────────┬─────┬────────┬───────────┬──────────┐ │ index ┆ id ┆ values ┆ group_len ┆ n_unique │ │ --- ┆ --- ┆ --- ┆ --- ┆ --- │ │ u32 ┆ i64 ┆ str ┆ u32 ┆ u32 │ ╞════════╪═════╪════════╪═══════════╪══════════╡ │ 0 ┆ 1 ┆ A ┆ 2 ┆ 1 │ │ 0 ┆ 1 ┆ B ┆ 2 ┆ 2 │ # B 出现在 index 0 和 1 │ 1 ┆ 1 ┆ B ┆ 2 ┆ 2 │ │ 1 ┆ 1 ┆ C ┆ 2 ┆ 1 │ │ 2 ┆ 2 ┆ A ┆ 2 ┆ 1 │ │ 2 ┆ 2 ┆ B ┆ 2 ┆ 2 │ # B 出现在 index 2 和 3 │ 3 ┆ 2 ┆ B ┆ 2 ┆ 2 │ │ 3 ┆ 2 ┆ C ┆ 2 ┆ 1 │ │ 4 ┆ 3 ┆ A ┆ 2 ┆ 1 │ │ 4 ┆ 3 ┆ B ┆ 2 ┆ 2 │ # B 出现在 index 4 和 5 │ 5 ┆ 3 ┆ B ┆ 2 ┆ 2 │ │ 5 ┆ 3 ┆ C ┆ 2 ┆ 1 │ └────────┴─────┴────────┴───────────┴──────────┘
从中间结果可以看出，对于 id=1，values=”B” 的 n_unique 是 2，而 group_len 也是 2，这表明 “B” 出现在 id=1 组的所有原始行中。同样，对于 id=2 和 id=3，”B” 也满足条件。而 “A” 和 “C” 的 n_unique 都是 1，不等于 group_len，因此它们不会被包含在交集中。
最终交集结果：
shape: (3, 2) ┌─────┬───────────┐ │ id ┆ values │ │ --- ┆ --- │ │ i64 ┆ list[str] │ ╞═════╪═══════════╡ │ 1 ┆ ["B"] │ │ 2 ┆ ["B"] │ │ 3 ┆ ["B"] │ └─────┴───────────┘
这与我们期望的输出完全一致。
总结与注意事项
这种方法巧妙地绕过了 Polars 中 reduce 在处理复杂列表集合操作时的限制，转而利用了其强大的窗口函数和聚合能力。
关键概念回顾：
- explode: 将列表列展平，将列表中的每个元素转换为单独的行。
- with_row_index(): 为DataFrame添加一个唯一的行索引，这对于追踪原始数据行至关重要。
- over() (窗口函数): 允许在指定的分组（或窗口）内执行计算，例如 pl.len().over(“id”) 计算每个 id 组的长度，pl.col.index.n_unique().over(“id”, “values”) 计算每个 (id, value) 对的唯一行索引数。
- filter(): 根据条件选择行。
- group_by().agg(): 重新分组并聚合数据。
注意事项：
- 性能考量: 对于非常大的数据集，explode 操作可能会显著增加DataFrame的行数，从而增加内存消耗和计算时间。然而，Polars 的优化引擎通常能高效处理这些操作。
- 数据类型: 本文示例是针对 list[str] 类型。对于其他可哈希（hashable）的元素类型（如整数、浮点数），此方法同样适用。
- 替代方案: 对于某些特定场景，如果数据量不大或逻辑允许，也可以先将数据导出到 python 列表，使用 Python 的集合操作进行计算，再导回 Polars。但通常而言，在 Polars 内部完成操作能更好地利用其并行计算能力。
通过理解和应用这种基于行索引和窗口函数的策略，您可以高效地解决 Polars 中按组计算列表交集的复杂问题，充分发挥 Polars 在大数据处理方面的优势。
