UDF: Python UDF и их цена - когда стоит, когда не стоит

from pyspark.sql import functions as F

# Строки
F.upper, F.lower, F.trim, F.regexp_replace, F.split, F.concat_ws

# Числа
F.round, F.abs, F.log, F.sqrt, F.pow, F.greatest, F.least

# Дата/время
F.to_date, F.date_format, F.datediff, F.date_add, F.year, F.month

# Условная логика
F.when, F.coalesce, F.nullif, F.isnull, F.isnan

# Массивы и структуры
F.array_contains, F.explode, F.collect_list, F.struct, F.map_keys

# Высокоуровневые функции над массивами (HOF)
F.transform, F.filter, F.aggregate, F.exists

Container killed by YARN for exceeding memory limits.
12.0 GB of 12.0 GB physical memory used.
Consider boosting spark.yarn.executor.memoryOverhead or
disabling yarn.nodemanager.vmem-check-enabled.

# Catalyst НЕ МОЖЕТ протолкнуть filter через UDF
df.withColumn("score", my_python_udf(col("text"))) \
  .filter(col("score") > 0.5)

# Physical plan:
# Filter (score > 0.5)
#   ↑
# BatchEvalPython [my_python_udf]   ← UDF выполнится для ВСЕХ строк
#   ↑
# FileScan parquet (читает ВСЕ данные)

# Catalyst компилирует как THREE отдельных блока - не один:

# Блок 1 (JVM Codegen): filter + project перед UDF
# ↓
# Barrier: сериализация в Python, выполнение UDF, десериализация
# ↓
# Блок 2 (JVM Codegen): операции после UDF

df.withColumn("cleaned", my_udf(col("raw_text"))) \
  .filter(col("cleaned").isNotNull()) \
  .explain(extended=True)

# Physical Plan:
# *(2) Project [...]
# +- *(2) Filter (isnotnull(cleaned#45))
#    +- BatchEvalPython [my_udf(raw_text#12)], [raw_text#12, cleaned#45]
#       +- *(1) FileScan parquet [...] PushedFilters: [], ...

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType, DoubleType, IntegerType

# Способ 1: декоратор
@udf(returnType=StringType())
def normalize_phone(phone: str) -> str:
    if phone is None:
        return None
    digits = "".join(c for c in phone if c.isdigit())
    if len(digits) == 11 and digits[0] == "7":
        return f"+{digits}"
    if len(digits) == 10:
        return f"+7{digits}"
    return None

# Способ 2: явная регистрация
def parse_category(code: str) -> str:
    mapping = {"A": "premium", "B": "standard", "C": "budget"}
    return mapping.get(code, "unknown")

parse_category_udf = udf(parse_category, StringType())

# Использование в DataFrame
df.withColumn("phone_clean", normalize_phone("phone")) \
  .withColumn("category_name", parse_category_udf("category_code"))

# Регистрация для SQL
spark.udf.register("normalize_phone", normalize_phone)
spark.sql("SELECT normalize_phone(phone) FROM customers")

@udf(returnType=IntegerType())  # объявлено: Int
def bad_return(s: str) -> str:
    return "not_an_int"  # возвращает: str

# При вызове - ошибка в Task:
# pyspark.sql.utils.PythonException: An exception was thrown from the Python
# worker. Please see the stack trace below.
# Traceback (most recent call last):
#   ...
# TypeError: an integer is required (got type str)

# Опасно: AttributeError при None
@udf("string")
def bad_upper(s: str) -> str:
    return s.upper()  # AttributeError: 'NoneType' object has no attribute 'upper'

# Безопасно: явная проверка на None
@udf("string")
def safe_upper(s: str) -> str:
    return s.upper() if s is not None else None

# Ещё лучше: встроенная функция без UDF вообще
from pyspark.sql.functions import upper
df.withColumn("name_upper", upper(col("name")))

from pyspark.sql.types import StructType, StructField, StringType, DoubleType

@udf(DoubleType())
def extract_amount(payment):
    # payment - это Row объект или None
    if payment is None:
        return None
    # payment.amount тоже может быть None!
    return float(payment.amount) if payment.amount is not None else 0.0

