Shuffle и перекос данных: диагностика и приёмы

Задание в Spark UI выглядит странно: из двухсот задач стадии сто девяносто девять закончились за минуту, а последняя идёт уже час. Кластер простаивает, а отчёт Отто не готов. Почти всегда это перекос данных: после shuffle одной партиции досталось в десятки раз больше строк, чем остальным. Сегодня научимся его находить и лечить.

Очередь в одну кассу

В супермаркете десять касс. Девять кассиров обслуживают покупателей с корзинками, а к десятой подъехал оптовик с двадцатью тележками. Пока он не расплатится, магазин не может закрыться, сколько бы касс ни простаивало.

Shuffle раскладывает строки по партициям по хешу ключа: все строки одного ключа попадают в одну партицию. Если один ключ встречается в миллионы раз чаще других, его партиция становится тем оптовиком.

Откуда берётся перекос

Ключи в жизни распределены неровно:

• одна большая зона, город или категория занимает половину заказов;
• служебные значения: тестовый пользователь, NULL вместо неизвестного клиента, пустая строка;
• популярный товар или сервис, на который приходится основная нагрузка.

Перекос бьёт по операциям, которым нужны все строки ключа в одной партиции: соединениям, оконным функциям, groupBy с collect_list. Обычной группировке с count или sum он почти не страшен: частичные итоги считаются до shuffle, и от каждой задачи по сети едет одна строка на ключ.

Как найти

1. Spark UI. В стадии видно распределение длительности задач: если максимальная в разы больше медианы, это перекос. В песочнице такая задача подсвечена красным.
2. Распределение ключей. df.groupBy("ключ").count().orderBy(F.desc("count")) покажет, кто наверху. Если первый ключ занимает десятки процентов строк, вы нашли виновника.

Как лечить

AQE. В Spark 3 и 4 адаптивное исполнение умеет само делить перекошенные партиции в sort-merge join (spark.sql.adaptive.skewJoin.enabled, включено по умолчанию). Это первое, что стоит проверить: не выключен ли AQE. Песочница этот режим не моделирует, поэтому ниже AQE выключен, чтобы увидеть перекос как есть.

Broadcast. Если одна сторона соединения маленькая, broadcast join обходится без shuffle большой стороны, и перекосу просто негде появиться.

Соль. Если обе стороны большие, ключ разбавляют солью. К каждой строке большой стороны добавляют число от 0 до N-1, например остаток от деления номера строки. Вторую сторону размножают N раз, по копии на каждое значение соли, и соединяют по паре «ключ, соль». Строки горячего ключа расходятся по N партициям.

Отдельная обработка. Иногда проще отделить горячий ключ: его строки обработать отдельно (часто через broadcast), остальные обычным способом, и сложить результаты.