本文详细介绍了如何在Polars中对包含字符串列表的列进行分组聚合，以找出每个组内所有列表的交集元素。通过巧妙地结合explode、with_row_index、over以及条件过滤等操作，我们将复杂的列表交集问题转化为高效的扁平化数据处理，最终实现精确的分组交集聚合。
1. 问题描述与挑战
在数据处理中，我们经常会遇到需要对包含列表（list）类型数据的列进行聚合操作的场景。一个常见的需求是，对于某个分组键（如 id），我们希望找到该组内所有列表的共同元素，即它们的交集。
考虑以下Polars 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 列中所有列表的交集。例如，对于 id=1，期望的输出是 [“B”]，因为 [“A”, “B”] 和 [“B”, “C”] 的交集是 [“B”]。
尝试直接使用 pl.reduce 和 list.set_intersection 可能不会立即得到期望的结果，因为 reduce 通常在聚合时处理的是列的 Series 而非内部列表元素的交集逻辑，或者需要更复杂的初始化和累积函数。例如：
# 尝试1：直接reduce # df.group_by("id").agg( # pl.reduce(function=lambda acc, x: acc.list.set_intersection(x), # exprs=pl.col("values")) # ) # 结果是 list[list[str]]，不是预期的单个列表交集 # 尝试2：explode后reduce # df.group_by("id").agg( # pl.reduce(function=lambda acc, x: acc.list.set_intersection(x), # exprs=pl.col("values").explode()) # ) # 结果是 list[str] 但包含了所有元素，不是交集
这些尝试表明，直接在列表列上进行集合操作可能需要更精细的控制。
2. Polars解决方案：基于扁平化与计数的方法
解决此问题的关键在于将列表数据扁平化，然后利用Polars强大的分组和窗口函数来识别那些在组内所有原始列表中都存在的元素。
以下是实现这一目标的详细步骤和代码：
步骤 1：计算每个分组的行数 (group_len)
首先，我们需要知道每个 id 分组中有多少行。这对于后续判断一个元素是否在所有列表中都出现至关重要。
df_processed = df.with_columns( pl.len().over("id").alias("group_len") ) print("Step 1: Add group_len") print(df_processed)
输出：
shape: (6, 3) ┌─────┬───────────┬───────────┐ │ id ┆ values ┆ group_len │ │ --- ┆ --- ┆ --- │ │ i64 ┆ list[str] ┆ u32 │ ╞═════╪═══════════╪═══════════╡ │ 1 ┆ ["A", "B"]│ 2 │ │ 1 ┆ ["B", "C"]│ 2 │ │ 2 ┆ ["A", "B"]│ 2 │ │ 2 ┆ ["B", "C"]│ 2 │ │ 3 ┆ ["A", "B"]│ 2 │ │ 3 ┆ ["B", "C"]│ 2 │ └─────┴───────────┴───────────┘
步骤 2：添加行索引 (index)
为了在扁平化后区分原始行，我们需要为每行添加一个唯一的索引。
df_processed = df_processed.with_row_index() print("nStep 2: Add row index") print(df_processed)
输出：
shape: (6, 4) ┌───────┬─────┬───────────┬───────────┐ │ index ┆ id ┆ values ┆ group_len │ │ --- ┆ --- ┆ --- ┆ --- │ │ u32 ┆ i64 ┆ list[str] ┆ u32 │ ╞═══════╪═════╪═══════════╪═══════════╡ │ 0 ┆ 1 ┆ ["A", "B"]│ 2 │ │ 1 ┆ 1 ┆ ["B", "C"]│ 2 │ │ 2 ┆ 2 ┆ ["A", "B"]│ 2 │ │ 3 ┆ 2 ┆ ["B", "C"]│ 2 │ │ 4 ┆ 3 ┆ ["A", "B"]│ 2 │ │ 5 ┆ 3 ┆ ["B", "C"]│ 2 │ └───────┴─────┴───────────┴───────────┘
步骤 3：扁平化 values 列 (explode)
使用 explode(“values”) 将 values 列中的每个列表元素展开成独立的行。这样，我们就可以对单个字符串元素进行操作。
df_processed = df_processed.explode("values") print("nStep 3: Explode values column") print(df_processed)
输出：
shape: (12, 4) ┌───────┬─────┬────────┬───────────┐ │ index ┆ id ┆ values ┆ group_len │ │ --- ┆ --- ┆ --- ┆ --- │ │ u32 ┆ i64 ┆ str ┆ u32 │ ╞═══════╪═════╪════════╪═══════════╡ │ 0 ┆ 1 ┆ A ┆ 2 │ │ 0 ┆ 1 ┆ B ┆ 2 │ │ 1 ┆ 1 ┆ B ┆ 2 │ │ 1 ┆ 1 ┆ C ┆ 2 │ │ 2 ┆ 2 ┆ A ┆ 2 │ │ 2 ┆ 2 ┆ B ┆ 2 │ │ 3 ┆ 2 ┆ B ┆ 2 │ │ 3 ┆ 2 ┆ C ┆ 2 │ │ 4 ┆ 3 ┆ A ┆ 2 │ │ 4 ┆ 3 ┆ B ┆ 2 │ │ 5 ┆ 3 ┆ B ┆ 2 │ │ 5 ┆ 3 ┆ C ┆ 2 │ └───────┴─────┴────────┴───────────┘
4. 计算每个 (id, value) 组合的唯一原始行数 (n_unique)
这是核心步骤。对于每个 id 和 values 的组合，我们计算其对应的原始 index 的唯一数量。如果 n_unique 等于 group_len，则表示这个 values 元素在 id 组内的所有原始行中都存在。
df_processed = df_processed.with_columns( pl.col("index").n_unique().over("id", "values").alias("n_unique") ) print("nStep 4: Calculate n_unique for each (id, value) pair") print(df_processed)
输出（部分关键行）：
shape: (12, 5) ┌───────┬─────┬────────┬───────────┬──────────┐ │ index ┆ id ┆ values ┆ group_len ┆ n_unique │ │ --- ┆ --- ┆ --- ┆ --- ┆ --- │ │ u32 ┆ i64 ┆ str ┆ u32 ┆ u32 │ ╞═══════╪═════╪════════╪═══════════╪══════════╡ │ 0 ┆ 1 ┆ A ┆ 2 ┆ 1 │ # 'A' for id=1 appears only in original row 0 │ 0 ┆ 1 ┆ B ┆ 2 ┆ 2 │ # 'B' for id=1 appears in original rows 0 and 1 │ 1 ┆ 1 ┆ B ┆ 2 ┆ 2 │ │ 1 ┆ 1 ┆ C ┆ 2 ┆ 1 │ # 'C' for id=1 appears only in original row 1 ...
从上面的输出可以看出，对于 id=1，values=’B’ 的 n_unique 是 2，这与 group_len (2) 相等，说明 ‘B’ 在 id=1 的所有原始列表中都存在。而 ‘A’ 和 ‘C’ 的 n_unique 是 1，不等于 group_len，说明它们只存在于部分列表中。
5. 过滤并重建列表
最后，我们过滤出那些 n_unique 等于 group_len 的行，然后按 id 分组，并将 values 列聚合回列表，同时去除重复项 (unique()) 以确保结果的正确性。
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") ) .group_by("id", maintain_order=True) .agg(pl.col("values").unique()) ) print("nFinal Result:") print(final_result)
最终输出：
shape: (3, 2) ┌─────┬───────────┐ │ id ┆ values │ │ --- ┆ --- │ │ i64 ┆ list[str] │ ╞═════╪═══════════╡ │ 1 ┆ ["B"] │ │ 2 ┆ ["B"] │ │ 3 ┆ ["B"] │ └─────┴───────────┘
这正是我们期望的交集结果。
3. 注意事项与总结
- 扁平化策略：此方法的核心思想是将列表数据扁平化，将列表操作转换为更通用的 DataFrame 行操作。这使得Polars能够高效地利用其向量化和并行处理能力。
- 窗口函数 (over) 的强大：over() 方法在这里发挥了关键作用，它允许我们在特定分组（id, values）内执行聚合操作（n_unique），而无需显式的 group_by。
- 内存考虑：对于非常大的列表或大量分组，explode 操作可能会显著增加 DataFrame 的行数，从而增加内存消耗。在极端情况下，可能需要考虑分块处理或使用其他内存优化技术。
- 通用性：这种模式不仅适用于字符串列表的交集，也可以扩展到其他类型元素的列表，只要能够定义其“存在”和“唯一性”。
通过上述步骤，我们成功地解决了Polars中列表字符串列的交集聚合问题，展示了Polars在处理复杂数据转换和聚合方面的灵活性和强大功能。这种方法避免了直接在列表上进行复杂的 reduce 操作，转而利用了DataFrame的扁平化能力和窗口函数，提供了一个高效且易于理解的解决方案。
