Shuffle Partitions: как выбрать spark.sql.shuffle.partitions вручную

Почему Shuffle - самая дорогая операция в Spark¶

Spark - распределённый вычислительный движок. Данные разбиты на партиции и распределены по узлам кластера. Пока трансформации "узкие" (narrow) - filter, map, withColumn - каждая партиция обрабатывается независимо, без взаимодействия с другими узлами. Это быстро, потому что нет сетевого обмена.

Но как только вы применяете широкую (wide) трансформацию - groupBy, join, distinct, orderBy, window - Spark вынужден перераспределить данные между узлами. Строки, которые должны попасть в одну группу агрегации или одну сторону JOIN, могут лежать на разных машинах. Их нужно "собрать вместе". Этот процесс и называется Shuffle.

Shuffle дорог по трём причинам:

1. Сетевой ввод-вывод (Network I/O): данные физически перемещаются по сети между executor'ами. Даже в быстрой сети это порядки медленнее, чем чтение из оперативной памяти. Для больших датасетов передача гигабайт или терабайт данных занимает значительное время.

2. Запись на диск (Disk I/O): в процессе Shuffle Spark записывает промежуточные shuffle-файлы на локальные диски executor'ов (shuffle write), а затем другие executor'ы их читают (shuffle read). Даже на SSD это значительно медленнее, чем работа в памяти. При нехватке памяти данные дополнительно сбрасываются на диск (disk spill) - это второй уровень замедления.

3. Перераспределение (Redistribution): каждая строка данных должна быть хешированна, назначена определённой партиции и записана в соответствующий файл. Это вычислительная работа, пропорциональная размеру данных и количеству партиций.

Именно поэтому количество shuffle-партиций - один из ключевых параметров настройки производительности Spark. Неправильное значение может сделать задачу в 10–50 раз медленнее оптимальной.

Анатомия Shuffle: Exchange оператор и stage boundary¶

Exchange оператор в физическом плане¶

Когда Spark компилирует логический запрос в физический план выполнения, оптимизатор Catalyst определяет, где нужен shuffle. В этих местах в план вставляется специальный оператор - Exchange.

Посмотрим на физический план простого GROUP BY:

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, sum as spark_sum, count

spark = SparkSession.builder \
    .appName("Shuffle Demo") \
    .config("spark.sql.shuffle.partitions", "4") \
    .getOrCreate()

df = spark.createDataFrame([
    ("2024-01-01", "search", 100),
    ("2024-01-01", "display", 200),
    ("2024-01-02", "search", 150),
    ("2024-01-02", "email", 80),
], ["date", "channel", "clicks"])

aggregated = df.groupBy("channel").agg(
    spark_sum("clicks").alias("total_clicks"),
    count("*").alias("events")
)

aggregated.explain("formatted")

В выводе explain вы увидите что-то похожее на:

== Physical Plan ==
AdaptiveSparkPlan (1)
+- HashAggregate (2)
   +- Exchange (3)          ← ВОТ ОН, Exchange оператор
      hashPartitioning(channel, 4)
      +- HashAggregate (4)
         +- Scan (5)

Оператор Exchange - это граница между двумя stage'ами. Всё, что выше Exchange (в данном примере - HashAggregate (4) и Scan (5)), выполняется в Stage 0. Всё, что ниже (финальный HashAggregate (2)), - в Stage 1. Stage 1 не может начаться, пока все задачи Stage 0 не завершили запись своих shuffle-файлов.

Stage boundary как барьер синхронизации¶

Граница между stage'ами - это точка полной синхронизации. Даже если 199 из 200 задач первого stage завершились за 1 минуту, а одна зависла на 10 минут, следующий stage не начнётся до завершения этой одной задачи. Именно поэтому так важно избегать перекоса (skew) в партициях.

Число 4 в hashPartitioning(channel, 4)¶

Значение 4 - это и есть spark.sql.shuffle.partitions. Exchange создаёт ровно столько выходных партиций, сколько задано этим параметром. В Stage 1 будет запущено ровно 4 задачи - по одной на каждую выходную партицию. Вот почему этот параметр напрямую управляет параллелизмом после shuffle.

Визуализация жизненного цикла Shuffle¶

