spark-optimization

Apache Spark の最適化

パーティショニング戦略、メモリ管理、シャッフルの最適化、パフォーマンスチューニングなど、Apache Spark ジョブを最適化するための本番環境向けパターンです。

このスキルを使用しない場合

• タスクが Apache Spark の最適化に関係しない場合
• 対象範囲外の別のドメインやツールが必要な場合

手順

• 目標、制約、必要な入力を明確にします。
• 関連するベストプラクティスを適用し、結果を検証します。
• 実行可能な手順と検証方法を提示します。

このスキルを使用する場合

• 遅い Spark ジョブを最適化する場合
• メモリやエグゼキューターの設定を調整する場合
• 効率的なパーティショニング戦略を実装する場合
• Spark のパフォーマンス問題をデバッグする場合
• 大規模なデータセット向けに Spark パイプラインをスケーリングする場合
• シャッフルやデータの偏りを減らす場合

主な概念

1. Spark の実行モデル

Driver Program
    ↓
Job (triggered by action)
    ↓
Stages (separated by shuffles)
    ↓
Tasks (one per partition)

2. 主なパフォーマンス要因

要因	影響	解決策
シャッフル	ネットワーク I/O、ディスク I/O	ワイド変換を最小限に抑える
データの偏り	タスクの実行時間の不均衡	ソルティング、ブロードキャスト結合
シリアライズ	CPU オーバーヘッド	Kryo、列指向フォーマットを使用する
メモリ	GC の負荷、スピル	エグゼキューターのメモリを調整する
パーティション	並列処理	パーティション数を適切に設定する

クイックスタート

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

# Create optimized Spark session
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())

# Read with optimized settings
df = (spark.read
    .format("parquet")
    .option("mergeSchema", "false")
    .load("s3://bucket/data/"))

# Efficient transformations
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：最適なパーティショニング

# Calculate optimal partition count
def calculate_partitions(data_size_gb: float, partition_size_mb: int = 128) -> int:
    """
    Optimal partition size: 128MB - 256MB
    Too few: Under-utilization, memory pressure
    Too many: Task scheduling overhead
    """
    return max(int(data_size_gb * 1024 / partition_size_mb), 1)

# Repartition for even distribution
df_repartitioned = df.repartition(200, "partition_key")

# Coalesce to reduce partitions (no shuffle)
df_coalesced = df.coalesce(100)

# Partition pruning with predicate pushdown
df = (spark.read.parquet("s3://bucket/data/")
    .filter(F.col("date") == "2024-01-01"))  # Spark pushes this down

# Write with partitioning for future queries
(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. Broadcast Join - Small table joins
# Best when: One side < 10MB (configurable)
small_df = spark.read.parquet("s3://bucket/small_table/")  # < 10MB
large_df = spark.read.parquet("s3://bucket/large_table/")  # TBs

# Explicit broadcast hint
result = large_df.join(
    F.broadcast(small_df),
    on="key",
    how="left"
)

# 2. Sort-Merge Join - Default for large tables
# Requires shuffle, but handles any size
result = large_df1.join(large_df2, on="key", how="inner")

# 3. Compatible bucketed tables may avoid a shuffle; inspect the actual plan
# Write bucketed tables
(df.write
    .bucketBy(200, "customer_id")
    .sortBy("customer_id")
    .mode("overwrite")
    .saveAsTable("bucketed_orders"))

# Join compatible bucketed tables (verify chosen physical plan)
orders = spark.table("bucketed_orders")
customers = spark.table("bucketed_customers")  # Same bucket count
result = orders.join(customers, on="customer_id")

# 4. Skew Join Handling
# Enable AQE skew join optimization
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")

# Manual salting for severe skew
def salt_join(df_skewed, df_other, key_col, num_salts=10):
    """Add salt to distribute skewed keys"""
    # Add salt to skewed side
    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"))
    )

    # Explode other side with all salts
    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"))
    )

    # Join on salted key
    return df_salted.join(df_exploded, on="salted_key", how="inner")

パターン 3：キャッシュと永続化

from pyspark import StorageLevel

