In Spark, moving data is almost always more expensive than processing it. Shuffle is where that cost hides.

Every Spark job that involves a groupBy, a join, or a distinct carries a hidden cost that rarely shows up until it’s too late: the cluster is running, resources are allocated, and yet the job crawls. In many cases, the culprit isn’t the volume of data or the complexity of the transformations — it’s shuffle, the mechanism Spark uses to redistribute data across partitions so that all records sharing the same key end up together on the same node.

Shuffle is triggered by every wide dependency, and unlike narrow transformations that execute in isolation within each partition, it forces the engine to serialize data, move it across the network, and deserialize it on the other end. It’s an inherently expensive operation — and in most distributed jobs, network I/O is the tightest bottleneck of all, well below the speed of local CPU processing.

The problem isn’t that shuffle exists. It’s that it’s frequently unmanaged: aggregations executed after moving all the data instead of before, joins on non-selective keys, unnecessary repartitions inherited from Spark versions that lacked today’s intelligence, or configurations left at their defaults regardless of workload. The result is jobs that consume more time, more memory, and more cluster cost than they should.

This article breaks down shuffle in detail — what it is, why it dominates performance in distributed jobs, and how to reduce its impact across both batch and Structured Streaming workloads. We’ll cover configuration parameters, engine-native features like Adaptive Query Execution, and practical strategies validated on Databricks environments, always with one central premise: managing shuffle isn’t about eliminating it, but understanding when and how it happens in order to act with precision instead of guesswork.

Wide vs. Narrow Dependencies

In Apache Spark and Databricks, the concepts of narrow and wide dependencies directly describe how data moves between input and output partitions during a transformation.

  1. Narrow Dependencies (No Shuffle):
    Each output partition depends on data from one parent partition. Spark can process it without moving data across the cluster. These are incredibly fast, since Spark executes these in-memory within the same stage (pipelining).
    Its main benefits are no shuffle, fewer stages, lower latency, and generally better performance.
    Examples include map, filter, select, flatMap, withColumn, coalesce, mapPartitions.

  2. Wide Dependencies (Triggers Shuffle):
    An output partition depends on data from multiple parent partitions. Spark must redistribute data across all executors, creating a shuffle and usually a new stage boundary where Spark must write data to disk and create a new execution stage.
    The main costs are network transfer, disk spill, serialization, more tasks, and possible data skew or shuffle fetch failures.
    Wide transformations include groupBy, join (non-broadcast), repartition, distinct, sortBy / orderBy, cogroup, and *ByKey operations.

Narrow Dependency (No Shuffle)           Wide Dependency (Triggers Shuffle)
   [Part 1] ───> [Part 1]                   [Part 1] ───┐ ┌───> [Part 1]
   [Part 2] ───> [Part 2]                                ╳ 
   [Part 3] ───> [Part 3]                   [Part 2] ───┘ └───> [Part 2]

The Spark DAG scheduler divides a job into stages at every shuffle boundary. Within a stage, narrow transformations are pipelined and execute in a single pass. Across a stage boundary, data must be physically transferred between executors, with a hard rule: every task in Stage N must finish before any task in Stage N+1 can start.

The Two Phases of Every Shuffle

Phase 1 — Shuffle Write (Map Side)
Each map task in the upstream stage processes its input partition and partitions every output record by destination, computed as partition_id = hash(key) % numPartitions. Records accumulate in an in-memory buffer — the PartitionedAppendOnlyMap when aggregation is needed — keyed by (partition_id, key). This buffer is what enables map-side combine.

To better understand how this buffer can become a problem, let’s consider a common example of associative operations like sum or count, where Spark consolidates millions of input rows into thousands of partial sums before any data leaves the executor. When the in-memory buffer exceeds its memory budget (controlled by spark.shuffle.file.buffer, default 32 KB), the task spills a sorted run to local disk using TimSort. Multiple spill files accumulate. When the task finishes processing all input, the ExternalSorter executes a final merge pass: it opens all spill files simultaneously, uses a min-heap to interleave them in partition order, and produces:

  • One shuffle data file (a single sorted file containing all rows for all reduce partitions)
  • One index file (the byte offset for each destination partition within the data file)

The executor registers the file pair with its BlockManager, which reports the location to the driver’s MapOutputTracker. Spill files are deleted after the merge.

Spark also offers two alternative write paths:

  • Bypass Merge Sort Shuffle Writer — for jobs with fewer output partitions than spark.shuffle.sort.bypassMergeThreshold (default 200) and no map-side combine. It opens one file per reduce partition and writes directly, skipping the sort and ExternalSorter. Faster for small partition counts.
  • Tungsten Unsafe Shuffle Writer — for DataFrame / SQL operations. Operates on off-heap UnsafeRow binary pages, sorts a long[] of pointers (24-bit partition_id + 40-bit page offset) instead of Java objects. Achieves 2–3× faster shuffle writes by avoiding GC pressure.

Phase 2 — Shuffle Read (Reduce Side):
Each reduce task in the downstream stage queries the driver’s MapOutputTracker for the location of all shuffle blocks assigned to its partition. It then fetches those blocks in parallel from the remote executors via Netty / HTTP through the BlockTransferService.

Incoming bytes are bounded by spark.reducer.maxSizeInFlight to prevent flooding the network. Streams from different sources are merged on the fly using a min-heap — this is what distinguishes Spark from Hadoop MapReduce, which materializes a fully merged file on disk before the reduce begins.

If the merged data exceeds the executor’s remaining memory, Spark spills the reduce-side merge to disk before completing. A job that spills on both sides writes the same data to disk four times: write (spill) → write (final file) → read (fetch) → write (reduce spill) → read (reduce merge).

Why Shuffles Are Expensive

Let’s break this down by considering its five cost layers:

Cost Layers Description
Disk I/O (Write + Read) Every byte of shuffled data is written to disk and read back. Keep in mind that reduce-side reads are random across many small files, while writes are mostly sequential. For example, for a 100 GB shuffle on HDDs (~100 MB/s sequential), the disk I/O alone can dominate runtime, especially when random reads are involved.
Network I/O Shuffle is an all-to-all communication pattern. Each reduce task fetches its slice from every map output. With M mappers and R reducers, there are up to M × R total reads. The ratio of useful network bandwidth to total bytes transferred is poor because every connection’s overhead (TCP setup, Netty RPC, deserialization on the receiver) is paid per block.
File Management Overhead The OS “charges” for each file creation, open, seek, and close, even with a more optimized sort shuffle than the hash shuffle produced. With bypass-merge sort, the cost grows linearly with the number of reduce partitions.
Serialization / Deserialization Every record is serialized before being written to disk and deserialized after being fetched. This is CPU-intensive and creates JVM heap pressure. PySpark pays an extra tax: every RDD record is pickled to a Python worker and unpickled back, which can be 17× slower than the JVM equivalent.
Stage Synchronization Barrier A stage completes only when its slowest task finishes. One straggler holds the entire cluster hostage. A single skewed partition, GC pause, slow disk, or network hiccup on one executor doubles the wall-clock time of the whole stage. This is the most architecturally important fact about shuffles.

How to Identify and Check Shuffle Issues

In the Databricks Spark UI, wide dependencies typically appear as shuffle read/write metrics and a new stage.

On the physical plan of your query, look for any Exchange to see the shuffle nodes and understand how much data is moving and whether it is necessary.

# Run one of your typical queries
df = spark.sql("""
    SELECT bio_id, COUNT(*) as alert_count
    FROM etl_data_delivered_silver.heart_rate_threshold_alert
    WHERE entry_date = '2024-01-01'
    GROUP BY bio_id
""")

# Check physical plan
df.explain("formatted")

# Count exchanges (shuffles)
plan = df._jdf.queryExecution().executedPlan().toString()
shuffle_count = plan.count("Exchange")
print(f"Number of shuffles: {shuffle_count}")

With the result count, Exchange nodes and check their input/output sizes. Each one is a potential target for elimination — e.g., replacing a shuffle-heavy sort-merge join with a broadcast join, or restructuring the query to avoid an unnecessary repartition.

Common Symptoms & What They Mean:

Spark UI Signal Likely Cause
Shuffle Spill (Disk) > 0 across all tasks Partition count too low — increase spark.sql.shuffle.partitions
Shuffle Spill (Disk) on reduce side only Reduce-side memory exhausted — increase executor memory or reduce fetch size
GC Time > 20% of task duration Too many Java objects — enable Kryo or migrate RDD → DataFrame
Max task duration / median > 5× Data skew — apply salting or rely on AQE skew join
Shuffle Read Fetch Wait Time > 10s Network saturation or BlockManager contention
Small shuffle blocks (< 100 KB) Too many partitions — raise advisoryPartitionSizeInBytes or lower shuffle.partitions

Handling Shuffling

Here, we will present a comprehensive overview of configuration parameters and methods for handling each of the main challenges to managing shuffles in modern projects like Spark/Databricks.

However, our focus in this article is only on programmatic/runtime methods at the code level and their configuration tuning, some of which we will discuss in detail later.

If you are looking for more details about the data layout and storage-level methods, as well as their configuration parameters for Delta Lake, I invite you to consult the article Maintaining your Delta Lake.

