UDF и pandas UDF: почему обычные UDF медленные

Финансы «Гонца» поменяли формулу сбора за доставку: 49 рублей плюс 12 рублей за километр. Первое желание любого, кто знает Python, написать функцию fee(km) и применить её к колонке. Spark это позволяет: такая функция называется UDF. Но у неё есть цена, и сегодня мы её измерим.

Переводчик посередине

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

Spark исполняет DataFrame на JVM. Функция на Python живёт в отдельном процессе Python. Чтобы применить её, Spark должен переслать значения из JVM в Python, вызвать функцию и переслать ответ обратно. Это и есть переводчик посередине.

UDF

UDF это функция пользователя. Её объявляют декоратором @F.udf(тип_результата) и вызывают на колонке, как встроенную.

У UDF две беды:

• Пересылка и вызов на каждую строку. В Spark 4 значения пересылаются пачками через формат Arrow, это заметно быстрее, чем в Spark 3. Но функция Python всё равно вызывается отдельно для каждой строки, а интерпретатор Python медленный.
• Чёрный ящик для оптимизатора. Catalyst не знает, что внутри функции. Фильтр по результату UDF нельзя отдать источнику, а саму функцию нельзя склеить с соседними операторами.

В плане UDF видна как отдельный оператор ArrowEvalPython с кодом 101.

pandas UDF

pandas UDF получает не значение, а целую пачку значений колонки в виде pandas.Series, и возвращает Series той же длины. Функция вызывается один раз на пачку, а внутри работают векторные операции pandas, написанные на C. Объявляют её через @F.pandas_udf(тип) с подсказками типов pd.Series -> pd.Series.

В плане это тот же ArrowEvalPython, но с кодом 200. Пересылка осталась, а вызовов функции стало в тысячи раз меньше.