在 Spark SQL 中，Shuffle 操作是分布式数据处理的核心环节，合理的 Shuffle 优化能显著提升查询性能。

1. Shuffle 操作基础理解

Shuffle 在查询计划中的角色

触发 Shuffle 的常见操作

-- 以下操作都会触发 Shuffle
SELECT department, AVG(salary) FROM employees GROUP BY department  -- Group By
SELECT * FROM table1 JOIN table2 ON table1.id = table2.id          -- Join
SELECT * FROM table ORDER BY salary DESC                          -- Order By
SELECT DISTINCT department FROM employees                         -- Distinct

2. Shuffle 优化核心策略

分区数优化

// 动态调整 Shuffle 分区数（默认200）
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.minPartitionSize", "16MB")
spark.conf.set("spark.sql.adaptive.advisoryPartitionSizeInBytes", "64MB")

// 手动设置合理分区数
spark.conf.set("spark.sql.shuffle.partitions", 
  Math.max(2, spark.sparkContext.defaultParallelism * 2).toString)

AQE（自适应查询执行）优化

// 启用 AQE 全套优化
spark.conf.set("spark.sql.adaptive.enabled", "true")

// Shuffle 分区合并
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.initialPartitionNum", "1000")

// 数据倾斜处理
spark.conf.set("spark.sql.adaptive.skewedJoin.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewedPartitionFactor", "5")
spark.conf.set("spark.sql.adaptive.skewedPartitionThresholdInBytes", "256MB")

// Join 策略调整
spark.conf.set("spark.sql.adaptive.join.enabled", "true")

3. Join 操作的 Shuffle 优化

广播小表避免 Shuffle

// 自动广播阈值设置（默认10MB）
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "50MB")

// 手动广播提示
val largeTable = spark.table("large_table")
val smallTable = spark.table("small_table")

val result = largeTable.join(broadcast(smallTable), "id")

// SQL 中使用广播提示
spark.sql("""
SELECT /*+ BROADCAST(small_table) */ *
FROM large_table 
JOIN small_table ON large_table.id = small_table.id
""")

Sort Merge Join 优化

// 确保数据预排序减少 Shuffle 开销
spark.conf.set("spark.sql.join.preferSortMergeJoin", "true")

// 使用分桶表优化 Sort Merge Join
spark.sql("""
CREATE TABLE bucketed_table
USING parquet
CLUSTERED BY (id) SORTED BY (id) INTO 32 BUCKETS
AS SELECT * FROM source_table
""")

4. Aggregation 聚合优化

两阶段聚合解决数据倾斜

import org.apache.spark.sql.functions._

// 原始聚合（可能产生数据倾斜）
val simpleAgg = df.groupBy("user_id").agg(sum("amount").as("total_amount"))

// 两阶段聚合解决倾斜
val twoPhaseAgg = df
  .withColumn("salted_key", concat(col("user_id"), lit("_"), (rand() * 10).cast("int")))
  .groupBy("salted_key")
  .agg(sum("amount").as("partial_sum"))
  .groupBy(substring(col("salted_key"), 1, 10).as("user_id"))  // 提取原始key
  .agg(sum("partial_sum").as("total_amount"))

部分聚合下推

-- Spark 会自动将部分聚合下推到数据源
SELECT customer_id, SUM(amount) 
FROM sales 
WHERE sale_date > '2023-01-01'
GROUP BY customer_id

/* 执行计划可能包含：
Exchange hashpartitioning(customer_id, 200)  -- Shuffle
+- HashAggregate(keys=[customer_id], functions=[partial_sum(amount)])
   +- Project [customer_id, amount]
      +- Filter (sale_date > 2023-01-01)
         +- TableScan sales
*/

5. 数据序列化与压缩优化

Shuffle 序列化配置

// 使用 Kryo 序列化（比 Java 序列化更高效）
spark.conf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer")

// Shuffle 压缩配置
spark.conf.set("spark.shuffle.compress", "true")
spark.conf.set("spark.shuffle.spill.compress", "true")