import random
from pyspark.sql.functions import udf
from pyspark.sql.types import DoubleType

# Без объявления non-deterministic: Spark может выполнить UDF один раз
# и переиспользовать результат - нарушая ожидаемое поведение
@udf(DoubleType())
def random_score(user_id: int) -> float:
    return random.random()

# Правильно: явно объявить non-deterministic
random_score_nd = udf(lambda uid: random.random(), DoubleType()).asNondeterministic()

df.withColumn("score", random_score_nd(col("user_id")))

# Опасно: UDF вызывается дважды в одном плане
expensive_udf_result = expensive_udf(col("raw"))

df.withColumn("a", expensive_udf_result) \
  .withColumn("b", expensive_udf_result * 2)

# Physical plan может содержать два BatchEvalPython узла
# - Spark не всегда умеет понять, что это один и тот же вызов

# Безопасно: UDF вызывается один раз, результат сохраняется в колонке
df_with_score = df.withColumn("score_raw", expensive_udf(col("raw")))

df_with_score \
  .withColumn("score_norm", col("score_raw") / 100) \
  .withColumn("score_flag", col("score_raw") > 80)

# Ищем BatchEvalPython в физическом плане
df.withColumn("result", my_udf(col("input"))).explain()

# BatchEvalPython [my_udf(input#5)], [input#5, result#12]
# Это сигнал - здесь Python worker

# На executor node
ps aux | grep python
# pyspark.worker - отдельный Python process на каждый executor

org.apache.spark.SparkException: Task not serializable
Caused by: java.io.NotSerializableException: org.postgresql.jdbc.PgConnection
    at java.io.ObjectOutputStream.writeObject0(ObjectOutputStream.java:1184)
    at ...

# ОШИБКА: conn создан на Driver, не может быть сериализован в Task
import psycopg2

conn = psycopg2.connect("host=pg-host dbname=mydb user=etl password=secret")

@udf("string")
def enrich_from_db(user_id: int) -> str:
    cursor = conn.cursor()  # conn захвачен в closure → NotSerializableException
    cursor.execute("SELECT name FROM users WHERE id = %s", (user_id,))
    return cursor.fetchone()[0]

# РЕШЕНИЕ: создать ресурс внутри UDF (на каждую строку - неэффективно)
@udf("string")
def enrich_from_db(user_id: int) -> str:
    import psycopg2  # импорт внутри UDF гарантирует, что он доступен на executor
    if user_id is None:
        return None
    conn = psycopg2.connect("host=pg-host dbname=mydb user=etl password=secret")
    cursor = conn.cursor()
    cursor.execute("SELECT name FROM users WHERE id = %s", (user_id,))
    result = cursor.fetchone()
    conn.close()
    return result[0] if result else None

def enrich_partition(rows):
    import psycopg2

    # Ресурс создаётся ОДИН РАЗ на партицию (не на строку!)
    conn = psycopg2.connect("host=pg-host dbname=mydb user=etl password=secret")
    cursor = conn.cursor()

    try:
        for row in rows:
            user_id = row["user_id"]
            if user_id is None:
                yield row.asDict() | {"user_name": None}
                continue
            cursor.execute("SELECT name FROM users WHERE id = %s", (user_id,))
            result = cursor.fetchone()
            yield row.asDict() | {"user_name": result[0] if result else None}
    finally:
        conn.close()  # Соединение закрывается после обработки всей партиции

from pyspark.sql.types import StructType, StructField, StringType, LongType, DoubleType

result_schema = StructType([
    StructField("user_id", LongType()),
    StructField("amount", DoubleType()),
    StructField("user_name", StringType()),
])

enriched_df = df.rdd.mapPartitions(enrich_partition).toDF(result_schema)

def predict_partition(rows):
    import joblib
    import numpy as np

    # Модель загружается ОДИН РАЗ на партицию
    model = joblib.load("/shared/models/fraud_detector_v2.pkl")

    batch = list(rows)
    if not batch:
        return

    features = np.array([[r["amount"], r["hour"], r["merchant_code"]] for r in batch])
    predictions = model.predict_proba(features)[:, 1]

    for row, score in zip(batch, predictions):
        yield row.asDict() | {"fraud_score": float(score)}