1. Configuration / Feature Tuning (properties)

  1. Tune spark.sql.shuffle.partitions (default 200, almost always wrong, set to auto for auto-optimized shuffle on Databricks)
  2. AQE coalescing post-shuffle partitions (spark.sql.adaptive.coalescePartitions.enabled)
  3. AQE advisory partition size (spark.sql.adaptive.advisoryPartitionSizeInBytes)
  4. AQE minimum partition size (spark.sql.adaptive.coalescePartitions.minPartitionSize)
  5. AQE initial partition number (spark.sql.adaptive.coalescePartitions.initialPartitionNum)
  6. AQE parallelism-first coalescing (spark.sql.adaptive.coalescePartitions.parallelismFirst)
  7. AQE local shuffle reader (spark.sql.adaptive.localShuffleReader.enabled) — avoids extra shuffle when downstream stage doesn’t re-shuffle
  8. External Shuffle Service (spark.shuffle.service.enabled) – Set it to true; it is required for dynamic allocation
  9. Tune spark.sql.adaptive.enabled (master switch for all AQE optimizations)
  10. Tune file scan partitioning (spark.sql.files.maxPartitionBytes, spark.sql.files.openCostInBytes, spark.sql.files.minPartitionNum, spark.sql.files.maxPartitionNum)
  11. Tune spark.sql.adaptive.rebalancePartitionsSmallPartitionFactor for rebalanced shuffle
  12. Tune spark.sql.adaptive.skewJoin.skewedPartitionFactor / skewedPartitionThresholdInBytes to control skew-driven shuffle splits
  13. Tune spark.sql.adaptive.maxShuffledHashJoinLocalMapThreshold for SMJ → SHJ conversion
  14. Enable spark.sql.adaptive.nonEmptyPartitionRatioForBroadcastJoin to avoid unnecessary shuffles on empty partitions
  15. spark.sql.join.preferSortMergeJoin (false can avoid large shuffle)
  16. spark.sql.shuffle.compress / spark.shuffle.compress to reduce shuffle data volume
  17. spark.sql.adaptive.autoBroadcastJoinThreshold (default 10m) — AQE may switch to broadcast instead of shuffling
  18. Storage Partition Join configs (spark.sql.sources.v2.bucketing.*) to eliminate shuffles entirely
  19. spark.sql.requireAllClusterKeysForCoPartition — control shuffle elimination
  20. spark.reducer.maxSizeInFlight for reduce-side buffering (default 48 MB)
  21. spark.shuffle.file.buffer (default 32k) / spark.shuffle.spill.compress (LZ4, default true) for shuffle I/O tuning
  22. Run ANALYZE TABLE so the optimizer can size shuffles correctly
  23. Push-based shuffle (spark.shuffle.push.enabled, spark.shuffle.push.maxBlockSizeToPush) – large workloads, Spark 3.2+

2. Programmatic / Runtime Code-level Methods

  1. Avoid generic joins. Ensure your join keys are highly specific (like an ID) rather than generic strings (like a status or date), which can cause massive cross-shuffling.
  2. Reduce shuffle scope with filter / where / select early (predicate pushdown, column pruning)
  3. Combine multiple narrow transformations before a wide one to reduce shuffle stages
  4. Avoid Left Joins – isolating Nulls/Defaults and union then back
  5. Pre-aggregate before the shuffle (reduce row count crossing the network)
  6. Use DataFrame APIs over RDD to benefit from Tungsten / whole-stage codegen and improved shuffle
  7. Cache / persist before repeated actions to avoid re-shuffling
  8. coalesce(n) — narrow dependency, avoids shuffle when reducing partitions
  9. repartition(n) — triggers a shuffle but balances partition sizes evenly
  10. repartition(cols) — hash-based shuffle on a key
  11. repartitionByRange(cols) — range-based shuffle for ordered joins
  12. Partitioning hints (COALESCE, REPARTITION, REPARTITION_BY_RANGE, REBALANCE)
  13. Broadcast hint (hint("broadcast") / /*+ BROADCAST */) — eliminates shuffle for small side
  14. Join strategy hints (MERGE, SHUFFLE_HASH, SHUFFLE_REPLICATE_NL)
  15. .bucketBy(n, cols) writer + matching reader — persist bucketed tables to enable shuffle-free bucketed joins / group-bys on subsequent reads (classic Hive-style)
  16. .sortBy(cols) combined with .bucketBy() — sort-merge bucketed joins avoid shuffle AND avoid the per-partition sort at read time
  17. Use mapPartitions to do per-partition processing without re-shuffling
  18. Use approx_count_distinct, sketches, Bloom filters to shrink shuffled data
  19. Use DISTRIBUTE BY / CLUSTER BY in SQL for explicit shuffle key control
  20. Replace groupByKey with reduceByKey / aggregateByKey to do map-side combines (RDD-only)
  21. Use flatMapGroupsWithState / structured streaming stateful operators that minimize cross-partition shuffle

3. Data Layout / Storage-level Methods

  1. Pre-partition / pre-bucket data on join keys at write time (eliminates shuffle on read)
  2. Use Delta CLUSTER BY / Liquid Clustering on join/aggregation columns – For new tables on Databricks and plan the partition & Z-Order migration to cluster — replaces BUCKETED BY, adapts over time, supports data skipping on high-cardinality columns
  3. PARTITIONED BY with low-cardinality columns (date, region) for partition pruning
  4. Combine PARTITIONED BY (low cardinality) + CLUSTER BY (high cardinality) for the modern Delta pattern
  5. .bucketBy(n, cols) at write time for non-Delta / external sources (Parquet, CSV, JDBC) where Liquid Clustering doesn’t apply
  6. Avoid high-cardinality partition columns to prevent the small-files / many-directories problem
  7. Use ANALYZE TABLE ... COMPUTE STATISTICS to feed the optimizer with accurate bucket/partition sizes
  8. Use Iceberg bucket / partition transforms for Storage Partition Join
  9. Write data already sorted by join key (sortBy / sortWithinPartitions)
  10. Compact small files with OPTIMIZE to reduce shuffle block fragmentation
  11. Z-ORDER on commonly-shuffled columns to improve data locality
  12. Use VORDER (Databricks) for optimized write layout, reducing shuffle data volume
  13. Size files to match target partition size (spark.sql.adaptive.advisoryPartitionSizeInBytes)
  14. Use columnar formats (Parquet, Delta, Iceberg) — efficient compression lowers shuffle bytes
  15. Tune file sizes (spark.sql.files.maxPartitionBytes) to avoid overly small/large shuffle inputs
  16. Materialize pre-shuffled / pre-aggregated intermediate tables for repeated workloads
  17. Use range partitioning instead of hash partitioning when range joins / merges are common
  18. Enable predicate pushdown and partition pruning at the source to reduce data shuffled
  19. Use Dynamic File Pruning (DFP) on Databricks to skip irrelevant partitions before shuffle
  20. Use Photon-accelerated shuffle (Databricks) for vectorized, high-throughput shuffle
  21. Adopt shuffle-cleanup-aware file layouts (e.g., avoid many tiny files in hot partitions)
  22. Pre-shard data across the same key space if you control the write path
  23. Use external shuffle service for large/many shuffles
  24. Leverage disk-cache / SSD-backed nodes to speed shuffle spill
  25. Split large tables physically so each shard can be shuffled independently (sharded join pattern)

Programmatic / Runtime Code-level Methods to Reduce Shuffling

1. Avoid Generic Joins

This is one of the cornerstones of high-performance Spark engineering. This point is about avoiding inefficient work caused by poor logic.

When you join on “generic” keys (low-cardinality fields), you trigger two specific technical failures in Spark:

  • Data Skew: Spark uses Hash Partitioning to distribute data across the cluster. When you join on a key, Spark hashes that key to determine which executor the data should live on. So, on generic keys, you’ll likely see data skew, and the skewed data will all be sent to the same executor.
  • Explosive Joins: It has the risk of a “Many-to-Many” join that creates a massive amount of duplicate data.

If you need a “generic” value (like status or date) in your final report, use it on the filters and include it in your SELECT statement, but never use it as the primary key in your JOIN clause.

Always join on the most granular, unique identifier available (Primary Keys, UUIDs, or unique IDs) to ensure the Spark engine can distribute the workload evenly across the cluster.

2-3. Reduce Shuffle by Combining Multiple Narrow Transformations Before a Wide One, and Applying Predicate Pushdown/Column Pruning Early

This is one of the higher-leverage Spark optimizations because it targets both the amount of data shuffled and the number of shuffle boundaries. Here’s how it works and why.

When you chain several narrow transformations together, Catalyst’s whole-stage codegen fuses them into a single tight loop over each partition, with no intermediate materialization. That fused block runs right before the shuffle write of the next wide operation.

# Less efficient — filter happens after the shuffle
df.groupBy("user_id").agg(F.sum("amount")).filter("sum(amount) > 100")

# Better — filter/select before the wide op
df.filter("status = 'active'") \
  .select("user_id", "amount") \
  .groupBy("user_id").agg(F.sum("amount"))

The second version shuffles far fewer rows and fewer bytes per row, because the filter and column pruning happen inside the narrow stage, before the shuffle exchange.

How to verify it’s actually happening:

df.filter("status = 'active'").select("user_id", "amount").explain("formatted")

Check the physical plan for:

  • PushedFilters: [...] on the scan node — confirms predicate pushdown
  • ReadSchema: struct<user_id:...,amount:...> — confirms only the needed columns are read, not the full schema

Predicate pushdown

Pushing filter/where early lets Catalyst move the predicate down to the data source scan itself (Parquet, ORC, Delta, JDBC). This means rows are excluded before they’re even read off disk, not just dropped in-memory afterward.

  • For Parquet/ORC/Delta: row-group / file-level statistics let Spark skip entire files or row groups that can’t match the predicate (Delta’s data-skipping min/max stats do this especially well).
  • For JDBC sources: the filter gets translated into the WHERE clause of the SQL sent to the database, so filtering happens remotely.

Caveat: predicate pushdown breaks if the filter depends on a UDF or a non-deterministic expression, since Catalyst can’t reason about opaque code. Prefer built-in pyspark.sql.functions over UDFs when the filter is on a performance-critical path.

Column pruning

Calling select early (just referencing only needed columns) lets Spark skip reading unused columns entirely for columnar formats like Parquet/ORC/Delta; this is huge for wide tables where you only need a few of n columns. It also shrinks row width going into any downstream shuffle, since fewer/smaller columns get serialized to shuffle files.

Does Catalyst already do this for me?

Partially, the optimizer has rules like PushDownPredicates, ColumnPruning, CombineFilters that reorder logical plans automatically, even if you write the filter/select after the join in your code. So in simple cases, writing df.join(...).filter(...) vs df.filter(...).join(...) can produce the same physical plan.

But you shouldn’t rely on this:

  • It breaks with UDFs, certain non-deterministic functions, or when a cache()/checkpoint sits between operations (the optimizer won’t push predicates across a materialization boundary).
  • Explicit ordering makes your intent unambiguous and is more portable across Spark versions/configs.
  • It’s simply better practice for readability and predictable performance.

One more layer: AQE

Adaptive Query Execution (Spark 3.x+, on by default in Databricks) helps after the shuffle happens — coalescing small partitions, handling skew, switching join strategies at runtime. It doesn’t reduce what gets shuffled in the first place. So AQE and early filter/select are complementary, not substitutes: minimize shuffle input yourself, let AQE clean up what’s left.

