Dynamic Partition Pruning: механизм star-schema и условия срабатывания

# Таблица sales партиционирована по колонке event_date
# Например: s3a://data/sales/event_date=2026-01-01/...
#                            s3a://data/sales/event_date=2026-01-02/...
#                            ... и так далее

# Запрос с литеральным фильтром
result = spark.table("sales").filter("event_date = '2026-01-01'")
result.explain(mode="simple")
# == Physical Plan ==
# FileScan parquet [amount, region_id, event_date]
#   PartitionFilters: [isnotnull(event_date#12), (event_date#12 = 2026-01-01)]
#   PushedFilters: []
#   ReadSchema: struct<amount:double,region_id:int>

SELECT s.region_id, sum(s.amount) AS total
FROM   sales s
JOIN   promotions p ON s.event_date = p.promo_date
WHERE  p.campaign = 'black_friday_2026'

promos = spark.table("promotions").filter("campaign = 'black_friday_2026'")
result = spark.table("sales").join(promos, "event_date").groupBy("region_id").sum("amount")

# Сохранение fact-таблицы с правильным партиционированием
fact_df.write \
    .mode("overwrite") \
    .partitionBy("event_date", "region_id") \  # по тем же ключам, по которым будут join
    .saveAsTable("sales")

# В результате получаем структуру:
# s3a://warehouse/sales/
#   event_date=2026-01-01/
#     region_id=1/   part-00000.parquet ...
#     region_id=2/   part-00001.parquet ...
#   event_date=2026-01-02/
#     region_id=1/   ...
#   ... (1095 папок для 3 лет данных)

# Запрос: суммарные продажи в регионах, где прошла кампания "black_friday"
promos = spark.table("promotions").filter(col("campaign") == "black_friday_2026")
result = (
    spark.table("sales")              # fact: 10 ТБ, partition by event_date
    .join(promos, "event_date")       # join по partition-колонке
    .groupBy("region_id")
    .agg(F.sum("amount").alias("total_amount"))
)

# ✅ DPP сработает: sales партиционирован по той же колонке, что используется в join
fact_df.write.partitionBy("event_date").saveAsTable("sales")
result = spark.table("sales").join(dim_df.filter(...), "event_date")

# ❌ DPP не сработает: join по product_id, а партиционирование по event_date
fact_df.write.partitionBy("event_date").saveAsTable("sales")
result = spark.table("sales").join(products.filter(...), "product_id")
# Spark прочтёт все партиции event_date, потому что product_id не является
# partition-колонкой - нет файловой структуры для отсечения

# ❌ DPP не сработает: таблица вообще не партиционирована
fact_df.write.saveAsTable("sales_no_partitioning")
result = spark.table("sales_no_partitioning").join(dim_df.filter(...), "event_date")

# ✅ DPP сработает: join по event_date = partition-колонка
spark.table("sales").join(dim_dates.filter(...), "event_date")

# ❌ DPP не сработает: join по выражению, а не по колонке напрямую
from pyspark.sql.functions import date_trunc
spark.table("sales").join(
    dim_dates.filter(...),
    date_trunc("month", col("event_date")) == date_trunc("month", col("date"))
    # Вычисляемое выражение - Catalyst не может свести его к partition-колонке
)

# ❌ DPP не сработает: join по transformed-ключу
spark.table("sales").join(
    dim_dates.filter(...),
    col("event_date").cast("string") == col("promo_date")
    # Cast создаёт новое выражение, не совпадающее с partition-колонкой
)

# ✅ DPP сработает: promotions < autoBroadcastJoinThreshold → BHJ автоматически
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "100m")
# promotions = 200 КБ < 100 МБ → Catalyst выберет BHJ → DPP активируется

# ✅ DPP сработает: явный broadcast-хинт
from pyspark.sql.functions import broadcast
spark.table("sales").join(broadcast(promotions.filter(...)), "event_date")

# ❌ DPP не сработает: BHJ отключён
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "-1")
# Spark выберет SortMergeJoin даже для маленьких таблиц → DPP отключится

# ❌ DPP не сработает: принудительный SMJ-хинт
spark.table("sales").hint("merge").join(promotions.filter(...), "event_date")

# Оба параметра важны вместе
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", str(100 * 1024 * 1024))  # 100 МБ
spark.conf.set("spark.sql.optimizer.dynamicPartitionPruning.enabled", "true")   # по умолчанию true с Spark 3.0