fraud_df = transactions.rdd.mapPartitions(predict_partition).toDF(result_schema)

WARN TaskSetManager: Stage 3 contains a task of very large size (1.2 MB).
The maximum recommended task size is 100 KB. If init time is short,
consider using a broadcast variable to reduce the task size.

# Плохо: большой словарь захвачен в closure - копируется в каждый Task
city_population = {city: pop for city, pop in spark.table("cities").collect()}
# Допустим: 500 городов × 1 KB = 500 KB на Task
# 200 Tasks × 500 KB = 100 MB трафика от Driver к Executor'ам

@udf("long")
def get_population(city: str) -> int:
    return city_population.get(city, 0)  # city_population → в closure!

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

# Broadcast variable: рассылается ОДИН РАЗ на Executor, не на Task
city_pop_dict = {city: pop for city, pop in spark.table("cities").collect()}
city_pop_bv = spark.sparkContext.broadcast(city_pop_dict)

@udf(LongType())
def get_population(city: str) -> int:
    # .value получает объект на Executor без повторной сериализации
    return city_pop_bv.value.get(city, 0)

df.withColumn("population", get_population(col("city")))

# Освободить когда больше не нужен
city_pop_bv.unpersist()

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.functions import udf
from pyspark.sql.types import StringType
import time

spark = SparkSession.builder \
    .appName("UDF-Benchmark") \
    .config("spark.sql.shuffle.partitions", "8") \
    .getOrCreate()

# Генерируем тестовые данные: 10 миллионов строк
df = spark.range(10_000_000).select(
    F.concat(
        F.lit("  User_"),
        F.col("id").cast("string"),
        F.lit("@Example.COM  ")
    ).alias("email_raw")
).cache()
df.count()  # материализуем кэш

# ── Вариант 1: Python UDF ─────────────────────────────────────────────────────
@udf(StringType())
def extract_domain_udf(email: str) -> str:
    if email is None:
        return None
    clean = email.strip().lower()
    if "@" not in clean:
        return None
    return clean.split("@")[1]

start = time.time()
df.withColumn("domain_udf", extract_domain_udf(F.col("email_raw"))).count()
print(f"Python UDF: {time.time() - start:.1f}s")

# ── Вариант 2: встроенные функции ─────────────────────────────────────────────
start = time.time()
df.withColumn("domain_builtin",
    F.split(F.lower(F.trim(F.col("email_raw"))), "@").getItem(1)
).count()
print(f"Built-in functions: {time.time() - start:.1f}s")

Python UDF:       45.3s   (BatchEvalPython, pickle overhead)
Built-in functions: 4.1s  (WholeStageCodeGen, JVM native)

# Всё через UDF - типичная ошибка junior разработчика
@udf("string")
def extract_year(date_str: str) -> str:
    if date_str is None:
        return None
    return date_str[:4]

@udf("string")
def clean_phone(phone: str) -> str:
    if phone is None:
        return None
    return "".join(c for c in phone if c.isdigit())

@udf("string")
def categorize_amount(amount) -> str:
    if amount is None:
        return "unknown"
    if amount < 100:
        return "small"
    if amount < 10000:
        return "medium"
    return "large"

@udf("string")
def upper_name(name: str) -> str:
    return name.upper() if name else None

result = df \
    .withColumn("year", extract_year(col("created_at"))) \
    .withColumn("phone_clean", clean_phone(col("phone"))) \
    .withColumn("amount_cat", categorize_amount(col("amount"))) \
    .withColumn("name_upper", upper_name(col("name")))

# Physical plan: 4× BatchEvalPython - 4 барьера JVM↔Python

from pyspark.sql import functions as F

