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+), la forma en que escribís las operaciones group-by afecta la performance dramáticamente. 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

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")

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

Seteá spark.sql.shuffle.partitions según el tamaño de los datos. El default de 200 sirve para jobs grandes, pero usá menos para datos pequeños. 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; evita 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.

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.