Spark 3.2中配置PySpark Join识别分桶数据集并避免Shuffle

问题背景 #

在Apache Spark 3.2环境中，我拥有两个以相同列visitor_id、相同桶数（1024）分桶的Parquet格式数据集，执行Join操作后，执行计划仍出现BroadcastExchange和BroadcastHashJoin步骤，未按预期跳过Shuffle。

代码示例 #

from py4j.java_gateway import java_import
java_import(spark._sc._jvm, "org.apache.spark.sql.api.python.*")
path_1="s3://PATH1/*.parquet"
path_2="s3://PATH2/*.parquet"
df_1 = spark.read.parquet(path_1)
df_2 = spark.read.parquet(path_2)
joined = df_1.join(df_2, 'visitor_id')
joined.explain('formatted')

执行计划相关步骤 #

(4) Scan parquet 
(5) Filter
(6) Project
(7) BroadcastExchange
(8) BroadcastHashJoin
(9) Project

解决方案 #

1. 确保Spark加载分桶元数据 #

Spark默认不会自动识别Parquet文件的分桶信息，需显式开启分桶扫描或通过分桶表读取数据：

• 开启分桶扫描选项：读取数据时添加bucketedScan参数

  df_1 = spark.read.option("bucketedScan", "true").parquet(path_1)
  df_2 = spark.read.option("bucketedScan", "true").parquet(path_2)
  

• 注册为临时分桶表：如果数据已分桶存储，将其注册为临时表可以让Spark明确获取分桶元数据

  # 若数据未持久化成分桶表，先执行写入（仅需一次）
  # df_1.write.bucketBy(1024, "visitor_id").mode("overwrite").saveAsTable("temp_bucketed_table_1")
  # df_2.write.bucketBy(1024, "visitor_id").mode("overwrite").saveAsTable("temp_bucketed_table_2")
  
  # 从分桶表读取数据
  df_1 = spark.table("temp_bucketed_table_1")
  df_2 = spark.table("temp_bucketed_table_2")
  

2. 配置Spark优化器参数 #

调整以下配置，让Spark优先选择分桶Join而非广播Join：

# 启用分桶Join优化规则
spark.conf.set("spark.sql.optimizer.bucketedJoin.enabled", "true")
# 禁用自动广播（若其中一个表被误判为小表），可根据实际场景设置合理阈值（单位：字节）
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "-1")

提示：如果存在需要广播的小表，建议将阈值设置为实际小表的大小（比如10485760代表10MB），而非完全禁用。

3. 验证分桶一致性 #

确认两个数据集的分桶配置完全一致：

• 分桶列visitor_id的数据类型完全相同（如均为String或Int）
• 桶数严格相等（均为1024）
• 分桶算法一致（Spark默认使用哈希分桶，无需额外配置）

4. 避免分桶元数据丢失 #

如果在Join前执行了Filter、Project等操作，可能导致Spark丢失分桶信息：

• 优先执行Join操作，再进行过滤和投影
• 若必须提前处理，确保操作后仍保留分桶元数据（比如基于分桶表进行操作）

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