PySpark复杂关联查询优化求助：倾斜键与拖慢任务问题

一、针对高频倾斜key的精准拆分处理 #

• 提取出现次数最多的N个倾斜main_key，将这部分数据与非倾斜数据分离处理：
  1. 对倾斜key对应的df1数据添加随机盐值（如0-9的整数），实现数据拆分；
  2. 对df2中对应倾斜key的数据，复制多份并分别匹配盐值，扩容后与df1拆分的数据关联；
  3. 最后将倾斜部分的关联结果与非倾斜部分的常规关联结果合并。
• 代码示例：

  from pyspark.sql.functions import rand, IntegerType
  
  # 提前统计得到的倾斜key列表
  skew_keys = ['top_key_1', 'top_key_2']
  
  # 处理df1的倾斜数据，添加盐值
  df1_skew = df1.filter(df1.main_key.isin(skew_keys))\
                .withColumn("salt", (rand() * 10).cast(IntegerType()))
  
  # 处理df2的倾斜数据，复制10份并匹配盐值
  df2_skew = df2.filter(df2.main_key.isin(skew_keys))\
                .crossJoin(spark.range(0, 10).withColumnRenamed("id", "salt"))
  
  # 关联倾斜部分数据
  skew_join_result = df1_skew.join(df2_skew, 
                                   on=["main_key", "salt"] + [f"(df1.col{i}_is_null OR df1.col{i} = df2.col{i})" for i in range(1, 31)],
                                   how="inner")
  
  # 关联非倾斜部分数据
  non_skew_join_result = df1.filter(~df1.main_key.isin(skew_keys))\
                            .join(df2.filter(~df2.main_key.isin(skew_keys)),
                                  on=["main_key"] + [f"(df1.col{i}_is_null OR df1.col{i} = df2.col{i})" for i in range(1, 31)],
                                  how="inner")
  
  # 合并最终结果
  final_result = skew_join_result.union(non_skew_join_result).select("df1.id", "df2.id")
  

二、前置过滤减少关联数据量 #

• 针对每一列的匹配规则，提前过滤掉无匹配可能的数据：
  1. 对于df1中colX_is_null为false的行，df2中只有colX等于df1对应值的行才可能匹配，因此可以提前从df2中过滤出这些有效值；
  2. 对所有30列依次执行该过滤，大幅减少关联时的数据规模。
• 代码示例：

  df2_filtered = df2
  for i in range(1, 31):
      # 获取df1中该列非null的所有有效值（去重）
      valid_col_values = df1.filter(~df1[f"col{i}_is_null"]).select(f"col{i}").distinct()
      # 保留df2中该列在有效值范围内的行
      df2_filtered = df2_filtered.join(broadcast(valid_col_values), on=f"col{i}", how="left_semi")
  # 用过滤后的df2执行原关联逻辑
  join_result = df1.join(df2_filtered, 
                         on=["main_key"] + [f"(df1.col{i}_is_null OR df1.col{i} = df2.col{i})" for i in range(1, 31)],
                         how="inner")
  

三、调整Spark倾斜关联参数 #

• 调大倾斜分区检测阈值：默认64MB，可根据实际数据量调整为更大值（如256MB），避免误判：

  spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "256mb")
  

• 强制开启倾斜分区拆分：

  spark.conf.set("spark.sql.adaptive.skewJoin.splitSkewedPartition", "true")
  

• 开启动态分区合并，优化后续任务的分区数量：

  spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
  

四、拆分广播策略，规避AQE替换 #

• 如果df2整体数据量适中，但倾斜key导致AQE自动切换为SortMergeJoin，可以拆分df2为倾斜和非倾斜两部分：
  1. 对非倾斜key的df2数据执行广播哈希关联；
  2. 倾斜key部分单独用加盐方式关联；
  3. 合并两部分结果，既利用广播的高效性，又解决倾斜问题。

五、预计算匹配哈希，简化关联逻辑 #

• 将30列的匹配规则转化为哈希值对比，减少关联时的逐列判断开销：
  1. 对df1，将每列按规则转换（非null则取原值，null则标记为固定字符串），拼接后计算哈希值；
  2. 对df2，直接拼接30列的值计算相同哈希值；
  3. 关联时先按main_key + match_hash关联，哈希冲突概率极低，可大幅提升关联效率。
• 代码示例：

  from pyspark.sql.functions import concat_ws, sha2, when
  
  # 为df1生成匹配哈希
  df1_with_hash = df1.withColumn("match_hash",
      sha2(concat_ws("|", *[when(df1[f"col{i}_is_null"], "NULL").otherwise(df1[f"col{i}"]) for i in range(1, 31)]), 256))
  
  # 为df2生成匹配哈希
  df2_with_hash = df2.withColumn("match_hash",
      sha2(concat_ws("|", *[df2[f"col{i}"] for i in range(1, 31)]), 256))
  
  # 按main_key和哈希值关联
  join_result = df1_with_hash.join(df2_with_hash, on=["main_key", "match_hash"], how="inner")\
                             .select("df1.id", "df2.id")
  

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