# SparkSessionの初期化

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

# SparkSessionを作成
spark = SparkSession.builder \
    .appName("Partitioning") \
    .getOrCreate()

print("SparkSession準備完了")

# サンプルデータを作成（10万件）

data = [(i, f"name_{i}", i % 100, i * 10) for i in range(100000)]
df = spark.createDataFrame(data, ["id", "name", "category", "value"])

print("サンプルデータ（最初の5件）:")
df.show(5)
print(f"パーティション数: {df.rdd.getNumPartitions()}")

# ========================================
# STEP 19 サンプルデータ一括作成スクリプト
# ========================================

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

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

# サンプルデータ（10万件）
data = [(i, f"name_{i}", i % 100, i * 10) for i in range(100000)]
df = spark.createDataFrame(data, ["id", "name", "category", "value"])

print("✅ サンプルデータを作成しました")
print("  - df: 10万件のデータ（id, name, category, value）")

# パーティション数を確認

print(f"パーティション数: {df.rdd.getNumPartitions()}")

# 各パーティションのデータ件数を確認
partition_counts = df.rdd.mapPartitions(lambda x: [sum(1 for _ in x)]).collect()
print(f"各パーティションの件数: {partition_counts}")

# パーティション数を変更（増やす or 減らす）

print(f"元のパーティション数: {df.rdd.getNumPartitions()}")

# repartition: 増やす
df_repartitioned = df.repartition(16)
print(f"repartition(16)後: {df_repartitioned.rdd.getNumPartitions()}")

# repartition: 減らす
df_repartitioned2 = df.repartition(4)
print(f"repartition(4)後: {df_repartitioned2.rdd.getNumPartitions()}")

# パーティション数を減らす（シャッフルなし、効率的）

print(f"元のパーティション数: {df.rdd.getNumPartitions()}")

# coalesce: 減らす（効率的）
df_coalesced = df.coalesce(4)
print(f"coalesce(4)後: {df_coalesced.rdd.getNumPartitions()}")

# 注意：coalesceで増やすことはできない
df_coalesced2 = df.coalesce(16)
print(f"coalesce(16)後: {df_coalesced2.rdd.getNumPartitions()}")  # 増えない！

import time

# 大きめのデータ
large_data = [(i, i * 10) for i in range(1000000)]
df_large = spark.createDataFrame(large_data, ["id", "value"])

# パーティション数を変えて処理時間を比較
for num_partitions in [1, 4, 8, 16]:
    df_test = df_large.repartition(num_partitions)
    
    start = time.time()
    df_test.groupBy(F.col("id") % 100).count().count()
    end = time.time()
    
    print(f"パーティション数 {num_partitions:2d}: {end - start:.2f}秒")

# ファイル出力時、パーティション数 = 出力ファイル数

# 多すぎるパーティションのまま出力すると、小さいファイルが大量にできる
print(f"現在のパーティション数: {df.rdd.getNumPartitions()}")

# 出力前にcoalesceで減らす（ベストプラクティス）
df_output = df.coalesce(4)  # 4ファイルに出力
print(f"出力用パーティション数: {df_output.rdd.getNumPartitions()}")

# Parquetで出力（4ファイルになる）
# df_output.write.parquet("/output/data")

# フィルター後、データが減ってパーティションがスカスカになることがある

# 1%のデータだけ残るフィルター
df_filtered = df.filter(F.col("category") == 0)

print(f"フィルター前: {df.count()}件, {df.rdd.getNumPartitions()}パーティション")
print(f"フィルター後: {df_filtered.count()}件, {df_filtered.rdd.getNumPartitions()}パーティション")

# パーティションを減らす（効率化）
df_optimized = df_filtered.coalesce(2)
print(f"最適化後: {df_optimized.count()}件, {df_optimized.rdd.getNumPartitions()}パーティション")