Apache Spark 优化

优化 Apache Spark 作业的生产模式，包括分区策略、内存管理、Shuffle 优化和性能调优。

不使用此技能的场景

• 任务与 Apache Spark 优化无关
• 需要此范围之外的领域或工具

使用说明

• 明确目标、约束条件和所需输入
• 应用相关最佳实践并验证结果
• 提供可操作的步骤和验证方法
• 如需详细示例，打开 resources/implementation-playbook.md

使用此技能的场景

• 优化慢速 Spark 作业
• 调整内存和执行器配置
• 实现高效的分区策略
• 调试 Spark 性能问题
• 为大数据集扩展 Spark 管道
• 减少 Shuffle 和数据倾斜

核心概念

1. Spark 执行模型

Driver Program
    ↓
Job（由 action 触发）
    ↓
Stages（由 shuffle 分隔）
    ↓
Tasks（每个分区一个）

2. 关键性能因素

快速开始

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

# 创建优化的 Spark 会话
spark = (SparkSession.builder
    .appName("OptimizedJob")
    .config("spark.sql.adaptive.enabled", "true")
    .config("spark.sql.adaptive.coalescePartitions.enabled", "true")
    .config("spark.sql.adaptive.skewJoin.enabled", "true")
    .config("spark.serializer", "org.apache.spark.serializer.KryoSerializer")
    .config("spark.sql.shuffle.partitions", "200")
    .getOrCreate())

# 使用优化设置读取数据
df = (spark.read
    .format("parquet")
    .option("mergeSchema", "false")
    .load("s3://bucket/data/"))

# 高效的数据转换
result = (df
    .filter(F.col("date") >= "2024-01-01")
    .select("id", "amount", "category")
    .groupBy("category")
    .agg(F.sum("amount").alias("total")))

result.write.mode("overwrite").parquet("s3://bucket/output/")

模式

模式 1：优化分区

# 计算最佳分区数
def calculate_partitions(data_size_gb: float, partition_size_mb: int = 128) -> int:
    """
    最佳分区大小：128MB - 256MB
    过少：利用率不足，内存压力大
    过多：任务调度开销大
    """
    return max(int(data_size_gb * 1024 / partition_size_mb), 1)

# 重新分区以实现均匀分布
df_repartitioned = df.repartition(200, "partition_key")

# 合并减少分区数（无 shuffle）
df_coalesced = df.coalesce(100)

# 使用谓词下推的分区裁剪
df = (spark.read.parquet("s3://bucket/data/")
    .filter(F.col("date") == "2024-01-01"))  # Spark 会下推此条件

# 写入时分区以便后续查询
(df.write
    .partitionBy("year", "month", "day")
    .mode("overwrite")
    .parquet("s3://bucket/partitioned_output/"))

模式 2：连接优化

from pyspark.sql import functions as F
from pyspark.sql.types import *

# 1. 广播连接 - 小表连接
# 最佳场景：一侧 < 10MB（可配置）
small_df = spark.read.parquet("s3://bucket/small_table/")  # < 10MB
large_df = spark.read.parquet("s3://bucket/large_table/")  # TB 级

# 显式广播提示
result = large_df.join(
    F.broadcast(small_df),
    on="key",
    how="left"
)

# 2. 排序合并连接 - 大表默认方式
# 需要 shuffle，但可处理任意大小
result = large_df1.join(large_df2, on="key", how="inner")

# 3. 分桶连接 - 预排序，连接时无 shuffle
# 写入分桶表
(df.write
    .bucketBy(200, "customer_id")
    .sortBy("customer_id")
    .mode("overwrite")
    .saveAsTable("bucketed_orders"))

# 连接分桶表（无 shuffle！）
orders = spark.table("bucketed_orders")
customers = spark.table("bucketed_customers")  # 相同桶数
result = orders.join(customers, on="customer_id")

# 4. 倾斜连接处理
# 启用 AQE 倾斜连接优化
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionFactor", "5")
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "256MB")

# 严重倾斜的手动加盐处理
def salt_join(df_skewed, df_other, key_col, num_salts=10):
    """为倾斜键添加盐值以分散分布"""
    # 为倾斜侧添加盐值
    df_salted = df_skewed.withColumn(
        "salt",
        (F.rand() * num_salts).cast("int")
    ).withColumn(
        "salted_key",
        F.concat(F.col(key_col), F.lit("_"), F.col("salt"))
    )

    # 展开另一侧与所有盐值
    df_exploded = df_other.crossJoin(
        spark.range(num_salts).withColumnRenamed("id", "salt")
    ).withColumn(
        "salted_key",
        F.concat(F.col(key_col), F.lit("_"), F.col("salt"))
    )

    # 在加盐键上连接
    return df_salted.join(df_exploded, on="salted_key", how="inner")

