PySpark下保留value字段的重复复合键Spark SQL高效查询重写

原查询的性能痛点 #

原窗口函数实现的核心问题在于shuffle阶段的无效数据传输：

• 窗口函数COUNT(*) OVER (PARTITION BY key0, key1)要求Spark在shuffle时将原表所有列（包括你未用到的value、以及表内其他所有业务字段）全量按复合键分发
• 哪怕最终确认全表没有任何重复键，整个流程依然要完成全表数据的shuffle、逐行窗口值计算，资源浪费极其严重

等价高效实现 #

采用「先筛重复键、再回捞明细」的两步逻辑，在完全保留原查询输出规则（保留value字段、排序逻辑、重复行返回规则）的前提下，将shuffle成本降到最低：

WITH dup_key_stats AS (
    -- 第一步：仅针对复合键做轻量聚合，shuffle时只传输key0、key1两个字段
    SELECT 
        key0, 
        key1, 
        COUNT(1) AS dups_count
    FROM code_adm1
    GROUP BY key0, key1
    -- 直接筛出存在重复的复合键，生成的结果集通常极小
    HAVING COUNT(1) > 1
)
SELECT 
    t.key0, 
    t.key1, 
    t.value,
    -- 可选：带上重复次数，满足调试需求
    d.dups_count
FROM code_adm1 t
-- 第二步：用极小的重复键集合关联原表，Spark会自动广播这个小结果集，无额外shuffle开销
INNER JOIN dup_key_stats d
    ON t.key0 = d.key0 AND t.key1 = d.key1
ORDER BY t.key0 ASC, t.key1 ASC, t.value ASC

为什么这个写法比原实现快 #

• shuffle数据量大幅降低：第一步聚合阶段，shuffle仅传输key0和key1两个轻量字段，相比原窗口函数全列shuffle，数据量可降低一个数量级以上
• 无重复场景下资源消耗趋近于0：如果全表没有重复键，第一步聚合完成后直接返回空结果，不需要执行后续的全表关联逻辑，完全避免无意义的计算
• 语法完全兼容：规避了直接GROUP BY查询value字段的语法报错，不需要用FIRST()之类的聚合函数强制裁剪字段，完整保留所有重复行的原始value值

等价简化写法 #

逻辑和上面完全一致，Spark会自动生成相同的执行计划：

SELECT key0, key1, value
FROM code_adm1 t
WHERE EXISTS (
    SELECT 1 
    FROM code_adm1 s
    WHERE s.key0 = t.key0 AND s.key1 = t.key1
    GROUP BY s.key0, s.key1
    HAVING COUNT(1) > 1
)
ORDER BY key0, key1, value

关于COUNT(DISTINCT)的认知澄清 #

同事提到的COUNT(DISTINCT (key0, key1))方式，本身并没有比普通COUNT(1)在分组场景下的性能优势，反而因为去重逻辑会带来额外计算开销。之前出现的语法报错，本质是因为在聚合查询中直接返回非分组字段（value）不符合SQL语法规范，和用不用DISTINCT无关。

Spark SQL性能调优学习方向 #

• 核心基础：重点掌握Spark Shuffle的执行机制、AQE自适应查询执行的触发逻辑、广播Join的适用场景，这些是判断SQL执行效率的核心底层逻辑
• 实践方法：养成通过Spark UI的Stage页面对比不同SQL写法的shuffle读写数据量、任务执行时长的习惯，快速定位慢查询的瓶颈
• 权威资料：优先学习Spark官方文档中关于SQL性能调优的章节内容，是最贴合Spark执行逻辑的一手参考资料

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