Pandas UDF с Arrow: @pandas_udf, Series-батч вместо строки, скорость vs Row-at-a-time UDF

Для каждой строки в батче (например, 8192 строк):
  1. Executor берёт значения нужных колонок из UnsafeRow
  2. Конвертирует JVM-типы в Python-типы (int → PyLong, str → PyUnicode, ...)
  3. Сериализует Python-объекты через Pickle (бинарный формат)
  4. Записывает pickled байты в pipe / socket к Python-воркеру
  5. Python-воркер читает байты, десериализует через Pickle
  6. Вызывает вашу функцию с одним значением
  7. Получает результат, сериализует обратно через Pickle
  8. Отправляет обратно в JVM
  9. JVM десериализует результат, конвертирует в JVM-тип
  10. Записывает в выходной UnsafeRow

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.types import DoubleType
import time
import math

spark = SparkSession.builder \
    .master("local[4]") \
    .config("spark.driver.memory", "4g") \
    .getOrCreate()

# Создаём тестовый DataFrame: 5 миллионов строк
df = spark.range(5_000_000) \
    .withColumn("lat1", (F.rand() * 180 - 90).cast("double")) \
    .withColumn("lon1", (F.rand() * 360 - 180).cast("double")) \
    .withColumn("lat2", (F.rand() * 180 - 90).cast("double")) \
    .withColumn("lon2", (F.rand() * 360 - 180).cast("double")) \
    .cache()

df.count()  # форсируем кэширование

# ──────────────────────────────────────────────────────────
# ВАРИАНТ 1: классический @udf - row-at-a-time
# ──────────────────────────────────────────────────────────
# Формула гаверсинуса вычисляет расстояние между двумя
# точками на сфере по их географическим координатам.
# Это типичная "нестандартная математика" - нет встроенной
# функции Spark SQL для этого расчёта.
@F.udf(returnType=DoubleType())
def haversine_udf(lat1, lon1, lat2, lon2):
    """Расстояние в км между двумя точками на Земле."""
    R = 6371.0  # радиус Земли в км
    phi1, phi2 = math.radians(lat1), math.radians(lat2)
    dphi = math.radians(lat2 - lat1)
    dlambda = math.radians(lon2 - lon1)
    a = math.sin(dphi/2)**2 + math.cos(phi1) * math.cos(phi2) * math.sin(dlambda/2)**2
    return R * 2 * math.asin(math.sqrt(a))

start = time.time()
result_udf = df.withColumn(
    "distance_km",
    haversine_udf("lat1", "lon1", "lat2", "lon2")
)
count = result_udf.count()  # форсируем выполнение
elapsed_udf = time.time() - start

print(f"Классический @udf: {elapsed_udf:.1f} сек на {count:,} строк")
# Ожидаемо: 45–120 секунд в зависимости от железа

import numpy as np
import time

data = np.random.randn(10_000_000)

# Python-цикл: 10 миллионов Python-вызовов math.sin
start = time.time()
result_loop = [math.sin(x) for x in data]
print(f"Python loop: {time.time() - start:.2f}s")   # ~8–12 секунд

# NumPy: один вызов C-кода над всем массивом
start = time.time()
result_numpy = np.sin(data)
print(f"NumPy vectorized: {time.time() - start:.2f}s")  # ~0.05–0.1 секунды

# ──────────────────────────────────────────────────────────
# Конфигурация Arrow для Pandas UDF
# ──────────────────────────────────────────────────────────
spark = SparkSession.builder \
    .master("local[4]") \
    .config("spark.driver.memory", "4g") \
    # Arrow enabled by default в Spark 3.x, но лучше явно:
    .config("spark.sql.execution.arrow.pyspark.enabled", "true") \
    # Максимальное число строк в одном Arrow-батче.
    # Default: 10000. Можно настраивать (см. раздел tuning).
    .config("spark.sql.execution.arrow.maxRecordsPerBatch", "10000") \
    # Python-воркер переиспользуется между задачами.
    # Выключить только если воркер накапливает состояние.
    .config("spark.python.worker.reuse", "true") \
    .getOrCreate()

# Проверяем, что Arrow действительно включён:
arrow_enabled = spark.conf.get("spark.sql.execution.arrow.pyspark.enabled")
print(f"Arrow enabled: {arrow_enabled}")  # true

import pandas as pd
import numpy as np
from pyspark.sql.functions import pandas_udf

# ──────────────────────────────────────────────────────────
# Простой пример: математическая трансформация
# ──────────────────────────────────────────────────────────
# Тайп-хинт `pd.Series -> pd.Series` говорит Spark:
#   - принять одну колонку как pd.Series
#   - вернуть pd.Series той же длины
# Декоратор @pandas_udf нужен для регистрации в Spark.
@pandas_udf("double")
def normalize_zscore(series: pd.Series) -> pd.Series:
    """
    Z-score нормализация: (x - mean) / std.
    Это вычисляется в NumPy - C-код с SIMD.
    """
    # series.mean() и series.std() - NumPy операции над всем батчем
    # (series - mean) / std - vectorized вычитание и деление
    mean = series.mean()
    std = series.std()
    # Защита от деления на ноль: если std == 0, возвращаем 0
    if std == 0:
        return pd.Series(np.zeros(len(series)))
    return (series - mean) / std