# Cache when reusing DataFrame multiple times
df = spark.read.parquet("s3://bucket/data/")
df_filtered = df.filter(F.col("status") == "active")

# DataFrame cache default is version-dependent (modern Spark uses MEMORY_AND_DISK_DESER)
df_filtered.cache()

# Or with specific storage level
df_filtered.persist(StorageLevel.MEMORY_AND_DISK)

# Force materialization
df_filtered.count()

# Use in multiple actions
agg1 = df_filtered.groupBy("category").count()
agg2 = df_filtered.groupBy("region").sum("amount")
# Materialize consumers before unpersist (use bounded results or writes as appropriate).
agg1.show(20)
agg2.show(20)

# Unpersist when done
df_filtered.unpersist()

# Storage levels explained:
# MEMORY_ONLY - Fast, but may not fit
# MEMORY_AND_DISK - Spills to disk if needed (recommended)
# MEMORY_ONLY_SER - Serialized, less memory, more CPU
# DISK_ONLY - When memory is tight
# OFF_HEAP - Tungsten off-heap memory

# Checkpoint for complex lineage
spark.sparkContext.setCheckpointDir("s3://bucket/checkpoints/")
df_complex = (df
    .join(other_df, "key")
    .groupBy("category")
    .agg(F.sum("amount")))
df_complex = df_complex.checkpoint()  # Use the returned DataFrame with truncated lineage

パターン 4：メモリの調整

# Executor memory configuration
# spark-submit --executor-memory 8g --executor-cores 4

# Memory breakdown (8GB executor):
# - spark.memory.fraction = 0.6 (60% = 4.8GB for execution + storage)
#   - spark.memory.storageFraction = 0.5 (50% of 4.8GB = 2.4GB for cache)
#   - Remaining 2.4GB for execution (shuffles, joins, sorts)
# - 40% = 3.2GB for user data structures and internal metadata

spark = (SparkSession.builder
    .config("spark.executor.memory", "8g")
    .config("spark.executor.memoryOverhead", "2g")  # For non-JVM memory
    .config("spark.memory.fraction", "0.6")
    .config("spark.memory.storageFraction", "0.5")
    .config("spark.sql.shuffle.partitions", "200")
    # For memory-intensive operations
    .config("spark.sql.autoBroadcastJoinThreshold", "50MB")
    # Prevent OOM on large shuffles
    .config("spark.sql.files.maxPartitionBytes", "128MB")
    .getOrCreate())

# Monitor memory usage
def print_memory_usage(spark):
    """Print current memory usage"""
    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：シャッフルの最適化

# Reduce shuffle data size
spark.conf.set("spark.sql.shuffle.partitions", "200")  # Integer in Apache Spark; AQE may coalesce
spark.conf.set("spark.shuffle.compress", "true")
spark.conf.set("spark.shuffle.spill.compress", "true")

# Pre-aggregate before shuffle
df_optimized = (df
    # Local aggregation first (combiner)
    .groupBy("key", "partition_col")
    .agg(F.sum("value").alias("partial_sum"))
    # Then global aggregation
    .groupBy("key")
    .agg(F.sum("partial_sum").alias("total")))

# Avoid shuffle with map-side operations
# BAD: Shuffle for each distinct
distinct_count = df.select("category").distinct().count()

# Approximate distinct reduces aggregation state; it can still involve an exchange
approx_count = df.select(F.approx_count_distinct("category")).collect()[0][0]

# Use coalesce instead of repartition when reducing partitions
df_reduced = df.coalesce(10)  # No shuffle

# Optimize shuffle with compression
spark.conf.set("spark.io.compression.codec", "lz4")  # Fast compression

パターン 6：データフォーマットの最適化

# Parquet optimizations
(df.write
    .option("compression", "snappy")  # Fast compression
    .option("parquet.block.size", 128 * 1024 * 1024)  # 128MB row groups
    .parquet("s3://bucket/output/"))

# Column pruning - only read needed columns
df = (spark.read.parquet("s3://bucket/data/")
    .select("id", "amount", "date"))  # Spark only reads these columns

# Predicate pushdown - filter at storage level
df = (spark.read.parquet("s3://bucket/partitioned/year=2024/")
    .filter(F.col("status") == "active"))  # Pushed to Parquet reader

