パーティションとシャッフルの最適化：Spark チューニングの核心

# 現在のパーティション数
df.rdd.getNumPartitions()

# パーティション数を増やす（シャッフルあり）
df.repartition(200)

# パーティション数を減らす（シャッフルなし、既存パーティションを結合）
df.coalesce(50)

# キーでパーティション分割（シャッフルあり）
df.repartition("region")
df.repartition(100, "region")

# 典型的なパターン：処理後に書き込み前でファイル数を調整
df.filter(...) \
  .groupBy(...).agg(...) \
  .coalesce(10) \
  .write.parquet("s3://output/")
# → 10個の Parquet ファイルが生成される

# デフォルト: 200（多くの場合、調整が必要）
spark.conf.set("spark.sql.shuffle.partitions", 200)

from pyspark.sql.functions import broadcast

# 明示的にブロードキャスト指定
result = large_df.join(broadcast(small_df), "key")

# 閾値の調整（デフォルト 10MB）
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", 50 * 1024 * 1024)  # 50MB

# 書き込み時にバケット分割
df.write.bucketBy(100, "product_id").sortBy("product_id") \
  .saveAsTable("bucketed_sales")

# バケット同士の Join → シャッフルなし
sales = spark.table("bucketed_sales")
products = spark.table("bucketed_products")  # 同じキー・同じバケット数
result = sales.join(products, "product_id")

# 悪い例：全データをシャッフルしてから集約
df.groupBy("region").agg(sum("amount"))

# Spark は自動的に partial aggregation を行う：
# 1. 各パーティション内で部分集約（Map側）
# 2. 部分結果だけをシャッフル（データ量が大幅に減る）
# 3. シャッフル後に最終集約（Reduce側）
# → ユーザーが意識しなくても Catalyst が最適化

# パーティションごとのサイズを確認
df.groupBy(spark_partition_id()).count().show()

# Spark UI の Stage 詳細で Task の実行時間のばらつきを確認
# → 1つだけ極端に遅い Task があればスキュー

from pyspark.sql.functions import lit, rand, concat, col

# 問題：region = "東京" のデータが全体の80%を占める
# → groupBy("region") で東京のパーティションだけ巨大になる

# Salt Key で分散
salt_num = 10
salted = df.withColumn("salt", (rand() * salt_num).cast("int"))

# 1段階目：salt 付きで集約
partial = salted.groupBy("region", "salt").agg(sum("amount").alias("partial_sum"))

# 2段階目：salt を除いて最終集約
result = partial.groupBy("region").agg(sum("partial_sum").alias("total"))

# キャッシュなし：df を2回計算する
result1 = df.filter(...).groupBy(...).agg(...)
result2 = df.filter(...).groupBy(...).agg(...)

# キャッシュあり：df を1回だけ計算してメモリに保持
df.cache()  # または df.persist()
result1 = df.filter(...).groupBy(...).agg(...)
result2 = df.filter(...).groupBy(...).agg(...)
df.unpersist()  # 不要になったら解放

from pyspark import StorageLevel

df.persist(StorageLevel.MEMORY_ONLY)        # メモリのみ（デフォルト）
df.persist(StorageLevel.MEMORY_AND_DISK)    # メモリに入りきらなければディスク
df.persist(StorageLevel.DISK_ONLY)          # ディスクのみ
df.persist(StorageLevel.MEMORY_ONLY_SER)    # シリアライズしてメモリ（省メモリ）

# パーティション数 = 出力ファイル数
df.repartition(10).write.parquet("s3://output/")  # 10ファイル
df.coalesce(1).write.parquet("s3://output/")       # 1ファイル（小データ向け）

# パーティションカラム指定時
df.repartition("date").write.partitionBy("date").parquet("s3://output/")
# → 日付ごとにディレクトリ分割、各ディレクトリ内のファイル数はパーティション数に依存