Dataset<Row> df = spark.read().parquet("data.parquet");
int partitions = df.rdd().getNumPartitions();
System.out.println("파티션 수: " + partitions);

// 각 파티션의 데이터 수
df.mapPartitions(
    iter -> {
        int count = 0;
        while (iter.hasNext()) { iter.next(); count++; }
        return Collections.singletonList(count).iterator();
    },
    Encoders.INT()
).collectAsList().forEach(System.out::println);

// repartition - 파티션 수 변경 (셔플 발생)
Dataset<Row> repartitioned = df.repartition(200);

// coalesce - 파티션 수 줄이기 (셔플 없음)
Dataset<Row> coalesced = df.coalesce(100);

// 키 기반 재파티셔닝
Dataset<Row> keyPartitioned = df.repartition(col("department"));
Dataset<Row> keyWithCount = df.repartition(100, col("department"));

// 파일 쓰기 전 파티션 줄이기 (파일 수 제어)
df.coalesce(10)
  .write()
  .parquet("output");
// → 10개의 parquet 파일 생성

// 스큐 해결을 위한 재파티셔닝
df.repartition(200, col("key"))
  .groupBy("key")
  .agg(sum("value"));

// groupBy - 같은 키를 같은 파티션으로
df.groupBy("key").agg(sum("value"));

// join - 같은 키를 같은 파티션으로
df1.join(df2, "key");

// distinct - 중복 제거를 위해 재분배
df.distinct();

// repartition - 명시적 재분배
df.repartition(100);

// orderBy/sort - 전역 정렬
df.orderBy("column");

// 셔플 파티션 수 설정 (기본값: 200)
spark.conf().set("spark.sql.shuffle.partitions", "400");

// 또는 SparkSession 생성 시
SparkSession spark = SparkSession.builder()
    .appName("App")
    .config("spark.sql.shuffle.partitions", "400")
    .getOrCreate();

// AQE 활성화 (Spark 3.2+에서는 기본 활성화)
spark.conf().set("spark.sql.adaptive.enabled", "true");

// 파티션 병합 활성화
spark.conf().set("spark.sql.adaptive.coalescePartitions.enabled", "true");

// 최소 파티션 크기 (병합 기준)
spark.conf().set("spark.sql.adaptive.advisoryPartitionSizeInBytes", "64MB");

// 스큐 조인 최적화
spark.conf().set("spark.sql.adaptive.skewJoin.enabled", "true");

// Hash Partitioning
df.repartition(100, col("user_id"));

// 해시 함수: partition = hash(key) % numPartitions

// Range Partitioning
df.repartitionByRange(100, col("timestamp"));

// 예: timestamp 0-100 → Partition 0
//     timestamp 100-200 → Partition 1

import org.apache.spark.Partitioner;

public class RegionPartitioner extends Partitioner {
    private int numPartitions;

    public RegionPartitioner(int numPartitions) {
        this.numPartitions = numPartitions;
    }

    @Override
    public int numPartitions() {
        return numPartitions;
    }

    @Override
    public int getPartition(Object key) {
        String region = (String) key;
        switch (region) {
            case "ASIA": return 0;
            case "EUROPE": return 1;
            case "AMERICA": return 2;
            default: return 3;
        }
    }
}

// 사용
JavaPairRDD<String, Integer> partitioned =
    pairRDD.partitionBy(new RegionPartitioner(4));

// 파티션별 데이터 수 확인
df.groupBy(spark_partition_id().alias("partition"))
  .count()
  .orderBy(col("count").desc())
  .show();

// 출력:
// +----------+--------+
// |partition |   count|
// +----------+--------+
// |        5 |10000000|  ← 스큐!
// |        3 |   50000|
// |        1 |   45000|
// ...

import java.util.Random;

// 핫 키에 랜덤 솔트 추가
int saltBuckets = 10;
Random rand = new Random();

