PySpark用sort-merge join解决数据倾斜的完整案例

假设有两个大表 table1 和 table2 ，并通过 sort-merge join 来解决可能的数据倾斜问题。

python 复制代码

from pyspark.sql import SparkSession
from pyspark.sql.functions import col

# 初始化SparkSession
spark = SparkSession.builder.appName("SortMergeJoinExample").getOrCreate()

# 加载数据，假设数据来自parquet文件
table1 = spark.read.parquet("path/to/table1.parquet")
table2 = spark.read.parquet("path/to/table2.parquet")

# 查看表的大小
print("table1 size: ", table1.count())
print("table2 size: ", table2.count())

# 为了演示数据倾斜，假设我们直接使用join，这里用inner join举例
joined = table1.join(table2, table1["id"] == table2["id"], "inner")

# 先对连接键进行排序，为sort-merge join做准备

sorted_table1 = table1.sortWithinPartitions("id")
sorted_table2 = table2.sortWithinPartitions("id")

# 使用sort-merge join进行连接
joined = sorted_table1.join(sorted_table2, sorted_table1["id"] == sorted_table2["id"], "inner")

# 触发Action，查看执行计划，此时可以去Spark WebUI查看任务执行情况
joined.count()

# 停止SparkSession
spark.stop()

代码解释

初始化SparkSession：创建一个SparkSession对象，这是与Spark交互的入口。

python 复制代码

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

加载数据并查看表大小：从Parquet文件加载两张表，并打印出它们的行数，以此来了解表的规模。

python 复制代码

table1 = spark.read.parquet("path/to/table1.parquet")
table2 = spark.read.parquet("path/to/table2.parquet")

print("table1 size: ", table1.count())
print("table2 size: ", table2.count())

数据预处理：在进行 sort-merge join 之前，对两个表按照连接键 id 在每个分区内进行排序。

python 复制代码

sorted_table1 = table1.sortWithinPartitions("id")
sorted_table2 = table2.sortWithinPartitions("id")

执行sort-merge join：利用排序后的表，执行 sort-merge join 操作，这里选择的是内连接。

python 复制代码

joined = sorted_table1.join(sorted_table2, sorted_table1["id"] == sorted_table2["id"], "inner")

触发Action并查看执行情况：调用 count() 方法触发一个Action，此时Spark会真正执行整个计算流程。与此同时，可以打开Spark WebUI（通常是 http://your-spark-master:4040 ），在 Stages 页面查看任务执行计划，尤其是查看各个阶段的数据分布情况，确认数据倾斜是否得到解决。

python 复制代码

joined.count()

停止SparkSession：任务完成后，关闭SparkSession释放资源。

python 复制代码

spark.stop()

要在Spark WebUI中查看数据倾斜：

在执行 joined.count() 后，迅速打开浏览器访问Spark WebUI。进入 Stages 标签页，找到正在执行的 join 相关阶段。查看每个任务的处理数据量，如果之前存在数据倾斜，经过 sort-merge join 处理后，各个任务处理的数据量应该相对均匀。