模式 3：缓存与持久化

from pyspark import StorageLevel

# 在多次重用 DataFrame 时缓存
df = spark.read.parquet("s3://bucket/data/")
df_filtered = df.filter(F.col("status") == "active")

# 缓存到内存（默认 MEMORY_AND_DISK）
df_filtered.cache()

# 或使用特定存储级别
df_filtered.persist(StorageLevel.MEMORY_AND_DISK_SER)

# 强制物化
df_filtered.count()

# 在多个 action 中使用
agg1 = df_filtered.groupBy("category").count()
agg2 = df_filtered.groupBy("region").sum("amount")

# 使用完成后取消持久化
df_filtered.unpersist()

# 存储级别说明：
# MEMORY_ONLY - 速度快，但可能放不下
# MEMORY_AND_DISK - 必要时溢出到磁盘（推荐）
# MEMORY_ONLY_SER - 序列化，占用内存少，CPU 开销大
# DISK_ONLY - 内存紧张时使用
# OFF_HEAP - Tungsten 堆外内存

# 复杂血缘的检查点
spark.sparkContext.setCheckpointDir("s3://bucket/checkpoints/")
df_complex = (df
    .join(other_df, "key")
    .groupBy("category")
    .agg(F.sum("amount")))
df_complex.checkpoint()  # 打断血缘，物化数据

模式 4：内存调优

# 执行器内存配置
# spark-submit --executor-memory 8g --executor-cores 4

# 内存分解（8GB 执行器）：
# - spark.memory.fraction = 0.6（60% = 4.8GB 用于执行 + 存储）
#   - spark.memory.storageFraction = 0.5（4.8GB 的 50% = 2.4GB 用于缓存）
#   - 剩余 2.4GB 用于执行（shuffle、连接、排序）
# - 40% = 3.2GB 用于用户数据结构和内部元数据

spark = (SparkSession.builder
    .config("spark.executor.memory", "8g")
    .config("spark.executor.memoryOverhead", "2g")  # 用于非 JVM 内存
    .config("spark.memory.fraction", "0.6")
    .config("spark.memory.storageFraction", "0.5")
    .config("spark.sql.shuffle.partitions", "200")
    # 针对内存密集型操作
    .config("spark.sql.autoBroadcastJoinThreshold", "50MB")
    # 防止大 shuffle 时的 OOM
    .config("spark.sql.files.maxPartitionBytes", "128MB")
    .getOrCreate())

# 监控内存使用
def print_memory_usage(spark):
    """打印当前内存使用情况"""
    sc = spark.sparkContext
    for executor in sc._jsc.sc().getExecutorMemoryStatus().keySet().toArray():
        mem_status = sc._jsc.sc().getExecutorMemoryStatus().get(executor)
        total = mem_status._1() / (1024**3)
        free = mem_status._2() / (1024**3)
        print(f"{executor}: {total:.2f}GB total, {free:.2f}GB free")

模式 5：Shuffle 优化

# 减少 shuffle 数据大小
spark.conf.set("spark.sql.shuffle.partitions", "auto")  # 使用 AQE
spark.conf.set("spark.shuffle.compress", "true")
spark.conf.set("spark.shuffle.spill.compress", "true")

# 在 shuffle 前预聚合
df_optimized = (df
    # 先本地聚合（combiner）
    .groupBy("key", "partition_col")
    .agg(F.sum("value").alias("partial_sum"))
    # 再全局聚合
    .groupBy("key")
    .agg(F.sum("partial_sum").alias("total")))

# 使用 map 端操作避免 shuffle
# 不好：每次 distinct 都触发 shuffle
distinct_count = df.select("category").distinct().count()

# 好：近似去重计数（无 shuffle）
approx_count = df.select(F.approx_count_distinct("category")).collect()[0][0]

# 减少分区时使用 coalesce 而非 repartition
df_reduced = df.coalesce(10)  # 无 shuffle

# 使用压缩优化 shuffle
spark.conf.set("spark.io.compression.codec", "lz4")  # 快速压缩