# ──────────────────────────────────────────────────────────
# Pandas UDF с несколькими входными колонками
# ──────────────────────────────────────────────────────────
# Для нескольких колонок нужно использовать F.struct() или
# аргументы в порядке колонок
@pandas_udf("double")
def haversine_pandas_udf(
    lat1: pd.Series,
    lon1: pd.Series,
    lat2: pd.Series,
    lon2: pd.Series
) -> pd.Series:
    """
    Формула гаверсинуса - расстояние в км между точками.

    Ключевое отличие от @udf:
    lat1, lon1, lat2, lon2 - это pd.Series с 10000 значений,
    а не одно float-значение. Все numpy операции работают
    над целым массивом одновременно.
    """
    R = 6371.0  # радиус Земли в км

    # np.radians применяется ко всему Series сразу - C-код
    phi1 = np.radians(lat1)
    phi2 = np.radians(lat2)
    dphi = np.radians(lat2 - lat1)
    dlambda = np.radians(lon2 - lon1)

    # np.sin, np.cos - векторизованные операции с SIMD
    # Каждая операция обрабатывает весь массив батча
    a = (np.sin(dphi / 2) ** 2
         + np.cos(phi1) * np.cos(phi2) * np.sin(dlambda / 2) ** 2)

    # np.arcsin, np.sqrt - тоже векторизованы
    return pd.Series(R * 2 * np.arcsin(np.sqrt(a)))

# Применяем Pandas UDF к DataFrame
result = df.withColumn(
    "distance_km",
    haversine_pandas_udf("lat1", "lon1", "lat2", "lon2")
)

start = time.time()
count = result.count()
elapsed_pandas_udf = time.time() - start

print(f"@pandas_udf: {elapsed_pandas_udf:.1f} сек")
# Ожидаемо: 2–8 секунд - в 10–50× быстрее классического @udf

from typing import Iterator

# ──────────────────────────────────────────────────────────
# Iterator-паттерн: загрузка модели один раз на partition
# ──────────────────────────────────────────────────────────
# Тайп-хинт Iterator[pd.Series] -> Iterator[pd.Series]
# говорит Spark: эта функция получает итератор батчей
# и возвращает итератор результатов.
@pandas_udf("double")
def ml_score_iterator(
    iterator: Iterator[pd.Series]
) -> Iterator[pd.Series]:
    """
    ML-скоринг с загрузкой модели один раз на partition.

    Структура:
    1. Код ДО цикла for выполняется один раз при старте partition
    2. Цикл for итерирует по батчам
    3. Каждый yield отправляет батч результатов в Spark
    """
    # ── ИНИЦИАЛИЗАЦИЯ (один раз на partition) ──────────────
    # Здесь можно загружать:
    # - ML-модель из файла или object storage
    # - Lookup-словари (тысячи записей)
    # - Тяжёлые коэффициенты для бизнес-правил
    # - Подключения к внешним сервисам (с осторожностью)
    import joblib
    import numpy as np

    # Загрузка модели выполняется ОДИН РАЗ для всей partition.
    # Если partition содержит 1M строк и батч 10k строк,
    # модель загружается 1 раз, а не 100 раз.
    model = joblib.load("/models/fraud_detection_v2.pkl")

    # Любая другая тяжёлая инициализация:
    # lookup_dict = load_large_dictionary()
    # scaler = load_feature_scaler()

    # ── ОБРАБОТКА БАТЧЕЙ ───────────────────────────────────
    # Теперь итерируем по батчам - каждый батч это pd.Series
    for batch in iterator:
        # batch - pd.Series с 10000 значений (или меньше в последнем батче)
        # Преобразуем в формат, который ожидает модель
        X = batch.values.reshape(-1, 1)  # (10000, 1) NumPy array

        # Инференс: модель работает над всем батчем
        # (внутри scikit-learn тоже NumPy/BLAS - быстро)
        scores = model.predict_proba(X)[:, 1]  # вероятность фрода

        # yield отправляет батч результатов в Spark
        # и функция "засыпает" до следующего батча
        yield pd.Series(scores)

# Для нескольких входных колонок:
@pandas_udf("double")
def ml_score_multi_cols(
    iterator: Iterator[tuple[pd.Series, pd.Series, pd.Series]]
) -> Iterator[pd.Series]:
    """ML-скоринг с несколькими признаками."""
    import joblib

    model = joblib.load("/models/risk_model.pkl")

    for amount_batch, merchant_id_batch, hour_batch in iterator:
        # Собираем признаки в матрицу
        X = np.column_stack([
            amount_batch.values,
            merchant_id_batch.values,
            hour_batch.values
        ])
        scores = model.predict_proba(X)[:, 1]
        yield pd.Series(scores)

