合同关联行分组优化求助：百万级数据高效分组方法

核心思路 #

你的问题本质是找出图中的连通分量：把每个合同ID（无论是ID1还是ID2）看作图的节点，每一行的ID1-ID2关系看作节点间的边，所有通过边直接/间接连通的节点属于同一组。迭代Join的方式会因为重复关联导致数据爆炸，而Union-Find（并查集）算法专门解决这类问题，时间复杂度接近O(n)，完全适配百万行规模的数据。

实现方案 #

方案一：离线脚本处理（推荐，性能最优） #

把数据导出到本地，用脚本实现并查集，处理完后再更新回数据库。以Python为例：

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):
        # 初始化节点的父节点为自身
        if x not in self.parent:
            self.parent[x] = x
        if y not in self.parent:
            self.parent[y] = y
        # 合并两个连通分量
        root_x = self.find(x)
        root_y = self.find(y)
        if root_x != root_y:
            self.parent[root_y] = root_x

# 1. 分块读取数据库数据，避免内存溢出
import pandas as pd
chunk_size = 100000
uf = UnionFind()

for chunk in pd.read_sql("SELECT ID1, ID2 FROM your_table", your_db_conn, chunksize=chunk_size):
    for _, row in chunk.iterrows():
        uf.union(row['ID1'], row['ID2'])

# 2. 生成分组映射：用根节点的唯一标识作为分组ID
group_map = {}
current_group = 1
for node in uf.parent:
    root = uf.find(node)
    if root not in group_map:
        group_map[root] = current_group
        current_group += 1

# 3. 批量更新回数据库
update_data = []
for chunk in pd.read_sql("SELECT ID1 FROM your_table", your_db_conn, chunksize=chunk_size):
    chunk['Group'] = chunk['ID1'].apply(lambda x: group_map[uf.find(x)])
    update_data.append(chunk)

pd.concat(update_data).to_sql("temp_group_update", your_db_conn, if_exists="replace", index=False)
your_db_conn.execute("""
    UPDATE your_table t
    SET "Group" = tu."Group"
    FROM temp_group_update tu
    WHERE t.ID1 = tu.ID1
""")

方案二：数据库内直接处理 #

如果不想导出数据，可在数据库中用递归CTE+并查集思路实现（以PostgreSQL为例，其他数据库可调整语法）：

1. 初始化并查集临时表

-- 收集所有唯一合同ID，初始化父节点为自身
CREATE TEMP TABLE uf_nodes AS
SELECT DISTINCT id AS node FROM (
    SELECT ID1 AS id FROM your_table
    UNION ALL
    SELECT ID2 AS id FROM your_table
) t;

ALTER TABLE uf_nodes ADD COLUMN parent_id VARCHAR(64); -- 按实际ID类型调整
UPDATE uf_nodes SET parent_id = node;

-- 加索引加速后续关联
CREATE INDEX idx_uf_node ON uf_nodes(node);
CREATE INDEX idx_uf_parent ON uf_nodes(parent_id);

2. 迭代合并连通分量（重复执行直到返回0行更新）

WITH cte_relations AS (
    SELECT ID1, ID2 FROM your_table
),
to_merge AS (
    SELECT
        u1.parent_id AS root1,
        u2.parent_id AS root2
    FROM cte_relations
    JOIN uf_nodes u1 ON cte_relations.ID1 = u1.node
    JOIN uf_nodes u2 ON cte_relations.ID2 = u2.node
    WHERE u1.parent_id != u2.parent_id
    LIMIT 10000 -- 批量合并，提升效率
)
UPDATE uf_nodes u
SET parent_id = (SELECT root1 FROM to_merge LIMIT 1)
WHERE parent_id IN (SELECT root2 FROM to_merge);

3. 路径压缩+生成分组映射

-- 路径压缩：让所有节点直接指向根节点
WITH RECURSIVE find_root AS (
    SELECT node, parent_id, 1 AS depth FROM uf_nodes
    UNION ALL
    SELECT fr.node, u.parent_id, fr.depth + 1
    FROM find_root fr
    JOIN uf_nodes u ON fr.parent_id = u.node
    WHERE fr.parent_id != u.parent_id
)
UPDATE uf_nodes u
SET parent_id = fr.parent_id
FROM find_root fr
WHERE fr.node = u.node
AND fr.depth = (SELECT MAX(depth) FROM find_root WHERE node = u.node);

-- 生成分组ID并更新原表
CREATE TEMP TABLE group_mapping AS
SELECT
    node,
    DENSE_RANK() OVER (ORDER BY parent_id) AS group_id
FROM uf_nodes;

UPDATE your_table t
SET "Group" = gm.group_id
FROM group_mapping gm
WHERE t.ID1 = gm.node;

关键优化点 #

• 路径压缩：并查集的核心优化，减少后续查询的层级，大幅提升效率。
• 批量处理：不管是脚本还是数据库操作，分块/批量处理避免内存溢出或锁表。
• 索引优化：数据库内处理时，给临时表的节点和父节点字段加索引，加速关联查询。
• 分组ID简化：可以直接用根节点的ID作为Group值，无需额外分配自增ID，进一步提升效率。

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