How to set Spark SQL shuffle partitions

Spark SQL moves rows between tasks when a query groups, joins, sorts, or repartitions data. The number of shuffle partitions sets the initial task layout for that exchange, so an unsuitable value can leave each task with too much data or create more scheduling work than the query needs.

The spark.sql.shuffle.partitions property applies to Spark SQL and DataFrame shuffles and defaults to 200. Setting it through spark.conf.set() changes the active SparkSession without placing an application-specific tuning value in a shared cluster defaults file.

Current Spark releases enable Adaptive Query Execution, which can combine small post-shuffle partitions at runtime. The configured value still appears on the Exchange operators as the initial count, while the effective result can contain fewer partitions after adaptive coalescing.

Steps to set Spark SQL shuffle partitions:

Start shuffle_partitions_check.py with the Spark session setup.

shuffle_partitions_check.py

from pyspark.sql import SparkSession, functions as F
 
 
spark = (
    SparkSession.builder
    .appName("shuffle-partitions-check")
    .config("spark.ui.enabled", "false")
    .config("spark.ui.showConsoleProgress", "false")
    .getOrCreate()
)
spark.sparkContext.setLogLevel("ERROR")

Set the session shuffle-partition count to six after the log-level line.
shuffle_partitions_check.py
```
spark.conf.set("spark.sql.shuffle.partitions", "6")
```
The value 6 represents an initial count suited to this small application. A production value should reflect its data volume and available executor cores because AQE cannot recover parallelism that an initial count makes too low.

A stateful Structured Streaming query cannot restart from the same checkpoint with a different spark.sql.shuffle.partitions value. Spark uses this count to partition state, so a different value requires a new checkpoint.

Define the input DataFrame after the partition setting.

shuffle_partitions_check.py

orders = spark.range(0, 60, 1, numPartitions=3).select(
    (F.col("id") % 6).alias("bucket")
)

Build the grouped result after the input DataFrame.
shuffle_partitions_check.py
```
totals = orders.groupBy("bucket").count().orderBy("bucket")
```
The grouped count creates a hash shuffle, and the global ordering creates a range shuffle. Both exchanges use the session's configured initial count.

Append the query execution block after the grouped result.

shuffle_partitions_check.py

rows = totals.collect()
 
print(
    "spark.sql.shuffle.partitions = "
    + spark.conf.get("spark.sql.shuffle.partitions")
)
print(f"result partitions after AQE = {totals.rdd.getNumPartitions()}")
print("Physical plan")
totals.explain(mode="simple")
print("rows = " + ", ".join(f"{row.bucket}:{row['count']}" for row in rows))
 
spark.stop()

Calling collect() completes the query before explain() prints the final adaptive plan. The tiny result is safe to return to the driver; do not collect an unbounded production DataFrame.

Submit the completed application to a local Spark master.

$ spark-submit --master local[2] shuffle_partitions_check.py
spark.sql.shuffle.partitions = 6
result partitions after AQE = 1
Physical plan
== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=true
+- == Final Plan ==
   ResultQueryStage 3
   +- *(3) Sort [bucket#1L ASC NULLS FIRST], true, 0
      +- AQEShuffleRead coalesced
         +- ShuffleQueryStage 1
            +- Exchange rangepartitioning(bucket#1L ASC NULLS FIRST, 6), ENSURE_REQUIREMENTS, [plan_id=60]
##### snipped #####
                        +- Exchange hashpartitioning(bucket#1L, 6), ENSURE_REQUIREMENTS, [plan_id=36]
##### snipped #####
rows = 0:10, 1:10, 2:10, 3:10, 4:10, 5:10

The configuration readback and both Exchange lines show an initial count of six. result partitions after AQE = 1 shows that AQE combined the small shuffle output, while the six row counts confirm that the grouped query completed.