# Использование с F.struct для множественных колонок
result = df.withColumn(
    "fraud_score",
    ml_score_multi_cols("amount", "merchant_id", "hour_of_day")
)

# ──────────────────────────────────────────────────────────
# applyInPandas: произвольная логика для каждой группы
# ──────────────────────────────────────────────────────────
# Пример: нормализация продаж по регионам.
# Для каждого region вычисляем z-score относительно
# среднего и стандартного отклонения ТОЛЬКО этого региона.
from pyspark.sql.types import StructType, StructField, StringType, DoubleType, LongType

# Схема результирующего DataFrame должна быть объявлена явно.
# Это не опционально - Spark не может вывести схему из Python-кода.
output_schema = StructType([
    StructField("region", StringType(), True),
    StructField("order_id", LongType(), True),
    StructField("amount", DoubleType(), True),
    StructField("amount_normalized", DoubleType(), True),
    StructField("amount_rank", DoubleType(), True),
])

def normalize_by_region(pdf: pd.DataFrame) -> pd.DataFrame:
    """
    Функция принимает pd.DataFrame для ОДНОЙ группы (одного region).
    Возвращает pd.DataFrame с той же или другой схемой.

    Здесь мы можем использовать весь арсенал Pandas:
    - произвольные вычисления по всей группе
    - операции, которые требуют видеть все строки группы
    - интеграцию с любыми Python-библиотеками
    """
    # pdf - это обычный pd.DataFrame только для одного region
    # Нормализуем amount внутри группы
    mean = pdf["amount"].mean()
    std = pdf["amount"].std()

    pdf["amount_normalized"] = (pdf["amount"] - mean) / (std if std > 0 else 1)

    # Ранг внутри группы (0.0 = минимум, 1.0 = максимум)
    pdf["amount_rank"] = pdf["amount"].rank(pct=True)

    return pdf

# groupBy указывает ключ разбиения на группы
# applyInPandas применяет функцию к каждой группе
result = df.groupBy("region").applyInPandas(
    normalize_by_region,
    schema=output_schema
)

# ──────────────────────────────────────────────────────────
# Анализ физического плана
# ──────────────────────────────────────────────────────────
# Применяем pandas_udf к DataFrame
result = df.withColumn(
    "distance_km",
    haversine_pandas_udf("lat1", "lon1", "lat2", "lon2")
)

# explain() показывает физический план
result.explain(extended=True)

# Ключевые части вывода:
#
# == Physical Plan ==
# *(2) Project [id#0L, lat1#4, lon1#6, lat2#8, lon2#10, distance_km#42]
# +- ArrowEvalPython [haversine_pandas_udf(lat1#4, lon1#6, lat2#8, lon2#10)#41],
#    [id#0L, lat1#4, lon1#6, lat2#8, lon2#10, distance_km#41],
#    evalType=SQL_SCALAR_PANDAS_UDF
#    +- *(1) Project [id#0L, ...]
#       +- *(1) Range (0, 5000000, step=1, splits=4)
#
# Что означают ключевые части:
#
# ArrowEvalPython - оператор для Pandas UDF через Arrow.
#   Если бы это был классический @udf, было бы BatchEvalPython.
#
# evalType=SQL_SCALAR_PANDAS_UDF - тип Pandas UDF:
#   SQL_SCALAR_PANDAS_UDF = Series → Series
#   SQL_SCALAR_PANDAS_ITER_UDF = Iterator of Series
#   SQL_GROUPED_MAP_PANDAS_UDF = applyInPandas
#
# *(1) и *(2) - номера Whole-Stage Codegen стадий.
# Обратите внимание: ArrowEvalPython разрывает WSCG!
# Стадия *(1) компилируется в JVM bytecode.
# ArrowEvalPython - граница, данные уходят в Python.
# Стадия *(2) - снова JVM bytecode.
#
# Это означает: нативные функции Spark SQL (F.sum, F.filter)
# внутри WSCG работают со скоростью нативного JVM-кода.
# Pandas UDF - всегда разрыв WSCG, всегда slower.

# Для сравнения - классический @udf в explain():
# BatchEvalPython [haversine_udf(...)], evalType=SQL_BATCHED_UDF
# Вместо ArrowEvalPython - BatchEvalPython.
# Принципиальная разница: BatchEvalPython использует Pickle,
# ArrowEvalPython - Arrow IPC.

# ──────────────────────────────────────────────────────────
# Полный бенчмарк: @udf vs @pandas_udf vs native SQL
# ──────────────────────────────────────────────────────────
import time
import math
import numpy as np
import pandas as pd
from pyspark.sql import functions as F
from pyspark.sql.functions import pandas_udf, udf
from pyspark.sql.types import DoubleType

