Agregaciones a Gran Escala con PySpark

Overview

PySpark es la API de Python para Apache Spark, un engine de procesamiento de datos distribuido. Las agregaciones group-by son una de las operaciones más comunes en pipelines de datos: sumar revenue por customer, contar eventos por día, promediar métricas por región. En datasets grandes (100GB+), cómo escribís las operaciones group-by cambia la performance un montón. Esta receta cubre agregaciones básicas, window functions, UDAFs, broadcast joins y tuning de particiones.

When to Use

• Datasets más grandes que 100GB que no caben en una sola máquina
• Agregaciones a través de billones de rows (clickstream, IoT, transaction logs)
• Cuando necesitas procesamiento distribuido a través de un cluster
• Pipelines que leen de/escriben a storage distribuido (S3, HDFS, GCS)
• Cuando pandas/Polars se quedan sin memoria

When NOT to Use

• Datasets bajo 10GB: pandas o Polars son más rápidos por no tener overhead de serialización
• Análisis interactivo en data pequeña: pandas es más ergonómico
• Procesamiento real-time: usa Structured Streaming o Flink
• Transformaciones simples sin agregación: un SQL query en el warehouse es más simple

Requisitos Previos

• PySpark 3.5+: instala con pip install pyspark o usa un runtime de cluster (Databricks, EMR, Dataproc).
• Java 8 u 11: Spark corre sobre la JVM, apuntá JAVA_HOME a tu instalación de JDK.
• Python 3.8+: PySpark 3.5 necesita Python 3.8, 3.9, 3.10 o 3.11.
• Un backend de storage distribuido (S3, HDFS, GCS, Azure Blob) cuando tu dataset supera un solo nodo.

Solution

Agregación group-by básica

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

spark = SparkSession.builder \
    .appName("aggregations") \
    .config("spark.sql.adaptive.enabled", "true") \
    .config("spark.sql.adaptive.coalescePartitions.enabled", "true") \
    .getOrCreate()

df = spark.read.parquet("s3://data-lake/orders/")

# Group-by básico
result = (
    df
    .groupBy("customer_id")
    .agg(
        F.sum("amount").alias("total_spent"),
        F.count("order_id").alias("order_count"),
        F.avg("amount").alias("avg_order_value"),
        F.max("order_date").alias("last_order_date"),
        F.min("order_date").alias("first_order_date"),
    )
    .orderBy(F.desc("total_spent"))
)

result.show(20)

Group by múltiples columnas

result = (
    df
    .groupBy("customer_id", F.date_format("order_date", "yyyy-MM").alias("month"))
    .agg(
        F.sum("amount").alias("monthly_spent"),
        F.countDistinct("order_id").alias("unique_orders"),
    )
    .orderBy("customer_id", "month")
)

Window functions

from pyspark.sql import Window

# Definir window specification
window_spec = Window.partitionBy("customer_id").orderBy(F.desc("order_date"))

# Row number — el order más reciente gets 1
df_with_rank = df.withColumn(
    "order_rank",
    F.row_number().over(window_spec)
)

# Running total per customer ordered por date
running_total_window = (
    Window
    .partitionBy("customer_id")
    .orderBy("order_date")
    .rowsBetween(Window.unboundedPreceding, Window.currentRow)
)

df_with_running = df.withColumn(
    "running_total",
    F.sum("amount").over(running_total_window)
)

# Lag — amount del order anterior
df_with_lag = df.withColumn(
    "prev_amount",
    F.lag("amount", 1).over(window_spec)
)

# Percentile dentro del group
df_with_pct = df.withColumn(
    "amount_percentile",
    F.percent_rank().over(Window.partitionBy("category").orderBy("amount"))
)

Múltiples agregaciones con diferentes groupings

from pyspark.sql import DataFrame

def aggregate_multiple_ways(df: DataFrame) -> DataFrame:
    """Realizar múltiples agregaciones en un solo pass."""
    return (
        df
        .groupBy("customer_id")
        .agg(
            F.sum("amount").alias("total_spent"),
            F.sum(F.when(F.col("status") == "completed", F.col("amount")).otherwise(0)).alias("completed_amount"),
            F.sum(F.when(F.col("status") == "cancelled", F.col("amount")).otherwise(0)).alias("cancelled_amount"),
            F.count(F.when(F.col("amount") > 100, 1)).alias("large_orders"),
            F.collect_set("category").alias("categories"),
            F.expr("percentile(amount, 0.95)").alias("p95_amount"),
        )
    )

