Large-Scale Aggregation with PySpark

Overview

PySpark is the Python API for Apache Spark, a distributed data processing engine. Group-by aggregations are one of the most common operations in data pipelines: summing revenue by customer, counting events by day, averaging metrics by region. On large datasets (100GB+), how you write group-by operations changes performance a lot. This recipe covers basic aggregations, window functions, UDAFs, broadcast joins, and partition tuning.

When to Use

• Datasets larger than 100GB that don’t fit on a single machine
• Aggregations across billions of rows (clickstream, IoT, transaction logs)
• When you need distributed processing across a cluster
• Pipelines that read from/write to distributed storage (S3, HDFS, GCS)
• When pandas/Polars run out of memory

When NOT to Use

• Datasets under 10GB: pandas or Polars are faster due to no serialization overhead
• Interactive analysis on small data: pandas is more ergonomic
• Real-time processing: use Structured Streaming or Flink
• Simple transformations without aggregation: a SQL query on the warehouse is simpler

Prerequisites

• PySpark 3.5+: install with pip install pyspark or use a cluster runtime (Databricks, EMR, Dataproc).
• Java 8 or 11: Spark sits on top of the JVM, so point JAVA_HOME at your JDK install.
• Python 3.8+: PySpark 3.5 needs Python 3.8, 3.9, 3.10, or 3.11.
• A distributed storage backend (S3, HDFS, GCS, Azure Blob) when your dataset outgrows a single node.

Solution

Basic group-by aggregation

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

# Basic group-by
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 multiple columns

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

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

# Row number — latest order gets 1
df_with_rank = df.withColumn(
    "order_rank",
    F.row_number().over(window_spec)
)

# Running total per customer ordered by 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 — previous order amount
df_with_lag = df.withColumn(
    "prev_amount",
    F.lag("amount", 1).over(window_spec)
)

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

Multiple aggregations with different groupings

from pyspark.sql import DataFrame

def aggregate_multiple_ways(df: DataFrame) -> DataFrame:
    """Perform multiple aggregations in a single 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 for small dimension tables

# Small dimension table — broadcast to all executors
customers = spark.read.parquet("s3://data-lake/customers/")
orders = spark.read.parquet("s3://data-lake/orders/")

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

# Without broadcast — triggers a shuffle (slow for large tables)
# joined = orders.join(customers, on="customer_id", how="left")

Aggregation with 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 for custom aggregation (vectorized — faster than 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")
    )
)

Partition tuning

# Set shuffle partitions (default is 200 — often too many for small data)
spark.conf.set("spark.sql.shuffle.partitions", "50")

# Repartition before group-by to avoid skew
df_repartitioned = df.repartition(100, "customer_id")

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

# Coalesce after aggregation to reduce small files
result = result.coalesce(10)
result.write.parquet("s3://data-lake/aggregated/")

Handling data skew

# Salting technique for skewed keys
from pyspark.sql.functions import concat, lit, rand, floor, explode, array

# Add a salt key to split large groups
df_salted = df.withColumn(
    "salt",
    floor(rand() * 10).cast("int")
)

# Group by with salt — splits large groups across partitions
partial = (
    df_salted
    .groupBy("customer_id", "salt")
    .agg(F.sum("amount").alias("partial_sum"))
)

# Second aggregation without salt to combine
final = (
    partial
    .groupBy("customer_id")
    .agg(F.sum("partial_sum").alias("total_sum"))
)

Saving results

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

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

# Write to Hive table
result.write \
    .mode("overwrite") \
    .saveAsTable("analytics.orders_summary")

Explanation

A group-by aggregation in Spark flows through several stages. Understanding the pipeline helps you tune each step:

The shuffle is the most expensive step. repartition() triggers a full shuffle across the cluster, so you want to keep how often it happens to a minimum. When you broadcast a small dimension table, Spark sends a copy to every executor instead of shuffling the large fact table. Adaptive Query Execution (AQE) can coalesce small partitions and split skewed ones at runtime, which is why enabling it’s the single highest-impact tuning change for most workloads.

For skewed keys, the salting technique splits a large group into 10 sub-groups, aggregates each separately, then combines them in a second pass. This adds one extra stage but stops a single executor from holding the entire large group. For a deeper comparison of distributed DataFrame engines, see Polars vs PySpark for large datasets and Pandas ETL pipeline patterns.

Variants

Using Spark SQL

# Register DataFrame as a 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
""")

Streaming aggregation with 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 for iterative workloads

# Cache a DataFrame used multiple times
df_cached = df.filter(F.col("status") == "completed").cache()

# First action — materializes cache
result1 = df_cached.groupBy("customer_id").agg(F.sum("amount").alias("total"))

# Second action — uses cache
result2 = df_cached.groupBy("category").agg(F.avg("amount").alias("avg"))

# Unpersist when done
df_cached.unpersist()

Best Practices

Tune spark.sql.shuffle.partitions to match your data size. The default 200 works for large jobs, but drop it for smaller data. Use broadcast() for dimension tables under 10MB to avoid expensive shuffles.

Use Pandas UDFs instead of regular UDFs; they’re vectorized and 10-100x faster. Filter early and push filters before joins and aggregations to reduce data volume.

Use coalesce() instead of repartition() when cutting down partitions; it skips a full shuffle. Enable Adaptive Query Execution (spark.sql.adaptive.enabled=true) so Spark can tune the plan at runtime.

Use partitionBy when writing to enable predicate pushdown for downstream reads, and avoid collect() on large DataFrames because it brings all data to the driver. If you aren’t running at cluster scale, Dask parallel DataFrame operations can be a simpler alternative. For datasets that fit in memory on a single node, Polars fast DataFrame operations are often faster than Spark. For orchestrating PySpark jobs in production pipelines, see Airflow DAG scheduling patterns.

Common Mistakes

Not setting shuffle partitions is a frequent issue. The default 200 creates tiny tasks for small aggregations; set it to 20-50 for small data.

Using regular UDFs instead of Pandas UDFs is slower because regular UDFs serialize each row individually. Pandas UDFs process in batches. A 5MB dimension table shuffled across 200 partitions is wasteful, so broadcast small tables instead letting Spark shuffle them.

Calling collect() on large results brings all data to the driver and can crash it. Use show(), take(), or write to storage instead. If you reuse a DataFrame three or more times, cache it; otherwise Spark recomputes the lineage each time.

Summary

• Set spark.sql.shuffle.partitions based on data size; the default 200 is too many for small aggregations.
• Use broadcast() for dimension tables under 10MB to avoid expensive shuffles.
• Prefer Pandas UDFs over regular UDFs: they’re vectorized and 10-100x faster.
• Use coalesce() to reduce partitions without a full shuffle; repartition() when you need to increase or fix skew.
• Enable Adaptive Query Execution (spark.sql.adaptive.enabled=true) so Spark can coalesce and skew-correct at runtime.
• Handle skewed keys with the salting technique: split, aggregate, then combine in a second stage.
• Write partitioned output (partitionBy) to enable predicate pushdown for downstream reads.
• Full companion code with tests: stack-practices-resources.

See Also

• PySpark SQL documentation — Official PySpark SQL API reference with examples for every function.
• Spark SQL functions — Built-in SQL functions available in spark.sql() and F.expr().
• Adaptive Query Execution guide — How AQE coalesces partitions and handles skew at runtime.
• Databricks blog on salting — Practical guide to handling skewed joins and aggregations.
• Spark performance tuning guide — Official tuning recommendations for shuffle partitions, broadcast, and caching.