# Дополнительная настройка: использовать статистику для оценки выгодности DPP
spark.conf.set("spark.sql.optimizer.dynamicPartitionPruning.useStats", "true")
# Если статистики нет, Catalyst может ошибочно решить, что DPP невыгоден
# Установка в true заставляет использовать доступные table stats

spark = SparkSession.builder \
    .appName("DPP-Production") \
    .config("spark.sql.optimizer.dynamicPartitionPruning.enabled", "true") \
    # ↑ главный включатель DPP (по умолчанию true начиная с Spark 3.0)

    .config("spark.sql.adaptive.enabled", "true") \
    # ↑ AQE улучшает DPP: в runtime может переключить SMJ → BHJ
    # и тогда DPP тоже активируется там, где статически не мог

    .config("spark.sql.autoBroadcastJoinThreshold", str(100 * 1024 * 1024)) \
    # ↑ порог, до которого dimension автоматически broadcastится
    # По умолчанию 10 МБ - часто слишком мало для реальных dimensions

    .getOrCreate()

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

spark = SparkSession.builder \
    .appName("DPP-Demo") \
    .config("spark.sql.optimizer.dynamicPartitionPruning.enabled", "true") \
    .config("spark.sql.adaptive.enabled", "true") \
    .config("spark.sql.autoBroadcastJoinThreshold", str(50 * 1024 * 1024)) \
    .config("spark.sql.warehouse.dir", "/tmp/spark-warehouse") \
    .enableHiveSupport() \
    .getOrCreate()

# ─── 1. Dimension-таблица: регионы (маленькая, 4 строки) ────────────────────
dim_regions = spark.createDataFrame(
    [
        (1, "North",  "Russia",  "active"),
        (2, "South",  "Russia",  "active"),
        (3, "East",   "China",   "active"),
        (4, "West",   "Germany", "inactive"),
    ],
    schema=StructType([
        StructField("region_id",   IntegerType(), False),
        StructField("region_name", StringType(),  True),
        StructField("country",     StringType(),  True),
        StructField("status",      StringType(),  True),
    ])
)

# ─── 2. Fact-таблица: продажи (большая, партиционирована по region_id) ───────
# Генерируем 10 млн строк с равномерным распределением по 4 регионам
fact_df = (
    spark.range(1, 10_000_001)
    .withColumn("region_id",  ((F.col("id") % 4) + 1).cast("int"))
    .withColumn("amount",     (F.rand(seed=42) * 10000).cast("double"))
    .withColumn("product_id", ((F.col("id") % 100) + 1).cast("int"))
    .withColumn("event_date", (
        F.date_add(F.lit(datetime.date(2025, 1, 1)), (F.col("id") % 365).cast("int"))
    ))
)

# Сохраняем с партиционированием по region_id - ключ для DPP
fact_df.write \
    .mode("overwrite") \
    .partitionBy("region_id") \
    .saveAsTable("fact_sales")
    # Результат: 4 папки region_id=1, region_id=2, region_id=3, region_id=4

spark.sql("ANALYZE TABLE fact_sales COMPUTE STATISTICS FOR ALL COLUMNS")
# Собираем статистику - Catalyst будет принимать лучшие решения

# Фильтруем dimension: оставляем только АКТИВНЫЕ регионы в России
# region_id=1 (North, Russia, active) и region_id=2 (South, Russia, active)
active_ru_regions = dim_regions.filter(
    (F.col("country") == "Russia") & (F.col("status") == "active")
)
# active_ru_regions содержит region_id IN (1, 2)

# Join с fact по partition-колонке region_id
result = (
    spark.table("fact_sales")
    .join(active_ru_regions, "region_id")            # join по partition-колонке
    .groupBy("region_name", "country")
    .agg(
        F.sum("amount").alias("total_amount"),
        F.count("*").alias("order_count"),
        F.avg("amount").alias("avg_amount"),
    )
)

result.explain(mode="formatted")

== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- HashAggregate(keys=[region_name#7, country#8], ...)
   +- Exchange hashpartitioning(region_name#7, country#8, 200), ...)
      +- HashAggregate(keys=[region_name#7, country#8], ...)
         +- Project [region_name#7, country#8, amount#3]
            +- BroadcastHashJoin [region_id#2], [region_id#0], Inner, ...
               :- FileScan parquet default.fact_sales [region_id#2, amount#3, ...]
               :    PartitionFilters: [isnotnull(region_id#2),
               :                       dynamicpruning#42 [region_id#2]]
               :              ^^^^^^^^^^^^^^^^^^^^^^^^^^
               :              ВОТ ОН - МАРКЕР УСПЕШНОГО DPP
               :    ReadSchema: struct<amount:double,product_id:int,event_date:date>
               +- BroadcastExchange HashedRelationBroadcastMode(...), ...
                  +- Filter ((country#8 = Russia) AND (status#9 = active))
                     +- Scan OneRowRelation [region_id#0, region_name#7, ...]
                             (dimension data - сканируется один раз)

== Subqueries ==
Subquery:0 Hosting operator id = 1 for dynamicpruning#42
+- BroadcastExchange HashedRelationBroadcastMode(List(cast(input[0,int,false]...))
   +- Filter ((country#8 = Russia) AND (status#9 = active))
      +- Scan OneRowRelation [region_id#0, ...]

# Явно отключаем broadcast для этого запроса
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "-1")

result_no_dpp = (
    spark.table("fact_sales")
    .join(active_ru_regions, "region_id")
    .groupBy("region_name", "country")
    .agg(F.sum("amount").alias("total_amount"))
)

result_no_dpp.explain(mode="formatted")

+- SortMergeJoin [region_id#2], [region_id#0], Inner
   :- Sort [region_id#2 ASC NULLS FIRST], ...
   :  +- Exchange hashpartitioning(region_id#2, 200), ...
   :     +- FileScan parquet default.fact_sales [region_id#2, amount#3, ...]
   :          PartitionFilters: [isnotnull(region_id#2)]
   :                            ^^^^^^^^^^^^^^^^^^^^^^^^^^
   :                            НЕТ dynamicpruning# - DPP не сработал!
   :          ReadSchema: struct<amount:double,...>
   +- Sort [region_id#0 ASC NULLS FIRST], ...
      +- Exchange hashpartitioning(region_id#0, 200), ...
         +- Filter ((country#8 = Russia) AND (status#9 = active))
            +- Scan OneRowRelation [region_id#0, ...]

# Возвращаем broadcast обратно
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", str(50 * 1024 * 1024))

# join по product_id - НЕ partition-колонка
dim_products = spark.createDataFrame(
    [(1, "electronics"), (2, "clothing"), (3, "food")],
    ["product_id", "category"]
)

# fact_sales партиционирован по region_id, но join идёт по product_id
result_wrong_join = (
    spark.table("fact_sales")
    .join(broadcast(dim_products.filter(F.col("category") == "electronics")), "product_id")
    .groupBy("category")
    .agg(F.sum("amount").alias("total"))
)

result_wrong_join.explain(mode="formatted")
# В FileScan будет: PartitionFilters: [isnotnull(region_id#2)]
# НЕТ dynamicpruning# - BHJ есть, но join не по partition-колонке → DPP не сработал

# Без DPP: читаем все 4 партиции
spark.conf.set("spark.sql.optimizer.dynamicPartitionPruning.enabled", "false")
result.write.mode("overwrite").format("noop").save()  # запустить без реальной записи
# Смотрим в Spark UI: size of files read ≈ весь объём fact_sales

# С DPP: читаем только 2 партиции из 4
spark.conf.set("spark.sql.optimizer.dynamicPartitionPruning.enabled", "true")
result.write.mode("overwrite").format("noop").save()
# Смотрим в Spark UI: size of files read ≈ 50% от объёма (только region_id=1,2)

# DPP + Delta Lake
fact_df.write \
    .format("delta") \
    .mode("overwrite") \
    .partitionBy("region_id") \
    .saveAsTable("fact_sales_delta")

# Дополнительно: OPTIMIZE + ZORDER улучшает Data Skipping внутри файлов
spark.sql("OPTIMIZE fact_sales_delta ZORDER BY (event_date, product_id)")
# После OPTIMIZE: крупные Parquet файлы с хорошей статистикой → ещё меньше чтения

# Iceberg с hidden partitioning
# Физически партиционируем по месяцам, но join по event_date работает для DPP
spark.sql("""
    CREATE TABLE iceberg.sales (
        id BIGINT, region_id INT, amount DOUBLE, event_date DATE
    ) USING iceberg
    PARTITIONED BY (months(event_date), region_id)
""")
# DPP будет работать при join по event_date или region_id

# Перед выполнением
result.explain(mode="formatted")
# AdaptiveSparkPlan isFinalPlan=false  ← статический план

# Закешировать и посмотреть финальный план
result.cache()
result.count()  # trigger execution
result.explain(mode="formatted")
# AdaptiveSparkPlan isFinalPlan=true   ← финальный план после AQE-переключений
# Здесь dynamicpruning# уже видно если DPP сработал

# ❌ Партиционирование по user_id с кардинальностью 100 миллионов
fact_df.write.partitionBy("user_id").saveAsTable("fact_events")
# Результат: 100 млн папок в хранилище → directory listing занимает минуты
# Object Storage (S3, MinIO) имеет лимиты на list-операции

# ✅ Партиционирование по колонке с разумной кардинальностью
fact_df.write.partitionBy("event_date", "region_id").saveAsTable("fact_events")
# event_date: ~1000 уникальных дат за 3 года (разумно)
# region_id: ~100 регионов (разумно)
# Вместе: ~100 000 партиций - это уже на грани, но управляемо

def check_dpp(df):
    """Проверить, сработал ли DPP в физическом плане."""
    plan_str = df._sc._jvm.PythonSQLUtils.explainString(
        df._jdf.queryExecution(), "formatted"
    )
    if "dynamicpruning#" in plan_str:
        print("✅ DPP активирован")
    else:
        print("❌ DPP НЕ сработал - проверьте условия:")
        print("   1. Fact-таблица партиционирована по join-ключу?")
        print("   2. Join по partition-колонке (не по выражению)?")
        print("   3. Тип join совместим (INNER, LEFT SEMI, RIGHT OUTER)?")
        print("   4. Dimension достаточно мала для BroadcastHashJoin?")
        print("   5. spark.sql.optimizer.dynamicPartitionPruning.enabled = true?")

check_dpp(result)

from pyspark.sql.functions import udf

@udf("string")
def normalize_region(r):
    return r.strip().lower()

# ❌ UDF в join-условии: Catalyst не может выразить это через partition-фильтр
result = spark.table("fact_sales").join(
    dim_regions,
    normalize_region(col("region_name")) == normalize_region(col("dim_region_name"))
)
# DPP не сработает: UDF - чёрный ящик для Catalyst, нет связи с partition-колонкой

# ✅ Нормализация отдельно, до join
dim_clean = dim_regions.withColumn("region_id_clean", col("region_id"))
spark.table("fact_sales").join(dim_clean.filter(...), col("region_id") == col("region_id_clean"))

# ❌ Subquery возвращает не partition-ключ, а вычисляемое значение
subquery = spark.sql("""
    SELECT DISTINCT month(promo_date) AS promo_month FROM promotions WHERE campaign='BF'
""")
result = spark.table("fact_sales").join(
    subquery,
    F.month("event_date") == F.col("promo_month")
)
# DPP не сработает: `month(event_date)` - это выражение, не partition-колонка event_date

# ✅ Оставаться на уровне partition-ключа
promo_dates = spark.sql("SELECT DISTINCT promo_date FROM promotions WHERE campaign='BF'")
result = spark.table("fact_sales").join(promo_dates, "event_date")
# join по event_date = partition-колонка → DPP сработает

spark.conf.set("spark.sql.optimizer.dynamicPartitionPruning.enabled", "true")  # уже true по умолчанию
spark.conf.set("spark.sql.adaptive.enabled", "true")                           # уже true в Spark 3.2+
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", str(100 * 1024 * 1024)) # увеличить с 10МБ до 100МБ
spark.conf.set("spark.sql.optimizer.dynamicPartitionPruning.useStats", "true") # использовать статистику

# Таблицы
orders    = spark.table("orders")    # 500 ГБ, партиционирована по order_date
stores    = spark.table("stores")    # 2 МБ, 300 магазинов
products  = spark.table("products")  # 150 МБ, 200 000 товаров

# Запрос
result = (
    orders
    .join(stores.filter(col("country") == "RU"), col("store_id") == col("id"))
    .join(products.filter(col("category") == "electronics"), "product_id")
    .filter(col("order_date") >= "2026-01-01")
    .groupBy("country", "category")
    .agg(F.sum("amount").alias("total"))
)