def benchmark(name, df_func, df, n_runs=3):
    """Запускает функцию n_runs раз и возвращает медианное время."""
    times = []
    for i in range(n_runs):
        start = time.time()
        df_func(df).count()  # count() форсирует полное выполнение
        elapsed = time.time() - start
        times.append(elapsed)
        print(f"  Run {i+1}: {elapsed:.2f}s")
    median = sorted(times)[n_runs // 2]
    print(f"  Median: {median:.2f}s\n")
    return median

# ──────────────────────────────────────────────────────────
# Вариант 1: Классический @udf - row-at-a-time
# ──────────────────────────────────────────────────────────
@udf(returnType=DoubleType())
def haversine_row_udf(lat1, lon1, lat2, lon2):
    R = 6371.0
    phi1, phi2 = math.radians(lat1), math.radians(lat2)
    dphi = math.radians(lat2 - lat1)
    dlambda = math.radians(lon2 - lon1)
    a = math.sin(dphi/2)**2 + math.cos(phi1)*math.cos(phi2)*math.sin(dlambda/2)**2
    return R * 2 * math.asin(math.sqrt(a))

# ──────────────────────────────────────────────────────────
# Вариант 2: @pandas_udf - Arrow batch, Python loops
# ──────────────────────────────────────────────────────────
# Внимание: здесь умышленно написан НЕПРАВИЛЬНЫЙ pandas_udf
# с Python-циклом внутри. Это покажет, что батчирование
# само по себе без векторизации даёт меньший прирост.
@pandas_udf("double")
def haversine_pandas_loop(lat1, lon1, lat2, lon2):
    """Pandas UDF с Python-циклом ВНУТРИ (антипаттерн!)."""
    result = []
    for a, b, c, d in zip(lat1, lon1, lat2, lon2):
        R = 6371.0
        phi1, phi2 = math.radians(a), math.radians(c)
        dphi = math.radians(c - a)
        dlambda = math.radians(d - b)
        v = math.sin(dphi/2)**2 + math.cos(phi1)*math.cos(phi2)*math.sin(dlambda/2)**2
        result.append(R * 2 * math.asin(math.sqrt(v)))
    return pd.Series(result)

# ──────────────────────────────────────────────────────────
# Вариант 3: @pandas_udf - Arrow batch + NumPy vectorized
# ──────────────────────────────────────────────────────────
@pandas_udf("double")
def haversine_pandas_numpy(lat1, lon1, lat2, lon2):
    """Pandas UDF с NumPy векторизацией (правильный подход)."""
    R = 6371.0
    phi1 = np.radians(lat1)
    phi2 = np.radians(lat2)
    dphi = np.radians(lat2 - lat1)
    dlambda = np.radians(lon2 - lon1)
    a = np.sin(dphi/2)**2 + np.cos(phi1)*np.cos(phi2)*np.sin(dlambda/2)**2
    return pd.Series(R * 2 * np.arcsin(np.sqrt(a)))

# ──────────────────────────────────────────────────────────
# Вариант 4: Native Spark SQL (ground truth - самый быстрый)
# ──────────────────────────────────────────────────────────
# Spark SQL функции выполняются в Whole-Stage Codegen
# без разрыва WSCG, без Python, в нативном JVM байткоде.
# Формула гаверсинуса на Spark SQL:
def haversine_native_sql(df):
    return df.withColumn(
        "distance_km",
        F.lit(6371.0) * 2 * F.asin(F.sqrt(
            F.pow(F.sin(F.radians(F.col("lat2") - F.col("lat1")) / 2), 2) +
            F.cos(F.radians(F.col("lat1"))) *
            F.cos(F.radians(F.col("lat2"))) *
            F.pow(F.sin(F.radians(F.col("lon2") - F.col("lon1")) / 2), 2)
        ))
    )

# Запускаем бенчмарк
print("=== БЕНЧМАРК: 5M строк, формула гаверсинуса ===\n")

print("1. Классический @udf (row-at-a-time):")
t1 = benchmark("@udf", lambda df: df.withColumn("d", haversine_row_udf("lat1","lon1","lat2","lon2")), df)

print("2. @pandas_udf с Python-циклом ВНУТРИ (антипаттерн):")
t2 = benchmark("@pandas_udf loop", lambda df: df.withColumn("d", haversine_pandas_loop("lat1","lon1","lat2","lon2")), df)

print("3. @pandas_udf с NumPy векторизацией (правильно):")
t3 = benchmark("@pandas_udf numpy", lambda df: df.withColumn("d", haversine_pandas_numpy("lat1","lon1","lat2","lon2")), df)

print("4. Native Spark SQL (ground truth):")
t4 = benchmark("native SQL", haversine_native_sql, df)

print("\n=== ИТОГИ ===")
print(f"@udf (row-at-a-time):       {t1:.1f}s  (baseline)")
print(f"@pandas_udf (loop внутри):  {t2:.1f}s  ({t1/t2:.1f}× vs @udf)")
print(f"@pandas_udf (NumPy):        {t3:.1f}s  ({t1/t3:.1f}× vs @udf)")
print(f"Native Spark SQL:           {t4:.1f}s  ({t1/t4:.1f}× vs @udf)")

Память_батча ≈ maxRecordsPerBatch × sum(sizeof(column) for column in передаваемых колонок)

10 000 × (4 × 8 байт) = 320 000 байт ≈ 312 KB per batch

# ──────────────────────────────────────────────────────────
# Бенчмарк разных размеров батча
# ──────────────────────────────────────────────────────────
import psutil
import os

batch_sizes = [1_000, 5_000, 10_000, 25_000, 50_000, 100_000]
results = []

for batch_size in batch_sizes:
    # Пересоздаём SparkSession с новым размером батча
    spark.stop()
    spark = SparkSession.builder \
        .master("local[4]") \
        .config("spark.sql.execution.arrow.pyspark.enabled", "true") \
        .config("spark.sql.execution.arrow.maxRecordsPerBatch", str(batch_size)) \
        .getOrCreate()

    df = recreate_dataframe(spark)  # ваш тестовый DataFrame

    # Замер памяти Python-воркера
    process = psutil.Process(os.getpid())
    mem_before = process.memory_info().rss / 1024 / 1024

    start = time.time()
    df.withColumn("d", haversine_pandas_numpy("lat1","lon1","lat2","lon2")).count()
    elapsed = time.time() - start

    mem_after = process.memory_info().rss / 1024 / 1024

    results.append({
        "batch_size": batch_size,
        "time_s": elapsed,
        "mem_delta_mb": mem_after - mem_before
    })
    print(f"batch={batch_size:7d}: {elapsed:.2f}s, mem +{mem_after-mem_before:.0f}MB")

# Типичные результаты (5M строк, 4 double колонки):
# batch=  1_000: 8.3s, mem +12MB    - много overhead, мало пользы SIMD
# batch=  5_000: 4.1s, mem +25MB    - хороший баланс
# batch= 10_000: 3.0s, mem +45MB    - default, хорошо для простых UDF
# batch= 25_000: 2.7s, mem +90MB    - лучше для NumPy, больше памяти
# batch= 50_000: 2.5s, mem +175MB   - отличная SIMD, требует памяти
# batch=100_000: 2.4s, mem +340MB   - максимальная SIMD, риск OOM при больших строках

def safe_batch_size(
    python_worker_memory_mb: int,
    row_size_bytes: int,
    safety_factor: float = 0.3
) -> int:
    """
    Рассчитывает максимальный безопасный размер батча.

    safety_factor=0.3 означает использование не более 30%
    памяти Python-воркера для одного батча.
    Оставшиеся 70% нужны для:
    - Python-объектов и pandas overhead
    - Промежуточных результатов вычислений
    - ML-модели (если загружена в память)
    """
    available_bytes = python_worker_memory_mb * 1024 * 1024 * safety_factor
    return int(available_bytes / row_size_bytes)

# Пример:
# Python-воркер имеет 4 GB памяти
# Строка: 4 double колонки = 32 байта
# safe_batch_size(4096, 32) = 4096 * 1024 * 1024 * 0.3 / 32 = 40 миллионов
# Это слишком много - ограничиваем сверху
# На практике 10_000–50_000 - разумный диапазон

# Для ML-инференса с тяжёлой моделью (1 GB):
# Эффективная память = 4096 - 1024 = 3072 MB
# safe_batch_size(3072, 32) → ещё больше, но реальный ограничитель —
# время инференса на батч. Оптимальный размер для GPU-инференса
# определяется пропускной способностью GPU, а не RAM.

# ──────────────────────────────────────────────────────────
# NULL vs NaN: важная ловушка
# ──────────────────────────────────────────────────────────
@pandas_udf("double")
def safe_divide(numerator: pd.Series, denominator: pd.Series) -> pd.Series:
    """
    Деление с защитой от NULL и нуля.
    """
    # Когда Spark передаёт NULL через Arrow в Pandas,
    # NULL → float('nan') для числовых типов
    # NULL → None для строковых типов

    # ПРОБЛЕМА: denominator.isna() ловит И NaN И None,
    # но числовой 0 при делении даёт python ZeroDivisionError
    # или numpy RuntimeWarning (inf или nan)

    # ПРАВИЛЬНЫЙ подход: использовать numpy где возможно
    result = numerator / denominator  # numpy: 1/0 → inf, 0/0 → nan

    # Заменяем inf и nan на 0 (или другое sentinel value)
    result = result.replace([np.inf, -np.inf], np.nan)

    # Важно: возвращаем nan, а не None.
    # Arrow преобразует float nan обратно в Spark NULL корректно.
    return result

# Проверка:
test_df = spark.createDataFrame(
    [(1.0, 2.0), (3.0, 0.0), (None, 5.0), (6.0, None)],
    ["num", "den"]
)

test_df.withColumn("ratio", safe_divide("num", "den")).show()
# +----+----+-----+
# | num| den|ratio|
# +----+----+-----+
# | 1.0| 2.0|  0.5|
# | 3.0| 0.0| null|  ← 3/0 → inf → null
# |null| 5.0| null|  ← null numerator
# | 6.0|null| null|  ← null denominator
# +----+----+-----+

# ──────────────────────────────────────────────────────────
# Timestamp: настройка timezone
# ──────────────────────────────────────────────────────────
# Spark TimestampType → Arrow Timestamp[us, tz=None]
# или Timestamp[us, tz="UTC"] в зависимости от конфига

# Рекомендация: всегда работайте в UTC внутри Pandas UDF
@pandas_udf("long")
def extract_hour_utc(ts: pd.Series) -> pd.Series:
    """Извлекает час из timestamp в UTC."""
    # ts.dt.hour работает с timezone-naive или timezone-aware Series
    # Если Arrow передал без timezone (naive), Pandas считает UTC
    return ts.dt.hour.astype("int64")

# Если нужна конвертация timezone:
@pandas_udf("long")
def extract_hour_moscow(ts: pd.Series) -> pd.Series:
    """Извлекает час в московском времени."""
    # Делаем timestamp timezone-aware (UTC) и конвертируем
    ts_utc = ts.dt.tz_localize("UTC")
    ts_moscow = ts_utc.dt.tz_convert("Europe/Moscow")
    return ts_moscow.dt.hour.astype("int64")

# Настройте WARNING-уровень для Arrow fallback:
import logging
logging.getLogger("pyspark").setLevel(logging.WARNING)

# Или в SparkSession:
spark.sparkContext.setLogLevel("WARN")

# Сообщения, которые сигнализируют о fallback:
# "Arrow optimization failed, falling back to non-optimized method"
# Обычная причина: unsupported data type (MapType, некоторые nested types)

# Проверить, используется ли Arrow в конкретном плане:
result.explain(extended=True)
# Ищите: ArrowEvalPython (хорошо) vs BatchEvalPython (fallback)

# ──────────────────────────────────────────────────────────
# Шаг 1: наивный @udf (медленный baseline)
# ──────────────────────────────────────────────────────────
import math
from pyspark.sql.types import DoubleType

@udf(returnType=DoubleType())
def fraud_score_udf(
    amount,
    prev_lat, prev_lon,
    curr_lat, curr_lon,
    time_delta_minutes,
    category_risk  # float: коэффициент риска категории товара
):
    """
    Скоринг фрода: row-at-a-time UDF.
    Вычисляется для каждой строки отдельно.
    """
    if any(x is None for x in [amount, prev_lat, prev_lon,
                                curr_lat, curr_lon, time_delta_minutes]):
        return None

    # 1. Гео-скорость: км/ч между предыдущей и текущей транзакцией
    R = 6371.0
    phi1, phi2 = math.radians(prev_lat), math.radians(curr_lat)
    dphi = math.radians(curr_lat - prev_lat)
    dlambda = math.radians(curr_lon - prev_lon)
    a = math.sin(dphi/2)**2 + math.cos(phi1)*math.cos(phi2)*math.sin(dlambda/2)**2
    distance_km = R * 2 * math.asin(math.sqrt(a))
    speed_kmh = distance_km / max(time_delta_minutes / 60, 0.001)

    # 2. Компоненты скора
    # Скорость > 900 км/ч подозрительна (быстрее самолёта невозможно)
    speed_score = min(speed_kmh / 900.0, 1.0)

    # Аномальная сумма: нормализуем относительно порога 50 000 руб.
    amount_score = min(amount / 50_000.0, 1.0)

    # Категорийный риск передаётся как коэффициент 0.0–1.0
    category_score = category_risk

    # Итоговый скор: взвешенная сумма компонент
    return 0.4 * speed_score + 0.3 * amount_score + 0.3 * category_score

# Применяем к данным
result_baseline = transactions.withColumn(
    "fraud_score",
    fraud_score_udf(
        "amount", "prev_lat", "prev_lon",
        "curr_lat", "curr_lon",
        "time_delta_minutes", "category_risk"
    )
)

start = time.time()
flagged = result_baseline.filter(F.col("fraud_score") > 0.7).count()
print(f"@udf: {time.time() - start:.1f}s, flagged: {flagged:,}")

# ──────────────────────────────────────────────────────────
# Шаг 2: pandas_udf с NumPy - правильная векторизация
# ──────────────────────────────────────────────────────────
@pandas_udf("double")
def fraud_score_pandas(
    amount: pd.Series,
    prev_lat: pd.Series, prev_lon: pd.Series,
    curr_lat: pd.Series, curr_lon: pd.Series,
    time_delta_minutes: pd.Series,
    category_risk: pd.Series
) -> pd.Series:
    """
    Тот же скор фрода, но векторизованный.

    Каждый аргумент - pd.Series с 10000 значений (батч).
    Все операции работают над целым батчем через NumPy.

    Обработка None/NaN: np.where и np.nan_to_num
    позволяют избежать Python-циклов для null-handling.
    """
    # ── Гео-расстояние: векторизовано ──────────────────────
    R = 6371.0
    phi1 = np.radians(prev_lat)
    phi2 = np.radians(curr_lat)
    dphi = np.radians(curr_lat - prev_lat)
    dlambda = np.radians(curr_lon - prev_lon)

    a = np.sin(dphi/2)**2 + np.cos(phi1)*np.cos(phi2)*np.sin(dlambda/2)**2
    distance_km = R * 2 * np.arcsin(np.sqrt(np.clip(a, 0, 1)))

    # ── Скорость: векторизовано, защита от деления на 0 ────
    # np.maximum гарантирует denominator >= 0.001 (без if)
    time_hours = np.maximum(time_delta_minutes / 60.0, 0.001)
    speed_kmh = distance_km / time_hours

    # ── Компоненты скора: всё vectorized ───────────────────
    # np.minimum эквивалентно min() но для всего массива
    speed_score = np.minimum(speed_kmh / 900.0, 1.0)
    amount_score = np.minimum(amount / 50_000.0, 1.0)
    category_score = category_risk.values  # NumPy array напрямую

    # ── Финальный скор ─────────────────────────────────────
    result = 0.4 * speed_score + 0.3 * amount_score + 0.3 * category_score

    # ── Null handling: где были NaN во входе → NaN в выходе ─
    # NaN propagation автоматически работает в NumPy:
    # nan + anything = nan, так что строки с null входными
    # данными автоматически получат nan в результате,
    # который Spark преобразует обратно в NULL
    return pd.Series(result)

# Применяем pandas_udf
result_vectorized = transactions.withColumn(
    "fraud_score",
    fraud_score_pandas(
        "amount", "prev_lat", "prev_lon",
        "curr_lat", "curr_lon",
        "time_delta_minutes", "category_risk"
    )
)

start = time.time()
flagged = result_vectorized.filter(F.col("fraud_score") > 0.7).count()
print(f"@pandas_udf: {time.time() - start:.1f}s, flagged: {flagged:,}")

# ──────────────────────────────────────────────────────────
# Шаг 3: Iterator - загрузка тяжёлых данных один раз
# ──────────────────────────────────────────────────────────
# Реальный сценарий: категорийный риск хранится в большом
# словаре (50 000 SKU → коэффициент риска).
# Загружать этот словарь для каждого батча было бы неэффективно.

from typing import Iterator

@pandas_udf("double")
def fraud_score_with_lookup(
    iterator: Iterator[tuple[pd.Series, pd.Series, pd.Series,
                              pd.Series, pd.Series, pd.Series, pd.Series]]
) -> Iterator[pd.Series]:
    """
    Fraud score с предзагрузкой категорийного справочника.

    Справочник загружается ОДИН РАЗ при старте partition.
    Затем обрабатывается каждый батч по очереди.
    """
    # ── ИНИЦИАЛИЗАЦИЯ (один раз на partition) ──────────────
    # Загружаем справочник категорийных рисков
    # (в реальности - из S3, Redis, или broadcast переменной)
    import json
    with open("/data/category_risk_dict.json") as f:
        category_risk_dict = json.load(f)  # {sku_id: risk_score}

    # Конвертируем в NumPy-friendly структуру для быстрого lookup
    # defaultdict с 0.5 как риском для неизвестных категорий
    from collections import defaultdict
    risk_lookup = defaultdict(lambda: 0.5, category_risk_dict)

    # ── ОБРАБОТКА БАТЧЕЙ ───────────────────────────────────
    for (amount_s, prev_lat_s, prev_lon_s,
         curr_lat_s, curr_lon_s, time_s, sku_s) in iterator:

        # Lookup категорийного риска для всего батча
        # np.vectorize быстрее Python-цикла, но не так быстро как NumPy ufunc
        # Для справочного lookup это оптимальный подход
        category_risk = np.array([risk_lookup[sku] for sku in sku_s])

        # Остальная логика - та же NumPy векторизация
        R = 6371.0
        phi1 = np.radians(prev_lat_s)
        phi2 = np.radians(curr_lat_s)
        dphi = np.radians(curr_lat_s - prev_lat_s)
        dlambda = np.radians(curr_lon_s - prev_lon_s)
        a = np.sin(dphi/2)**2 + np.cos(phi1)*np.cos(phi2)*np.sin(dlambda/2)**2
        distance_km = R * 2 * np.arcsin(np.sqrt(np.clip(a, 0, 1)))

        time_hours = np.maximum(time_s / 60.0, 0.001)
        speed_kmh = distance_km / time_hours
        speed_score = np.minimum(speed_kmh / 900.0, 1.0)
        amount_score = np.minimum(amount_s / 50_000.0, 1.0)

        result = 0.4 * speed_score + 0.3 * amount_score + 0.3 * category_risk
        yield pd.Series(result)

# ❌ АНТИПАТТЕРН: Python-цикл внутри pandas_udf
@pandas_udf("double")
def bad_pandas_udf(series: pd.Series) -> pd.Series:
    result = []
    for value in series:           # ← Цикл! 10000 Python-итераций на батч
        result.append(value * 2.7183)
    return pd.Series(result)

# ✅ ПРАВИЛЬНО: NumPy операция над всем Series
@pandas_udf("double")
def good_pandas_udf(series: pd.Series) -> pd.Series:
    return series * 2.7183         # ← Одна NumPy операция над всем батчем

# ❌ АНТИПАТТЕРН: pandas_udf для операций, которые Spark умеет нативно
@pandas_udf("string")
def upper_case_udf(s: pd.Series) -> pd.Series:
    return s.str.upper()

result = df.withColumn("name_upper", upper_case_udf("name"))

# ✅ ПРАВИЛЬНО: встроенная Spark SQL функция
# F.upper работает в Whole-Stage Codegen без Python-границы
result = df.withColumn("name_upper", F.upper("name"))

# ❌ АНТИПАТТЕРН: многократный вызов pandas_udf в цикле на маленьких DataFrame
for user_id in user_ids:
    user_df = all_df.filter(F.col("user_id") == user_id)
    scored = user_df.withColumn("score", fraud_score_pandas(...))
    scored.write.mode("append").parquet(f"s3://bucket/scores/{user_id}/")

# Проблема: каждый filter создаёт новый план, каждый write - новый job.
# pandas_udf вызывается N раз для крошечных DataFrame.
# Overhead на инициализацию Python-воркера >> время вычислений.

# ✅ ПРАВИЛЬНО: обрабатываем всех пользователей одним DataFrame
all_scored = all_df.withColumn("score", fraud_score_pandas(...))
all_scored.write.mode("overwrite").partitionBy("user_id").parquet("s3://bucket/scores/")
# Или используем applyInPandas если нужна user-level логика

# Spark может молча отключить Arrow при unsupported типах.
# Если в логах есть:
# "Arrow optimization failed, ... falling back to non-optimized"
# - ваш pandas_udf работает как обычный @udf по производительности!

# Диагностика: проверьте explain()
df.withColumn("result", my_pandas_udf("col")).explain()
# Ищите: ArrowEvalPython (OK) или BatchEvalPython (fallback!)

# ──────────────────────────────────────────────────────────
# Матрица выбора в коде
# ──────────────────────────────────────────────────────────

# Задача: проверить, является ли email корректным

# Уровень 1 - Native Spark SQL (лучший вариант):
df.withColumn("email_valid",
    F.col("email").rlike(r"^[a-zA-Z0-9+_.-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$")
)

# Уровень 3 - pandas_udf (если нужна сложная валидация через Python-библиотеку):
import re
@pandas_udf("boolean")
def validate_email_pandas(emails: pd.Series) -> pd.Series:
    pattern = re.compile(r"^[a-zA-Z0-9+_.-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$")
    return emails.str.match(pattern)

# Уровень 4 - @udf (не нужен здесь, но для примера):
@udf("boolean")
def validate_email_udf(email):
    import re
    if email is None:
        return None
    return bool(re.match(r"^[a-zA-Z0-9+_.-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$", email))

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

# Медленная row-at-a-time реализация
@udf(returnType=IntegerType())
def levenshtein_slow(s1, s2):
    """Расстояние Левенштейна: минимальное число правок."""
    if s1 is None or s2 is None:
        return None
    m, n = len(s1), len(s2)
    dp = [[0] * (n + 1) for _ in range(m + 1)]
    for i in range(m + 1):
        dp[i][0] = i
    for j in range(n + 1):
        dp[0][j] = j
    for i in range(1, m + 1):
        for j in range(1, n + 1):
            if s1[i-1] == s2[j-1]:
                dp[i][j] = dp[i-1][j-1]
            else:
                dp[i][j] = 1 + min(dp[i-1][j], dp[i][j-1], dp[i-1][j-1])
    return dp[m][n]

# Применяется к 20 миллионам пар запросов
result = searches_df.withColumn(
    "edit_distance",
    levenshtein_slow("query", "reference")
)

from Levenshtein import distance as lev_distance
import numpy as np
# Подсказка: np.vectorize или list comprehension внутри pandas_udf