AQE: adaptive coalescing - автоматический подбор числа партиций

spark.conf.set("spark.sql.adaptive.enabled", "true")

spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")

spark.conf.set("spark.sql.adaptive.advisoryPartitionSizeInBytes", "128m")

spark.conf.set("spark.sql.adaptive.coalescePartitions.minPartitionNum", "1")

spark.conf.set("spark.sql.adaptive.coalescePartitions.parallelismFirst", "true")

spark.conf.set("spark.sql.shuffle.partitions", "2000")

# Статический explain - план до выполнения
df.groupBy("region").agg(sum("revenue")).explain("formatted")

== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- HashAggregate(keys=[region#12], functions=[sum(revenue#34)])
   +- Exchange hashpartitioning(region#12, 2000), ENSURE_REQUIREMENTS, [id=#89]
      +- ...

== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=true
+- == Final Plan ==
   HashAggregate(keys=[region#12], functions=[sum(revenue#34)])
   +- AQEShuffleRead coalesced
      +- ShuffleQueryStage 0
         +- Exchange hashpartitioning(region#12, 2000), ENSURE_REQUIREMENTS, [id=#89]
            +- ...

coalesced from 2000 to 14 partitions

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

def create_spark(aqe_enabled: bool, shuffle_partitions: int = 2000) -> SparkSession:
    """
    Создаёт SparkSession с заданными настройками AQE.
    Параметр aqe_enabled позволяет сравнить поведение до и после включения AQE.
    shuffle_partitions=2000 симулирует настройку "с запасом" для пиковой нагрузки.
    """
    builder = (
        SparkSession.builder
        .appName(f"AQE-Demo-{'enabled' if aqe_enabled else 'disabled'}")
        .master("local[4]")

        # Фиксируем высокое число shuffle партиций, как на продакшн-кластере
        # для пиковой нагрузки. Без AQE это будет катастрофой на малых данных.
        .config("spark.sql.shuffle.partitions", str(shuffle_partitions))

        # AQE: главный тумблер
        .config("spark.sql.adaptive.enabled", str(aqe_enabled).lower())

        # Adaptive Coalescing: целевой размер партиции 64 МБ
        .config("spark.sql.adaptive.coalescePartitions.enabled", "true")
        .config("spark.sql.adaptive.advisoryPartitionSizeInBytes", "64m")

        # Минимум партиций - не ниже числа доступных ядер (4 в данном случае)
        .config("spark.sql.adaptive.coalescePartitions.minPartitionNum", "4")

        # Включаем логирование адаптивного плана для анализа
        .config("spark.sql.adaptive.logLevel", "WARN")
    )
    return builder.getOrCreate()


def generate_small_dataset(spark: SparkSession, size_mb: float = 50.0):
    """
    Генерирует датасет ~50 МБ - типичный «выходной день» для нашего ETL.
    При spark.sql.shuffle.partitions=2000 это создаст партиции по ~25 КБ,
    что является экстремальным примером over-partitioning.
    """
    # Оцениваем число строк для нужного объёма данных.
    # Строка с 5 колонками ≈ 100 байт → 50 МБ ≈ 500_000 строк
    num_rows = int(size_mb * 1024 * 1024 / 100)

    df = spark.range(num_rows).select(
        F.col("id"),
        # Регион: 10 уникальных значений - типичный GROUP BY ключ
        (F.col("id") % 10).cast("string").alias("region"),
        # Категория: 50 значений - второй уровень группировки
        (F.col("id") % 50).cast("string").alias("category"),
        # Выручка: случайное значение от 1 до 1000
        (F.rand(seed=42) * 999 + 1).alias("revenue"),
        # Дата: последние 30 дней
        (F.current_date() - (F.col("id") % 30).cast("int")).alias("sale_date"),
    )
    return df.cache()


def run_aggregation(spark: SparkSession, df, label: str):
    """
    Выполняет типичную GROUP BY агрегацию и замеряет время.
    Именно после GROUP BY происходит Shuffle с adaptivePartitions или без.
    """
    print(f"\n{'='*60}")
    print(f"  {label}")
    print(f"{'='*60}")

    # Запрос с двухуровневой агрегацией - two-stage shuffle
    query = (
        df
        .groupBy("region", "category")
        .agg(
            F.sum("revenue").alias("total_revenue"),
            F.count("*").alias("order_count"),
            F.avg("revenue").alias("avg_revenue"),
            F.countDistinct("sale_date").alias("active_days"),
        )
        .orderBy("total_revenue", ascending=False)
    )

    start = time.time()
    result = query.collect()  # Действие, которое реально запускает выполнение
    elapsed = time.time() - start

    print(f"  Результат: {len(result)} строк")
    print(f"  Время выполнения: {elapsed:.2f}с")
    return elapsed

# Создаём сессию БЕЗ AQE
spark_no_aqe = create_spark(aqe_enabled=False, shuffle_partitions=2000)

# Генерируем 50 МБ данных - «выходной» объём
df_small = generate_small_dataset(spark_no_aqe, size_mb=50.0)
df_small.count()  # Прогреваем кэш

# Смотрим статический план - видим 2000 партиций
print("=== СТАТИЧЕСКИЙ ПЛАН (без AQE) ===")
df_small.groupBy("region").agg(F.sum("revenue")).explain("formatted")

== Physical Plan ==
*(2) HashAggregate(keys=[region#12], functions=[sum(revenue#34)])
+- Exchange hashpartitioning(region#12, 2000), ENSURE_REQUIREMENTS, [id=#89]
   +- *(1) HashAggregate(keys=[region#12], ...

# Запускаем и замеряем время
time_no_aqe = run_aggregation(
    spark_no_aqe,
    df_small,
    "Без AQE: spark.sql.shuffle.partitions=2000, данные=50МБ"
)

# После выполнения откройте Spark UI (обычно localhost:4040)
# Перейдите в Stages → найдите стадию с 2000 тасками
# Обратите внимание на Scheduler Delay в сводной таблице тасок

# Создаём новую сессию С AQE
spark_aqe = create_spark(aqe_enabled=True, shuffle_partitions=2000)

# Используем те же самые данные того же объёма
df_small_aqe = generate_small_dataset(spark_aqe, size_mb=50.0)
df_small_aqe.count()

# Статический план всё ещё покажет 2000 - AQE ещё не выполнялся
print("=== СТАТИЧЕСКИЙ ПЛАН (с AQE, до выполнения) ===")
df_small_aqe.groupBy("region").agg(F.sum("revenue")).explain("formatted")

== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false   ← Маркер: план будет адаптирован
+- HashAggregate(keys=[region#12], ...)
   +- Exchange hashpartitioning(region#12, 2000), ...

# Запускаем и замеряем время
time_aqe = run_aggregation(
    spark_aqe,
    df_small_aqe,
    "С AQE: advisoryPartitionSizeInBytes=64MB, данные=50МБ"
)

# Теперь смотрим финальный (адаптированный) план
print("=== ФИНАЛЬНЫЙ ПЛАН (с AQE, после выполнения) ===")
adapted_df = df_small_aqe.groupBy("region").agg(F.sum("revenue"))
adapted_df.explain("formatted")

== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=true
+- == Final Plan ==
   *(2) HashAggregate(...)
   +- AQEShuffleRead coalesced          ← Вот он!
      +- ShuffleQueryStage 0, Statistics(sizeInBytes=48.0 MiB, ...)
         +- Exchange hashpartitioning(region#12, 2000), ...

== Initial Plan ==                       ← Исходный план для сравнения
   HashAggregate(...)
   +- Exchange hashpartitioning(region#12, 2000), ...

def analyze_stage_metrics(spark: SparkSession, has_aqe: bool):
    """
    Выводит ключевые метрики для понимания эффекта AQE.
    В реальном коде эти данные смотрим в Spark UI → Stages.
    Здесь используем SparkContext.statusTracker для программного доступа.
    """
    sc = spark.sparkContext
    # Получаем информацию о всех завершённых стадиях
    stage_ids = sc.statusTracker().getJobIdsForGroup(None)

    aqe_label = "С AQE" if has_aqe else "Без AQE"
    print(f"\n{'='*60}")
    print(f"  Анализ метрик: {aqe_label}")
    print(f"{'='*60}")

    # Ключевые наблюдения для интерпретации результата
    if has_aqe:
        print("""
  Ожидаемые метрики в Spark UI (Stages → Reduce Stage):
  ─────────────────────────────────────────────────────
  • Number of Tasks:     ~4 (после coalescing с minPartitionNum=4)
  • Median Task Duration: 0.5-2с (больше полезной работы)
  • Scheduler Delay:     2-5% от Duration (норма)
  • Shuffle Read/task:   ~12 МБ (вместо 25 КБ)
  • Files written:       4-8 файлов (вместо 2000)

  В DAG-графе (SQL tab):
  • Оператор: AQEShuffleRead coalesced
  • Надпись: "coalesced from 2000 to 4 partitions"
        """)
    else:
        print("""
  Ожидаемые метрики в Spark UI (Stages → Reduce Stage):
  ─────────────────────────────────────────────────────
  • Number of Tasks:     2000 (все запущены!)
  • Median Task Duration: 8-20 мс (в основном overhead)
  • Scheduler Delay:     15-30% от Duration (красный флаг!)
  • Shuffle Read/task:   ~25 КБ (мизерная полезная работа)
  • Files written:       2000 мелких файлов

  Это признаки over-partitioning:
  ─────────────────────────────────────────────────────
  ⚠ Scheduler Delay > 10% → executor тратит время на orchestration
  ⚠ Task Duration < 100мс → задача тривиальная, overhead доминирует
  ⚠ Shuffle Read < 1MB → партиция почти пустая
        """)


# Запускаем анализ для обоих вариантов
analyze_stage_metrics(spark_no_aqe, has_aqe=False)
analyze_stage_metrics(spark_aqe, has_aqe=True)

# Итоговое сравнение
print(f"\n{'='*60}")
print(f"  ИТОГОВОЕ СРАВНЕНИЕ")
print(f"{'='*60}")
print(f"  Без AQE:   {time_no_aqe:.2f}с")
print(f"  С AQE:     {time_aqe:.2f}с")
print(f"  Ускорение: {time_no_aqe/time_aqe:.1f}x")

def benchmark_advisory_sizes(spark_base: SparkSession, df, sizes_mb: list):
    """
    Сравниваем влияние разных значений advisoryPartitionSizeInBytes.
    Понимание этого помогает найти баланс между параллелизмом и размером файлов.
    """
    results = {}

    for size_mb in sizes_mb:
        size_bytes = size_mb * 1024 * 1024

        # Создаём отдельную сессию для чистоты эксперимента
        spark = (
            SparkSession.builder
            .appName(f"AQE-Advisory-{size_mb}MB")
            .master("local[4]")
            .config("spark.sql.shuffle.partitions", "2000")
            .config("spark.sql.adaptive.enabled", "true")
            .config("spark.sql.adaptive.coalescePartitions.enabled", "true")
            .config("spark.sql.adaptive.advisoryPartitionSizeInBytes", str(int(size_bytes)))
            .config("spark.sql.adaptive.coalescePartitions.minPartitionNum", "1")
            # parallelismFirst=false → строго соблюдаем advisoryPartitionSizeInBytes
            .config("spark.sql.adaptive.coalescePartitions.parallelismFirst", "false")
            .getOrCreate()
        )

        df_local = df.repartition(2000)  # Принудительно 2000 исходных партиций

        start = time.time()
        query_result = (
            df_local
            .groupBy("region", "category")
            .agg(F.sum("revenue").alias("total"))
            .collect()
        )
        elapsed = time.time() - start

        results[size_mb] = elapsed
        print(f"  advisoryPartitionSizeInBytes={size_mb}MB: {elapsed:.2f}с")
        spark.stop()

    return results


# Тестируем разные целевые размеры партиций
print("\n=== Benchmark: advisoryPartitionSizeInBytes ===")
print("Данные: 50 МБ, shuffle.partitions=2000, parallelismFirst=false")
print()
# В реальном эксперименте запускаем с нашим датасетом
# benchmark_results = benchmark_advisory_sizes(spark_aqe, df_small_aqe, [16, 32, 64, 128, 256])

def simulate_incremental_etl(spark: SparkSession, day_type: str):
    """
    Симулирует реальный ETL-пайплайн с разным объёмом данных:
    - 'weekday': типичный будний день (500 МБ)
    - 'weekend': выходной день (5 МБ)
    - 'peak': пиковый день (5 ГБ, конец месяца)

    Именно для этого сценария AQE наиболее ценен:
    один и тот же код и конфигурация работают оптимально на всех трёх объёмах.
    """
    sizes = {
        "weekday": 500,    # 500 МБ → ~250 строк × 2 МБ блоки
        "weekend": 5,      # 5 МБ → мизерный объём, 2000 партиций = катастрофа без AQE
        "peak":    5000,   # 5 ГБ → нужен высокий параллелизм
    }

    size_mb = sizes[day_type]
    # Масштабируем число строк к нужному размеру
    num_rows = int(size_mb * 1024 * 1024 / 100)

    raw = spark.range(num_rows).select(
        F.col("id"),
        (F.col("id") % 100).cast("string").alias("product_id"),
        (F.col("id") % 50).cast("string").alias("store_id"),
        (F.rand(seed=42) * 999 + 1).alias("amount"),
        (F.current_date() - (F.col("id") % 7).cast("int")).alias("txn_date"),
    )

    # Типичный ETL: агрегация по продуктам и магазинам
    silver = (
        raw
        .groupBy("product_id", "store_id", "txn_date")
        .agg(
            F.sum("amount").alias("daily_revenue"),
            F.count("*").alias("txn_count"),
        )
    )

    # Вторая агрегация - Gold layer
    gold = (
        silver
        .groupBy("product_id", "txn_date")
        .agg(
            F.sum("daily_revenue").alias("total_revenue"),
            F.sum("txn_count").alias("total_txns"),
            F.countDistinct("store_id").alias("active_stores"),
        )
    )

    start = time.time()
    row_count = gold.count()
    elapsed = time.time() - start

    print(f"  Day type: {day_type:8s} | Size: {size_mb:5d} МБ | "
          f"Rows: {row_count:6d} | Time: {elapsed:.2f}с")
    return elapsed


print("\n=== Симуляция инкрементального ETL (с AQE) ===")
print("Конфигурация: shuffle.partitions=2000, advisoryPartitionSizeInBytes=64MB")
print()
for day in ["weekend", "weekday", "peak"]:
    simulate_incremental_etl(spark_aqe, day)

def production_aqe_config(
    spark_builder,
    peak_data_gb: float,
    executor_memory_gb: int,
    num_executors: int,
    target_file_size_mb: int = 128,
):
    """
    Задаёт production-готовые настройки AQE на основе характеристик кластера.

    Параметры:
        peak_data_gb:       Пиковый объём данных после shuffle в ГБ
        executor_memory_gb: RAM одного Executor'а
        num_executors:      Количество Executor'ов
        target_file_size_mb: Целевой размер Parquet-файла на выходе

    Логика расчёта:
        - shuffle_partitions: правило 4x от пикового объёма (с запасом под AQE)
        - advisoryPartitionSizeInBytes: целевой размер файла - хотим, чтобы
          каждая AQE-таска создала ровно один файл нужного размера
        - minPartitionNum: число ядер кластера - чтобы кластер не простаивал
    """
    total_cores = num_executors * 4  # Предполагаем 4 ядра на Executor

    # Число shuffle партиций с запасом (AQE схлопнет лишние)
    shuffle_partitions = max(200, int(peak_data_gb * 1024 / 64 * 4))

    # Целевой размер = целевой размер файла (для хорошего layout в Lakehouse)
    advisory_size_bytes = target_file_size_mb * 1024 * 1024

    # Нижний лимит = число ядер (чтобы все заняты)
    min_partitions = total_cores

    return (
        spark_builder
        # Включаем AQE
        .config("spark.sql.adaptive.enabled", "true")
        .config("spark.sql.adaptive.coalescePartitions.enabled", "true")

        # Ключевые параметры coalescing
        .config("spark.sql.shuffle.partitions", str(shuffle_partitions))
        .config("spark.sql.adaptive.advisoryPartitionSizeInBytes",
                str(advisory_size_bytes))
        .config("spark.sql.adaptive.coalescePartitions.minPartitionNum",
                str(min_partitions))

        # parallelismFirst=false: соблюдаем целевой размер строго
        # Это важно для Lakehouse: хотим файлы нужного размера, не максимальный параллелизм
        .config("spark.sql.adaptive.coalescePartitions.parallelismFirst", "false")

        # Skew Join - включаем вместе с coalescing (работают дополняя друг друга)
        .config("spark.sql.adaptive.skewJoin.enabled", "true")
        .config("spark.sql.adaptive.skewJoin.skewedPartitionFactor", "5")
        .config("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "256m")
    )


# Пример использования для кластера: 20 Executor'ов, 16 ГБ каждый, пик 2 ТБ данных
builder = SparkSession.builder.appName("Production-ETL")
configured_builder = production_aqe_config(
    spark_builder=builder,
    peak_data_gb=2048,      # 2 ТБ пиковый объём после shuffle
    executor_memory_gb=16,  # 16 ГБ на Executor
    num_executors=20,       # 20 Executor'ов
    target_file_size_mb=128 # Хотим Parquet-файлы по 128 МБ
)

# Расчётные значения для данного примера:
# shuffle_partitions = max(200, 2048 * 1024 / 64 * 4) = max(200, 131072) = 131072
# advisoryPartitionSizeInBytes = 128 МБ
# minPartitionNum = 20 × 4 = 80 ядер
# Итог: AQE может схлопнуть 131072 исходных партиций до ~(2TB/128MB) = ~16384 тасок

# НЕПРАВИЛЬНО: отключение AQE для детерминированного числа файлов
spark.conf.set("spark.sql.adaptive.enabled", "false")
spark.conf.set("spark.sql.shuffle.partitions", "500")
df.write.parquet("s3://bucket/output/")
# Результат: 500 файлов независимо от объёма данных.
# При 5 МБ данных - 500 файлов по 10 КБ (катастрофа для S3 listing).
# При 5 ТБ данных - 500 файлов по 10 ГБ (Parquet reader будет страдать).

# ПРАВИЛЬНО: доверяйте AQE + задайте advisoryPartitionSizeInBytes
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.shuffle.partitions", "10000")
spark.conf.set("spark.sql.adaptive.advisoryPartitionSizeInBytes", "128m")
spark.conf.set("spark.sql.adaptive.coalescePartitions.parallelismFirst", "false")
df.write.parquet("s3://bucket/output/")
# Результат: ~(total_data_size / 128MB) файлов.
# При 5 МБ → 1 файл. При 5 ТБ → ~40000 файлов.

# НЕПРАВИЛЬНО: слишком мало исходных партиций
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.shuffle.partitions", "50")  # Мало!
# Если данных окажется 10 ТБ → 50 партиций по 200 ГБ → огромные таски, spill, OOM.
# AQE может только СХЛОПЫВАТЬ, но не РАСШИРЯТЬ число партиций выше заданного значения.

# ПРАВИЛЬНО: задавать с запасом на пиковый объём
spark.conf.set("spark.sql.shuffle.partitions", "8000")  # Для пика 2 ТБ
# При 10 МБ данных AQE схлопнет до 1-4 партиций - всё нормально.
# При 2 ТБ данных AQE оставит ~16000 партиций по 128 МБ - всё нормально.

# Ситуация: хотим строго контролировать размер выходных файлов в Iceberg
spark.conf.set("spark.sql.adaptive.advisoryPartitionSizeInBytes", "128m")

# Но если оставить parallelismFirst=true (дефолт), Spark может создать
# больше партиций, чем нужно, чтобы «занять» все ядра кластера.
# В итоге файлы будут мельче 128 МБ - цель не достигнута.

# ПРАВИЛЬНО для write-heavy пайплайнов с Lakehouse:
spark.conf.set("spark.sql.adaptive.coalescePartitions.parallelismFirst", "false")
# Теперь Spark строго соблюдает advisoryPartitionSizeInBytes,
# даже если часть ядер кластера будет простаивать.

# Датасет: транзакции e-commerce за квартал
# Объём: ~1.2 ТБ (после shuffle - около 800 ГБ после агрегации)
# Кластер: 40 Executor'ов, 8 ядер каждый, 32 ГБ RAM

raw_txns = spark.read.parquet("s3://data-lake/bronze/transactions/2024-Q4/")

# Ваш Gold-layer пайплайн:
gold = (
    raw_txns
    .groupBy("merchant_id", "product_category", "txn_date")
    .agg(
        F.sum("amount").alias("daily_gmv"),
        F.count("*").alias("txn_count"),
        F.approx_count_distinct("user_id").alias("unique_buyers"),
    )
    .groupBy("merchant_id", "product_category")
    .agg(
        F.sum("daily_gmv").alias("quarterly_gmv"),
        F.sum("txn_count").alias("quarterly_txns"),
        F.avg("unique_buyers").alias("avg_daily_buyers"),
    )
)

gold.write.format("parquet").mode("overwrite").save("s3://data-lake/gold/merchant_summary/")