模式 6：数据格式优化

# Parquet 优化
(df.write
    .option("compression", "snappy")  # 快速压缩
    .option("parquet.block.size", 128 * 1024 * 1024)  # 128MB 行组
    .parquet("s3://bucket/output/"))

# 列裁剪 - 只读取需要的列
df = (spark.read.parquet("s3://bucket/data/")
    .select("id", "amount", "date"))  # Spark 只读取这些列

# 谓词下推 - 在存储层过滤
df = (spark.read.parquet("s3://bucket/partitioned/year=2024/")
    .filter(F.col("status") == "active"))  # 下推到 Parquet 读取器

# Delta Lake 优化
(df.write
    .format("delta")
    .option("optimizeWrite", "true")  # 装箱优化
    .option("autoCompact", "true")  # 合并小文件
    .mode("overwrite")
    .save("s3://bucket/delta_table/"))

# 多维查询的 Z-ordering
spark.sql("""
    OPTIMIZE delta.`s3://bucket/delta_table/`
    ZORDER BY (customer_id, date)
""")

模式 7：监控与调试

# 启用详细指标
spark.conf.set("spark.sql.codegen.wholeStage", "true")
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")

# 查看查询计划
df.explain(mode="extended")
# 模式：simple、extended、codegen、cost、formatted

# 获取物理计划统计信息
df.explain(mode="cost")

# 监控任务指标
def analyze_stage_metrics(spark):
    """分析最近的阶段指标"""
    status_tracker = spark.sparkContext.statusTracker()

    for stage_id in status_tracker.getActiveStageIds():
        stage_info = status_tracker.getStageInfo(stage_id)
        print(f"Stage {stage_id}:")
        print(f"  Tasks: {stage_info.numTasks}")
        print(f"  Completed: {stage_info.numCompletedTasks}")
        print(f"  Failed: {stage_info.numFailedTasks}")

# 识别数据倾斜
def check_partition_skew(df):
    """检查分区倾斜"""
    partition_counts = (df
        .withColumn("partition_id", F.spark_partition_id())
        .groupBy("partition_id")
        .count()
        .orderBy(F.desc("count")))

    partition_counts.show(20)

    stats = partition_counts.select(
        F.min("count").alias("min"),
        F.max("count").alias("max"),
        F.avg("count").alias("avg"),
        F.stddev("count").alias("stddev")
    ).collect()[0]

    skew_ratio = stats["max"] / stats["avg"]
    print(f"Skew ratio: {skew_ratio:.2f}x (>2x indicates skew)")

配置速查表

# 生产环境配置模板
spark_configs = {
    # 自适应查询执行（AQE）
    "spark.sql.adaptive.enabled": "true",
    "spark.sql.adaptive.coalescePartitions.enabled": "true",
    "spark.sql.adaptive.skewJoin.enabled": "true",

    # 内存
    "spark.executor.memory": "8g",
    "spark.executor.memoryOverhead": "2g",
    "spark.memory.fraction": "0.6",
    "spark.memory.storageFraction": "0.5",

    # 并行度
    "spark.sql.shuffle.partitions": "200",
    "spark.default.parallelism": "200",

    # 序列化
    "spark.serializer": "org.apache.spark.serializer.KryoSerializer",
    "spark.sql.execution.arrow.pyspark.enabled": "true",

    # 压缩
    "spark.io.compression.codec": "lz4",
    "spark.shuffle.compress": "true",

    # 广播
    "spark.sql.autoBroadcastJoinThreshold": "50MB",

    # 文件处理
    "spark.sql.files.maxPartitionBytes": "128MB",
    "spark.sql.files.openCostInBytes": "4MB",
}

最佳实践

建议

• 启用 AQE - 自适应查询执行可处理许多问题
• 使用 Parquet/Delta - 带压缩的列式格式
• 广播小表 - 避免小连接时的 shuffle
• 监控 Spark UI - 检查倾斜、溢出、GC
• 合理设置分区大小 - 每个分区 128MB - 256MB

禁忌

• 不要 collect 大数据 - 保持数据分布式
• 不要不必要地使用 UDF - 使用内置函数
• 不要过度缓存 - 内存是有限的
• 不要忽略数据倾斜 - 它会主导作业时间
• 不要用 .count() 检查存在性 - 使用 .take(1) 或 .isEmpty()

资源

• Spark Performance Tuning
• Spark Configuration
• Databricks Optimization Guide