Spark SQL, группировки и оконные функции

В команде «Гонца» не все пишут на Python. Аналитики думают на SQL, и для них цепочка filter(...).groupBy(...).agg(...) выглядит чужой. Хорошая новость: Spark понимает SQL. И не как надстройку, которая работает медленнее, а как второй способ записать тот же план. Сегодня запустим один запрос двумя способами, убедимся, что Spark выполняет их одинаково, и разберём оконные функции.

Два языка, один план

Если вы проходили оконные функции в курсе аналитика, почти весь SQL из этого урока вам знаком. Spark SQL поддерживает обычный SELECT: WHERE, GROUP BY, HAVING, JOIN, ORDER BY, подзапросы и WITH.

Чтобы обратиться к DataFrame из SQL, ему дают имя таблицы: orders.createOrReplaceTempView("orders"). Это временное представление: данные никуда не копируются, имя просто ссылается на план DataFrame. Живёт оно, пока жива сессия. Запрос выполняется через spark.sql("...") и возвращает обычный DataFrame.

Почему планы одинаковые? И SQL, и методы DataFrame сначала превращаются в одно и то же дерево операций, логический план. Дальше с ним работает оптимизатор Catalyst из урока про DataFrame, ему всё равно, на каком языке вы писали. Выбирайте тот, на котором запрос читается легче.

Группировки

groupBy в DataFrame и GROUP BY в SQL работают как в любой базе. Под капотом Spark считает группировку в две стадии, как мы видели в первых уроках: частичные итоги внутри каждой партиции, shuffle по ключу группировки и окончательные итоги.

Оконные функции

Оконная функция считает значение для каждой строки по её соседям. В отличие от группировки, строки не схлопываются: каждый день остаётся отдельной строкой, но рядом с ним можно посчитать нарастающий итог недели или разницу со вчерашним днём.

В DataFrame окно описывают объектом Window:

w = Window.partitionBy("zone_id").orderBy("day")

partitionBy делит строки на группы (здесь по зонам), orderBy задаёт порядок внутри группы. Дальше любая функция с .over(w) считается по этому окну:

• F.row_number(), F.rank(), F.dense_rank(): место строки в группе;
• F.lag("orders"), F.lead("orders"): значение из предыдущей или следующей строки;
• F.sum("orders").over(w): сумма от начала группы до текущей строки, то есть нарастающий итог.

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