Broadcast join para tablas de dimensiones pequeñas

# Tabla de dimensiones pequeña — broadcast a todos los executors
customers = spark.read.parquet("s3://data-lake/customers/")
orders = spark.read.parquet("s3://data-lake/orders/")

# Broadcast join — evita shuffle
joined = orders.join(
    F.broadcast(customers),
    on="customer_id",
    how="left"
)

# Sin broadcast — triggerea un shuffle (lento para tablas grandes)
# joined = orders.join(customers, on="customer_id", how="left")

Agregación con pivot

# Pivot: rows=customer, columns=month, values=sum(amount)
pivoted = (
    df
    .groupBy("customer_id")
    .pivot("month")  # or .pivot("month", ["2025-01", "2025-02", "2025-03"])
    .agg(F.sum("amount"))
)

User Defined Aggregate Function (UDAF)

from pyspark.sql.types import DoubleType
from pyspark.sql.functions import udf
from pyspark.sql.functions import struct

# Pandas UDF para agregación custom (vectorized — más rápido que regular UDF)
@F.pandas_udf(DoubleType())
def custom_metric(amounts: pd.Series, quantities: pd.Series) -> float:
    """Weighted average price."""
    total_qty = quantities.sum()
    if total_qty == 0:
        return 0.0
    return (amounts * quantities).sum() / total_qty

result = (
    df
    .groupBy("customer_id")
    .agg(
        custom_metric(F.col("amount"), F.col("quantity")).alias("weighted_avg_price")
    )
)

Tuning de particiones

# Setear shuffle partitions (default es 200 — a menudo demasiado para data pequeña)
spark.conf.set("spark.sql.shuffle.partitions", "50")

# Repartition antes de group-by para evitar skew
df_repartitioned = df.repartition(100, "customer_id")

result = (
    df_repartitioned
    .groupBy("customer_id")
    .agg(F.sum("amount").alias("total"))
)

# Coalesce después de agregación para reducir small files
result = result.coalesce(10)
result.write.parquet("s3://data-lake/aggregated/")

Manejar data skew

# Técnica de salting para keys skewed
from pyspark.sql.functions import concat, lit, rand, floor, explode, array

# Agregar salt key para splittear grupos grandes
df_salted = df.withColumn(
    "salt",
    floor(rand() * 10).cast("int")
)

# Group by con salt — splittea grupos grandes a través de particiones
partial = (
    df_salted
    .groupBy("customer_id", "salt")
    .agg(F.sum("amount").alias("partial_sum"))
)

# Segunda agregación sin salt para combinar
final = (
    partial
    .groupBy("customer_id")
    .agg(F.sum("partial_sum").alias("total_sum"))
)

Guardar resultados

# Escribir como Parquet partitioned por date
result.write \
    .partitionBy("year", "month") \
    .mode("overwrite") \
    .parquet("s3://data-lake/aggregated/orders_by_month/")

# Escribir como CSV
result.write \
    .mode("overwrite") \
    .option("header", "true") \
    .csv("s3://data-lake/aggregated/orders_csv/")

# Escribir a Hive table
result.write \
    .mode("overwrite") \
    .saveAsTable("analytics.orders_summary")

Explicación

Una agregación group-by en Spark fluye por varias etapas. Entender el pipeline ayuda a tunear cada paso:

El shuffle es el paso más costoso. repartition() dispara un full shuffle a través del cluster, así que conviene mantener al mínimo cuántas veces ocurre. Cuando broadcasteás una tabla de dimensiones pequeña, Spark manda una copia a cada executor en lugar de shufflear la tabla de hechos grande. Adaptive Query Execution (AQE) puede coalescear particiones pequeñas y splitear las skewed en runtime, por eso habilitarlo es el cambio de tuning de mayor impacto para la mayoría de los workloads.

Para keys skewed, la técnica de salting splitea un grupo grande en 10 sub-grupos, agrega cada uno por separado, y los combina en un segundo pass. Esto agrega una etapa extra pero frena a un solo executor de retener todo el grupo grande. Para una comparación más profunda de engines de DataFrames distribuidos, véase Polars vs PySpark para datasets grandes y patrones de ETL con pandas.

Variants

Usar Spark SQL

# Registrar DataFrame como temp view
df.createOrReplaceTempView("orders")

result = spark.sql("""
    SELECT
        customer_id,
        SUM(amount) AS total_spent,
        COUNT(DISTINCT order_id) AS unique_orders,
        PERCENTILE(amount, 0.95) AS p95_amount
    FROM orders
    WHERE status = 'completed'
    GROUP BY customer_id
    ORDER BY total_spent DESC
""")

