How to configure Spark adaptive query execution

Spark SQL can change a physical plan after shuffle and join statistics are available, which matters when a query plan does not match the data shape seen at runtime. Adaptive Query Execution, or AQE, lets Spark re-optimize SQL and DataFrame work while the job runs instead of relying only on the first compiled plan.

AQE is controlled by Spark SQL configuration, with spark.sql.adaptive.enabled as the umbrella switch. Apache Spark enables it by default in Spark 3.2.0 and later, but setting it explicitly is useful when a cluster profile, notebook session, or application wrapper may have changed the default.

Set AQE at the layer that starts the application, then verify it inside the same session that runs the query. The local PySpark check uses a small shuffle-heavy query so the active settings, adaptive plan marker, and final row result are visible from one command.

1. Choose the configuration surface that owns the Spark application.

   Use the SparkSession builder when application code creates the session, spark-submit --conf when an operator starts the job, or spark-defaults.conf when every job from the same client should inherit the setting.

2. Add the AQE settings before the SparkSession is created.

   aqe-check.py

   from pyspark.sql import SparkSession
   from pyspark.sql.functions import col
    
   spark = (
       SparkSession.builder.appName("sg-aqe-check")
       .config("spark.sql.adaptive.enabled", "true")
       .config("spark.sql.adaptive.coalescePartitions.enabled", "true")
       .config("spark.sql.shuffle.partitions", "8")
       .getOrCreate()
   )
   spark.sparkContext.setLogLevel("ERROR")
    
   orders = spark.range(0, 1000).select(
       (col("id") % 8).alias("customer_id"),
       (col("id") * 10).alias("amount"),
   )
   customers = spark.range(0, 8).select(col("id").alias("customer_id"))
    
   result = (
       orders.join(customers, "customer_id")
       .groupBy("customer_id")
       .count()
       .orderBy("customer_id")
   )
    
   rows = result.collect()
    
   print(f"spark.sql.adaptive.enabled = {spark.conf.get('spark.sql.adaptive.enabled')}")
   print(
       "spark.sql.adaptive.coalescePartitions.enabled = "
       f"{spark.conf.get('spark.sql.adaptive.coalescePartitions.enabled')}"
   )
   print(f"spark.sql.shuffle.partitions = {spark.conf.get('spark.sql.shuffle.partitions')}")
   result.explain()
   print(f"row_count = {len(rows)}")
   print("first_rows = " + ", ".join(f"{row.customer_id}:{row['count']}" for row in rows[:3]))
    
   spark.stop()

   Keep spark.sql.shuffle.partitions=8 only for this local check. Use a workload-sized shuffle partition value for production jobs.

3. Run the AQE check script.

   $ python aqe-check.py
   spark.sql.adaptive.enabled = true
   spark.sql.adaptive.coalescePartitions.enabled = true
   spark.sql.shuffle.partitions = 8
   == Physical Plan ==
   AdaptiveSparkPlan isFinalPlan=true
   +- == Final Plan ==
      ResultQueryStage 3
      +- *(4) Sort [customer_id#1L ASC NULLS FIRST], true, 0
         +- AQEShuffleRead coalesced
   ##### snipped #####
   row_count = 8
   first_rows = 0:125, 1:125, 2:125

4. Confirm the explain output shows an adaptive final plan after the action runs.

   AdaptiveSparkPlan isFinalPlan=true means Spark printed the final adaptive physical plan after collect() materialized the result. AQEShuffleRead coalesced shows AQE coalesced shuffle reads for this small check. If an ordinary physical plan appears instead, set spark.sql.adaptive.enabled before creating the SparkSession and rerun the query action before calling explain().