Dataset<Row> salted = df.withColumn(
    "salted_key",
    concat(col("key"), lit("_"), lit(rand.nextInt(saltBuckets)))
);

// 솔팅된 키로 집계
Dataset<Row> partialAgg = salted
    .groupBy("salted_key")
    .agg(sum("value").alias("partial_sum"));

// 원래 키로 최종 집계
Dataset<Row> finalResult = partialAgg
    .withColumn("original_key", split(col("salted_key"), "_").getItem(0))
    .groupBy("original_key")
    .agg(sum("partial_sum").alias("total"));

import static org.apache.spark.sql.functions.broadcast;

// 작은 테이블을 브로드캐스트
Dataset<Row> result = largeTable.join(
    broadcast(smallTable),
    "key"
);

// Spark 3.0+ AQE 스큐 조인 활성화
spark.conf().set("spark.sql.adaptive.enabled", "true");
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");

// 1단계: 파티션 내 부분 집계
Dataset<Row> partial = df
    .repartition(1000)  // 많은 파티션으로 분산
    .groupBy("key", spark_partition_id().alias("part"))
    .agg(sum("value").alias("partial_sum"));

// 2단계: 최종 집계
Dataset<Row> result = partial
    .groupBy("key")
    .agg(sum("partial_sum").alias("total"));

// 1. Broadcast Hash Join (작은 테이블)
// 자동 또는 명시적으로 브로드캐스트
spark.conf().set("spark.sql.autoBroadcastJoinThreshold", "100MB");
df1.join(broadcast(df2), "key");

// 2. Sort-Merge Join (큰 테이블)
// 양쪽 테이블을 키로 정렬 후 머지
// 기본 조인 전략

// 3. Shuffle Hash Join
// 한쪽만 해시 테이블 생성
df1.join(df2.hint("shuffle_hash"), "key");

// 조인 전 필터링 (셔플 데이터 감소)
Dataset<Row> filtered1 = df1.filter(col("status").equalTo("ACTIVE"));
Dataset<Row> filtered2 = df2.filter(col("type").equalTo("VALID"));
Dataset<Row> joined = filtered1.join(filtered2, "key");

// 필요한 컬럼만 선택
Dataset<Row> slim1 = df1.select("key", "needed_col1");
Dataset<Row> slim2 = df2.select("key", "needed_col2");
Dataset<Row> joined = slim1.join(slim2, "key");

// 사전 파티셔닝 (같은 키로 파티셔닝되면 셔플 회피)
Dataset<Row> part1 = df1.repartition(100, col("key"));
Dataset<Row> part2 = df2.repartition(100, col("key"));
Dataset<Row> joined = part1.join(part2, "key");
// → 이미 같은 키가 같은 파티션에 있으므로 셔플 불필요

// 파티션 컬럼으로 저장
df.write()
    .partitionBy("year", "month")
    .parquet("output/data");

// 생성되는 디렉토리 구조:
// output/data/
// ├── year=2024/
// │   ├── month=01/
// │   │   ├── part-00000.parquet
// │   │   └── part-00001.parquet
// │   ├── month=02/
// │   ...
// └── year=2025/
//     ...

// 파티션 프루닝 (해당 파티션만 읽음)
Dataset<Row> filtered = spark.read()
    .parquet("output/data")
    .filter(col("year").equalTo(2024).and(col("month").equalTo(1)));
// → year=2024/month=01/ 디렉토리만 스캔

// 버케팅으로 저장 (조인 최적화)
df.write()
    .bucketBy(100, "user_id")
    .sortBy("timestamp")
    .saveAsTable("user_events");

// 같은 버케팅으로 저장된 테이블 조인 시 셔플 없음
Dataset<Row> users = spark.table("users").bucketBy(100, "user_id");
Dataset<Row> events = spark.table("user_events");
Dataset<Row> joined = users.join(events, "user_id");
// → 셔플 없이 조인 (같은 user_id가 같은 버킷에 있음)