如何手动运行PySpark分区函数调试以排查数据倾斜问题

完全可以复现Spark真实分区逻辑生成自定义列，无需跑全量作业即可定位初始读取阶段的倾斜根因，这个排查思路比依赖WebUI等全量跑完再排查的效率高很多，尤其适配你已经确认倾斜发生在Hive表查询初始阶段、后续转换逻辑无影响的场景。

核心实现步骤 #

1. 精准复现对应场景的Spark真实分区逻辑 #

不要自己随手写散列算法，必须和Spark实际运行时的分区规则完全对齐，否则统计出的分布结果没有参考性，常见场景对应的复现逻辑：

• Shuffle阶段Hash分区倾斜（比如repartition、groupBy、join触发的shuffle）：直接复用Spark内置HashPartitioner的计算规则，即keyHash % spark.sql.shuffle.partitions，注意对null值的处理是固定映射到0号分区，不要把null的hash计算为随机值。可以用内置的UTF8String工具类计算key的hash值，保证和运行逻辑完全一致。
  示例代码：

  from pyspark.sql import SparkSession
  from pyspark.sql.functions import udf, col, spark_partition_id
  from pyspark.util import UTF8String
  
  # 初始化时必须配置和正式作业完全相同的shuffle分区数
  spark = SparkSession.builder \
      .config("spark.sql.shuffle.partitions", "200") \
      .getOrCreate()
  
  # 定义和HashPartitioner完全一致的分区计算函数
  def calc_real_partition(key_val, second_key=None):
      # 单key计算逻辑
      key_hash = UTF8String.fromString(str(key_val)).hashCode()
      # 多key场景的hash计算，和Spark的多列hash规则对齐
      if second_key is not None:
          key_hash = 31 * key_hash + UTF8String.fromString(str(second_key)).hashCode()
      return key_hash % int(spark.conf.get("spark.sql.shuffle.partitions"))
  
  partition_udf = udf(calc_real_partition, "int")
  

• Hive分桶表读取倾斜：Hive分桶的散列规则是hashCode & Integer.MAX_VALUE % numBuckets，和Spark默认Hash分区的取模逻辑多了一步无符号整数处理，不能直接用Spark的Hash逻辑计算，否则会得到错误的分桶分布。
• 文件读取阶段的切分倾斜（无shuffle场景）：如果倾斜根本没有触发shuffle，只是单个存储文件过大导致任务分配不均，不需要计算key的散列值，直接用input_file_name()列分组统计每个文件的行数、大小即可定位。

2. 交叉校验保证分区逻辑和真实运行一致 #

自定义UDF写完之后，必须做一轮逻辑对齐校验，避免计算逻辑和实际运行不符导致排查走偏：

• 把原始Hive读取的数据集，按照你正式作业里的分区规则做一次repartition，用Spark内置的spark_partition_id()统计每个真实分区的行数分布；
• 对比真实分布结果和你自定义UDF计算出来的assigned_partition列的分布结果，两边按分区号排序后的计数差值为0，就说明逻辑完全对齐。
  示例校验代码：

  hive_df = spark.sql("select * from your_target_hive_table")
  # 真实Spark分区的统计结果
  real_partition_stat = hive_df.repartition(200, col("test_partition_key")) \
      .withColumn("real_pid", spark_partition_id()) \
      .groupby("real_pid").count()
  # 自定义函数的统计结果
  custom_partition_stat = hive_df.withColumn("assigned_partition", partition_udf(col("test_partition_key"))) \
      .groupby("assigned_partition").count()
  

3. 基于分布结果定位倾斜根因 #

统计完各分区的 key 分布后，结合你已经确认的“倾斜发生在Hive查询初始阶段”的前提，重点排查两类根因：

• 少数热点key的占比远超平均水平：比如某个分区80%的行都对应同一个空值、默认值、无意义的系统枚举值，这种就是典型的Hash分区key分布不均问题，之前调整多列组合repartition无效，大概率是选的列组合没有打散这些热点值。
• 存储文件本身的大小分布不均：比如某个Parquet/Orc文件大小达几十GB，其余文件仅几百MB，这种即使没有shuffle，Spark也会把整个大文件分配给单个任务，导致任务时长出现数量级差异，这类倾斜靠计算key的分区逻辑是排查不出来的。

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