Medev/skills

Spark & PySpark

Casi todo trabajo lento de Spark lo es por una de tres razones: un shuffle evitable, una clave con sesgo (skew), o filas que pasan por Python una a una. El error costoso es ajustar el tamaño del clúster antes de leer el plan: pagar por más ejecutores para ejecutar el mismo DAG ineficiente más rápido.

Procedimiento operativo

Paso 1: Recopilar información

1. El código del trabajo y, para un trabajo lento, la vista de la UI de Spark de la etapa lenta (distribución de duración de tareas, tamaños de lectura/escritura de shuffle) y la salida de df.explain(True).
2. Tamaño y formato de los datos de entrada, y los tamaños de ambos lados de cada join.
3. Configuración del clúster: número de ejecutores, núcleos y memoria.
4. Si el trabajo es una nueva creación o un pase de ajuste; para ajuste, diagnosticar desde el plan antes de tocar el código.

Paso 2: Comenzar con una sesión adecuada

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = (
    SparkSession.builder
    .appName("etl")
    .config("spark.sql.shuffle.partitions", "200")
    .config("spark.sql.adaptive.enabled", "true")
    .getOrCreate()
)

Habilitar Adaptive Query Execution (AQE); fusiona particiones de shuffle y cambia estrategias de join en tiempo de ejecución. Luego dimensionar spark.sql.shuffle.partitions deliberadamente: apuntar a 100-200 MB por partición de shuffle, así que particiones ≈ tamaño de datos de shuffle ÷ 128 MB, y al menos 2-3 veces el total de núcleos de ejecutores para que ningún núcleo esté inactivo.

Paso 3: Permanecer en la API de DataFrame y fuera de Python

Las operaciones de DataFrame pasan por Catalyst y Tungsten, lo que proporciona optimización de consultas y memoria fuera del heap. Pasar a RDDs solo cuando no exista un primitivo de DataFrame. Las UDF de Python serializan cada fila al worker de Python y viceversa, rompiendo la generación de código. Orden de preferencia:

1. Funciones integradas de pyspark.sql.functions (ej. F.regexp_extract, F.when).
2. UDF de Pandas (vectorizadas) cuando Python sea inevitable.
3. UDF escalares de Python solo como último recurso.

from pyspark.sql.functions import pandas_udf

@pandas_udf("double")
def normalize(s):
    return (s - s.mean()) / s.std()

Paso 4: Elegir la estrategia de join por tamaño

• Hacer broadcast del lado más pequeño cuando quepa en memoria. Spark emite broadcast automáticamente por debajo de spark.sql.autoBroadcastJoinThreshold (por defecto 10 MB); hacer broadcast explícitamente de tablas de hasta unos pocos cientos de MB cuando la memoria del ejecutor lo permita, y considerar ~1 GB como el límite práctico: broadcast copia la tabla a cada ejecutor y provoca OOM en el trabajo si se juzga mal.

from pyspark.sql.functions import broadcast
result = large.join(broadcast(small), "key")

• Filtrar y proyectar columnas antes de unirse para reducir el shuffle; cada columna y fila eliminada son bytes de shuffle ahorrados.
• Dos tablas grandes → aceptar el sort-merge join, pero verificar primero si hay skew (Paso 5).

Paso 5: Diagnosticar y corregir el skew

Síntomas de skew: la etapa está completa excepto por una o dos tareas; la duración máxima de la tarea supera ~3-5 veces la mediana; una partición de lectura de shuffle es mucho más grande que el resto en la UI de Spark. El manejo de skew de AQE divide una partición cuando es más de 5 veces el tamaño de partición mediana y más de 256 MB (los valores predeterminados); verificar que se activó en el plan. Cuando AQE no puede ayudar (ej. agregación sesgada), aplicar salting a las claves calientes: agregar un sufijo aleatorio 0..N a la clave en el lado grande, explotar el lado pequeño en todos los N sufijos, unir y luego quitar la sal. Elegir N ≈ el factor de skew (una clave 20 veces la mediana obtiene N=20).

Paso 6: Controlar la partición y el diseño de salida

