Профилирование PySpark: Java Stack Traces, PySpark Profiler и WholeStageCodegen

Как читать Java-исключения, находить Python-узкие места через PySpark Profiler и трактовать наличие/отсутствие WholeStageCodegen в Spark UI SQL tab.

optimization platform

Чтение Java Stack Traces¶

Ошибки в Spark выглядят как многокилометровые Java-трейсы. Правило: прокрутить до самой нижней строчки Caused by: - там настоящая причина.

java.lang.RuntimeException: Exception while running task
  at org.apache.spark.executor.Executor$TaskRunner.run(Executor.scala:597)
  at java.util.concurrent.ThreadPoolExecutor.runWorker(...)
  ... 45 more
Caused by: org.apache.spark.SparkException: Task failed while writing rows.
  at org.apache.spark.sql.execution.datasources.FileFormatWriter$.write(...)
  ... 23 more
Caused by: org.apache.parquet.io.ParquetEncodingException: ...
  at org.apache.parquet.column.values.dictionary.DictionaryValuesWriter.getBytes(...)
  ... 7 more
Caused by: java.lang.OutOfMemoryError: Java heap space   ← НАСТОЯЩАЯ ПРИЧИНА

Паттерны частых ошибок¶

Caused by	Диагноз	Действие
`OutOfMemoryError: Java heap space`	Недостаточно heap	Увеличить `executor.memory` или партиций
`OutOfMemoryError: GC overhead limit exceeded`	GC не успевает	G1GC, уменьшить `memory.fraction`
`FetchFailed`	Shuffle-файл недоступен	Executor умер, проверить `shuffle.service.enabled`
`FileNotFoundException` (shuffle)	`spark.local.dir` переполнен	Добавить диски, NVMe
`ArrowInvalid` / `ArrowTypeError`	Тип Python ↔ Arrow несовместим	Проверить return-тип Pandas UDF
`PicklingError`	Closure захватил непикуемый объект	Не передавать SparkSession в UDF
`TaskNotSerializable`	JVM-объект не сериализуем	Kryo или реструктуризация кода
`Container killed... memory limit`	memoryOverhead мало	Увеличить `executor.memoryOverhead`

Разбор Python-исключения в UDF¶

Python-исключения внутри UDF оборачиваются в Java-обёртку:

org.apache.spark.SparkException: Job aborted due to stage failure:
  ...
Caused by: org.apache.spark.api.python.PythonException:
  Traceback (most recent call last):
    File ".../pyspark/worker.py", line 239, in main
      func(split_index, iterator)
    File "your_code.py", line 42, in transform
      return json.loads(s)["missing_key"]   ← здесь ваша ошибка
  KeyError: 'missing_key'

Ищите PythonException: → ниже будет обычный Python traceback.

WholeStageCodegen в Spark UI¶

Вкладка SQL → Details показывает физический план с цветовой разметкой:

Оператор обёрнут в WholeStageCodegen → Tungsten объединил несколько операций в один цикл, генерируя Java bytecode на лету. Эффективно.
Оператор вне WholeStageCodegen → оператор выполняется через интерпретируемый режим (медленнее)

# В EXPLAIN это видно явно:
df.explain(mode="formatted")

# Ищите:
# *(1) Project [...]          ← * означает WholeStageCodegen (номер в скобках - id)
# *(1) Filter [...]
# *(1) HashAggregate [...]
# vs
# Exchange [...]              ← без * - shuffle, всегда вне WholeStageCodegen
# Scan parquet [...]          ← без * - IO операции не кодогенерируются

Когда WholeStageCodegen НЕ работает¶

Операции, которые «разрывают» codegen и переводят в интерпретируемый режим:

Python UDF - данные уходят в Python Worker
Некоторые сложные Window функции
RDD.mapPartitions с Python-функцией

# Проверить включение codegen:
spark.conf.set("spark.sql.codegen.wholeStage", "true")   # по умолчанию включено

# Отключить для отладки (замедляет в 2-5×, но даёт точные row-level ошибки):
spark.conf.set("spark.sql.codegen.wholeStage", "false")

PySpark Profiler¶

Встроенный профайлер Python для поиска узких мест в Python UDF.

from pyspark import BasicProfiler

# Включить профилирование при создании SparkContext
sc = SparkContext(
    conf=SparkConf().setAppName("Profiled"),
    profiler_cls=BasicProfiler
)
spark = SparkSession(sc)

После выполнения задачи:

# Вывести статистику по всем Python Workers
sc.show_profiles()

# Сохранить в файлы (pstats формат)
sc.dump_profiles("/tmp/spark-profiles/")

Анализ через pstats:

import pstats, io

stream = io.StringIO()
ps = pstats.Stats("/tmp/spark-profiles/profile-0", stream=stream)
ps.sort_stats("cumulative")
ps.print_stats(20)   # Топ-20 функций по кумулятивному времени
print(stream.getvalue())

Вывод:

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
  1200000    8.234    0.000   15.891    0.000 your_udf.py:15(parse_json)
  1200000    4.102    0.000    4.102    0.000 {method 'loads' of 'json.decoder'}

Узкое место найдено: parse_json вызывается 1.2М раз и занимает 16 секунд.

Альтернатива: cProfile напрямую в UDF¶

import cProfile, pstats, io
from pyspark.sql.functions import udf

def make_profiled_udf(func):
    profiler = cProfile.Profile()
    def wrapper(*args):
        return profiler.runcall(func, *args)
    wrapper._profiler = profiler
    return wrapper

# После нескольких итераций вывести в драйвер:
# Только для отладки - не для production

Чек-лист диагностики производительности¶

Slow job:
├── Spark UI → Jobs: есть ли долгий Stage?
│   └── Stages → Tasks: есть ли skew (один task в 10× дольше)?
│       ├── Да → salting / AQE SkewJoin
│       └── Нет → смотреть Task Metrics
│
├── Spark UI → SQL → Details
│   ├── Много операторов без WholeStageCodegen?
│   │   └── Убрать Python UDF, заменить на built-ins
│   └── Большой Shuffle Size?
│       └── Уменьшить shuffle.partitions / Bucketing / Broadcast
│
└── Executor Memory Issues?
    ├── Caused by OutOfMemoryError?
    │   └── executor.memory ↑ или партиций ↑
    └── Container killed (YARN/K8s)?
        └── executor.memoryOverhead ↑