如何在Polars中高效处理连通ID分组？替代Python迭代的方案

问题本质 #

你的需求属于图论中的连通分量识别问题：每个子列表中的ID属于同一个连通图，需要将所有相互连通的ID归为一组并分配唯一标识。

现有方案的问题 #

纯Polars原生方案依赖频繁explode列表，会导致数据量急剧膨胀（尤其是大数据集），进而引发内存或性能瓶颈；而Python迭代+iter_rows的方式虽然性能较好，但需要逐行处理，没有充分利用Polars的向量化优势。

更优的实现思路 #

方案1：Polars + 并查集（Union-Find）高效实现 #

并查集是解决连通分量问题的最优数据结构之一，时间复杂度接近O(n)。我们可以将Polars的数据批量提取到Python中构建并查集，再将结果映射回Polars，避免频繁展开列表：

import polars as pl

# 示例数据
data = {"ConnectedIDs": [[1, 2], [6], [8], [2], [7], [2, 7], [7, 9], [6, 8]]}
df = pl.DataFrame(data)

# 实现并查集
class UnionFind:
    def __init__(self):
        self.parent = {}
    
    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]
    
    def union(self, x, y):
        x_root = self.find(x)
        y_root = self.find(y)
        if x_root != y_root:
            self.parent[y_root] = x_root

# 初始化并查集，处理所有连通对
uf = UnionFind()
for ids in df["ConnectedIDs"].to_list():
    if not ids:
        continue
    # 先将第一个ID加入并查集
    root = ids[0]
    if root not in uf.parent:
        uf.parent[root] = root
    # 合并当前列表中的所有ID
    for id_ in ids[1:]:
        if id_ not in uf.parent:
            uf.parent[id_] = id_
        uf.union(root, id_)

# 为每个根节点分配唯一ID
root_to_group = {root: idx for idx, root in enumerate(set(uf.find(x) for x in uf.parent))}

# 生成结果DataFrame
result_df = pl.DataFrame({
    "ConnectedID": list(uf.parent.keys()),
    "UnitedID": [root_to_group[uf.find(x)] for x in uf.parent.keys()]
}).sort("ConnectedID")

print(result_df)

这个方案的优势：

• 避免了Polars中大规模explode操作，内存占用更低
• 并查集的高效性保证了处理大数据集的性能
• 仅在数据提取和结果映射时进行Polars与Python的转换，中间计算完全基于高效的并查集

方案2：Polars原生优化（针对中小数据集） #

如果你的数据集规模不算极大，可以优化纯Polars方案，减少不必要的展开和重复计算：

import polars as pl

df = pl.DataFrame({"ConnectedIDs": [[1, 2], [6], [8], [2], [7], [2, 7], [7, 9], [6, 8]]})

# 1. 展开所有ID并记录所属的原始行索引
expanded = df.with_row_index("row_idx").explode("ConnectedIDs").rename({"ConnectedIDs": "id"})

# 2. 构建ID之间的关联关系，通过行索引进行分组合并
connected_pairs = expanded.join(expanded, on="row_idx", suffix="_other")
connected_pairs = connected_pairs.filter(pl.col("id") != pl.col("id_other")).select("id", "id_other")

# 3. 递归合并连通分量（用循环代替递归避免栈溢出）
groups = {}
for id1, id2 in connected_pairs.iter_rows():
    # 查找两个ID的所属组
    group1 = groups.get(id1, {id1})
    group2 = groups.get(id2, {id2})
    # 合并组
    merged = group1.union(group2)
    # 更新所有组内ID的映射
    for id_ in merged:
        groups[id_] = merged

# 4. 为每个组分配唯一ID并生成结果
group_ids = {id_: idx for idx, group in enumerate(set(frozenset(g) for g in groups.values())) for id_ in group}
# 处理孤立的ID（没有出现在任何关联对中的ID）
all_ids = expanded["id"].unique()
isolated_ids = all_ids.filter(~pl.col("id").is_in(group_ids.keys()))
for idx, id_ in enumerate(isolated_ids, start=len(group_ids)):
    group_ids[id_] = idx // len(isolated_ids)  # 孤立ID各自为一组

result_df = pl.DataFrame({
    "ConnectedID": list(group_ids.keys()),
    "UnitedID": list(group_ids.values())
}).sort("ConnectedID")

print(result_df)

这个方案纯用Polars完成，但对于超大数据集，join和iter_rows仍可能存在性能瓶颈，因此更适合中小规模数据。

结论 #

• 对于大规模数据集，最优方案是Polars + 并查集：既利用了Polars的数据处理优势，又借助并查集的高效算法避免了展开操作的性能问题，比纯Python迭代更高效。
• 纯Polars原生方案仅适合中小数据集，大规模场景下难以规避展开/join带来的性能问题。
• 如果追求极致性能，可以考虑用Rust编写自定义Polars扩展（利用Rust的并查集实现和Polars的原生API），但开发成本较高，一般情况下Python并查集方案已足够满足需求。

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