• repartition(n, col) hace un shuffle completo para equilibrar los datos; usar antes de escrituras amplias.
• coalesce(n) reduce las particiones sin shuffle; usar para evitar archivos de salida pequeños.
• Apuntar a archivos de salida de 128 MB-1 GB; miles de archivos pequeños perjudican a todos los lectores downstream.
• Particionar la salida solo por columnas de baja cardinalidad (fecha, región: docenas a cientos de valores, nunca IDs):

df.write.partitionBy("dt").mode("overwrite").parquet("/data/out")

Paso 7: Cachear solo lo que se reutiliza

Cachear solo cuando un DataFrame se use múltiples veces en el DAG, materializarlo y despersistir al terminar:

df.cache()
df.count()  # materializar
# ... reutilizar df ...
df.unpersist()

Cachear un DataFrame usado una sola vez desperdicia memoria y puede desalojar datos que sí importaban.

Paso 8: Verificar con el plan

Volver a ejecutar df.explain(True) y confirmar que aparezca la estrategia deseada (BroadcastHashJoin, filtros empujados, sin Exchange donde se eliminó uno). Leer formatos columnares (Parquet, ORC) para que funcione el pushdown de predicados, nunca hacer collect() de datos grandes al driver, y preferir referencias F.col sobre cadenas de columnas al encadenar.

Ejemplo práctico: join malo vs bueno

Malo: una tabla de eventos de 2 TB unida a una dimensión de 200 MB, ancho completo, UDF de Python para analizar:

result = (events.join(dims, "store_id")          # sort-merge: shufflea todos los 2 TB
    .withColumn("region", parse_region_udf("meta"))  # Python fila por fila
    .filter(F.col("dt") == "2024-06-01"))            # filtrar DESPUÉS del join

El plan muestra un Exchange en ambos lados (2 TB shuffleados), y la UDF bloquea la generación de código.

Bueno: filtrar y proyectar primero, hacer broadcast de la dimensión, usar una función integrada:

result = (events
    .filter(F.col("dt") == "2024-06-01")              # podar a ~30 GB primero
    .select("store_id", "amount", "meta")
    .join(broadcast(dims.select("store_id", "region_code")), "store_id")
    .withColumn("region", F.regexp_extract("meta", r"region=(\w+)", 1)))

El plan ahora muestra BroadcastHashJoin sin Exchange del lado grande: ~30 GB escaneados, cero shuffle de la tabla grande, sin límite de Python.

Entregable

Producir el código del trabajo revisado más una nota de ajuste que indique: la estrategia de join para cada join con la evidencia de tamaño, la configuración de las particiones de shuffle con su cálculo, cualquier skew encontrado (relación de tarea máxima vs mediana) y la corrección aplicada, decisiones de caché, expectativas de tamaño de archivos de salida y los deltas de explain antes y después para un pase de ajuste.

NO hacer

• NO hacer broadcast de una tabla que no se haya medido; un broadcast mal calculado provoca OOM en todos los ejecutores a la vez.
• NO escribir una UDF de Python antes de verificar pyspark.sql.functions; casi toda la lógica de cadenas, fechas y condicionales tiene una función integrada.
• NO hacer collect() o toPandas() de un DataFrame grande; canaliza los datos del clúster a través de un solo driver.
• NO particionar la salida por una columna de alta cardinalidad: millones de directorios, una fila cada uno.
• NO agregar ejecutores a un trabajo cuya etapa lenta es una tarea sesgada; más núcleos no pueden dividir una partición.
• NO cachear por reflejo; cachear solo DataFrames reutilizados downstream, y despersistirlos.

Nivel de calidad

Un trabajo terminado pasa solo cuando: cada join nombra su estrategia y la evidencia de tamaño para ella; spark.sql.shuffle.partitions se deriva del tamaño de los datos, no se deja en un valor por costumbre; la UI de Spark no muestra ninguna tarea por encima de ~3x la duración mediana en etapas de shuffle (o el skew está explicado y corregido); no sobrevive ninguna UDF escalar de Python donde exista una función integrada; los archivos de salida caen en el rango de 128 MB-1 GB; y el explain final confirma el plan previsto.