基于列值分组行：BigQuery/PySpark大数据量性能优化求助

原始数据 #

| ID | A  | B  | C  |
|----|----|----|----|
| 1  | 10 | 15 | 20 |
| 2  | 10 | 14 | 21 |
| 3  | 9  | 14 | 22 |
| 4  | 8  | 13 | 23 |
| 5  | 7  | 12 | 24 |

需求说明 #

识别任意A、B、C列值相同的行，将它们归为同一ID_REF，示例输出如下：

| ID | A  | B  | C  | ID_REF |
|----|----|----|----|--------|
| 1  | 10 | 15 | 20 | 1      |
| 2  | 10 | 14 | 21 | 1      |
| 3  | 9  | 14 | 22 | 1      |
| 4  | 8  | 13 | 22 | 1      |
| 5  | 7  | 12 | 23 | 2      |

遇到的问题 #

使用BigQuery或PySpark处理时，尝试过在JOIN的ON子句中用OR条件关联，小数据量下可行，但在2亿行的全表中运行速度极慢。

尝试过的代码（BigQuery） #

WITH
tb AS(
SELECT 1 AS ID, 10 AS A, 15 AS B, 20 AS C UNION ALL
SELECT 2 AS ID, 10 AS A, 14 AS B, 21 AS C UNION ALL
SELECT 3 AS ID, 9 AS A, 14 AS B, 22 AS C UNION ALL
SELECT 4 AS ID, 8 AS A, 13 AS B, 23 AS C UNION ALL
SELECT 5 AS ID, 7 AS A, 12 AS B, 24 AS C UNION ALL
SELECT 6 AS ID, 81 AS A, 23 AS B, 22 AS C
)

SELECT t1.ID, MIN(t2.ID) AS ID_REF
FROM tb AS t1
LEFT JOIN tb AS t2
ON t1.A = t2.A
   OR t1.B = t2.B
   OR t1.C = t2.C
GROUP BY 1

优化方案 #

直接用OR关联会产生大量笛卡尔积，导致性能爆炸。这个问题本质是连通分量查找——只要两行在A/B/C任意列有相同值，就属于同一连通分量，以下是针对BigQuery和PySpark的高效实现：

BigQuery 实现 #

通过递归CTE构建值与ID的映射，查找连通分量并取分量内最小ID作为ID_REF：

WITH tb AS (
  SELECT 1 AS ID, 10 AS A, 15 AS B, 20 AS C UNION ALL
  SELECT 2 AS ID, 10 AS A, 14 AS B, 21 AS C UNION ALL
  SELECT 3 AS ID, 9 AS A, 14 AS B, 22 AS C UNION ALL
  SELECT 4 AS ID, 8 AS A, 13 AS B, 23 AS C UNION ALL
  SELECT 5 AS ID, 7 AS A, 12 AS B, 24 AS C UNION ALL
  SELECT 6 AS ID, 81 AS A, 23 AS B, 22 AS C
),
-- 展开每个值对应的ID，建立值-ID映射
value_id_map AS (
  SELECT ID, 'A' AS col, A AS val FROM tb UNION ALL
  SELECT ID, 'B' AS col, B AS val FROM tb UNION ALL
  SELECT ID, 'C' AS col, C AS val FROM tb
),
-- 递归查找连通的ID，追踪每个ID所在分量的最小ID
recursive_cte AS (
  SELECT 
    ID, 
    ID AS min_id,
    CAST(ID AS STRING) AS visited_ids
  FROM tb
  UNION ALL
  SELECT 
    r.ID,
    LEAST(r.min_id, m.ID) AS min_id,
    CONCAT(r.visited_ids, ',', CAST(m.ID AS STRING)) AS visited_ids
  FROM recursive_cte r
  JOIN value_id_map vm ON r.ID = vm.ID
  JOIN value_id_map m ON vm.val = m.val AND m.ID != r.ID
  WHERE NOT REGEXP_CONTAINS(r.visited_ids, CONCAT('^', CAST(m.ID AS STRING), '$|,', CAST(m.ID AS STRING), '$'))
),
-- 对每个ID取最小的分量ID作为ID_REF
final_groups AS (
  SELECT ID, MIN(min_id) AS ID_REF FROM recursive_cte GROUP BY ID
)
-- 关联原表输出完整结果
SELECT t.*, f.ID_REF FROM tb t JOIN final_groups f ON t.ID = f.ID

PySpark 实现 #

使用GraphFrames库处理连通分量，适合大规模数据：

1. 先安装依赖：

pip install graphframes

2. 代码实现：

from pyspark.sql import SparkSession
from graphframes import GraphFrame

spark = SparkSession.builder.appName("ConnectedComponents").getOrCreate()

# 加载原始数据
data = [
    (1, 10, 15, 20),
    (2, 10, 14, 21),
    (3, 9, 14, 22),
    (4, 8, 13, 23),
    (5, 7, 12, 24),
    (6, 81, 23, 22)
]
df = spark.createDataFrame(data, ["ID", "A", "B", "C"])

# 生成边表：同一值（A/B/C）的ID两两相连
def generate_edges(col_name):
    return df.select("ID", col_name)\
             .join(df.select("ID", col_name).withColumnRenamed("ID", "dst"), col_name)\
             .filter(df.ID != df.dst)\
             .select("ID", "dst")

edges_a = generate_edges("A")
edges_b = generate_edges("B")
edges_c = generate_edges("C")
edges = edges_a.union(edges_b).union(edges_c).distinct()

# 构建图并计算连通分量
vertices = df.select("ID")
g = GraphFrame(vertices, edges)
connected_components = g.connectedComponents()

# 获取每个连通分量的最小ID作为ID_REF
min_id_per_component = connected_components.groupBy("component").agg({"ID": "min"}).withColumnRenamed("min(ID)", "ID_REF")
result = df.join(connected_components, on="ID").join(min_id_per_component, on="component").drop("component")

# 展示结果
result.show()

内容的提问来源于stack exchange，提问作者José Carlos Bezerra Filho