Agregación streaming con Structured Streaming

streaming_df = spark \
    .readStream \
    .format("kafka") \
    .option("kafka.bootstrap.servers", "localhost:9092") \
    .option("subscribe", "orders") \
    .load()

aggregated = (
    streaming_df
    .selectExpr("CAST(value AS STRING) AS json")
    .selectExpr("json_tuple(json, 'customer_id', 'amount') AS (customer_id, amount)")
    .groupBy("customer_id")
    .agg(F.sum("amount").alias("total"))
)

query = aggregated \
    .writeStream \
    .outputMode("complete") \
    .format("console") \
    .start()

Caching para workloads iterativos

# Cachear un DataFrame usado múltiples veces
df_cached = df.filter(F.col("status") == "completed").cache()

# Primera action — materializa el cache
result1 = df_cached.groupBy("customer_id").agg(F.sum("amount").alias("total"))

# Segunda action — usa el cache
result2 = df_cached.groupBy("category").agg(F.avg("amount").alias("avg"))

# Unpersist cuando termines
df_cached.unpersist()

Best Practices

Tuneá spark.sql.shuffle.partitions para matchear el tamaño de tus datos. El default 200 sirve para jobs grandes, pero bajalo para data pequeña. Usá broadcast() para tablas de dimensiones menores a 10MB y evitá shuffles costosos.

Usá Pandas UDFs en lugar de UDFs regulares; son vectorizados y 10-100x más rápidos. Filtrá temprano y empujá los filtros antes de joins y agregaciones para reducir el volumen de datos.

Usá coalesce() en lugar de repartition() cuando reducís particiones; saltea un full shuffle. Habilitá Adaptive Query Execution (spark.sql.adaptive.enabled=true) para que Spark optimice el plan en runtime. Usá partitionBy al escribir para habilitar predicate pushdown en lecturas downstream, y evitá collect() en DataFrames grandes porque trae toda la data al driver. Si no estás a escala de cluster, las operaciones paralelas con DataFrames en Dask pueden ser una alternativa más simple. Para datasets que caben en memoria en un solo nodo, las operaciones rápidas con Polars suelen ser más rápidas que Spark. Para orquestar jobs de PySpark en pipelines de producción, véase patrones de scheduling de DAGs con Airflow.

Common Mistakes

No setear shuffle partitions es un problema frecuente. El default de 200 crea tareas pequeñas para agregaciones chicas; setealo a 20-50 para datos pequeños.

Usar UDFs regulares en lugar de Pandas UDFs es más lento porque las UDFs regulares serializan cada fila individualmente. Las Pandas UDFs procesan en batches. Una tabla de dimensión de 5MB shufflada a través de 200 particiones es un desperdicio, así que broadcasteá tablas pequeñas en lugar de dejar que Spark las shufflee.

Llamar collect() en resultados grandes trae toda la data al driver y puede crashearlo. Usá show(), take() o escribí a storage. Si reusás un DataFrame tres o más veces, cachéalo; si no, Spark recomputa el lineage cada vez.

Resumen

• Seteá spark.sql.shuffle.partitions según el tamaño de los datos; el default 200 es demasiado para agregaciones chicas.
• Usá broadcast() para tablas de dimensiones menores a 10MB y evitá shuffles costosos.
• Preferí Pandas UDFs sobre UDFs regulares: son vectorizados y 10-100x más rápidos.
• Usá coalesce() para reducir particiones sin full shuffle; repartition() cuando necesitás aumentar o arreglar skew.
• Habilitá Adaptive Query Execution (spark.sql.adaptive.enabled=true) para que Spark coalescee y corrija skew en runtime.
• Manejá keys skewed con la técnica de salting: spliteá, agregá, combiná en una segunda etapa.
• Escribí output particionado (partitionBy) para habilitar predicate pushdown en lecturas downstream.
• Código companion completo con tests: stack-practices-resources.

See Also

• Documentación de PySpark SQL — Referencia oficial de la API PySpark SQL con ejemplos para cada función.
• Funciones de Spark SQL — Funciones SQL built-in disponibles en spark.sql() y F.expr().
• Guía de Adaptive Query Execution — Cómo AQE coalescea particiones y maneja skew en runtime.
• Blog de Databricks sobre salting — Guía práctica para manejar joins y agregaciones skewed.
• Guía de tuning de performance de Spark — Recomendaciones oficiales para shuffle partitions, broadcast y caching.