4. Isolating Nulls/Defaults and union then back is better than Left Joins

A standard Left Join forces the system to send all rows from the left table across the network. If your join column contains many NULL (empty) values, or default values instead, engines like Apache Spark or SQL group all those NULL/Default keys onto a single computer. This creates a massive “hot spot” that slows down processing or causes out-of-memory errors.

Instead of joining everything at once, you split the work into simpler steps:

  1. Filter out nulls/dfaults: Separate the left table into two groups: rows with valid keys and rows where the key is NULL/default.
  2. Run only on the smaller group that has valid keys (the inner join data): This avoids clumping data on one machine.
  3. Union back: Combine the joined results with your separated NULL rows using a UNION command, filling in blank values for the right side.

5. Pre-aggregate before the shuffle

Pre-aggregate before the shuffle is one of the most impactful optimizations in distributed data processing (Spark, Flink, Presto/Trino, BigQuery, etc.). The idea is to do partial aggregation locally on each node before sending data across the network, so the shuffle only moves already-reduced data instead of raw rows.

Aggregate within each partition first (map-side / local combine), so each node only sends out one row per group instead of a huge amount. To make it clear, let’s take an example of counting events per user across 1,000 partitions with 1M users**:

  • Naive: shuffle all raw rows byuser_id, then aggregate on the receiving side. If you have 1 billion rows, you shuffle 1 billion rows.
  • Pre-aggregated: each partition first computes local (user_id, count) pairs — collapsing its local rows down to at most 1M rows (one per user seen locally). Then shuffle those partial sums, and do a final merge (sum-of-sums) on the receiving side.

If each partition originally has ~1M rows but only touches, say, 10K distinct users, you just cut that partition’s shuffle output by 100x.

It works great for associative/commutative aggregations: sum, count, min, max, count distinct, with sketches like HyperLogLog.

It works less well or not at all for aggregations that can’t be partially combined without keeping all the data, like exact median/percentiles or collect_list of everything; those still need the full data in one place (though approximations can help).

The benefit scales with cardinality reduction: if every row already has a unique key, pre-aggregation does nothing. There’s no reduction, and you’ve just added an extra combine step. It shines most when many rows map to relatively few groups.

How engines expose this:

  • Spark: reduceByKey / combineByKey do map-side combining automatically, unlike plaingroupByKey, which shuffles raw values. This is why “avoid groupByKey” is classic Spark advice.
  • Spark SQL / Catalyst: does partial aggregation automatically as part of the physical plan (HashAggregate with a partial + final phase) — you usually don’t need to hand-roll it.
  • SQL engines generally: look for a “partial aggregate” or “local aggregate” step in the query plan (EXPLAIN) before the exchange/shuffle step — that’s this optimization happening under the hood.
  • Flink: similar concept via pre-aggregation windows / AggregateFunction with local combining.

In any case, take care with some cases; for example, avoid using UDAFs / collect_list — aggregations that can’t pre-combine well. This is an expensive operation, since collect_list can’t meaningfully partial-aggregate, so most data still has to move.

df.groupBy("user_id").agg(F.collect_list("event_id"))

For unbounded “collect everything” aggregations like that, consider windowing, sampling, or restructuring (e.g., aggregate to counts/stats instead of raw lists) if volume is large.

On Databricks, let Adaptive Query Execution help

On Databricks (AQE on by default since Spark 3.x), you generally don’t need to hand-tune partition counts. It dynamically coalesces post-shuffle partitions and can handle skew automatically:

spark.conf.set("spark.sql.adaptive.enabled", "true")               # default true on Databricks
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")       # auto-splits skewed shuffle partitions
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")

Combined with Photon (Databricks’ vectorized engine), the partial-aggregation step itself runs faster in native code, but the shuffle-reduction principle is identical. Check the query profile UI (Spark UI → SQL tab → the aggregate node’s “shuffle bytes written” before vs after the partial agg) to confirm it’s working as expected.

result.explain("formatted")   # or check the Spark UI's SQL/DataFrame tab

Look for a HashAggregate (partial) sitting directly above the scan and before the Exchange — if you instead see the Exchange immediately above the scan with no partial-aggregate step, something (usually a UDAF or non-associative function) is preventing the optimization.

Two-phase Manual Pre-Aggregation over Heavy Data Skew Cases

Consider it only when auto-optimization isn’t enough (e.g., skew). Sometimes you want to force an explicit local combine, especially to fight against data skew on a few hot keys:

from pyspark.sql import functions as F

# Salt the key to spread a hot key across multiple shuffle partitions,
# pre-aggregate on the salted key, then combine the salted results
salted = df.withColumn("salt", (F.rand() * 8).cast("int")) \
            .withColumn("salted_key", F.concat_ws("_", "user_id", "salt"))

stage1 = salted.groupBy("salted_key", "user_id").agg(F.sum("amount").alias("partial_sum"))

final = stage1.groupBy("user_id").agg(F.sum("partial_sum").alias("total_amount"))

This forces two shuffle stages, but each is much smaller and avoids one partition getting overwhelmed by a skewed key.

6. DataFrame/Dataset APIs outperform RDDs

Catalyst optimizer sees structure; RDDs are opaque

RDD transformations are arbitrary JVM lambdas — Spark can’t look inside them. It just executes your map/filter functions as black boxes in the order you wrote them.

DataFrame/SQL operations are an expression tree (a logical plan) that Catalyst can analyze, reorder, and rewrite before anything runs — predicate pushdown, column pruning, join reordering, constant folding, etc.

# RDD: Spark has no idea what this lambda does internally.
# It can't push the filter into the file scan, can't prune columns.
rdd = sc.textFile("events.csv") \
        .map(lambda line: line.split(",")) \
        .filter(lambda x: x[2] == "purchase") \
        .map(lambda x: (x[0], float(x[3])))

# DataFrame: this is a plan Catalyst can optimize before execution
df = spark.read.csv("events.csv", header=True, inferSchema=True)
result = df.filter(F.col("event_type") == "purchase") \
           .select("user_id", "amount")

Run result.explain(True) and you’ll see the filter and column-pruning pushed all the way down to the file scan (PushedFilters, ReadSchema only including needed columns) — parquet/columnar sources can even skip reading unneeded columns/row groups entirely. The RDD version reads and parses every column of every row, always.

