From 6ad74b7e32965465a912eeca817b033b2d0e7214 Mon Sep 17 00:00:00 2001
From: Chang chen <baibaichen@gmail.com>
Date: Thu, 25 Dec 2025 20:14:30 +0800
Subject: [PATCH] [SPARK-53413][SQL][FOLLOWUP] Fix `extractShuffleIds` to find
 `AdaptiveSparkPlanExec` anywhere in plan tree

---
 .../org/apache/spark/sql/execution/SQLExecution.scala      | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala
index 1cab0f8d35af5..19bafeb196122 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala
@@ -71,11 +71,12 @@ object SQLExecution extends Logging {
   }
 
   private def extractShuffleIds(plan: SparkPlan): Seq[Int] = {
-    plan match {
+    val shuffleIdsOption = plan.collectFirst {
       case ae: AdaptiveSparkPlanExec =>
         ae.context.shuffleIds.asScala.keys.toSeq
-      case nonAdaptivePlan =>
-        nonAdaptivePlan.collect {
+    }
+    shuffleIdsOption.getOrElse {
+        plan.collect {
           case exec: ShuffleExchangeLike => exec.shuffleId
         }
     }