如何基于多列匹配条件对Polars DataFrame行进行分组？

问题描述 #

给定如下Polars DataFrame模拟数据：

data = {'col1': [1, 1, 1, 1, 1, 2, 2], 'col2': ['br', 'bra', 'bra', 'col', 'col', 'br', 'b'], 'col3': ['brazil', 'brazil', 'brasil', 'collombia', 'columbia', 'brazil', 'brazil']}
df = pl.DataFrame(data)

需按以下规则对行进行分组：

• 同一组内的行必须满足col1值相等；
• 行的col2值或col3值，需与组内至少另一行的对应值相等（即通过col2或col3的关联关系，将有共同值的行归为一组）。

分组示例：

• 前3行归为一组：col1=1，且col2='bra'、col3='brazil'存在重复关联；
• 第4-5行归为一组：col1=1，且col2='col'存在重复值；
• 第6-7行归为一组：col1=2，且col3='brazil'存在重复值。

预期输出结果：

{'col1': [1, 1, 1, 1, 1, 2, 2], 'col2': ['br', 'bra', 'bra', 'col', 'col', 'br', 'b'], 'col3': ['brazil', 'brazil', 'brasil', 'collombia', 'columbia', 'brazil', 'brazil'], 'group_idx': [1, 1, 1, 2, 2, 3, 3],}

解决方案 #

可以利用图论连通分量的思路实现：将每行的col2和col3看作两个节点，行本身作为连接这两个节点的边，这样有共同节点（即col2或col3值相同）的行会属于同一个连通分量。再结合col1的分组限制，最终生成每个行的组索引。

代码实现 #

import polars as pl
import networkx as nx

# 初始化数据
data = {'col1': [1, 1, 1, 1, 1, 2, 2], 'col2': ['br', 'bra', 'bra', 'col', 'col', 'br', 'b'], 'col3': ['brazil', 'brazil', 'brasil', 'collombia', 'columbia', 'brazil', 'brazil']}
df = pl.DataFrame(data)

def assign_group_ids(df_group):
    # 构建图：col2和col3的值为节点，每行作为边连接对应节点
    G = nx.Graph()
    for row in df_group.iter_rows(named=True):
        G.add_edge(row['col2'], row['col3'])
    
    # 为每个节点分配连通分量ID
    node_to_group = {}
    for group_idx, component in enumerate(nx.connected_components(G), 1):
        for node in component:
            node_to_group[node] = group_idx
    
    # 为每行映射对应的组ID（col2和col3属于同一分量，取任意一个即可）
    return df_group.with_columns(
        pl.col('col2').map_elements(lambda x: node_to_group[x], return_dtype=pl.Int32).alias('group_idx')
    )

# 按col1分组处理连通分量
temp_df = df.group_by('col1', maintain_order=True).apply(assign_group_ids)

# 生成全局唯一的group_idx（避免不同col1分组内的ID重复）
offset_df = temp_df.group_by('col1')['group_idx'].max().with_columns(
    pl.col('group_idx').cumsum().shift(1).fill_null(0).alias('offset')
)
result_df = temp_df.join(offset_df, on='col1').with_columns(
    (pl.col('group_idx') + pl.col('offset')).alias('group_idx')
).drop('offset')

# 输出结果
print(result_df.to_dict())

代码说明 #

1. 分组构建图结构：针对每个col1分组，将行的col2和col3值连接为边，通过图的连通分量找到所有关联的行；
2. 映射组ID：每个连通分量对应一个组ID，将行的col2值映射到所属的连通分量ID；
3. 全局ID去重：通过累加前序col1分组的最大组ID，为每个分组的组ID添加偏移量，保证全局group_idx唯一；
4. 最终输出结果与预期完全一致。

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