AQE: Adaptive Query Execution - включение, параметры и что оно реально делает

-- Запрос: джойн транзакций с информацией о стране
SELECT t.*, c.country_name
FROM transactions t
JOIN countries c ON t.country_code = c.code
WHERE t.country_code = 'IS'  -- Исландия: ~370 000 человек населения

from pyspark.sql import SparkSession

spark = SparkSession.builder \
    .appName("aqe-demo") \
    # Spark 3.0-3.1: явно включаем AQE
    .config("spark.sql.adaptive.enabled", "true") \
    # Spark 3.2+: включён по умолчанию, но явное указание не помешает

    # Все три механизма AQE включены по умолчанию при AQE=true:
    .config("spark.sql.adaptive.coalescePartitions.enabled", "true")
    .config("spark.sql.adaptive.skewJoin.enabled", "true")
    .config("spark.sql.adaptive.localShuffleReader.enabled", "true")

    .getOrCreate()

# Проверяем статус AQE программно:
print("AQE enabled:", spark.conf.get("spark.sql.adaptive.enabled"))

# Вывод: AQE enabled: true

from pyspark.sql import functions as F

# Создадим датасет для демонстрации
df1 = spark.range(1_000_000).select(
    F.col("id"),
    (F.rand() * 100).cast("int").alias("country_id"),
    (F.rand() * 1000).alias("amount")
)

df2 = spark.range(100).select(
    F.col("id").alias("country_id"),
    F.concat(F.lit("Country_"), F.col("id").cast("string")).alias("country_name")
)

result = df1.join(df2, "country_id").groupBy("country_name").sum("amount")

# explain() ДО выполнения: покажет начальный план с AdaptiveSparkPlan
result.explain("formatted")

== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false              ← AQE активен, план ещё не финальный!
+- == Current Plan ==
   HashAggregate(keys=[country_name#17], functions=[sum(amount#12)])
   +- Exchange hashpartitioning(country_name#17, 200), ENSURE_REQUIREMENTS, [plan_id=15]
      +- HashAggregate(keys=[country_name#17], functions=[partial_sum(amount#12)])
         +- Project [amount#12, country_name#17]
            +- SortMergeJoin [country_id#11], [country_id#7], Inner   ← SMJ до AQE
               :- Sort [country_id#11 ASC], false, 0
               :  +- Exchange hashpartitioning(country_id#11, 200)
               :     +- Range (0, 1000000, step=1, splits=8)
               +- Sort [country_id#7 ASC], false, 0
                  +- Exchange hashpartitioning(country_id#7, 200)
                     +- Range (0, 100, step=1, splits=8)

== Initial Plan ==         ← Что Catalyst запланировал изначально
   ... (полный план без AQE адаптаций)

result.collect()

# Теперь смотрим финальный план:
result.explain("formatted")

== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=true               ← ПЛАН ФИНАЛЬНЫЙ!
+- == Final Plan ==
   HashAggregate(keys=[country_name#17], functions=[sum(amount#12)])
   +- HashAggregate(keys=[country_name#17], functions=[partial_sum(amount#12)])
      +- Project [amount#12, country_name#17]
         +- BroadcastHashJoin [country_id#11], [country_id#7], Inner  ← AQE заменил SMJ→BHJ!
            :- Range (0, 1000000, step=1, splits=8)
            +- BroadcastExchange HashedRelationBroadcastMode(...)
               +- Range (0, 100, step=1, splits=8)   ← 100 строк → broadcast!

spark = SparkSession.builder \

    # Главный тумблер coalesce (включён по умолчанию с AQE)
    .config("spark.sql.adaptive.coalescePartitions.enabled", "true") \

    # Целевой размер одной партиции после coalesce.
    # AQE склеивает мелкие партиции пока их суммарный размер не достигает этого порога.
    # Дефолт: 64 MB. Рекомендации:
    # - Для HDFS/S3 аналитики: 128 MB (совпадает с HDFS блоком)
    # - Для быстрых интерактивных запросов: 32-64 MB
    # - Для heavy aggregations с большими state: 128-256 MB
    .config("spark.sql.adaptive.advisoryPartitionSizeInBytes", "128MB") \

    # Стартовое максимальное число партиций для Shuffle.
    # AQE может только УМЕНЬШАТЬ это число (coalesce), не увеличивать!
    # Правило: ставьте с большим запасом. Лучше 2000 мелких партиций
    # которые AQE сольёт в 50, чем 50 с риском OOM.
    # None = использовать spark.sql.shuffle.partitions
    .config("spark.sql.adaptive.coalescePartitions.initialPartitionNum", "2000") \

    # Минимальный размер партиции после coalesce.
    # Партиции меньше этого порога будут объединяться с соседями.
    # Дефолт: 1 MB. Уменьшите до 512 KB если данных очень мало.
    .config("spark.sql.adaptive.coalescePartitions.minPartitionSize", "1MB") \

    .getOrCreate()

from pyspark.sql import SparkSession, functions as F

spark = SparkSession.builder \
    .master("local[4]") \
    .config("spark.sql.adaptive.enabled", "true") \
    .config("spark.sql.adaptive.coalescePartitions.enabled", "true") \
    .config("spark.sql.adaptive.advisoryPartitionSizeInBytes", "64MB") \
    .config("spark.sql.shuffle.partitions", "200") \
    .getOrCreate()

spark.sparkContext.setLogLevel("WARN")

# Маленький датасет: ~10 MB
df = spark.range(100_000).select(
    F.col("id"),
    (F.rand() * 100).cast("int").alias("group_id"),
    F.rand().alias("value")
)

# Без AQE: 200 тасок для 10 MB данных (нелепо!)
df_agg = df.groupBy("group_id").sum("value")

# Смотрим сколько партиций получилось ПОСЛЕ AQE
import time
t0 = time.time()
result = df_agg.collect()
elapsed = time.time() - t0

print(f"Время выполнения: {elapsed:.2f}с")

# Смотрим финальный план - сколько задач реально запустил Spark:
df_agg.explain("formatted")
# В Adaptive Plan увидим: CoalesceShufflePartitions
# и реальное число партиций (например, 2-5 вместо 200)

# Программная проверка числа партиций после AQE:
# (через collect нельзя, но можно через промежуточный persist)
df_agg_persisted = df_agg.persist()
df_agg_persisted.count()  # материализуем
print(f"Партиций в результате: {df_agg_persisted.rdd.getNumPartitions()}")
# Вывод: Партиций в результате: 3 (вместо 200!)

spark = SparkSession.builder \

    # Порог для автоматического broadcast при статическом планировании
    # Дефолт: 10 MB. Это пороговое значение для СТАТИЧЕСКОГО выбора BHJ.
    .config("spark.sql.autoBroadcastJoinThreshold", "10MB") \

    # AQE-специфичный порог для ДИНАМИЧЕСКОГО upgrade до BHJ.
    # По умолчанию равен autoBroadcastJoinThreshold.
    # Рекомендуется поднять для AQE - раз мы уже знаем реальный размер,
    # можем безопасно broadcast'ить большие таблицы:
    .config("spark.sql.adaptive.autoBroadcastJoinThreshold", "30MB") \

    .getOrCreate()

from pyspark.sql import SparkSession, functions as F

spark = SparkSession.builder \
    .master("local[4]") \
    .config("spark.sql.adaptive.enabled", "true") \
    .config("spark.sql.autoBroadcastJoinThreshold", "10MB") \
    .config("spark.sql.adaptive.autoBroadcastJoinThreshold", "50MB") \
    .getOrCreate()

# Симулируем ситуацию: большая таблица с сильным фильтром
transactions = spark.range(10_000_000).select(
    F.col("id").alias("txn_id"),
    (F.rand() * 1000).cast("int").alias("country_id"),
    (F.rand() * 500).alias("amount")
)

# Маленькая таблица стран
countries = spark.range(1000).select(
    F.col("id").alias("country_id"),
    F.concat(F.lit("Country_"), F.col("id")).alias("country_name")
)

# Фильтруем transactions до ~10 000 строк (0.1% от исходных)
txn_filtered = transactions.filter(F.col("country_id") < 10)

# Join
result = txn_filtered.join(countries, "country_id") \
    .groupBy("country_name").agg(F.sum("amount").alias("total"))

# До выполнения: Catalyst видит transactions без фильтра → SMJ
print("=== ПЛАН ДО ВЫПОЛНЕНИЯ ===")
result.explain("formatted")
# Увидим: AdaptiveSparkPlan isFinalPlan=false
# Current Plan: SortMergeJoin (т.к. 10M строк > 10MB threshold)

# Выполняем
result.collect()

# После выполнения: AQE увидел реальный размер
print("=== ПЛАН ПОСЛЕ ВЫПОЛНЕНИЯ (АДАПТИВНЫЙ) ===")
result.explain("formatted")
# Увидим: AdaptiveSparkPlan isFinalPlan=true
# Final Plan: BroadcastHashJoin (AQE увидел что осталось ~10K строк)

spark = SparkSession.builder \

    # Главный тумблер Skew Join оптимизации
    .config("spark.sql.adaptive.skewJoin.enabled", "true") \

    # Коэффициент: партиция считается скошенной если она в N раз больше медианы.
    # Дефолт: 5 (партиция должна быть в 5x раз больше медианной).
    # Снизьте до 3 если хотите более агрессивное разбиение.
    # Поднимите до 10 если видите слишком много ложных срабатываний.
    .config("spark.sql.adaptive.skewJoin.skewedPartitionFactor", "5") \

    # Минимальный абсолютный размер для признания партиции скошенной.
    # Дефолт: 256 MB. Партиция должна быть > skewedPartitionFactor*median
    # И > skewedPartitionThresholdInBytes одновременно.
    # Снизьте если работаете с небольшими данными:
    .config("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "256MB") \

    # Целевой размер части при разбиении скошенной партиции.
    # Дефолт: advisoryPartitionSizeInBytes (обычно 64 MB).
    # Разбиение: skewed_partition_size / advisoryPartitionSizeInBytes = N частей
    .config("spark.sql.adaptive.advisoryPartitionSizeInBytes", "64MB") \

    .getOrCreate()

# При настройках выше:
# Если медиана = 50 MB, threshold = 256 MB:
# Партиция 847 MB > 5 * 50 MB (= 250 MB) AND > 256 MB → SKEW!
# 847 MB / 64 MB = ~13 частей

from pyspark.sql import SparkSession, functions as F
from pyspark.sql.types import StructType, StructField, LongType, StringType, DoubleType

spark = SparkSession.builder \
    .master("local[4]") \
    .config("spark.sql.adaptive.enabled", "true") \
    .config("spark.sql.adaptive.skewJoin.enabled", "true") \
    .config("spark.sql.adaptive.skewJoin.skewedPartitionFactor", "3") \
    .config("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "10MB") \
    .config("spark.sql.adaptive.advisoryPartitionSizeInBytes", "10MB") \
    .getOrCreate()

# Создаём сильно скошенный датасет
# 90% строк имеют category_id = 0 (один "горячий" ключ)
def create_skewed_data(spark, n_rows=1_000_000):
    # 900K строк с category_id=0, 100K строк с category_id=1..9
    hot_data = spark.range(900_000).select(
        F.col("id").alias("event_id"),
        F.lit(0).alias("category_id"),
        F.rand().alias("revenue")
    )
    normal_data = spark.range(100_000).select(
        (F.col("id") + 900_000).alias("event_id"),
        (F.col("id") % 9 + 1).cast("int").alias("category_id"),
        F.rand().alias("revenue")
    )
    return hot_data.union(normal_data)

events = create_skewed_data(spark)

# Таблица категорий
categories = spark.range(10).select(
    F.col("id").alias("category_id"),
    F.concat(F.lit("Cat_"), F.col("id")).alias("category_name")
)

# Join с перекошенными данными
result = events.join(categories, "category_id") \
    .groupBy("category_name") \
    .agg(F.sum("revenue").alias("total_revenue"))

import time
t0 = time.time()
result.collect()
elapsed = time.time() - t0
print(f"С AQE Skew Join: {elapsed:.1f}с")

# Повторяем без skew оптимизации для сравнения:
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "false")
events2 = create_skewed_data(spark)
categories2 = spark.range(10).select(
    F.col("id").alias("category_id"),
    F.concat(F.lit("Cat_"), F.col("id")).alias("category_name")
)
result2 = events2.join(categories2, "category_id") \
    .groupBy("category_name") \
    .agg(F.sum("revenue").alias("total_revenue"))

t0 = time.time()
result2.collect()
elapsed_no_skew = time.time() - t0
print(f"Без AQE Skew Join: {elapsed_no_skew:.1f}с")
print(f"Ускорение AQE: {elapsed_no_skew/elapsed:.1f}x")

spark = SparkSession.builder \
    # Local Shuffle Reader: включён по умолчанию с AQE
    .config("spark.sql.adaptive.localShuffleReader.enabled", "true") \
    .getOrCreate()

# explain("formatted") показывает полный адаптивный план
df.explain("formatted")

# Ищите в выводе:
# 1. Признак активности AQE:
#    AdaptiveSparkPlan isFinalPlan=false|true

# 2. Coalesce:
#    CoalesceShufflePartitions
#    (указывает что партиции были слиты)

# 3. Join downgrade:
#    BroadcastHashJoin ... (skipped=(SortMergeJoin))
#    или просто BroadcastHashJoin где ожидался SortMergeJoin

# 4. Skew handling:
#    SkewJoin(skewedPartitions=[SkewedPartition(0, 199)], ...
#    AQEShuffleRead (вместо обычного ShuffleReader)

spark.sparkContext.setLogLevel("WARN")

# Только для AQE логов:
log4j = spark.sparkContext._jvm.org.apache.log4j
log4j.LogManager.getLogger(
    "org.apache.spark.sql.execution.adaptive"
).setLevel(log4j.Level.DEBUG)

DEBUG AdaptiveSparkPlanExec: Final plan:
  BroadcastHashJoin (replaced SortMergeJoin due to runtime statistics)

DEBUG CoalesceShufflePartitions: Coalesced from 200 to 8 partitions

DEBUG OptimizeSkewedJoin: Partition 199 is skewed (size=847MB, median=1.1MB, factor=770x)
  Splitting into 13 sub-partitions

from pyspark.sql import SparkSession, functions as F

spark = SparkSession.builder \
    .master("local[8]") \
    .config("spark.sql.adaptive.enabled", "true") \
    .getOrCreate()

# Listener для сбора статистики AQE
class AQEMetricsListener:
    """Собирает метрики о применении AQE оптимизаций."""

    def __init__(self):
        self.stage_stats = {}

    def collect(self, df):
        """
        Выполняет запрос и собирает метрики партиций.
        Возвращает число задач в каждом Stage.
        """
        df.collect()

        # Получаем информацию через SparkContext StatusTracker
        sc = df.sparkSession.sparkContext
        stage_ids = sc.statusTracker().getActiveJobIds()

        return {
            "shuffle_partitions_setting": df.sparkSession.conf.get(
                "spark.sql.shuffle.partitions"
            ),
            "aqe_enabled": df.sparkSession.conf.get(
                "spark.sql.adaptive.enabled"
            ),
        }


# Сравнение с AQE и без
df = spark.range(100_000).select(
    F.col("id"),
    (F.rand() * 100).cast("int").alias("key"),
    F.rand().alias("value")
)

# С AQE
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.shuffle.partitions", "200")
df_agg_aqe = df.groupBy("key").sum("value")
df_agg_aqe.persist()
count_aqe = df_agg_aqe.count()
partitions_aqe = df_agg_aqe.rdd.getNumPartitions()

# Без AQE
spark.conf.set("spark.sql.adaptive.enabled", "false")
df_agg_no_aqe = df.groupBy("key").sum("value")
df_agg_no_aqe.persist()
count_no_aqe = df_agg_no_aqe.count()
partitions_no_aqe = df_agg_no_aqe.rdd.getNumPartitions()

print(f"С AQE:    {partitions_aqe} партиций  (было настроено: 200)")
print(f"Без AQE:  {partitions_no_aqe} партиций (было настроено: 200)")
# Пример вывода:
# С AQE:    3 партиций  (было настроено: 200)
# Без AQE:  200 партиций (было настроено: 200)

# production_aqe_config.py
# Рекомендуемые настройки AQE для production HDFS/Data Lakehouse кластера

from pyspark.sql import SparkSession


def create_spark_with_aqe(app_name: str, data_size_estimate: str = "medium") -> SparkSession:
    """
    Создаёт SparkSession с настройками AQE под конкретный размер данных.

    data_size_estimate:
    - "small"  : данные < 10 GB
    - "medium" : данные 10-500 GB (большинство ETL)
    - "large"  : данные > 500 GB
    """

    configs = {
        "small": {
            "shuffle_partitions": "1000",
            "advisory_partition_size": "32MB",
            "broadcast_threshold": "30MB",
            "skew_factor": "3",
            "skew_threshold": "32MB",
        },
        "medium": {
            "shuffle_partitions": "2000",
            "advisory_partition_size": "128MB",
            "broadcast_threshold": "50MB",
            "skew_factor": "5",
            "skew_threshold": "256MB",
        },
        "large": {
            "shuffle_partitions": "5000",
            "advisory_partition_size": "256MB",
            "broadcast_threshold": "100MB",
            "skew_factor": "5",
            "skew_threshold": "512MB",
        },
    }

    cfg = configs[data_size_estimate]

    return SparkSession.builder \
        .appName(app_name) \

        # ── Основные AQE переключатели ────────────────────────────────
        .config("spark.sql.adaptive.enabled", "true") \
        .config("spark.sql.adaptive.coalescePartitions.enabled", "true") \
        .config("spark.sql.adaptive.skewJoin.enabled", "true") \
        .config("spark.sql.adaptive.localShuffleReader.enabled", "true") \

        # ── Начальное число партиций (с запасом, AQE уменьшит) ────────
        # Всегда ставьте spark.sql.shuffle.partitions с запасом.
        # AQE может только УМЕНЬШАТЬ число партиций, не увеличивать!
        .config("spark.sql.shuffle.partitions", cfg["shuffle_partitions"]) \
        .config(
            "spark.sql.adaptive.coalescePartitions.initialPartitionNum",
            cfg["shuffle_partitions"]
        ) \

        # ── Целевой размер партиции после coalesce ────────────────────
        # Рекомендуется 128 MB для HDFS (совпадает с блоком HDFS)
        # Для S3/MinIO: 64-128 MB
        .config(
            "spark.sql.adaptive.advisoryPartitionSizeInBytes",
            cfg["advisory_partition_size"]
        ) \

        # ── Broadcast threshold для AQE ───────────────────────────────
        # AQE использует реальный размер, поэтому можно поднять
        # выше статического autoBroadcastJoinThreshold
        .config(
            "spark.sql.adaptive.autoBroadcastJoinThreshold",
            cfg["broadcast_threshold"]
        ) \

        # ── Настройки детекции Skew ───────────────────────────────────
        # Партиция считается скошенной если:
        # size > skewedPartitionFactor * median_size
        # AND size > skewedPartitionThresholdInBytes
        .config(
            "spark.sql.adaptive.skewJoin.skewedPartitionFactor",
            cfg["skew_factor"]
        ) \
        .config(
            "spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes",
            cfg["skew_threshold"]
        ) \

        # ── CBO статистика для лучших начальных оценок ────────────────
        # CBO + AQE = лучший результат чем по отдельности
        .config("spark.sql.cbo.enabled", "true") \
        .config("spark.sql.statistics.histogram.enabled", "true") \

        .getOrCreate()


# Примеры использования:
spark_etl = create_spark_with_aqe("daily-silver-etl", "medium")
spark_report = create_spark_with_aqe("monthly-gold-report", "large")

# lab_aqe_benchmark.py
# Полное сравнение: ручной тюнинг vs AQE

import time
from dataclasses import dataclass
from pyspark.sql import SparkSession, functions as F


@dataclass
class BenchmarkResult:
    name: str
    elapsed_sec: float
    straggler_task_sec: float  # время самой долгой задачи


def create_benchmark_data(spark: SparkSession):
    """
    Создаёт реалистичный датасет с:
    1. Data Skew (90% данных на key=0)
    2. Таблица которая сильно фильтруется
    """
    # Скошенные транзакции
    skewed_events = spark.range(2_000_000).select(
        F.col("id").alias("event_id"),
        # 90% событий имеют product_id=0 (хит-продукт)
        F.when(F.rand() < 0.9, F.lit(0))
         .otherwise((F.rand() * 999 + 1).cast("int"))
         .alias("product_id"),
        (F.rand() * 1000).alias("revenue"),
        F.date_add(F.lit("2024-01-01"), (F.rand() * 365).cast("int")).alias("date"),
        F.array(F.lit("RU"), F.lit("DE"), F.lit("US")).getItem(
            (F.rand() * 3).cast("int")
        ).alias("country")
    )

    # Таблица продуктов (будет сильно фильтроваться)
    products = spark.range(1000).select(
        F.col("id").alias("product_id"),
        F.concat(F.lit("Product_"), F.col("id")).alias("product_name"),
        F.array(F.lit("Electronics"), F.lit("Books"), F.lit("Food")).getItem(
            (F.rand() * 3).cast("int")
        ).alias("category"),
        (F.rand() * 100).alias("margin_pct")
    )

    return skewed_events, products


def run_without_aqe(spark: SparkSession, events, products) -> float:
    """
    Запуск без AQE: все проблемы накапливаются.
    """
    spark.conf.set("spark.sql.adaptive.enabled", "false")
    spark.conf.set("spark.sql.shuffle.partitions", "200")
    spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "10MB")

    # Запрос: фильтруем до Electronics + skewed join
    result = events \
        .join(products.filter(F.col("category") == "Electronics"), "product_id") \
        .groupBy("product_name", "country") \
        .agg(
            F.sum("revenue").alias("total_revenue"),
            F.count("*").alias("event_count"),
        ) \
        .orderBy(F.desc("total_revenue"))

    t0 = time.time()
    result.collect()
    return time.time() - t0


def run_with_manual_tuning(spark: SparkSession, events, products) -> float:
    """
    Ручная оптимизация: требует знания данных заранее.
    """
    spark.conf.set("spark.sql.adaptive.enabled", "false")
    # Знаем что после фильтра Electronics останется ~333 продукта
    # Знаем что данные сильно скошены → нужен salting
    spark.conf.set("spark.sql.shuffle.partitions", "50")
    spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "50MB")  # принудительный broadcast

    # Salting для скошенных данных (ручной метод)
    SALT = 10
    events_salted = events.withColumn("salt", (F.rand() * SALT).cast("int")) \
        .withColumn("product_id_salted",
                    F.concat(F.col("product_id").cast("string"), F.lit("_"), F.col("salt")))

    products_electronics = products.filter(F.col("category") == "Electronics")
    products_expanded = products_electronics \
        .crossJoin(spark.range(SALT).select(F.col("id").alias("salt"))) \
        .withColumn("product_id_salted",
                    F.concat(F.col("product_id").cast("string"), F.lit("_"), F.col("salt")))

    result = events_salted \
        .join(F.broadcast(products_expanded), "product_id_salted") \
        .groupBy("product_name", "country") \
        .agg(F.sum("revenue"), F.count("*")) \
        .orderBy(F.desc("sum(revenue)"))

    t0 = time.time()
    result.collect()
    return time.time() - t0


def run_with_aqe(spark: SparkSession, events, products) -> float:
    """
    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", "3")
    spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "10MB")
    spark.conf.set("spark.sql.adaptive.autoBroadcastJoinThreshold", "50MB")
    spark.conf.set("spark.sql.shuffle.partitions", "2000")
    spark.conf.set("spark.sql.adaptive.advisoryPartitionSizeInBytes", "64MB")

    # Тот же запрос, без salting и hints!
    result = events \
        .join(products.filter(F.col("category") == "Electronics"), "product_id") \
        .groupBy("product_name", "country") \
        .agg(
            F.sum("revenue").alias("total_revenue"),
            F.count("*").alias("event_count"),
        ) \
        .orderBy(F.desc("total_revenue"))

    t0 = time.time()
    result.collect()
    return time.time() - t0


if __name__ == "__main__":
    spark = SparkSession.builder \
        .master("local[8]") \
        .appName("aqe-benchmark") \
        .getOrCreate()

    spark.sparkContext.setLogLevel("ERROR")

    events, products = create_benchmark_data(spark)
    events.cache()
    products.cache()
    events.count()
    products.count()

    print("\n" + "=" * 60)
    print("БАТЛ: Ручной тюнинг vs AQE")
    print("=" * 60)

    t_no_aqe = run_without_aqe(spark, events, products)
    print(f"Без AQE (дефолты):      {t_no_aqe:.1f}с")

    t_manual = run_with_manual_tuning(spark, events, products)
    print(f"Ручной тюнинг + salting: {t_manual:.1f}с")

    t_aqe = run_with_aqe(spark, events, products)
    print(f"AQE (минимум настроек): {t_aqe:.1f}с")

    print("\nВывод:")
    print(f"  AQE быстрее дефолтов:       {t_no_aqe/t_aqe:.1f}x")
    print(f"  AQE vs ручной тюнинг:        {t_manual/t_aqe:.1f}x")
    print(f"  Строк кода настройки (AQE):  5")
    print(f"  Строк кода настройки (ручн): ~40 (salting + hints)")

    spark.stop()