Диаграмма иллюстрирует полный цикл shuffle с тремя executor'ами и тремя выходными партициями. Stage 0 (Map Phase): каждый executor обрабатывает свою входную партицию и хеширует каждую строку, определяя, в какую выходную партицию она попадёт. В примере с groupBy("channel") все строки с channel="search" получат одинаковый хеш и окажутся в одной партиции. Shuffle Write: результат записывается на локальный диск executor'а в виде отдельных файлов (slice'ов) для каждой выходной партиции. Executor 1 создаёт 3 slice'а - по одному для каждого будущего reducer'а. Network I/O: reducer'ы читают свои slice'ы со всех узлов по сети. Stage 1 (Reduce Phase): каждый reducer получает все данные для своей партиции и выполняет итоговую агрегацию.

Что такое spark.sql.shuffle.partitions¶

spark.sql.shuffle.partitions - это конфигурационный параметр, определяющий количество выходных партиций после любой операции, требующей shuffle:

• groupBy().agg()
• join() (SortMergeJoin, ShuffleHashJoin)
• distinct()
• orderBy() / sort()
• window() (функции окна)
• DataFrame.repartition(n) с числовым аргументом

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

spark.conf.get("spark.sql.shuffle.partitions")

Важно понимать: этот параметр не влияет на количество партиций при чтении данных (это управляется другими параметрами: spark.sql.files.maxPartitionBytes и spark.default.parallelism). Он управляет только тем, что происходит после Exchange оператора.

Почему 200 - это дефолт¶

Значение 200 появилось в ранних версиях Spark как компромисс для типичных кластеров того времени: несколько десятков узлов, датасеты порядка гигабайт. Для задач такого масштаба 200 партиций давали разумный баланс.

Сегодня реальность другая: данные исчисляются терабайтами, кластеры - тысячами ядер, а данные могут быть как мегабайтовыми (небольшие справочники) так и многотерабайтными (event logs). Дефолт в 200 почти никогда не является оптимальным:

• Для датасета в 100 МБ: 200 партиций по 512 КБ каждая - тысячи мелких задач, огромный overhead на планирование
• Для датасета в 10 ТБ: 200 партиций по 51 ГБ каждая - OOM (Out of Memory) у каждого executor'а

HashPartitioning: как данные распределяются по партициям¶

При groupBy или join Spark использует HashPartitioning: хеш от значения ключевых столбцов определяет номер партиции.

partition_id = hash(row["channel"]) % num_shuffle_partitions

Это детерминированная функция: одинаковое значение ключа всегда попадает в одну и ту же партицию. Благодаря этому все строки с channel="search" гарантированно оказываются у одного reducer'а - это необходимо для корректной агрегации.

RoundRobinPartitioning: для операций без ключа¶

Когда нет явного ключа распределения (например, DataFrame.repartition(n) без указания колонок, или DISTRIBUTE BY), Spark использует RoundRobinPartitioning - строки распределяются по партициям по кругу, обеспечивая равномерное распределение по количеству строк. Однако при этом строки с одинаковым значением могут оказаться в разных партициях - такое распределение подходит только для балансировки нагрузки, а не для агрегаций.

Связь: Partition = Task = Parallelism¶

Количество shuffle-партиций напрямую определяет степень параллелизма после shuffle. Каждая партиция - это одна задача (task) в Stage 1, и каждая задача выполняется на одном ядре одного executor'а.

Количество shuffle-партиций = Количество задач в Stage 1

Это означает:

• Если у вас 100 ядер в кластере и 50 shuffle-партиций - половина ядер простаивает
• Если у вас 100 ядер и 10 000 shuffle-партиций - Spark последовательно запускает 100 задач в каждой "волне", итого 100 волн × overhead на планирование каждой задачи
• Если у вас 100 ядер и 200 shuffle-партиций - 2 волны, каждая на 100 ядер - близко к оптимуму

Вот почему говорят, что количество партиций должно быть кратным числу ядер в кластере (или несколько кратным - 2–3×, чтобы обеспечить динамическую балансировку: когда одна задача завершится раньше, свободное ядро сразу берёт следующую).

Проблемы неправильного количества партиций¶

Проблема: слишком мало партиций¶

Когда партиций мало, каждая задача получает огромный объём данных. Рассмотрим конкретный пример: у вас 10 ГБ данных после shuffle и только 10 партиций - каждая задача получает по 1 ГБ.

Disk Spill (сброс на диск): Spark выделяет executor'у, например, 4 ГБ памяти для обработки данных. Если задача получила 1 ГБ данных, которые нужно отсортировать или сагрегировать - это может вписаться в память. Но если joined датасет раздувается до 8 ГБ в памяти (без сжатия) - Spark начинает сбрасывать часть данных на диск. Это называется Memory Spill (→ Disk Spill). Каждое обращение к диску вместо памяти - замедление в 10–100 раз.

OOM (Out of Memory): если memory spill невозможен или данные слишком велики даже для disk spill, executor падает с java.lang.OutOfMemoryError. Spark попытается перезапустить задачу на другом executor'е - но это время, которое уже потеряно.

Straggler задача: при малом числе партиций одна "особенно тяжёлая" партиция (hot key в агрегации) становится straggler'ом - задачей, которая выполняется значительно дольше остальных. Пока она работает, весь следующий stage ждёт.

Проблема: слишком много партиций¶

Когда партиций слишком много, каждая задача получает ничтожно малый объём данных - возможно, несколько килобайт.

Task scheduling overhead: каждая задача в Spark - это объект с метаданными, которые нужно сериализовать, отправить executor'у, запустить JVM-контекст (для PySpark - ещё и Python-процесс), и после завершения - отправить результат обратно драйверу. Этот overhead составляет несколько миллисекунд на задачу. При 10 000 задач это уже десятки секунд накладных расходов на планирование - и это время не выполняет никакой полезной работы.

Small Files Problem: количество выходных файлов напрямую связано с количеством shuffle-партиций. Если вы пишете результат в S3/MinIO/HDFS, каждая партиция создаёт отдельный файл. 10 000 партиций → 10 000 файлов. Чтение такого датасета в будущем потребует 10 000 обращений к хранилищу только для листинга файлов, плюс overhead на открытие каждого файла отдельно. В S3 листинг тысяч файлов сам по себе занимает секунды.

Фазы Shuffle в деталях¶

Shuffle Write (Map-сторона)¶

В процессе shuffle write каждый mapper (задача Stage 0) проходит по всем своим строкам и для каждой строки:

1. Вычисляет хеш ключевых колонок
2. Определяет целевую партицию: hash % num_shuffle_partitions
3. Записывает строку в соответствующий буфер
4. Когда буфер заполняется - сбрасывает его на диск

Итого, один mapper создаёт N файлов (по одному на каждую выходную партицию). Если у вас 500 mapper'ов и 200 shuffle-партиций, это 500 × 200 = 100 000 промежуточных shuffle-файлов на дисках executor'ов. Именно это называется "shuffle file explosion" - взрывной рост количества временных файлов.

Для оптимизации Spark использует Sort-Based Shuffle (по умолчанию с Spark 1.4): данные сортируются по partition id перед записью, что позволяет записать один файл данных и один индексный файл для каждого mapper'а, вместо N отдельных файлов. Это значительно снижает нагрузку на файловую систему.

Shuffle Read (Reduce-сторона)¶

После завершения всех mapper'ов каждый reducer (задача Stage 1) читает свой slice данных с каждого mapper'а. Для reducer'а №0 это означает: подключиться к 500 mapper'ам, скачать по одному slice'у с каждого, объединить всё в один поток, опционально отсортировать (для SortMergeJoin или ORDER BY) и выполнить агрегацию.

Сетевые запросы при shuffle read - это сотни или тысячи TCP-соединений. Spark делает это параллельно, но управление таким количеством соединений само по себе создаёт overhead.

Disk Spill: когда памяти не хватает¶

Disk Spill - один из главных индикаторов того, что shuffle-партиции слишком большие. Разберёмся в деталях.

Memory Spill - объём данных, который был вытеснен из JVM heap на промежуточное хранение. Это может быть запись в памяти (OffHeap) или на диск.

Disk Spill - объём данных, физически записанных на локальный диск executor'а из-за нехватки памяти. Это самое медленное из возможных: данные записываются на диск, потом читаются обратно.

spark.conf.set("spark.sql.shuffle.partitions", "5")
spark.conf.set("spark.executor.memory", "2g")

large_df = spark.range(0, 10_000_000).selectExpr(
    "id",
    "cast(id % 100 as string) as category",
    "rand() * 1000 as amount",
    "array(rand(), rand(), rand(), rand()) as features"
)

result = large_df.groupBy("category").agg(
    spark_sum("amount").alias("total"),
    count("*").alias("cnt")
)

result.explain("formatted")
result.count()

После выполнения откройте Spark UI → вкладку Stages → нажмите на второй stage (после shuffle). В таблице Summary Metrics найдите строки:

• Shuffle Read Size / Records - сколько данных прочитал каждый reducer
• Spill (Memory) - объём данных, временно вытесненных в промежуточное хранение
• Spill (Disk) - объём, записанный на диск из-за нехватки памяти

Если Spill (Disk) > 0 - это прямой сигнал: партиции слишком большие, нужно увеличить spark.sql.shuffle.partitions.

Формула расчёта оптимального числа партиций¶

Откуда берутся числа¶

Практическое правило, проверенное на production-кластерах: одна shuffle-партиция должна занимать 128–512 МБ в памяти executor'а в несжатом виде (или 100–200 МБ в виде сжатых данных, которые пишутся на диск/S3).

Почему именно этот диапазон?

• Меньше 128 МБ: накладные расходы на планирование задачи начинают доминировать над полезной работой
• Больше 512 МБ: риск disk spill при агрегации (реальный объём данных в памяти может в 2–5 раз превышать размер на диске из-за деserialization и накладных расходов структур данных JVM)

Математика расчёта¶

Полная формула включает несколько шагов:

Шаг 1: Определите объём shuffle write

Это метрика "Shuffle Write" в Spark UI на вкладке Stages для stage'а, который выполняет широкую трансформацию. Обозначим её как S (в МБ).

Шаг 2: Оцените коэффициент раздувания (expansion factor)

Данные на диске сжаты (Snappy/Zstd для Parquet), но в памяти executor'а они несжатые. Типичный коэффициент раздувания: 2–5× (зависит от типов данных: числовые данные - ближе к 2×, строковые JSON - ближе к 5×).

Шаг 3: Рассчитайте базовое количество партиций

N_base = (S × expansion_factor) / target_partition_size

Где target_partition_size = 200–400 МБ (целевой размер партиции в памяти).

Шаг 4: Округлите до числа, кратного числу ядер кластера

N_final = ceil(N_base / total_cores) × total_cores × parallelism_factor

Где parallelism_factor = 2–3 (для лучшей балансировки нагрузки).

Практический пример расчёта¶

def calculate_shuffle_partitions(
    shuffle_write_mb: float,
    total_cluster_cores: int,
    expansion_factor: float = 3.0,
    target_partition_mb: int = 256,
    parallelism_factor: int = 2,
) -> int:
    """
    Рассчитывает оптимальное количество shuffle-партиций.

    Args:
        shuffle_write_mb: Объём Shuffle Write в МБ (из Spark UI)
        total_cluster_cores: Общее число ядер в кластере
        expansion_factor: Коэффициент раздувания данных (диск → память)
        target_partition_mb: Целевой размер партиции в памяти (МБ)
        parallelism_factor: Множитель для обеспечения балансировки

    Returns:
        Рекомендуемое значение spark.sql.shuffle.partitions
    """
    data_in_memory_mb = shuffle_write_mb * expansion_factor

    n_base = data_in_memory_mb / target_partition_mb

    n_rounded_to_cores = (
        (int(n_base) + total_cluster_cores - 1) // total_cluster_cores
    ) * total_cluster_cores

    n_final = n_rounded_to_cores * parallelism_factor

    n_final = max(n_final, total_cluster_cores)

    print(f"Shuffle Write:        {shuffle_write_mb:.0f} MB")
    print(f"Data in memory (~):   {data_in_memory_mb:.0f} MB "
          f"(expansion ×{expansion_factor})")
    print(f"Base partitions:      {n_base:.1f}")
    print(f"Rounded to cores:     {n_rounded_to_cores}")
    print(f"Final (×{parallelism_factor} factor): {n_final}")
    print(f"Expected task size:   "
          f"{data_in_memory_mb / n_final:.0f} MB/task")

    return n_final

calculate_shuffle_partitions(
    shuffle_write_mb=80_000,
    total_cluster_cores=40,
    expansion_factor=3.0,
    target_partition_mb=256,
    parallelism_factor=2,
)

Вывод для сценария "80 ГБ shuffle write, 40 ядер":

Shuffle Write:        80000 MB
Data in memory (~):   240000 MB (expansion ×3.0)
Base partitions:      937.5
Rounded to cores:     960
Final (×2 factor):    1920
Expected task size:   125 MB/task

Итоговый ответ: spark.sql.shuffle.partitions = 1920. Это кратно 40 ядрам и даёт разумный размер задачи ~125 МБ.

Пример для трёх сценариев¶

scenarios = [
    ("Малый датасет (справочники)", 500, 40),
    ("Средний датасет (дневные агрегаты)", 50_000, 40),
    ("Большой датасет (терабайтный clickstream)", 2_000_000, 40),
]

for name, shuffle_mb, cores in scenarios:
    print(f"\n{'='*50}")
    print(f"Сценарий: {name}")
    n = calculate_shuffle_partitions(shuffle_mb, cores)
    print(f"→ Рекомендуется: spark.sql.shuffle.partitions = {n}")

Вывод:

==================================================
Сценарий: Малый датасет (справочники)
Shuffle Write:        500 MB
Data in memory (~):   1500 MB (expansion ×3.0)
Base partitions:      5.9
Rounded to cores:     40
Final (×2 factor):    80
Expected task size:   19 MB/task
→ Рекомендуется: spark.sql.shuffle.partitions = 80

==================================================
Сценарий: Средний датасет (дневные агрегаты)
Shuffle Write:        50000 MB
Data in memory (~):   150000 MB (expansion ×3.0)
Base partitions:      585.9
Rounded to cores:     600
Final (×2 factor):    1200
Expected task size:   125 MB/task
→ Рекомендуется: spark.sql.shuffle.partitions = 1200

==================================================
Сценарий: Большой датасет (терабайтный clickstream)
Shuffle Write:        2000000 MB
Data in memory (~):   6000000 MB (expansion ×3.0)
Base partitions:      23437.5
Rounded to cores:     23440
Final (×2 factor):    46880
Expected task size:   128 MB/task
→ Рекомендуется: spark.sql.shuffle.partitions = 46880

Важное замечание по малому датасету: n_base = 5.9, но мы не ставим 5 партиций. Мы поднимаем до минимума - числа, кратного числу ядер (40). Задачи будут маленькими (19 МБ), но хотя бы кластер будет утилизирован полностью. Реальный оптимум для малых датасетов - использовать Broadcast Join вместо shuffle (см. раздел про broadcast).

Анализ метрик в Spark UI¶

Где найти объём Shuffle Write¶

В Spark UI откройте вкладку Jobs, выберите нужный Job, затем нажмите на Stage, который содержит вашу широкую трансформацию.

На странице Stage найдите секцию Summary Metrics - она показывает агрегаты по всем задачам этого stage'а. Ключевые метрики:

Shuffle Write Size, Total - суммируйте значения по всем задачам. Это и есть S в вашей формуле.

Признаки неправильного размера партиций¶

Слишком мало партиций:

Task Duration: Min=45s, Median=3min 20s, Max=18min
Shuffle Read Size: Min=10MB, Median=800MB, Max=4.5GB
Spill (Disk): Min=0, Median=2.1GB, Max=8.7GB

Огромный Max Duration, колоссальный Spill (Disk) - классический признак. Partition size 4.5 ГБ при executor memory 8 ГБ - неминуемо приведёт к spill.

Слишком много партиций:

Task Duration: Min=10ms, Median=25ms, Max=200ms
Shuffle Read Size: Min=0B, Median=512KB, Max=2MB

Задачи выполняются за миллисекунды - overhead на планирование, запуск и завершение задачи превышает полезную работу.

Оптимальный размер:

Task Duration: Min=15s, Median=45s, Max=75s
Shuffle Read Size: Min=80MB, Median=150MB, Max=320MB
Spill (Disk): 0B (нет spill)

Задачи выполняются десятки секунд, нет spill, Max/Median < 2× - отличный результат.

Как получить Shuffle Write из кода¶

from pyspark.sql import SparkSession

def get_shuffle_write_mb(spark: SparkSession, stage_id: int) -> float:
    """
    Получает объём Shuffle Write для конкретного stage из SparkContext.

    Внимание: данные доступны только после завершения stage.
    """
    sc = spark.sparkContext
    status = sc.statusTracker()

    stage_info = status.getStageInfo(stage_id)
    if stage_info is None:
        print(f"Stage {stage_id} not found or not completed")
        return 0.0

    total_shuffle_write = stage_info.shuffleWriteBytes()

    return total_shuffle_write / (1024 ** 2)


result = df.groupBy("channel").agg(spark_sum("clicks"))
result.collect()

for stage_id in range(0, 5):
    mb = get_shuffle_write_mb(spark, stage_id)
    if mb > 0:
        print(f"Stage {stage_id}: Shuffle Write = {mb:.1f} MB")

Это даёт программный доступ к той же метрике, что видна в Spark UI, что позволяет автоматически рассчитывать и устанавливать оптимальное число партиций для следующего запуска.

Специфика PySpark: JVM vs Python¶

PySpark - это мост между Python и JVM-процессом Apache Spark. Когда вы работаете с DataFrame API или SQL-запросами, данные остаются в JVM и сериализуются в Python только при явном вызове .collect(), .toPandas(), или при использовании Python UDF.

Python UDF и shuffle partitions¶

Python UDF (User-Defined Function) - самая дорогая операция в PySpark с точки зрения памяти. Каждая строка данных для UDF сериализуется из JVM в Python через механизм Py4J (сокет), обрабатывается Python-функцией, и результат сериализуется обратно в JVM.

Это означает:

• Каждая задача, обрабатывающая данные через Python UDF, держит в памяти две копии данных: одну в JVM, одну в Python-процессе
• Реальное потребление памяти на executor'е при использовании Python UDF в 2 раза выше расчётного

При использовании Python UDF стоит уменьшать целевой размер партиции до 64–128 МБ вместо 200–400 МБ, чтобы компенсировать двойное потребление памяти.

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType

@udf(StringType())
def python_udf_example(s: str) -> str:
    return s.upper()

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

df.withColumn("upper_channel", python_udf_example(col("channel"))) \
  .groupBy("upper_channel") \
  .agg(spark_sum("clicks")) \
  .show()

При наличии Python UDF в pipeline рекомендуется мониторить потребление памяти executor'ов в Spark UI (вкладка Executors → колонка Memory Used).

Pandas UDF (Vectorized UDF) и память¶

Pandas UDF (также называемые vectorized UDF или Arrow UDF) значительно эффективнее обычных Python UDF - они передают данные блоками (Arrow batches) через Apache Arrow, минуя построчную сериализацию Py4J.

Однако весь Arrow batch при этом оказывается в памяти Python-процесса целиком. Размер batch контролируется параметром:

spark.conf.set("spark.sql.execution.arrow.maxRecordsPerBatch", "10000")

При большом размере shuffle-партиций Arrow batch может занять значительный объём памяти Python-процесса. Для Pandas UDF с тяжёлыми вычислениями (ML-inference, сложная обработка текста) уменьшайте maxRecordsPerBatch.

AQE: Adaptive Query Execution¶

Начиная с Spark 3.0, существует механизм AQE (Adaptive Query Execution) - адаптивная оптимизация запросов, которая может автоматически корректировать план выполнения на основе статистик времени выполнения.

Включение AQE¶

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

В Spark 3.2+ AQE включён по умолчанию.

Coalesce Shuffle Partitions¶

Главная функция AQE для shuffle-партиций - Coalesce Shuffle Partitions: если несколько маленьких выходных партиций можно объединить без превышения целевого размера, AQE делает это автоматически.

Пример: вы задали spark.sql.shuffle.partitions = 2000, но после shuffle оказалось, что данные занимают всего 400 МБ. AQE автоматически объединит 2000 партиций в ~4 (по 100 МБ каждая), и Stage 1 запустит 4 задачи вместо 2000.

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

result = df.groupBy("channel").agg(spark_sum("clicks"))
result.explain("formatted")

В плане появится AQEShuffleRead вместо Exchange:

== Physical Plan ==
AdaptiveSparkPlan (1)
+- == Final Plan ==
   HashAggregate (2)
   +- AQEShuffleRead (3)       ← AQE скоалесцировал партиции
      coalesced_partition_specs: [(0, 0, 49), ...]
      +- ShuffleQueryStage (4)
         +- Exchange (5)
            hashPartitioning(channel, 2000)

coalesced_partition_specs: [(0, 0, 49), ...] означает, что AQE объединил исходные партиции 0–49 в одну задачу.

Параметр advisoryPartitionSizeInBytes¶

Это "советуемый" размер одной партиции для AQE. Если текущая партиция меньше этого значения, AQE попытается объединить её с соседними:

spark.conf.set(
    "spark.sql.adaptive.advisoryPartitionSizeInBytes",
    "134217728"
)

134217728 байт = 128 МБ. Рекомендуемое значение для production: 128–256 МБ.

Когда AQE недостаточно¶

AQE - мощный инструмент, но он имеет ограничения:

1. AQE не умеет разделять партиции (кроме skew join): если вы задали spark.sql.shuffle.partitions = 50, а данных оказалось 100 ГБ - AQE не увеличит количество партиций до нужного. Он может только уменьшать количество партиций (coalesce). Поэтому лучше задавать с запасом (большее число партиций), чтобы AQE мог уменьшить.

2. Сильные колебания объёма данных: если объём данных варьируется в 100× от запуска к запуску (например, исторический бэкфилл vs дневной инкремент), AQE с одним фиксированным advisory size может не справиться оптимально в обоих случаях.

3. Сложные многошаговые пайплайны: в DAG с десятками stage'ей AQE применяет оптимизацию для каждого stage независимо, что не всегда даёт глобально оптимальный результат.

4. Structured Streaming: AQE не поддерживается в streaming режиме (micro-batch). Для стриминга партиции нужно настраивать вручную для каждого micro-batch.

spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
spark.conf.set("spark.sql.adaptive.advisoryPartitionSizeInBytes", "134217728")
spark.conf.set("spark.sql.adaptive.coalescePartitions.initialPartitionNum", "400")

initialPartitionNum - начальное число партиций для AQE (аналог spark.sql.shuffle.partitions при включённом AQE). AQE начнёт с этого числа и уменьшит при необходимости.

Связь с проблемой мелких файлов¶

Почему shuffle partitions создают мелкие файлы¶

Когда вы записываете DataFrame в Parquet/Delta/Iceberg, Spark создаёт ровно столько выходных файлов, сколько партиций в записываемом DataFrame в момент записи. А количество партиций после последнего shuffle равно spark.sql.shuffle.partitions.

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

result = events_df \
    .groupBy("date", "channel") \
    .agg(spark_sum("clicks")) \
    .write \
    .partitionBy("date") \
    .parquet("/silver/events_agg")

Если date содержит 30 дней, а shuffle.partitions = 200, то каждый день получит до 200 файлов-фрагментов. Итого: 30 × 200 = 6 000 файлов вместо 30 (по одному на дату).

При следующем чтении этого датасета Spark откроет 6 000 файлов. В S3 это 6 000 API-вызовов только для чтения метаданных файлов (ListObjects + HeadObject) - это занимает секунды и стоит денег.

Как настроить shuffle.partitions для контроля выходных файлов¶

Простой подход: установить shuffle.partitions равным желаемому числу выходных файлов (с учётом партиционирования по дате):

output_days = 30
files_per_day = 4

spark.conf.set(
    "spark.sql.shuffle.partitions",
    str(output_days * files_per_day)
)

result = events_df \
    .groupBy("date", "channel") \
    .agg(spark_sum("clicks")) \
    .write \
    .partitionBy("date") \
    .parquet("/silver/events_agg")

Это создаст 4 файла на каждый день.

repartition vs coalesce для управления выходными файлами¶

Более гибкий подход - использовать явный repartition() перед записью, независимо от shuffle.partitions:

result = events_df \
    .groupBy("date", "channel") \
    .agg(spark_sum("clicks"))

files_per_date = result \
    .repartition(4, "date") \
    .write \
    .partitionBy("date") \
    .mode("overwrite") \
    .parquet("/silver/events_agg")

repartition(4, "date") создаёт 4 партиции, гарантируя, что все строки с одинаковым date окажутся в одной из 4 партиций. При записи partitionBy("date") это даст ровно 4 файла на каждую дату.

Разница между repartition и coalesce:

coalesce(n) объединяет существующие партиции без перемешивания данных. Это дёшево, но может дать неравномерные файлы, если изначальные партиции были неравномерными.

Применение настроек в коде¶

Установка в SparkSession при запуске¶

from pyspark.sql import SparkSession

spark = SparkSession.builder \
    .appName("Production ETL") \
    .config("spark.sql.shuffle.partitions", "800") \
    .config("spark.sql.adaptive.enabled", "true") \
    .config("spark.sql.adaptive.advisoryPartitionSizeInBytes", "134217728") \
    .getOrCreate()

Динамическое изменение в runtime¶

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

small_result = customers_df \
    .groupBy("segment") \
    .agg(count("*").alias("cnt"))
small_result.cache()
small_result.count()

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

large_result = events_df \
    .join(customers_df, "user_id") \
    .groupBy("segment", "date") \
    .agg(spark_sum("revenue"))

Динамическое изменение работает только для новых операций - уже запланированные DataFrame не пересчитываются. Это мощный приём для пайплайнов, где разные части имеют принципиально разный объём данных.

Автоматическая подстройка по объёму данных¶

def auto_shuffle_partitions(
    df,
    total_cluster_cores: int,
    target_partition_mb: int = 200,
) -> int:
    """
    Оценивает размер DataFrame и рассчитывает оптимальное число партиций.

    Использует размер DataFrame в сжатом виде как прокси для shuffle write.
    Менее точно, чем реальные метрики Spark UI, но применимо до первого запуска.
    """
    size_in_bytes = df.rdd.map(lambda row: len(str(row))).sum()
    size_in_mb = size_in_bytes / (1024 ** 2)

    n_partitions = max(
        total_cluster_cores,
        int(size_in_mb * 3 / target_partition_mb / total_cluster_cores)
        * total_cluster_cores
        * 2
    )

    print(f"Estimated size: {size_in_mb:.0f} MB → partitions: {n_partitions}")
    return n_partitions

Предупреждение: df.rdd.map(lambda row: len(str(row))).sum() - очень грубая оценка, которая сама вызывает Action и может быть медленной. В production лучше использовать реальные метрики Spark UI из предыдущего запуска или оценивать объём файлов на входе.

Лабораторная практика: тюнинг пайплайна по метрикам¶

Бизнес-кейс¶

Ежедневный ETL pipeline рассчитывает агрегаты по clickstream-данным за предыдущий день. Данные хранятся в Parquet на MinIO, объём - около 80 ГБ сжатых данных. Кластер: 10 воркеров × 4 ядра = 40 ядер. Исходный код оставлен с дефолтными настройками.

from pyspark.sql import SparkSession
from pyspark.sql.functions import (
    col, sum as spark_sum, count, avg, max as spark_max,
    date_format, to_date, hour
)

spark = SparkSession.builder \
    .appName("Clickstream Daily Aggregation") \
    .getOrCreate()

events_df = spark.read.parquet("s3a://datalake/clickstream/2024-01-15/")

result = events_df \
    .withColumn("event_hour", hour("event_ts")) \
    .groupBy("campaign_id", "channel", "date_part", "event_hour") \
    .agg(
        spark_sum("clicks").alias("total_clicks"),
        spark_sum("impressions").alias("total_impressions"),
        spark_sum("revenue").alias("total_revenue"),
        count("*").alias("event_count"),
        avg("session_duration").alias("avg_session"),
        spark_max("bid_price").alias("max_bid"),
    )

result.write \
    .partitionBy("date_part") \
    .mode("overwrite") \
    .parquet("s3a://datalake/aggregated/daily/")

Шаг 1: Запуск с дефолтными настройками (200 партиций)¶

После запуска откройте Spark UI. Обнаружим:

Stage 1 (HashAggregate после shuffle):
  Task count: 200
  Duration: Min=2s, Median=6min 45s, Max=18min
  Shuffle Read Size: Min=18MB, Median=1.2GB, Max=4.5GB
  Spill (Memory): 89.3 GB total
  Spill (Disk): 54.7 GB total

Диагноз: 200 партиций явно недостаточно для 80 ГБ данных. Медианный размер shuffle read 1.2 ГБ при типичном executor memory 8 ГБ приводит к massовому disk spill. Максимальная задача 18 минут - straggler.

Шаг 2: Расчёт оптимального числа партиций¶

Из Spark UI: Shuffle Write Total = 80 ГБ = 81 920 МБ (сжатые данные).

n_optimal = calculate_shuffle_partitions(
    shuffle_write_mb=81_920,
    total_cluster_cores=40,
    expansion_factor=3.0,
    target_partition_mb=256,
    parallelism_factor=2,
)
print(f"\nОптимальное значение: spark.sql.shuffle.partitions = {n_optimal}")

Shuffle Write:        81920 MB
Data in memory (~):   245760 MB (expansion ×3.0)
Base partitions:      960.0
Rounded to cores:     960
Final (×2 factor):    1920
Expected task size:   128 MB/task
→ Оптимальное значение: spark.sql.shuffle.partitions = 1920

Шаг 3: Применение рассчитанного значения¶

spark = SparkSession.builder \
    .appName("Clickstream Daily Aggregation - Tuned") \
    .config("spark.sql.shuffle.partitions", "1920") \
    .config("spark.sql.adaptive.enabled", "true") \
    .config("spark.sql.adaptive.advisoryPartitionSizeInBytes", "134217728") \
    .getOrCreate()

events_df = spark.read.parquet("s3a://datalake/clickstream/2024-01-15/")

result = events_df \
    .withColumn("event_hour", hour("event_ts")) \
    .groupBy("campaign_id", "channel", "date_part", "event_hour") \
    .agg(
        spark_sum("clicks").alias("total_clicks"),
        spark_sum("impressions").alias("total_impressions"),
        spark_sum("revenue").alias("total_revenue"),
        count("*").alias("event_count"),
        avg("session_duration").alias("avg_session"),
        spark_max("bid_price").alias("max_bid"),
    )

result.write \
    .partitionBy("date_part") \
    .mode("overwrite") \
    .parquet("s3a://datalake/aggregated/daily/")

Шаг 4: Результаты после тюнинга¶

Stage 1 (HashAggregate после shuffle):
  Task count: 1920 (AQE coalesced → ~480 effective tasks)
  Duration: Min=20s, Median=42s, Max=68s
  Shuffle Read Size: Min=95MB, Median=148MB, Max=310MB
  Spill (Memory): 0 B
  Spill (Disk): 0 B
  Total time: 6 min 15 sec  (было: 25 min)

Ускорение в 4× достигнуто за счёт:

• Полного устранения disk spill (54.7 ГБ → 0)
• Равномерного распределения нагрузки (Max/Median = 1.6× вместо 2.7×)
• AQE автоматически скоалесцировал 1920 партиций в ~480 эффективных

Конфигурация для анализа входных партиций¶

spark.sql.shuffle.partitions влияет только на shuffle output. Входные партиции при чтении файлов управляются другим параметром:

spark.conf.set(
    "spark.sql.files.maxPartitionBytes",
    str(128 * 1024 * 1024)
)

maxPartitionBytes = 128 МБ по умолчанию - максимальный размер входной партиции при чтении Parquet/ORC файлов. Это определяет количество задач в Stage 0 (Map Phase). Если у вас 1 ТБ Parquet-файлов и maxPartitionBytes=128MB, Stage 0 получит ~8 000 задач.

Полная таблица конфигураций shuffle¶

Broadcast Join как способ избежать shuffle¶

Если одна из сторон JOIN маленькая (справочник, dimension table), Spark может использовать Broadcast Join (BHJ - Broadcast Hash Join) вместо shuffle. В этом случае маленькая таблица рассылается на все executor'ы, и JOIN выполняется локально, без shuffle вообще.

from pyspark.sql.functions import broadcast

large_events = spark.read.parquet("s3a://events/")
small_campaigns = spark.read.parquet("s3a://campaigns/")

result = large_events.join(
    broadcast(small_campaigns),
    "campaign_id"
)

result.explain()

В плане будет:

BroadcastHashJoin campaign_id, Inner, BuildRight

Вместо:

SortMergeJoin campaign_id, Inner
+- Exchange hashPartitioning(campaign_id, 200)
+- Exchange hashPartitioning(campaign_id, 200)

Broadcast Join полностью устраняет shuffle для этого JOIN, что делает spark.sql.shuffle.partitions нерелевантным для данной операции. Это самый эффективный способ JOIN в Spark при условии, что одна сторона достаточно мала.

Порог для автоматического broadcast:

spark.conf.set(
    "spark.sql.autoBroadcastJoinThreshold",
    str(20 * 1024 * 1024)
)

По умолчанию: 10 МБ. Таблицы меньше этого порога Spark broadcast'ит автоматически. Можно увеличить до 20–50 МБ на кластерах с достаточным объёмом памяти executor'ов.

Антипаттерны¶

Антипаттерн 1: Оставить дефолт 200 для всех задач¶

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

result_tiny = tiny_df.groupBy("type").count()
result_huge = huge_df.groupBy("date", "user", "session").agg(...)

Одно значение не может быть оптимальным для датасетов разного размера. Используйте разные значения для разных частей пайплайна или включите AQE.

Антипаттерн 2: Установить огромное число "на всякий случай"¶

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

Если данных немного, 10 000 задач по 10 КБ каждая создадут колоссальный overhead на планирование. AQE частично исправит это, но лучше задавать адекватное начальное значение.

Антипаттерн 3: repartition() без причины¶

df.repartition(200) \
  .groupBy("channel") \
  .agg(spark_sum("clicks"))

repartition(200) перед groupBy вызывает два shuffle: один для repartition, второй для groupBy. Это вдвое дороже. repartition нужен только когда нужно изменить число партиций для последующей записи, а не перед агрегациями.

Антипаттерн 4: coalesce() в середине вычислений¶

df.groupBy("channel").agg(spark_sum("clicks")) \
  .coalesce(10) \
  .join(other_df, "channel")

coalesce(10) уменьшит параллелизм до 10 партиций. Последующий join будет выполняться с 10-кратно сниженным параллелизмом. coalesce нужен только перед финальной записью, не в середине вычислений.

Антипаттерн 5: Не мониторить Spill¶

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

result = df.groupBy("user_id", "session_id", "date").agg(...)
result.write.parquet("/output/")

print("Job completed successfully!")

Job завершился, но занял 3 часа вместо 30 минут из-за disk spill. Без мониторинга Spark UI вы никогда не узнаете о проблеме. Всегда проверяйте метрики Spill после завершения job'а.

Антипаттерн 6: Игнорировать связь shuffle.partitions с выходными файлами¶

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

result.write.partitionBy("date").parquet("/silver/events/")

Результат: 5000 файлов на каждую дату. При 30 датах - 150 000 файлов. Следующее чтение этого датасета займёт минуты только на листинг файлов. Добавьте repartition перед записью:

result \
    .repartition(4, "date") \
    .write \
    .partitionBy("date") \
    .parquet("/silver/events/")

Чеклист¶

Диагностика:

• Открыть Spark UI → Stages → найти Stage с shuffle операцией
• Проверить Spill (Disk) - если > 0, нужно больше партиций
• Проверить Task Duration Max vs Median - если Max/Median > 3×, есть skew
• Проверить Shuffle Read Size - должно быть 100–400 МБ/задача

Расчёт:

• Найти Shuffle Write Total для целевого Stage
• Применить формулу: N = (Shuffle Write × 3) / target_mb → округлить до кратного ядрам × 2
• Учесть expansion factor (числа → 2×, строки/JSON → 4–5×)

Применение:

• Установить через spark.conf.set или builder .config()
• Включить AQE: spark.sql.adaptive.enabled = true
• Задать advisoryPartitionSizeInBytes = 128–256 МБ
• Контролировать выходные файлы: repartition(n, key_col) перед записью

Для PySpark UDF:

• Уменьшить target partition size до 64–128 МБ
• Мониторить память executor'ов в Spark UI → Executors

Домашнее задание¶

Задача. Предоставлен лог метрик из Spark UI для трёх production-задач. Кластер одинаковый для всех: 16 воркеров × 4 ядра = 64 ядра, executor memory = 8 ГБ каждый.

Сценарий A: Агрегация справочников

Stage 0 - Input:  Shuffle Write Total = 480 MB, Tasks = 100
Stage 1 - GroupBy: Tasks = 200, Duration: Median=180ms, Max=520ms
           Shuffle Read Size: Median=2.4MB, Max=7.1MB
           Spill (Disk): 0 B

Сценарий B: Дневная агрегация событий

Stage 0 - Input:  Shuffle Write Total = 48 GB, Tasks = 400
Stage 1 - GroupBy: Tasks = 200, Duration: Median=4min, Max=22min
           Shuffle Read Size: Median=982MB, Max=5.8GB
           Spill (Disk): 31.4 GB total

Сценарий C: Терабайтный join для бэкфилла

Stage 0 - Input:  Shuffle Write Total = 1.8 TB, Tasks = 15000
Stage 1 - Join:   Tasks = 200, Duration: Median=N/A (OOM failures)
           Spill (Disk): N/A (executor OOM)
           Executor failures: 47 out of 64

Задание:

1. Для каждого сценария рассчитайте оптимальное значение spark.sql.shuffle.partitions с подробным обоснованием (формула, промежуточные значения).

2. Объясните, что именно происходит в каждом сценарии (диагноз проблемы или подтверждение оптимальности).

3. Для Сценария B: напишите Python-код, который устанавливает рассчитанное значение и запускает агрегацию.

4. Для Сценария C: предложите дополнительные меры помимо shuffle.partitions, которые помогут избежать OOM (минимум 2 разных подхода).

5. Для Сценария A: объясните, почему текущая конфигурация уже оптимальна или почему её стоит изменить, даже несмотря на отсутствие Spill.