A quick way to audit a codebase is grep for RDD usage that could likely be DataFrame ops instead. Look for .rdd, .map(, .flatMap(, .groupByKey(, .reduceByKey( on raw RDDs. If you see .rdd being called on an existing DataFrame just to do a .map(), that’s usually a red flag; you’re dropping out of Catalyst/Tungsten/Photon for something a .withColumn() + UDF (or better, a built-in F.* function) could do while staying in the optimized path.

Tungsten’s binary row format avoids JVM object overhead

RDDs store data as actual JVM objects (e.g., Java String, boxed Integer, Tuple2). This means:

  • Per-object memory overhead (headers, pointers) — a 4-byte int can cost 16+ bytes as a boxed object
  • Garbage collection pressure — millions of short-lived objects means GC pauses
  • Java serialization (or Kryo) overhead when data crosses the wire

DataFrames use UnsafeRow — a binary, off-heap-friendly row format that packs data tightly (like a C struct), with fixed-offset fields. Spark operates directly on these byte arrays without deserializing to JVM objects, and can put much of it in off-heap memory managed manually, sidestepping GC almost entirely.

RDD[(String, Double)] -> JVM objects: String object + boxed Double + Tuple2 wrapper, on-heap, GC-tracked
DataFrame row -> packed bytes in UnsafeRow: [8-byte offset][string bytes][8-byte double], no object headers

You don’t write this — it’s automatic once you’re in DataFrame land — but it’s why the same logical operation uses less memory and less GC time.

Whole-stage code generation (WSCG)

For RDDs, each transformation is a separate closure call — map then filter then map means three virtual function calls per row, with intermediate object allocation between each.

For DataFrames, Catalyst compiles a whole chain of operators into a single Java function at runtime (via Janino) — collapsing filter→project→aggregate into one tight loop with no per-row virtual dispatch, similar to what a hand-written for-loop would look like.

df.filter(F.col("amount") > 100).select("user_id", "amount").explain("codegen")

You’ll see generated Java source in the output — look for *(1) markers in explain() output indicating operators fused into a single “whole-stage codegen” unit. RDD chains never get this treatment; each stage is interpreted, not compiled.

“Improved shuffle”

  • Serialization: DataFrame shuffles serialize UnsafeRow bytes directly (near-memcpy speed). RDD shuffles typically use Java serialization or Kryo, both slower and producing larger payloads for the same logical data.
  • Sort-based shuffle + Tungsten shuffle manager: Spark’s default shuffle (SortShuffleManager) is Tungsten-aware for DataFrames — it can sort/spill using the binary row format directly, avoiding repeated ser/deserialize round-trips that RDD shuffles incur.
  • Automatic partial aggregation: as covered earlier, groupBy().agg() on a DataFrame gets a partial-aggregate step before the shuffle for free. The RDD equivalent (reduceByKey) requires you to manually choose the right combiner — easy to accidentally use groupByKey and lose the optimization.
  • AQE (Adaptive Query Execution) only applies to DataFrame/SQL plans — dynamic shuffle partition coalescing, skew-join splitting, and runtime join-strategy switching all require the structured plan that RDDs don’t provide.

Databricks-specific: Photon only accelerates DataFrame/SQL

Photon, Databricks’ native vectorized execution engine (C++, not JVM), only kicks in for DataFrame/SQL operators — it can’t accelerate arbitrary RDD lambda code, since Photon needs to understand the operator semantics to generate vectorized native code.

To check if Photon is engaged for a given query, don’t rely on explain() text — use the Spark UI’s SQL/DataFrame tab: Photon operators are shown in orange in the query DAG, standard JVM Spark operators in blue. On SQL warehouses or serverless compute, the query profile’s Execution Details view goes further, reporting the percentage of total task time spent in Photon directly.

On a Photon-enabled Databricks cluster, this is often the single biggest gap: identical logic written as RDD vs DataFrame can differ by multiples, not percentages, because one path is eligible for native vectorized execution and the other is stuck in the JVM interpreter (or even Python, with extra serialization cost if using PySpark RDDs).

When RDDs are still legitimate to use

  • Truly unstructured data/custom binary formats before you can impose a schema
  • Fine-grained control over partitioning schemes Catalyst doesn’t expose
  • Certain iterative graph/ML algorithms (though GraphFrames/MLlib DataFrame APIs cover most of this now)
  • Legacy code you haven’t migrated yet

But as a default for anything tabular — reads, filters, joins, aggregations — DataFrame/Dataset APIs should be the starting point on modern Spark/Databricks, with RDDs reserved for cases the structured APIs genuinely can’t express.

7. Cache / persist before repeated actions to avoid re-shuffling

Spark’s execution is lazy. Nothing actually runs until an action (count, collect, write, etc.) triggers it. If you reference the same DataFrame in multiple actions without caching, Spark recomputes the entire lineage from scratch each time, including re-running any shuffle that produced it.

shuffled = df.groupBy("user_id").agg(F.sum("amount").alias("total"))

# Without caching, EACH action below re-executes the groupBy + shuffle from scratch
count = shuffled.count()                          # triggers shuffle #1
top_users = shuffled.orderBy(F.desc("total")).limit(10).collect()  # triggers shuffle #2 (recomputes groupBy!)
shuffled.write.parquet("/output/totals")           # triggers shuffle #3 (recomputes groupBy again!)

That groupBy — which involves a full shuffle — runs three separate times. Caching breaks this cycle.

shuffled = df.groupBy("user_id").agg(F.sum("amount").alias("total")).cache()

count = shuffled.count()                          # triggers shuffle, THEN materializes cache
top_users = shuffled.orderBy(F.desc("total")).limit(10).collect()  # reads from cache, no re-shuffle
shuffled.write.parquet("/output/totals")           # reads from cache, no re-shuffle

cache() is a hint, not immediate — the data is only materialized into the cache on the first action that touches it. Subsequent actions reuse it. In the example above, the orderBy(F.desc("total")).limit(10) still performs a sort operation, but it’s sorting the already-aggregated cached data (much smaller dataset) rather than re-aggregating from the raw source.

Verifying it’s actually helping:

shuffled.explain()
spark.catalog.isCached("some_temp_view")  # for registered views

Look for “InMemoryTableScan” instead of “Exchange” + upstream shuffle operators on the second/third action — that confirms it’s reading from cache, not re-shuffling.

The Spark UI’s Storage tab shows cached RDD/DataFrame partitions, size in memory vs disk, and fraction cached — useful for confirming it actually fit and didn’t silently spill/evict.

When it doesn’t help

1. Single action — no benefit, pure overhead

# BAD: cache() adds serialization overhead for zero reuse
df.groupBy("user_id").sum("amount").cache().write.parquet("/out")  # only one action — don't cache

2. Cache count() as a “trigger” — a common but wasteful pattern

df.cache()
df.count()  # people do this "to force caching" — it works, but costs a full extra pass

This isn’t wrong, but if you don’t need the count for anything, it’s a wasted action just to warm the cache. If the next action is going to be write anyway, the cache still helps, but you paid for two passes (count + write) instead of eagerly materializing during the write itself.

3. Data doesn’t fit in memory → silent partition eviction
With MEMORY_ONLY, partitions that don’t fit are simply dropped and recomputed from lineage on next access — which for a post-shuffle DataFrame means re-shuffling anyway. MEMORY_AND_DISK avoids this by spilling instead of dropping.

4. Caching too early in the lineage

# Less useful: caching before the filter means you cache more data than needed
raw = spark.table("events").cache()
filtered = raw.filter(F.col("event_type") == "purchase")

# Better: filter first, cache the smaller, post-shuffle/post-filter result
filtered = spark.table("events").filter(F.col("event_type") == "purchase")
shuffled = filtered.groupBy("user_id").sum("amount").cache()

Cache as late as possible in the pipeline — ideally right after the expensive shuffle/aggregation, not before it.

cache() vs persist()

For most Databricks workloads on the DataFrame/Dataset API, MEMORY_AND_DISK (the default) is a reasonable choice — it spills to disk instead of dropping partitions and forcing a recompute (which would mean re-shuffling) if memory runs out.

Note: the default storage level differs by API. Dataset/DataFrame.cache() defaults to MEMORY_AND_DISK; the older RDD.cache() defaults to MEMORY_ONLY. If you’re coming from RDD habits or mixing RDD and DataFrame code, don’t assume the same default applies to both.

Storage Level Memory Usage CPU Overhead Explanation
MEMORY_ONLY Very High Low Kept as raw Java objects. If it doesn’t fit, missing blocks are recomputed.
MEMORY_AND_DISK High Low Default for DataFrame/Dataset.cache(). Spills extra data to local disk if memory is full.
MEMORY_ONLY_SER Low High Serializes data into compact byte arrays. Saves memory but costs CPU to unpack.
MEMORY_AND_DISK_SER Very Low High Serialized in memory, spills serialized blocks to disk. Saves memory, but costs CPU to deserialize.
DISK_ONLY None Moderate–High Bypasses memory entirely, writing directly to local executor disk. Data is always stored serialized here too, so there’s real (de)serialization cost — but in practice, disk I/O latency dominates the overhead far more than CPU does. Avoids recompute and memory pressure, but generally the slowest read path.

How to apply it:

from pyspark import StorageLevel

df.cache()                                    # = persist(MEMORY_AND_DISK)
df.persist(StorageLevel.MEMORY_AND_DISK_SER)  # serialized, saves memory, costs CPU to deserialize

Cached data holds onto executor memory/disk. Leaving it cached after you’re done wastes resources and can even push out other cached data or cause spills.

shuffled.unpersist()

General rules for cache/persist:

  • The same DataFrame (especially one that required a shuffle — groupBy, join, distinct, repartition) is used in 2+ actions within the same job/session
  • It’s small enough to reasonably fit MEMORY_AND_DISK without evicting other important cached data
  • You remember to unpersist() when done, or let it fall out of scope naturally at session end
  • Don’t cache when there’s only one downstream action — that’s pure overhead with no reuse to amortize it against

Delta Lake table alternative: materialize instead of cache:
For pipelines where the same aggregation feeds many downstream jobs (not just multiple actions in one session), consider writing the shuffled result as a Delta table instead of (or in addition to) in-session caching — this persists across job/cluster restarts, which .cache() does not (cache is tied to the Spark session/cluster lifetime). This is a general Spark + Delta Lake pattern, not something exclusive to Databricks — any Spark environment running open-source Delta Lake gets the same durability benefit.

shuffled.write.format("delta").mode("overwrite").saveAsTable("user_totals")
# Other jobs/notebooks read this without re-running the shuffle at all
downstream_df = spark.table("user_totals")
Serialized storage for memory-constrained caching

Serialization is a memory-saving technique for RAM-constrained environments. When you cache a DataFrame with serialized storage, Spark converts each partition from raw Java objects (which include class metadata, pointers, and overhead) into compact binary byte arrays. This meaningfully reduces memory footprint, because serialized data strips away object headers, padding, and reference pointers — instead of storing many individual Row objects with their internal structures, Spark stores a single compact byte buffer per partition. The exact compression ratio depends heavily on schema (wide, string-heavy schemas tend to compress more than narrow numeric ones), so treat any specific multiplier as workload-dependent rather than a fixed rule of thumb.

Serialization isn’t free — it trades CPU cycles for memory space. Every time Spark reads cached data, it must deserialize the byte arrays back into usable objects, adding real per-query overhead. This is still typically faster than re-reading from disk or recomputing the DataFrame, but the exact overhead depends on row width and access pattern rather than a single fixed percentage. On Databricks Serverless with Photon, the vectorized runtime partially mitigates deserialization costs through efficient columnar processing.

If you have enough RAM for MEMORY_ONLY, skip serialization — the CPU overhead usually isn’t worth it. Also, for Delta Lake tables, aggressive caching is often less necessary than it seems: Delta’s disk cache (below) accelerates repeated file scans, and file-level statistics power data skipping at query planning time — two different mechanisms that, combined, often reduce the need to lean on .cache() for read-heavy workloads.

Only use serialized caching (MEMORY_ONLY_SER or MEMORY_AND_DISK_SER) when:

  • Your cluster has limited executor memory (< 16 GB per executor)
  • You need to cache large DataFrames (> 10 GB) that don’t fit uncompressed
  • Your data is accessed repeatedly but infrequently enough that deserialization overhead is acceptable
  • You’re working on Serverless, where memory is fixed, and you can’t just add more executors

Databricks-specific angles

Disk cache is different from .cache() and often preferable for reads:
On Databricks, the disk cache (formerly “Delta cache”) automatically caches remote Parquet/Delta file data on local SSD after the first read — this operates below Spark’s DataFrame cache, at the file-scan layer, and doesn’t need you to call .cache() at all.

spark.conf.set("spark.databricks.io.cache.enabled", "true")

It’s on by default on instance types with local SSDs — Databricks automatically enables and sizes it for those instance types, using up to half the available local SSD space.

This helps repeated scans of the same table across queries/notebooks, but it does not avoid re-shuffling — a groupBy still re-shuffles every time even if the underlying file read is served from disk cache. For avoiding repeated shuffles specifically, you still need .cache()/.persist() on the shuffled DataFrame itself.

Photon + cache:
Photon can accelerate the operators that populate and scan the cache, but the fundamental cache-avoids-reshuffle logic is identical to open-source Spark — Photon just makes both the original shuffle and any cache scan faster in absolute terms.

8. coalesce(n)

This operator reduces the number of partitions without causing a full network shuffle. It simply merges adjacent partitions on the same executors.

It is much faster thanrepartition, but if you drastically reduce partitions (e.g., coalesce(1)), you force a single node to handle all the data, which can cause Out-Of-Memory (OOM) errors.

Use it right before exporting data to a legacy system or external API that requires a specific, fixed number of flat files (like a single CSV file via .coalesce(1).write.csv(...)).

Another use case is to reduce memory overhead. If a transformation drastically reduces the size of your data (for example, by filtering 99% of the rows), your pipeline will have thousands of empty partitions. Using coalesce efficiently consolidates them without the need to shuffle the data.

Take care to never do an unintentional serialization. Because coalesce avoids shuffles, it forces the entire upstream processing to run on only n partitions. If you call .coalesce(1) at the end of a heavy transformation pipeline, Spark may force the entire heavy transformation to run on a single CPU core, destroying your cluster’s parallelism.

Common anti-pattern: df.coalesce(1).write.csv(...) for “one output file” — almost always wrong. Use a controlled write with a small number of files (e.g., coalesce(8)) instead, or rely on Databricks Auto-Optimize / Optimize Writes.

9–10. repartition(n, *cols)

This operator forces a full network shuffle to redistribute data evenly across a specific number of partitions based on the chosen columns.

This is a very expensive network operation. On Databricks, manually running .repartition() is often discouraged unless for In-Memory compute operations, where you are specifically trying to fix data skew (where one worker node does all the work) or optimizing Stream Joins by forcing a repartition on a shared join key right before a heavy join — this prevents Spark from doing an expensive, unpredictable shuffle during the join execution itself.

Do not use it right before .write to control file sizes on Databricks Delta tables. Databricks features like Optimize Write and Liquid Clustering handle file sizes automatically. Manual repartitioning here just wastes expensive cluster compute time on a useless shuffle.

11. repartitionByRange(n, *cols)

A variant of repartition that distributes data using range partitioning instead of hash partitioning. Data is sampled, ordered ranges are computed, and each row is assigned to the partition whose range covers its key.

df.repartitionByRange(16, "event_date")           // 16 partitions, sorted by event_date
df.repartitionByRange(col("year"), col("month"))  // multi-column range

The reason to repartition here isn’t skew or parallelism, but preserving order for something downstream. Hash partitioning distributes keys uniformly but breaks any natural ordering. Range partitioning preserves approximate order, then rows with similar keys land in the same or adjacent partitions.

Cost profile:

  • The range partitioner samples the data first to discover distribution → small extra cost
  • Then performs a full shuffle like repartition(n, cols)
  • Within each partition, records are sorted by the range column

Then, repartitionByRange is strictly more expensive than a hash-based repartition, because it does everything a normal shuffle does plus an upfront sampling pass to learn the data distribution before it can even define the partition boundaries. It’s a good concrete illustration of “not all shuffles cost the same. Some carry extra pre-work!”

So, repartitionByRange is valid for improving ordering within a job:

  • You need sorted output within each partition (for downstream sort-merge join)
  • You’re writing time-series data and want date ranges clustered
  • You’re chaining into a window function or range-based aggregation
  • As preparation immediately before a write to a non-Delta table

Avoid repartitionByRange right before .write to a Delta table — sorted writes are better handled by OPTIMIZE ... ZORDER BY or Liquid Clustering. In other words, don’t pay for a shuffle-based sort if the table format can achieve the same physical layout benefit post-write, without an extra shuffle stage in your job’s critical path. That’s a direct extension of the “avoid the shuffle entirely when a cheaper mechanism exists” principle.

12. SQL Partitioning Hints (the /*+ */ family)

In SQL, the same operations are exposed as hints rather than DataFrame methods. These all map directly to DataFrame APIs:

Hint DataFrame Equivalent Shuffle? Notes
/*+ COALESCE(n) */ .coalesce(n) No Merge adjacent partitions
/*+ REPARTITION(n) */ .repartition(n) Yes Even round-robin
/*+ REPARTITION(cols) */ .repartition(cols) Yes Hash on cols
/*+ REPARTITION(n, cols) */ .repartition(n, cols) Yes Hash on cols into n
/*+ REPARTITION_BY_RANGE(cols) */ .repartitionByRange(cols) Yes Range on cols
/*+ REPARTITION_BY_RANGE(n, cols) */ .repartitionByRange(n, cols) Yes Range on cols into n
/*+ REBALANCE */ (no direct API) Yes Advisory size, AQE-aware
/*+ REBALANCE(n) */ (no direct API) Yes n partitions, AQE-aware
/*+ REBALANCE(cols) */ (no direct API) Yes Hash on cols, AQE-aware
/*+ REBALANCE(n, cols) */ (no direct API) Yes Hash on cols into n, AQE-aware

hint("REBALANCE", ...)

A hint-only operation (no corresponding DataFrame method) that tells Spark to redistribute data so that every output partition is of a reasonable size — neither too small nor too big.

Form Behavior
REBALANCE Round-robin repartition with advisory size
REBALANCE(n) Round-robin repartition into n partitions
REBALANCE(cols) Hash repartition on cols (similar to repartition(cols))
REBALANCE(n, cols) Hash repartition into n partitions on cols

REBALANCE is skew-aware: if skew is detected, Spark will split the skewed partitions to keep them within the advisory size. The hint documentation states explicitly:

Spark’s official documentation: “This is a best-effort: if there are skews, Spark will split the skewed partitions, to make these partitions not too big. This hint is useful when you need to write the result of the query to a table, to avoid too small/big files. This hint is ignored if AQE is not enabled.”

When to Use It:

  • Right before writing to a table, when you want reasonably-sized output files without choosing an exact partition count
  • When you’re not sure which column to hash by
  • When AQE is enabled (required)

When NOT to Use It:

  • AQE is disabled → hint is ignored
  • You need exact partitioning semantics → use repartition(n, cols) instead

Practical Heuristics:

  1. Filter before partitioning. A large DataFrame filtered down to smaller should be partitioned after the filter, not before.

  2. Prefer AQE before manual repartition. If AQE is enabled, the optimizer already handles:

    • Small-partition coalescing toward the advisory size (default 64 MB)
    • Skew splitting in joins (default threshold 256 MB)
    • Dynamic broadcast switching

    Add a manual repartition only when you can prove AQE didn’t help (and verify it in the SQL tab plan).

  3. One shuffle is better than two. If your pipeline is:

    df.filter(...).repartition(100, "key").join(other, "key")

    You’re forcing a shuffle even if the optimizer would have done one anyway. Trust the planner for the join — repartition only if you’ve measured the difference.

  4. Round numbers for partition counts. Aim for target_size = total_bytes / 128 MB. With AQE, overshoot is fine (it coalesces), but undershoot causes spills.

  5. Never .coalesce(1) a heavy pipeline. It collapses the entire upstream computation to one task. If you need one output file, write a small partition count (e.g., coalesce(8)) or use a custom writer that handles file merging.

Partitions Summary

Operator Target Layer Status in Modern Spark / Databricks
repartition Memory Recommended for fixing memory skew and preparing heavy joins.
repartition Storage Anti-pattern. Let Liquid Clustering or AQE handle disk layout.
coalesce Memory Recommended to collapse empty partitions after massive filters. Avoid coalesce(1) at the end of heavy pipelines — destroys parallelism.
coalesce Storage Anti-pattern. Let Auto-Optimize / Optimize Writes handle file sizes.
repartitionByRange Memory Use when you need sorted-by-key output within partitions (range joins, time-series windows). Avoid before .write to Delta — use ZORDER / Liquid Clustering instead.
hint("REBALANCE", ...) Memory / Write boundary AQE-aware skew splitting. Useful right before .write to avoid tiny/huge files without picking an exact partition count. Requires AQE enabled.

13. Use Broadcast When Joining a Massive Table with a Small Table

When joining a massive table (such as a fact table) with a small lookup/dimension table (typically under 100MB), force a Broadcast Join to copy the small table to all executors. This eliminates a cluster-wide shuffle.

from pyspark.sql.functions import broadcast

# Standard PySpark Broadcast Join Syntax
optimized_df = large_fact_df.join(broadcast(small_dim_df), "product_id")

14. Join Strategy Hints:

The MERGE, SHUFFLE_HASH, SHUFFLE_REPLICATE_NL hints tell Spark’s planner which physical join algorithm to use, overriding the default cost-based selection. They’re typically used when:

  • Statistics are stale or wrong (cost model picks a bad algorithm)
  • AQE is disabled
  • You know the data shape better than the optimizer does
  • You’re debugging which algorithm is actually running

With AQE enabled (default since 3.2), Spark can:

  • Upgrade a planned MERGE (SMJ) to a broadcast hash join (BHJ) at runtime if one side turns out small
  • Upgrade SMJ to shuffle hash join if all post-shuffle partitions fit the local map threshold
  • Split skewed partitions in SMJ (skew join handling)
  • Skip the second shuffle if the local shuffle reader can serve data without redistribution

This means with AQE, the hints are often unnecessary. The planner adapts at runtime using real measurements. Manual hints are mainly useful when AQE is disabled or when you need deterministic behavior for benchmarking/debugging.

But there are some legitimate cases to still use them, such as:

Scenario Hint to Use Why
Cost model picks BHJ, but you know the small side is actually just over
the threshold and will fail		 MERGE Force the safer SMJ
Cost model picks SMJ, but you know all partitions will be tiny post-filter SHUFFLE_HASH Skip the sort cost
Non-equi join between a small dimension and a large fact SHUFFLE_REPLICATE_NL The only viable option
Debugging join behavior in production Any Lock the algorithm to isolate a regression
AQE is disabled for some reason Any Manual control when the planner can’t adapt

The Sort-Merge Join is the default for two large inputs. Both sides are:

  1. Shuffled on the join key (wide dependency, full network exchange)
  2. Sorted within each partition
  3. Merged by walking both sorted iterators together

On the Shuffle Hash Join, both sides are shuffled on the join key, but instead of sorting, the build side (usually the smaller side) is loaded into an in-memory hash table per partition, and the probe side is iterated to look up matches. If spark.sql.adaptive.maxShuffledHashJoinLocalMapThreshold is set to a sensible value (e.g., matching your executor memory), AQE will upgrade a planned SMJ to SHJ at runtime once it measures the actual partition sizes. This means forcing SHJ manually is rarely necessary — AQE usually gets it right.

The Shuffle-and-Replicate Nested Loop Join is the most expensive join strategy, and the only one that handles non-equi joins efficiently under specific shapes. The algorithm:

  1. The probe side is shuffled (gets distributed)
  2. The build side is replicated in full to every executor
  3. A nested loop runs: every record of the probe side is compared against every record of the build side

This is the algorithm Spark uses as a fallback when no other strategy works

When SHUFFLE_REPLICATE_NL is the right choice:

  • One side is small enough to broadcast — but you can’t use BHJ for some reason (e.g., it’s a non-equi join, the size estimate is wrong, or the build side is just over the broadcast threshold)
  • Non-equi joins with very selective conditions (e.g., t1.range_start < t2.value AND t2.value < t1.range_end)
  • Cartesian-product-style joins when the build side is small

Quick Comparison

Strategy Shuffles Memory Cost Sort Cost Supports Non-Equi Scales To Summary
MERGE Both sides Low (streaming merge) Yes (both sides) No (equi only) Unlimited The safe, general-purpose large-large equi-join. Always works, scales infinitely, but expensive.
SHUFFLE_HASH Both sides High (per-partition hash table) No No (equi only) Limited by partition size Faster than SMJ when partitions fit in memory, but memory-bounded and vulnerable to skew. AQE usually handles this automatically.
SHUFFLE_REPLICATE_NL One side Very high (full build side replicated) No Yes Small build side only The fallback for non-equi joins with a small build side. Use only when BHJ isn’t an option.

Anti-Patterns

  • Forcing SHUFFLE_HASH on large/skewed data — invites OOM
  • Forcing MERGE on a clearly broadcastable side — wastes time and memory
  • Forcing SHUFFLE_REPLICATE_NL with a large build side — will OOM every executor
  • Adding these hints everywhere “just in case” — the cost model usually knows better than you do

How to use

df1.hint("MERGE").join(df2, "key")              // force sort-merge join
df1.hint("SHUFFLE_HASH").join(df2, "key")       // force shuffle hash join
df1.hint("SHUFFLE_REPLICATE_NL").join(df2, "k") // force shuffle-and-replicate nested loop

SQL equivalents:

SELECT /*+ MERGE(t1) */              * FROM t1 JOIN t2 ON t1.k = t2.k;
SELECT /*+ SHUFFLE_HASH(t1) */       * FROM t1 JOIN t2 ON t1.k = t2.k;
SELECT /*+ SHUFFLE_REPLICATE_NL(t1)*/ * FROM t1 JOIN t2 ON t1.k = t2.k;

The aliases are also accepted: SHUFFLE_MERGE / MERGEJOIN for MERGE; BROADCAST / BROADCASTJOIN / MAPJOIN for broadcast (which is a fourth, distinct strategy not in this list).

When multiple hints conflict, Spark follows a fixed priority:

BROADCAST > MERGE > SHUFFLE_HASH > SHUFFLE_REPLICATE_NL

If both sides of a join get a hint, the higher-priority one wins:

-- BROADCAST wins, MERGE is ignored
SELECT /*+ BROADCAST(t1), MERGE(t1, t2) */ * FROM t1 JOIN t2 ON t1.k = t2.k;

Spark logs a warning when a hint is overridden:

HintErrorLogger: Hint (strategy=merge) is overridden by another hint
                 and will not take effect.

15-16. Leverage Bucketing Over Massive Operations (joins, group-bys, and to Handling Non-Delta Sources)

Bucketing is highly effective for fixed-schema operations, such as joins and group-bys, where shuffle costs are a bottleneck. For example, if you must join two massive tables frequently on the same key (e.g., user_id), pre-sort and bucket the data when saving it to disk.

By aligning bucketed data with query patterns, Spark can read matching buckets directly and skip the shuffle for the join — provided both tables share the same bucket count and are bucketed on the same key(s) (spark.sql.sources.bucketing.enabled, on by default). If those conditions aren’t met, Spark silently falls back to a full shuffle join, so bucket alignment across tables matters as much as bucketing itself.

An important downside is that you must specify the bucket count upfront and can’t change it later without rewriting the table. This reliance on static configuration makes bucketing less adaptable to dynamic or evolving workloads. Use it when the performance gains from reduced shuffle outweigh the overhead of setup and maintenance.

While Databricks Liquid Clustering has effectively made physical table bucketing obsolete for most Delta use cases, it operates exclusively at the storage layer. It does not manage how data is distributed across your cluster’s executors while a pipeline job is running in memory — that’s a separate problem, solved by different tools.

Anyway, .bucketBy() and .sortBy() are Spark’s classic operations for the Hive-style bucketing API on DataFrameWriter. They are not supported by Delta Lake; Delta simply doesn’t persist bucketing metadata in its transaction log the way Hive’s metastore does, so Spark has no bucket-spec to detect on read.

This is exactly why Databricks pushes Liquid Clustering (and, for older tables, Z-ORDER) as the Delta-native alternative — it’s a completely different mechanism (clustering metadata + file-level stats) built to solve a similar layout problem in a way that fits Delta’s architecture.

Iceberg does have a concept of “bucketing,” but it’s not Spark’s .bucketBy() — it’s a partition transform: bucket(N, col), defined at table-creation time (CREATE TABLE ... PARTITIONED BY (bucket(16, user_id))). Iceberg tracks this in its own metadata layer, and Spark can then exploit it for Storage-Partitioned Joins (SPJ) — a newer Spark 3.3+ feature that lets two Iceberg tables bucketed/partitioned identically join without a shuffle, conceptually similar to what Hive bucketing achieved, but implemented through Iceberg’s partitioning system and DataSourceV2 rather than the legacy bucketing API.

Bucketing: a write-time construct

.bucketBy() is a DataFrameWriter method — it only applies when writing via .write.bucketBy(...).saveAsTable(...) against a Hive-compatible catalog table. It has no equivalent as an in-memory, mid-pipeline transformation, and it doesn’t apply to plain path-based writes (.parquet(), .save()). Bucketing produces a physical layout on disk that a future job can detect and exploit — it is not a lever you pull during active DataFrame processing.

Sorting within buckets: .sortBy() + .bucketBy()

Bucketing alone removes the shuffle for a matching join. Adding .sortBy(cols) on top removes the second cost that a sort-merge join normally pays at read time: the per-partition sort.

df.write
  .bucketBy(16, "user_id")
  .sortBy("user_id")
  .saveAsTable("events_bucketed")

Sort-merge join, Spark’s default strategy for large-to-large joins, requires both sides sorted by the join key before it can merge them row by row. Without.sortBy(), even a bucketed join still pays for an in-memory sort of each bucket at read time — bucketing skips the shuffle stage, but the sort stage still runs. With.sortBy(), the data is written pre-sorted within each bucket file, so Spark can detect this at read time and skip the sort entirely, going straight to the merge.

Requirements and gotchas:

  • .sortBy() is only valid combined with .bucketBy() — Spark throws an error if you call .sortBy() without also bucketing, since there’d be no stable per-bucket file boundary to sort within.
  • The sort order and bucketing key don’t have to be the same column, but for eliminating the sort-merge join’s sort phase specifically, sortBy should match the join key.
  • Same constraint as bucketing alone: both tables need matching bucket count and matching bucket/sort keys for Spark to recognize the layout and take the fast path. Check .explain("formatted") for the absence of both Exchange and Sort nodes directly above the join to confirm it actually landed.
  • This combination pays off most when the same join is executed repeatedly (e.g. a recurring batch job) — the one-time write-time cost of sorting is amortized across many reads. For a one-off join, it’s rarely worth the setup.

In short: .bucketBy() alone gets you shuffle-free joins; .bucketBy() + .sortBy() gets you shuffle-free and sort-free joins — the full elimination of what a sort-merge join normally costs.

Repartitioning: an execution-time construct

.repartition(num, cols) is the in-memory equivalent — it operates on a live DataFrame mid-pipeline and solves different, execution-time problems:

  • Eliminating shuffles in heavy joins: If you’re joining two massive DataFrames midway through a complex pipeline, pre-partitioning both by the join key ensures matching rows land on the same executors ahead of the join stage, avoiding a redundant shuffle when the join itself executes.
  • Handling non-Delta sources: If your pipeline ingests raw Parquet, CSV, or JDBC sources, Liquid Clustering and bucketing offer no protection during processing — .repartition() is the primary lever for shaping distribution over these raw, intermediate DataFrames.

On skew: hash-based repartitioning does not fix skew caused by a genuinely hot key. Sincepartition_id = hash(key) % n, every row for the same key always lands in the same partition regardless of n — a single oversized key stays oversized. Real mitigation requires either salting (appending a random suffix to the hot key and aggregating in two stages) or AQE’s skew join handling (spark.sql.adaptive.skewJoin.enabled), which detects and splits oversized partitions automatically at runtime.

Typical use cases

High-cardinality partitioning — the 100,000-directory problem
Question: your team debates partitioning the output bycustomer_id, which has 100,000 distinct values. Is this a good idea?
Answer: No. High-cardinality partition columns create 100,000 directories, most with tiny files — the opposite of what partitioning is meant to achieve. Partition by a lower-cardinality attribute like region or date, and reserve bucketing for customer_id at write time to optimize joins on that key instead.

In the pipeline (memory): use .repartition() to manage network traffic, control parallelism, and prepare for downstream joins during active transformations. Reach for salting or AQE skew handling — not repartition alone — when the underlying problem is a hot key.

17. Use mapPartitions to do per-partition processing without re-shuffling

The RDD operation mapPartitions applies a function once per partition, giving you an iterator over all rows in that partition, rather than calling your function once per row (like map/UDFs do). Crucially, this is a narrow transformation, so no shuffle, since it processes each partition independently, in place, using whatever partitioning already exists.

# map / row-wise UDF: function call overhead PER ROW
rdd.map(lambda row: expensive_setup_and_process(row))

# mapPartitions: expensive setup happens ONCE per partition, not once per row
def process_partition(rows):
    model = load_expensive_model()  # loaded once per partition, not per row
    for row in rows:
        yield model.predict(row)

rdd.mapPartitions(process_partition)

The mapPartitions doesn’t change partitioning; it’s a 1:1 partition-to-partition transformation, though the number of output rows from that partition can differ, since you control what you yield. No Exchange node appears in the plan for mapPartitions itself. It’s cheap for that reason — the data never crosses the network; it just gets processed locally, partition by partition, exactly where it already sits.

df.rdd.mapPartitions(process_partition).toDF()
result.explain()
# No "Exchange" for the mapPartitions step itself —
# only appears if something else in the plan (e.g. a later groupBy) needs one

The classic use case: amortizing expensive per-partition setup

def score_partition(rows):
    # DB connection, ML model load, regex compile, etc. — done ONCE per partition
    conn = create_db_connection()
    try:
        for row in rows:
            yield (row.id, lookup_enrichment(conn, row.id))
    finally:
        conn.close()

enriched = df.rdd.mapPartitions(score_partition).toDF(["id", "enrichment"])

Compare to a row-wise scalar Python UDF, which (depending on serialization) can re-instantiate connections/state per call or per batch far more often, and typically can’t hold long-lived state across the whole partition as cleanly as an explicit generator function can. A pandas_udf (vectorized) is a middle ground: state is already batched rather than per-row, but it still isn’t scoped to the whole partition the way mapPartitions/mapInPandas are.

Modern Spark equivalent: mapInPandas / mapInArrow

On modern Spark (3.x+) and Databricks, plain RDD mapPartitions is largely superseded by Arrow-based partition APIs on DataFrames — same “process a whole partition at once, no shuffle” idea, but vectorized and staying in DataFrame land (so you keep Catalyst/Tungsten benefits for everything else in the plan).

import pandas as pd

def process_batch(iterator):
    for pdf in iterator:  # each pdf is a pandas DataFrame = one partition's data (or a batch of it)
        pdf["score"] = model.predict(pdf[["feature1", "feature2"]])
        yield pdf

result = df.mapInPandas(process_batch, schema="id long, feature1 double, feature2 double, score double")

Advantages over .rdd.mapPartitions:

  • Stays in the DataFrame API: no drop to RDDs, so you don’t lose whole-stage codegen for surrounding operators
  • Vectorized via Arrow: batches of rows processed as columnar pandas DataFrames, not Python-object-per-row
  • Plays well with Photon-accelerated operators before/after it in the plan (the mapInPandas step itself runs in Python, but neighboring native operators aren’t penalized)

mapInArrow (Spark 3.3+) is the even lower-overhead version if you want raw PyArrow RecordBatches instead of pandas, skipping the pandas conversion cost entirely:

import pyarrow as pa

def process_batch(iterator):
    for batch in iterator:
        # process as pyarrow RecordBatch — no pandas materialization
        yield batch

result = df.mapInArrow(process_batch, schema="id long, score double")

mapInArrow (Spark 3.3+) has the same partition semantics as mapInPandas and mapPartitions: it’s a 1:1, narrow, partition-local transformation. The only difference between the three is the data representation handed to your function (Python-object iterator vs pandas DataFrame batches vs PyArrow RecordBatches) — none of that affects partitioning or shuffle behavior.

Quick check: it’s not shuffling:

result.explain()

They should NOT introduce an Exchange node on its own. If you see an Exchange right before/after it, that shuffle is coming from something else in your plan (a repartition call, a join, a groupBy), not from mapPartitions itself.

Since mapInArrow is often used specifically because it’s the lowest overhead of the three (no pandas conversion), it’s commonly placed in performance-sensitive pipelines where people are also more likely to be doing a repartition nearby to control partition sizing for the batch function. If you see an Exchange next to a mapInArrow call, it’s worth checking whether you added that repartition deliberately (e.g., to control batch size per partition) — since that’s a very natural pairing with mapInArrow even though the shuffle itself still isn’t coming from mapInArrow proper.

When to reach for this vs. built-in functions or a scalar UDF

Use mapPartitions/mapInPandas when:

  • You need an expensive per-partition setup (model loading, connections, compiled resources) that shouldn’t repeat per row
  • Your logic genuinely needs to see multiple rows together within a partition (e.g., stateful processing, custom batching, sequential dependencies between rows)
  • You’re calling out to an external library that only has a batch/dataframe-style interface (e.g., a pandas-based ML model, an Arrow-native library)

Don’t use it when:

  • The logic is expressible with built-in F.* functions: those get whole-stage codegen and Photon acceleration; a Python mapInPandas step is a black box to Catalyst/Photon, same problem as raw RDDs, just batched instead of per-row
  • A regular (scalar) pandas UDF (@pandas_udf) would do: those are also vectorized and Arrow-based but map more naturally to a single-column expression rather than manipulating a whole partition/DataFrame
# If this is achievable with built-ins, prefer it — stays fully in Catalyst/Photon:
df.withColumn("score", F.col("feature1") * 0.5 + F.col("feature2") * 0.3)

# vs mapInPandas for the same simple case — unnecessary overhead, opaque to the optimizer
def process_batch(iterator):
    for pdf in iterator:
        pdf["score"] = pdf["feature1"] * 0.5 + pdf["feature2"] * 0.3
        yield pdf
df.mapInPandas(process_batch, schema=...)

Interaction with partition count/skew

Since mapPartitions is 1:1 on partitions, the partitioning you had going in is the partitioning you’re stuck with — if upstream partitions are skewed (one partition has 10x the rows of others), mapPartitions inherits that skew directly; it does nothing to rebalance.

# If input partitions are skewed, repartition BEFORE mapPartitions —
# but note this costs a shuffle, trading it against the skew problem
df.repartition(200, "some_balanced_key").mapInPandas(process_batch, schema=...)

This is the tradeoff: mapPartitions avoids a shuffle, but if that means processing wildly uneven partitions, you might still want a shuffle first to rebalance — the “no shuffle” property is a benefit only when your existing partitioning is already reasonably balanced.

Databricks-specific notes

  • Photon does not accelerate mapPartitions/mapInPandas/UDF code — these are Python/JVM user code Photon can’t see into. Expect the plan to show a gap: Photon operators before/after, plain Python execution for the mapInPandas step itself.
  • For ML inference specifically, Databricks often recommends mapInPandas (or pandas_udf) over raw RDD mapPartitions specifically because it integrates with the DataFrame plan and Arrow serialization, which is materially faster than Python-object pickling used by classic RDD operations.
  • If you’re calling a model registered in Unity Catalog / MLflow, there’s often a purpose-built helper (mlflow.pyfunc.spark_udf) that wraps this pattern for you, handling batching and broadcasting the model efficiently — worth checking before hand-rolling mapInPandas for standard model-scoring workloads.

18. Use Approximate distinct counts

This is a good option when you can trade exactness for a smaller shuffle:

# Exact count(distinct) requires shuffling enough info to dedupe exactly
exact = df.groupBy("region").agg(F.countDistinct("user_id"))

# HyperLogLog sketch: merges as tiny fixed-size sketches instead of raw values
approx = df.groupBy("region").agg(F.approx_count_distinct("user_id", rsd=0.01))

approx_count_distinct is a textbook case of “pre-aggregate before the shuffle” — each partition builds a small HLL sketch locally, and only the sketches (not the raw distinct values) cross the network.

19. Use DISTRIBUTE BY / CLUSTER BY in SQL for explicit shuffle key control

Good one to dig into — there’s actually a naming collision on Databricks worth flagging up front: SQL’s CLUSTER BY clause is a completely different feature from Delta Lake’s CLUSTER BY table property (Liquid Clustering, a physical data layout feature). Same keyword, unrelated purposes. I’ll cover the SQL clause here and note the collision at the end.

The three keywords and what they map to:

SQL clause DataFrame equivalent Shuffles? Sorts?
DISTRIBUTE BY .repartition(cols) Yes No
SORT BY .sortWithinPartitions(cols) No Within-partition only
CLUSTER BY .repartition(cols).sortWithinPartitions(cols) Yes Within-partition
ORDER BY .orderBy(cols) Yes (range shuffle) Global, total order

CLUSTER BY is literally shorthand for DISTRIBUTE BY x SORT BY x, same column(s) for both partitioning and local sort.

Why use them instead of letting Spark decide

By default, a GROUP BY/JOIN picks its own shuffle partitioning based on the join/group keys automatically. DISTRIBUTE BY/CLUSTER BY let you explicitly control partitioning independent of any aggregation, useful when the thing you want partitioned isn’t itself the target of a GROUP BY.

-- Explicitly redistribute data by user_id before writing out,
-- so each output file/partition contains one user's data contiguously,
-- without needing a GROUP BY to force it
SELECT *
FROM events
DISTRIBUTE BY user_id

-- CLUSTER BY: same partitioning key AND sort within each partition,
-- useful for downstream range scans or predicate pushdown on user_id
SELECT *
FROM events
CLUSTER BY user_id

A very common reason to reach for this explicitly is when you want to control output file layout. You’re writing a table and want each partition/file to be organized by a key, without introducing an aggregation.

-- Without this, Spark's default shuffle partitioning for a plain SELECT
-- may not group rows by user_id at all — files could have scattered users
CREATE TABLE user_events_clustered AS
SELECT *
FROM raw_events
CLUSTER BY user_id

# DataFrame equivalent
df.repartition("user_id").sortWithinPartitions("user_id") \
  .write.format("delta").saveAsTable("user_events_clustered")

This matters for downstream jobs that filter or join on user_id. If data is already co-located and sorted by that key, later shuffles/joins on it can be cheaper (data skipping, fewer partitions touched), and file-level statistics (min/max per file) become more selective.

CLUSTER BY vs ORDER BY

If you don’t need a single, fully globally-ordered result (e.g., you just want good local ordering for compression/skipping within each output file), CLUSTER BY is significantly cheaper than ORDER BY because it avoids the expensive range-partitioning shuffle that global ordering requires.

-- CLUSTER BY: only sorts WITHIN each partition, not globally —
-- much cheaper, no global ordering guarantee across partitions
SELECT * FROM events CLUSTER BY event_time

Checking the effect:

EXPLAIN
SELECT * FROM events CLUSTER BY user_id

In the physical plan, look for RepartitionByExpression (or Exchange hashpartitioning) followed by a Sort operator scoped within each partition (not a global range Sort, as the ORDER BY giveaway).

DISTRIBUTE BY vs a plain repartition(N)

# repartition(N): round-robin/random-ish redistribution, no key semantics
df.repartition(200)

# repartition(cols): hash-partitions by the given column(s) — same semantics as DISTRIBUTE BY
df.repartition("user_id")

# repartition(N, cols): explicit partition count AND key
df.repartition(200, "user_id")

DISTRIBUTE BY user_id in SQL is equivalent to .repartition("user_id") — it hash-partitions by that key, meaning all rows with the same user_id land in the same partition. This is what you want before an operation that needs per-key locality (e.g., a subsequent mapPartitions, a windowed function partitioned by that key, or just controlling file layout at write time) without necessarily triggering a GROUP BY.

When to reach for the query-level clause

  • Writing output where downstream jobs benefit from key-based file co-location (without an aggregation forcing that shape)
  • Overriding default shuffle partitioning/count when you know your key cardinality better than Spark’s defaults
  • Wanting local (not global) sort ordering cheaply, without paying for a full ORDER BY range-shuffle
  • Preparing data for a subsequent operation that benefits from per-key partition locality (custom mapPartitions, certain window functions, join co-partitioning)

For most everyday aggregation/join work, you don’t need this, since Catalyst’s automatic shuffle partitioning is fine. Reach for DISTRIBUTE BY/CLUSTER BY when you need to decouple partitioning logic from aggregation logic, or explicitly shape output file layout.

20. Replace groupByKey with reduceByKey / aggregateByKey to do map-side combines (RDD)

Here are concrete examples, from the classic RDD pitfall to how it looks on modern Spark/Databricks with the DataFrame API and Adaptive Query Execution.

Below is the classic anti-pattern (RDD API) and how to fix it. reduceByKey/combineByKey/aggregateByKey All do a map-side combine before the shuffle. groupByKey does not — it’s almost always the wrong choice for aggregation.

# BAD: groupByKey ships every raw value across the network,
# then sums them at the destination
rdd = sc.parallelize([("user1", 1), ("user2", 1), ("user1", 1)] * 1_000_000)

bad = rdd.groupByKey().mapValues(sum)  # shuffles ALL raw (key, value) pairs

# GOOD: reduceByKey combines locally per-partition first,
# only partial sums cross the network
good = rdd.reduceByKey(lambda a, b: a + b)

The DataFrame API on Spark does this automatically. On modern Spark (3.x/4.x, which is what Databricks runs), you seldom write manual combiners — Catalyst’s HashAggregateExec inserts a partial aggregation stage before the shuffle automatically.

df = spark.table("events")

result = (
    df.groupBy("user_id")
      .agg(F.count("*").alias("event_count"), F.sum("amount").alias("total_amount"))
)

result.explain(True)

Look at the physical plan — you’ll see something like:

== Physical Plan ==
*(2) HashAggregate(keys=[user_id#x], functions=[count(1), sum(amount#x)])
+- Exchange hashpartitioning(user_id#x, 200)
   +- *(1) HashAggregate(keys=[user_id#x], functions=[partial_count(1), partial_sum(amount#x)])
      +- FileScan parquet events

As you can see, that’s the optimization happening as expected. The partial_count/partial_sum run per-partition before the Exchange (shuffle).

It’s important to note that pre-aggregation is only useful when many rows are grouped into a few sets. If the keys are nearly unique, you’ll still shuffle practically everything, so don’t expect great results on high-cardinality columns.

21. Structured Streaming Stateful Operators to Minimize Cross-Partition Shuffle

In batch Spark, groupBy shuffles every time it runs. In Structured Streaming, a stateful operator (groupByKey().mapGroupsWithState(...) / flatMapGroupsWithState(...), or SQL-level windowed aggregations) also shuffles by key, but only once per query’s lifetime in the sense that the same key always lands on the same partition thereafter, because HashPartitioning with a fixed spark.sql.shuffle.partitions is deterministic across microbatches.

That means:

  • State for key K is stored in the state store local to whichever partition hash(K) % numShufflePartitions maps to.
  • On the next microbatch, new events for key K shuffle to that same partition, so they’re processed alongside the already-resident state; no cross-partition lookup or repartitioning of state itself is needed.
  • The state store (in-memory HDFSBackedStateStoreProvider or RocksDB-backed on Databricks Runtime 17.3+) lives colocated with the executor/partition, checkpointed incrementally to cloud storage (DBFS/S3/ADLS).

So “minimizes cross-partition shuffle” really means: avoids re-shuffling state on every batch — only the new incoming micro-batch data shuffles by key; the accumulated state doesn’t move.

So pay attention: if you change spark.sql.shuffle.partitions between restarts of a stateful streaming query, keys remap to different partitions, and existing checkpointed state becomes orphaned/unreadable for those keys. This is one of the most common production bugs with stateful streaming. Databricks recommends fixing this value explicitly rather than relying on defaults, and never changing it without a full state migration plan.

API: mapGroupsWithState vs flatMapGroupsWithState

  • mapGroupsWithState: exactly one output row per group per trigger. Simpler, but forces a fixed output shape.
  • flatMapGroupsWithState: zero, one, or many output rows per group. Used for things like emitting intermediate session events, splitting a group into multiple derived records, or suppressing output until some condition is met.

In Scala/Java, use mapGroupsWithState/flatMapGroupsWithState directly on typed Dataset.

The PySpark equivalent is applyInPandasWithState (pandas UDF–based, vectorized).

from pyspark.sql.streaming.state import GroupState, GroupStateTimeout
import pandas as pd

def update_session(key, pdf_iter, state: GroupState):
    if state.hasTimedOut:
        # emit final result, clean up
        state.remove()
        return iter([])  # or yield final row(s)

    total = state.getOption[0] if state.exists else 0
    for pdf in pdf_iter:
        total += pdf["amount"].sum()

    state.update((total,))
    state.setTimeoutDuration("10 minutes")  # event-time or processing-time timeout
    yield pd.DataFrame({"user_id": [key[0]], "running_total": [total]})

result = (
    df.groupBy("user_id")
      .applyInPandasWithState(
          update_session,
          outputStructType="user_id string, running_total double",
          stateStructType="total double",
          outputMode="update",
          timeoutConf=GroupStateTimeout.ProcessingTimeTimeout,
      )
)

Watermarks: the shuffle-adjacent cleanup mechanism

Watermarks don’t reduce shuffle directly, but they bound how long state for a key stays resident before eviction. Without them, state accumulates forever and eventually the state store (and its checkpoint) becomes the bottleneck, degrading throughput on every batch as more state has to be scanned/serialized. This matters here because unbounded state defeats the whole point of the partition-stability optimization: shuffle cost per batch stays low, but state store I/O cost grows unboundedly instead.

When to reach for this vs. built-in stateful aggregations

Most streaming aggregation needs are covered by SQL-native stateful operations, which are simpler and let Catalyst optimize more:

df.withWatermark("event_time", "10 minutes") \
  .groupBy(F.window("event_time", "5 minutes"), "user_id") \
  .agg(F.sum("amount"))

Reach for flatMapGroupsWithState instead when you need arbitrary custom logic that doesn’t map to sum/count/avg-style aggregation:

  • Session windows with custom gap logic beyond session_window()
  • State machines (e.g., tracking an order through custom status transitions)
  • Emitting different output shapes/rows depending on accumulated state
  • Deduplication logic beyond dropDuplicates

The tradeoff: you lose some Catalyst-level optimization since it’s opaque user code (similar caveat to UDFs breaking predicate pushdown), and you own correctness for state eviction — you must call state.setTimeoutDuration/setTimeoutTimestamp yourself or state grows unbounded.

Configuration Parameters

By default, Spark uses HDFSBackedStateStoreProvider, which maintains the state as maps in memory, with snapshots saved to durable storage. This works well for small states, but full state snapshots at each checkpoint become expensive as the state grows.

Databricks defaults to a RocksDB state store, which keeps state on local SSD instead of the JVM heap, avoiding GC pressure and supporting much larger per-partition state.

Changelog checkpointing is another specific optimization from Databricks, complementary to the shuffle-partition-stability benefit, not a replacement for it. It writes only the delta of state changes per batch to durable storage instead of full snapshots, significantly cutting checkpoint I/O and, indirectly, end-to-end batch latency.

spark.conf.set("spark.sql.streaming.stateStore.providerClass",
                "org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider")
spark.conf.set("spark.sql.streaming.stateStore.rocksdb.changelogCheckpointing.enabled", "true")

How to verify it’s working

Check the Spark UI’s Structured Streaming tab per query, or in Databricks, the streaming query’s metrics:

  • stateOperators in the query progress JSON: shows numRowsTotal, numRowsUpdated, memoryUsedBytes per stateful operator.
  • Watch numRowsTotal growth over time: if it grows unbounded, your watermark/timeout isn’t evicting state.
  • Consistently high stateStore commit/load latency in the metrics is a sign your state store provider choice (HDFS-backed vs RocksDB) is undersized for your state volume.

Conclusion

As we saw, shuffle is expensive because it serializes data, sends it over the network, and deserializes it on the receiving node. Network I/O is almost always the bottleneck in distributed jobs — much slower than local CPU work.

But shuffle is not something to eliminate at all costs — it’s a fundamental mechanism of distributed computing that, when properly managed, allows Spark to process datasets far beyond the capacity of a single machine. The goal is not to avoid shuffle entirely, but to reduce its cost through informed decisions at each stage of the pipeline: filtering and projecting early, favoring map-side aggregations, choosing broadcast joins when applicable, and sizing partitions appropriately for the workload.

In modern versions of Spark, especially on Databricks, much of this optimization work has been absorbed by Adaptive Query Execution. This shifts the engineer’s role: rather than manually tuning partition counts or forcing repartitions “just in case,” the more effective practice is to trust AQE as the default and intervene only when a specific, measurable problem is identified — a confirmed skew, a small file problem, or a bottleneck evidenced by the execution plan itself.

This distinction matters. Many of the anti-patterns discussed here — indiscriminate use ofrepartition, unnecessary UDFs, or heavy reliance on RDDs when the DataFrame API would suffice — originate from practices developed in earlier versions of Spark, where the engine offered less native intelligence. Applying them today, without validation, can result in effort spent solving problems the engine already handles on its own.

Ultimately, shuffle optimization is a matter of diagnosis before action. Tools like.explain("formatted"), the Spark UI, and AQE metrics offer visibility into what’s actually happening in the execution plan. Combining that visibility with the mitigation strategies presented — from serialization and buffer configuration to structural choices like avoiding shuffle through broadcast joins — allows performance to be treated as an engineering process based on evidence, not assumptions.

At the end…

Every wide transformation triggers a shuffle. Every shuffle writes to disk, transfers over the network, and creates a stage barrier. As “Julius” would tell you, “The cheapest shuffle is the one you avoid.”. The next cheapest is the one that crosses with map-side combine applied. AQE, push-based shuffle, and the External Shuffle Service handle the rest, but the highest-value optimizations are always architectural: filter early, broadcast small tables, pre-aggregate, and pick the right primitive.
The right management of shuffles is where that you save cost and boost performance!