# Delta Lake optimizations
(df.write
    .format("delta")
    .option("optimizeWrite", "true")  # Bin-packing
    .option("autoCompact", "true")  # Compact small files
    .mode("overwrite")
    .save("s3://bucket/delta_table/"))

# Z-ordering for multi-dimensional queries
spark.sql("""
    OPTIMIZE delta.`s3://bucket/delta_table/`
    ZORDER BY (customer_id, date)
""")

パターン 7：モニタリングとデバッグ

# These enable execution features, not detailed telemetry
spark.conf.set("spark.sql.codegen.wholeStage", "true")
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")

# Explain query plan
df.explain(mode="extended")
# Modes: simple, extended, codegen, cost, formatted

# Show available optimizer estimates, not actual runtime statistics
df.explain(mode="cost")

# Monitor task metrics
def analyze_stage_metrics(spark):
    """Analyze recent stage metrics"""
    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}")

# Identify data skew
def check_partition_skew(df):
    """Check for partition skew"""
    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"] if stats["avg"] else None
    if skew_ratio is None:
        print("No rows; skew not assessed")
        return
    print(f"Skew ratio: {skew_ratio:.2f}x (>2x indicates skew)")

設定チートシート

# Production configuration template
spark_configs = {
    # Adaptive Query Execution (AQE)
    "spark.sql.adaptive.enabled": "true",
    "spark.sql.adaptive.coalescePartitions.enabled": "true",
    "spark.sql.adaptive.skewJoin.enabled": "true",

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

    # Parallelism
    "spark.sql.shuffle.partitions": "200",
    "spark.default.parallelism": "200",

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

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

    # Broadcast
    "spark.sql.autoBroadcastJoinThreshold": "50MB",

    # File handling
    "spark.sql.files.maxPartitionBytes": "128MB",
    "spark.sql.files.openCostInBytes": "4MB",
}

ベストプラクティス

推奨事項

• AQE を有効にする - 適応的クエリ実行により、多くの問題に対処できます
• Parquet/Delta を使用する - 圧縮に対応した列指向フォーマットです
• 小さなテーブルをブロードキャストする - 小規模な結合でシャッフルを回避できます
• Spark UI を監視する - 偏り、スピル、GC を確認します
• パーティション数を適切に設定する - パーティションあたり 128MB～256MB を目安にします

非推奨事項

• 大規模なデータを収集しない - データを分散したままにします
• 不要な UDF を使わない - 組み込み関数を使用します
• 過剰にキャッシュしない - メモリには限りがあります
• データの偏りを無視しない - ジョブ時間の大部分を占めることがあります
• 存在確認に .count() を使わない - .take(1) または .isEmpty() を使用します

リソース

• Spark のパフォーマンスチューニング
• Spark の設定
• Databricks 最適化ガイド

制限事項

• タスクが上記の説明に明確に合致する場合に限り、このスキルを使用してください。
• 出力を、環境固有の検証、テスト、専門家によるレビューの代わりとして扱わないでください。
• 必要な入力、権限、安全上の境界、成功基準が不足している場合は、作業を中断して確認を求めてください。

すべての値はワークロードに基づく仮説であり、あらゆる状況に通用する最適設定ではありません。Spark UI と正確性チェックを使って検証してください。上書き書き込みは既存の出力を削除する可能性があります。エグゼキューターの設定は、多くの場合、コンテキストの作成前に行う必要があります。メモリ状態を取得するヘルパーは非公開の JVM API を使用し、Scala の Option のアンラップが必要です。サポート対象のモニタリングを優先してください。ファイルの分割サイズはシャッフルパーティション数を制限せず、OOM も防ぎません。結合のソルティングでは、キーを文字列連結するのではなく、元の型付きキーとソルトを使用し、null、出力スキーマ、重複の多重度をテストしてください。Delta／Databricks 固有のオプションと OPTIMIZE を使用するには、互換性のあるランタイムがインストールされている必要があります。実際のアクションが完了するまでキャッシュを保持し、キャッシュストレージの状態と SQL 実行メモリを考慮してください。