result = df \
    .withColumn("year",
        F.year(F.col("created_at").cast("date"))
    ) \
    .withColumn("phone_clean",
        F.regexp_replace(F.col("phone"), "[^0-9]", "")
    ) \
    .withColumn("amount_cat",
        F.when(F.col("amount") < 100, "small")
         .when(F.col("amount") < 10000, "medium")
         .otherwise("large")
    ) \
    .withColumn("name_upper",
        F.upper(F.col("name"))
    )

# Physical plan: один WholeStageCodeGen блок - всё вместе

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf, col, upper, when, regexp_replace
from pyspark.sql.types import StringType

spark = SparkSession.builder.appName("UDF-Demo").getOrCreate()

df = spark.createDataFrame([
    (1, "john doe", "+7 (999) 123-45-67", "A"),
    (2, None, "8 800 555 3535", "B"),
    (3, "Jane Smith", None, "C"),
], ["id", "name", "phone", "tier"])

# ── ПЛОХО: UDF для того, что встроенные функции делают лучше ──────────────────
@udf("string")
def bad_upper(s):
    return s.upper() if s else None  # встроенный upper() быстрее в 5–10×

# ── ХОРОШО: встроенные функции ───────────────────────────────────────────────
df_good = df.withColumn("name_upper", upper(col("name")))  # JVM, Codegen


# ── ХОРОШО: UDF только для нестандартной логики (нормализация телефона) ───────
# Такой логики нет в встроенных функциях - UDF оправдан
@udf("string")
def normalize_phone(phone: str) -> str:
    if phone is None:
        return None
    digits = "".join(c for c in phone if c.isdigit())
    if len(digits) == 11:
        return f"+{digits}"
    if len(digits) == 10:
        return f"+7{digits}"
    return phone

df_phones = df.withColumn("phone_clean", normalize_phone(col("phone")))

# ── ХОРОШО: broadcast для справочника ────────────────────────────────────────
tier_labels = {"A": "Premium", "B": "Standard", "C": "Budget"}
tier_bv = spark.sparkContext.broadcast(tier_labels)

@udf("string")
def tier_label(code: str) -> str:
    return tier_bv.value.get(code, "Unknown") if code else None

df_result = df_phones.withColumn("tier_name", tier_label(col("tier")))
df_result.show()

tier_bv.unpersist()

# ПЛОХО
@udf("string")
def my_upper(s): return s.upper() if s else None

# ХОРОШО
F.upper(col("name"))

# ПЛОХО: UDF для условной логики
@udf("string")
def categorize(score):
    if score is None: return None
    if score > 0.8: return "high"
    if score > 0.5: return "medium"
    return "low"

# ХОРОШО: встроенная условная логика
F.when(col("score") > 0.8, "high") \
 .when(col("score") > 0.5, "medium") \
 .otherwise("low")

# ПЛОХО: regex через Python UDF
@udf("string")
def extract_code(text):
    import re
    m = re.search(r"CODE-(\d+)", text or "")
    return m.group(1) if m else None

# ХОРОШО: встроенный regexp_extract
F.regexp_extract(col("text"), r"CODE-(\d+)", 1)

# ПЛОХО: UDF для преобразования массива
@udf("array<double>")
def normalize_scores(scores):
    if scores is None: return None
    total = sum(scores) or 1
    return [s / total for s in scores]

# ХОРОШО: transform (HOF) - работает в JVM
from pyspark.sql.functions import transform, aggregate
total = aggregate(col("scores"), F.lit(0.0), lambda acc, x: acc + x)
F.transform(col("scores"), lambda x: x / total)

# ПЛОХО: 10M строк = 10M соединений с PostgreSQL
@udf("string")
def lookup(user_id):
    import psycopg2
    conn = psycopg2.connect(...)  # новое соединение на каждую строку!
    ...

# ХОРОШО: mapPartitions - одно соединение на партицию
df.rdd.mapPartitions(enrich_partition).toDF(schema)

# ПЛОХО: 10MB dict в closure каждого Task
big_dict = load_mapping()  # 10MB

@udf("string")
def lookup(key): return big_dict.get(key)  # big_dict → closure

# ХОРОШО: broadcast
bv = spark.sparkContext.broadcast(big_dict)

@udf("string")
def lookup(key): return bv.value.get(key)