Py4J Gateway и Python Workers: как Python-код переходит в JVM

# Что вы пишете
df = spark.read.parquet("/data/events")

# Что происходит за кулисами
# spark._jvm.org.apache.spark.sql.SparkSession → proxy SparkSession
# .read → proxy DataFrameReader (новый Java-объект, новый id в JVM)
# .parquet("/data/events") → вызов Java-метода, возвращает proxy DataFrame

# Антипаттерн: создание тысяч промежуточных DataFrame в цикле
for table in all_tables:
    df = spark.table(table)
    result = df.count()
    # df - proxy на JVM-объект, финализируется GC Python'а
    # при большом числе таблиц JVM heap заполняется proxy-объектами

# В SparkConf или spark-defaults.conf
spark.python.worker.reuse = true  # default: true

# Когда отключить worker reuse:
# - UDF использует global state, который нельзя сбрасывать
# - Memory leak в third-party библиотеке
# - Отладка (каждый запуск = чистое окружение)

import cloudpickle

# Этот UDF умеет захватывать замыкание:
threshold = 100.0

@udf("boolean")
def is_large(v):
    return v > threshold  # threshold попадает в байткод через cloudpickle

# cloudpickle сериализует:
# - байткод функции
# - глобальные переменные, захваченные в замыкании (threshold)
# - вложенные функции и классы

# НЕ сериализуется:
# - открытые файловые дескрипторы
# - database connections (psycopg2.connection, etc.)
# - thread.Lock
# - объекты C-extension без __reduce__

# Типичная ошибка:
db_conn = psycopg2.connect(...)  # глобальная переменная

@udf("string")
def lookup(id):
    return db_conn.execute(...)  # PicklingError при отправке на executor

@udf("string")
def lookup(id):
    # Соединение создаётся на Executor'е при первом вызове (worker reuse кэширует)
    import psycopg2
    conn = psycopg2.connect(host="...", database="...")
    return conn.execute("SELECT name FROM t WHERE id = %s", [id]).fetchone()[0]

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

# Scalar UDF: Series → Series (построчная логика над батчем)
@pandas_udf("double")
def normalize(s: pd.Series) -> pd.Series:
    return (s - s.mean()) / s.std()

df.withColumn("norm_age", normalize("age"))

# Grouped Map UDF (applyInPandas): каждая группа → pd.DataFrame → pd.DataFrame
def fill_gaps(key, pdf):
    # pdf - весь pandas DataFrame для данной группы
    return pdf.fillna(method="ffill")

df.groupBy("user_id").applyInPandas(fill_gaps, schema=df.schema)

# Grouped Aggregate UDF: Series → scalar (кастомные агрегаты)
@pandas_udf("double")
def trimmed_mean(s: pd.Series) -> float:
    from scipy.stats import trim_mean
    return trim_mean(s.dropna(), 0.1)

df.groupBy("category").agg(trimmed_mean("price"))

# Для тяжёлых ML-моделей (много памяти на батч):
spark.conf.set("spark.sql.execution.arrow.maxRecordsPerBatch", "1000")

# Для лёгких операций (накладные расходы важнее):
spark.conf.set("spark.sql.execution.arrow.maxRecordsPerBatch", "50000")

# Фильтр сработает ПОСЛЕ UDF (Catalyst не видит внутрь Python)
df.withColumn("result", my_udf("value")).filter(col("result") > 0)

# Если бизнес-логика позволяет, лучше фильтровать ДО UDF:
df.filter(col("value") > 0).withColumn("result", my_udf("value"))

# На каждом Executor'е при работающем job:
ps aux | grep pyspark.worker

# Типичный вывод:
# spark    12345  99.0  2.1  820MB python3 -m pyspark.worker
# spark    12346  98.5  2.0  815MB python3 -m pyspark.worker
# spark    12347  97.2  1.9  800MB python3 -m pyspark.worker
# - один процесс на партицию (при worker.reuse=true - переиспользуются)

# Включить детальный лог Worker'а:
spark.sparkContext.setLogLevel("DEBUG")

# Или через SparkConf:
# spark.python.worker.reuse=false  # каждый запуск = новый лог

yarn logs -applicationId application_XXXXX_YYYYY | grep "Python worker"

kubectl logs spark-exec-1 -c spark-kubernetes-driver | grep "PythonRunner"

try:
    df.write.mode("overwrite").parquet("/output")
except Exception as e:
    # Py4JJavaError содержит вложенное Java-исключение
    # Нужно искать "Caused by:" в трейсбеке
    print(str(e))  # первая строка - Python wrapper
    # Реальная причина - глубже в стеке

from pyspark.sql.functions import col, lower, trim, pandas_udf, udf
from pyspark.sql.types import StringType
import pandas as pd
import time

# Создаём тестовый датасет
df = spark.range(10_000_000).selectExpr(
    "CAST(id AS STRING) AS raw_name",
    "CONCAT('  USER_', CAST(id AS STRING), '  ') AS name"
)
df.cache()
df.count()  # прогрев

# Вариант 1: built-in функции
t0 = time.time()
df.select(lower(trim(col("name")))).count()
t_builtin = time.time() - t0

# Вариант 2: Python UDF (Pickle)
@udf(StringType())
def clean_udf(s):
    return s.strip().lower() if s else None

t0 = time.time()
df.select(clean_udf(col("name"))).count()
t_pyudf = time.time() - t0

# Вариант 3: Pandas UDF (Arrow)
@pandas_udf(StringType())
def clean_pandas(s: pd.Series) -> pd.Series:
    return s.str.strip().str.lower()

t0 = time.time()
df.select(clean_pandas(col("name"))).count()
t_pandas = time.time() - t0

print(f"Built-in:   {t_builtin:.1f}s (1.0×)")
print(f"Python UDF: {t_pyudf:.1f}s ({t_pyudf/t_builtin:.1f}×)")
print(f"Pandas UDF: {t_pandas:.1f}s ({t_pandas/t_builtin:.1f}×)")

# Посмотреть план с указанием типа UDF
df.select(clean_udf(col("name"))).explain(mode="extended")
# Ищем: BatchEvalPython или ArrowEvalPython

df.select(clean_pandas(col("name"))).explain(mode="extended")
# Ищем: ArrowEvalPython - хороший знак (Arrow путь)

import time

# Один вызов через Py4J
t0 = time.time()
for _ in range(1000):
    _ = spark.sparkContext.defaultParallelism  # одно свойство через TCP
print(f"1000 Py4J calls: {(time.time()-t0)*1000:.1f}ms")
# Обычно 50–200ms для 1000 вызовов - ~0.1ms на вызов

# Вывод: не вызывать Py4J в Python-цикле с большим числом итераций

import cloudpickle

# Проверить, сериализуется ли UDF до отправки на Executor
@udf("string")
def my_udf(s):
    return s.upper()

try:
    serialized = cloudpickle.dumps(my_udf.func)
    print(f"Размер сериализованного UDF: {len(serialized)} bytes")
    # Восстановить:
    restored = cloudpickle.loads(serialized)
    print(f"Тест: {restored('hello')}")
except Exception as e:
    print(f"Ошибка сериализации: {e}")
    # Это то, что увидит Executor при попытке выполнить задачу