UDF Pitfalls: closure, null без защиты и non-determinism

# ❌ UDF захватывает SparkSession - не сериализуется
spark = SparkSession.getOrCreate()
config = spark.conf.get("myapp.setting")   # OK пока не в UDF

@udf("string")
def process(s):
    return spark.read.parquet("ref_table/") \
        .filter(...).collect()[0][0]   # SparkSession в UDF → TaskNotSerializable

# ❌ UDF захватывает огромный объект (100 MB model)
model = load_big_model("model.pkl")   # 100 MB

@udf("float")
def predict(x):
    return model.predict([x])[0]
# 100 MB pickle на каждую задачу × N Tasks = огромный overhead

# ✅ Broadcast для тяжёлых объектов
model_bc = spark.sparkContext.broadcast(load_big_model("model.pkl"))

@udf("float")
def predict(x):
    return model_bc.value.predict([x])[0]
# model сериализуется один раз → передаётся через Block Manager

# ✅ Использовать Pandas UDF для батчевой обработки (Arrow, не Pickle)
@pandas_udf("float")
def predict_batch(x: pd.Series) -> pd.Series:
    return pd.Series(model.predict(x.values))

# ❌ UDF без null-проверки
@udf("string")
def clean(email):
    return email.strip().lower()   # AttributeError: 'NoneType' has no attribute 'strip'

# Одна null-строка → Task failure → Stage retry → Job failure

# ✅ Всегда проверяйте None в UDF
@udf("string")
def clean(email):
    if email is None:
        return None   # возвращаем null, а не падаем
    return email.strip().lower()

# ✅ Ещё лучше - заменить UDF на built-in (null-safe автоматически)
from pyspark.sql.functions import lower, trim
df.withColumn("email", lower(trim(col("raw_email"))))

# ❌ UDF с side effects или внешними вызовами
import random, datetime

@udf("string")
def add_timestamp(s):
    return f"{s}_{datetime.now()}"   # разные результаты при повторном вызове!

# Spark может вызвать UDF дважды (retry, спекулятивное выполнение)
# Результаты будут разные → нарушение consistency

# ✅ Явно пометить UDF как non-deterministic (отключает кэш и некоторые оптимизации)
from pyspark.sql.functions import udf
from pyspark.sql.types import StringType

@udf("string")
def external_api_call(s):
    return requests.get(f"http://api/{s}").text

external_api_call_nd = external_api_call.asNondeterministic()
# Теперь Spark знает - не кэшировать, не переупорядочивать

# ✅ Лучше - использовать built-in current_timestamp()
from pyspark.sql.functions import current_timestamp
df.withColumn("ts", current_timestamp())

# ❌ UDF в filter - Catalyst не может pushdown в Parquet
@udf("boolean")
def is_valid(s):
    return len(s) > 5

df = spark.read.parquet("data/") \
    .filter(is_valid("email"))
# Spark читает ВСЕ строки, потом фильтрует Python UDF

# ✅ Встроенный фильтр - Parquet row group pruning работает
df = spark.read.parquet("data/") \
    .filter(length("email") > 5)

# ❌ UDF возвращает сложный объект - медленная сериализация
from pyspark.sql.types import StructType, StructField, StringType, IntegerType

@udf(StructType([StructField("name", StringType()), StructField("age", IntegerType())]))
def parse(s):
    parts = s.split(",")
    return {"name": parts[0], "age": int(parts[1])}
# Pickle сериализует dict для каждой строки - очень медленно

# ✅ Встроенные функции для парсинга
from pyspark.sql.functions import split, col

df.withColumn("parts", split("raw", ",")) \
  .withColumn("name", col("parts")[0]) \
  .withColumn("age", col("parts")[1].cast("int"))

# Spark UI: Tasks где PythonTime >> JVM time - UDF overhead
# Вкладка Stages → Task Metrics → "Python Time" (если показывает)

# Сравнить время с UDF и без:
import time

start = time.time()
df.withColumn("x", my_udf("col")).count()
print(f"UDF: {time.time()-start:.1f}s")

start = time.time()
df.withColumn("x", lower(trim(col("col")))).count()
print(f"Built-in: {time.time()-start:.1f}s")

1. Built-in functions (pyspark.sql.functions.*)  → JVM, SIMD, нет Python
2. Pandas UDF (@pandas_udf)                      → Arrow batches, умеренные расходы
3. Python UDF (@udf)                             → Pickle, построчно, дорого
4. RDD.mapPartitions с Python функцией           → ещё медленнее, нет Catalyst