// 使用更高效的压缩算法
spark.conf.set("spark.io.compression.codec", "snappy")  // 或 lz4, zstd
spark.conf.set("spark.shuffle.compress", "true")

内存管理优化

// Shuffle 内存分配
spark.conf.set("spark.shuffle.memoryFraction", "0.2")
spark.conf.set("spark.shuffle.spill.numElementsForceSpillThreshold", "1000000")

// 堆外内存优化
spark.conf.set("spark.memory.offHeap.enabled", "true")
spark.conf.set("spark.memory.offHeap.size", "2g")

6. 数据倾斜专项优化

识别数据倾斜

// 分析 Key 分布识别倾斜
df.groupBy("join_key").count().orderBy(desc("count")).show(20)

// 通过 Spark UI 识别倾斜任务
// - 查看 Stages 页面中任务执行时间分布
// - 关注 Shuffle Read/Write 数据量异常的任务

倾斜处理技术

// 方法1：分离倾斜Key单独处理
val skewedKeys = Seq("key1", "key2", "key3")  // 已知的倾斜Key

val skewedData = df.filter(col("join_key").isin(skewedKeys: _*))
val normalData = df.filter(!col("join_key").isin(skewedKeys: _*))

// 对倾斜数据增加随机后缀
val saltedSkewed = skewedData
  .withColumn("salt", (rand() * 10).cast("int"))
  .withColumn("salted_key", concat(col("join_key"), lit("_"), col("salt")))

// 正常数据保持不变，最后合并结果

7. 物理执行计划优化

自定义分区器

import org.apache.spark.HashPartitioner
import org.apache.spark.sql.DataFrame

// 自定义分区策略
def optimizeShufflePartitioning(df: DataFrame, partitionColumns: Seq[String]): DataFrame = {
  val numPartitions = Math.max(df.rdd.partitions.length / 2, 1)
  df.repartition(numPartitions, partitionColumns.map(col): _*)
}

// 应用自定义分区
val optimizedDF = optimizeShufflePartitioning(df, Seq("department", "year"))

Z-Order 多维聚类

// 使用 Delta Lake 的 Z-Order 优化
df.write
  .format("delta")
  .option("dataChange", "false")
  .mode("overwrite")
  .saveAsTable("optimized_table")

// 对表进行 Z-Order 优化
spark.sql("OPTIMIZE optimized_table ZORDER BY (timestamp, user_id)")

8. 监控与调优验证

执行计划分析

val df = spark.sql("""
SELECT department, AVG(salary) 
FROM employees 
GROUP BY department
""")

// 查看优化后的物理计划
df.explain("formatted")

// 重点关注：
// - Exchange 操作的数量和类型
// - 分区数量是否合理
// - 是否存在不必要的 Shuffle

性能指标监控

# Spark UI 关键指标
# - Stages 页面的 Shuffle 读写大小
# - 任务执行时间的标准差（识别倾斜）
# - GC 时间和网络流量统计

# 命令行监控
./bin/spark-submit --conf spark.metrics.conf=metrics.properties ...

9. 最佳实践总结

优化检查清单

// Shuffle 优化配置模板
def setupShuffleOptimizations(spark: SparkSession): Unit = {
  // AQE 配置
  spark.conf.set("spark.sql.adaptive.enabled", "true")
  spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
  spark.conf.set("spark.sql.adaptive.skewedJoin.enabled", "true")
  
  // 分区配置
  spark.conf.set("spark.sql.shuffle.partitions", "200")
  spark.conf.set("spark.sql.adaptive.advisoryPartitionSizeInBytes", "64MB")
  
  // 序列化与压缩
  spark.conf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer")
  spark.conf.set("spark.shuffle.compress", "true")
  spark.conf.set("spark.io.compression.codec", "lz4")
  
  // 内存配置
  spark.conf.set("spark.memory.fraction", "0.6")
  spark.conf.set("spark.memory.storageFraction", "0.5")
}

根据数据特征选择策略

通过系统性地应用这些 Shuffle 优化策略，可以显著提升 Spark SQL 查询的性能，特别是在处理大规模数据和高并发场景时效果更为明显。