How to set Spark SQL shuffle partitions

Spark SQL creates shuffle partitions when a query has to redistribute rows for joins, aggregations, sorts, or partitioning operations. Setting spark.sql.shuffle.partitions gives a SQL or DataFrame workload a deliberate shuffle task target instead of accepting the default value for every job.

The setting is a Spark SQL configuration key, so it must be present before the session or submitted application builds the query plan. Use the SparkSession builder for application-owned code, spark-submit with --conf for one submitted job, or spark-defaults.conf when a shared client should apply the same default to many applications.

Adaptive Query Execution can coalesce shuffle partitions after runtime statistics are available, which means the configured value is the initial shuffle target rather than a promise that every final stage keeps that count. For checkpointed stateful streaming queries, treat the partition count as part of the state layout and start with a new checkpoint when the state partitioning must change.

1. Choose the partition count and the configuration layer that starts the Spark application.

   Use SparkSession.builder.config("spark.sql.shuffle.partitions", "...") for one application, spark-submit with --conf for one submitted job, or spark-defaults.conf for a shared client default. Set the value before the SparkSession or SQL session creates the workload.

2. Save the local shuffle-partition check as shuffle_partitions_check.py.

   shuffle_partitions_check.py

   from pyspark.sql import SparkSession, functions as F
    
   spark = (
       SparkSession.builder
       .appName("sg-shuffle-partitions-check")
       .master("local[2]")
       .config("spark.ui.enabled", "false")
       .config("spark.ui.showConsoleProgress", "false")
       .config("spark.sql.adaptive.enabled", "false")
       .config("spark.sql.shuffle.partitions", "6")
       .getOrCreate()
   )
   spark.sparkContext.setLogLevel("ERROR")
    
   orders = spark.range(0, 60, 1, numPartitions=3).select(
       (F.col("id") % 6).alias("bucket")
   )
    
   totals = orders.groupBy("bucket").count().orderBy("bucket")
   rows = totals.collect()
    
   print(f"spark.sql.shuffle.partitions = {spark.conf.get('spark.sql.shuffle.partitions')}")
   print(f"result partitions = {totals.rdd.getNumPartitions()}")
   print("Physical plan")
   totals.explain(mode="simple")
   print("rows = " + ", ".join(f"{row.bucket}:{row['count']}" for row in rows))
    
   spark.stop()

   The local check disables spark.sql.adaptive.enabled only so the plan and result partition count show the configured value directly. Keep or tune Adaptive Query Execution separately when production jobs should coalesce small shuffle partitions.
   Related: How to configure Spark adaptive query execution

3. Run the check with the Python interpreter that has PySpark installed.

   $ python3 shuffle_partitions_check.py
   spark.sql.shuffle.partitions = 6
   result partitions = 6
   Physical plan
   == Physical Plan ==
   *(3) Sort [bucket#1L ASC NULLS FIRST], true, 0
   +- Exchange rangepartitioning(bucket#1L ASC NULLS FIRST, 6)
      +- *(2) HashAggregate(keys=[bucket#1L], functions=[count(1)])
         +- Exchange hashpartitioning(bucket#1L, 6)
   ##### snipped #####
   rows = 0:10, 1:10, 2:10, 3:10, 4:10, 5:10

4. Confirm the configuration value, result partition count, and plan exchange count match the intended value.

   The hashpartitioning(bucket#1L, 6) and rangepartitioning(bucket#1L ASC NULLS FIRST, 6) plan lines show Spark planned six-way shuffle exchanges for this query. The grouped rows confirm that each bucket still returned the expected count.

5. Remove the temporary check script when it is not part of the application.

   $ rm shuffle_partitions_check.py