[SPARK-41752][SQL][UI] Group nested executions under the root execution

core/src/main/protobuf/org/apache/spark/status/protobuf/store_types.proto

@@ -417,6 +417,7 @@ message SQLExecutionUIData {
   repeated int64 stages = 11;
   bool metric_values_is_null = 12;
   map<int64, string> metric_values = 13;
+  optional int64 root_execution_id = 14;
 }
 
 message SparkPlanGraphNode {

core/src/main/resources/org/apache/spark/ui/static/webui.css

@@ -187,6 +187,18 @@ pre {
   display: none;
 }
 
+.sub-execution-list {
+  font-size: 0.9rem;
+}
+
+.sub-execution-list.collapsed {
+  display: none;
+}
+
+.table-striped .sub-execution-list table tr {
+  background-color: inherit;
+}
+
 .description-input {
   overflow: hidden;
   text-overflow: ellipsis;

core/src/main/scala/org/apache/spark/internal/config/UI.scala

@@ -229,4 +229,11 @@ private[spark] object UI {
     .stringConf
     .transform(_.toUpperCase(Locale.ROOT))
     .createWithDefault("LOCAL")
+
+  val UI_SQL_GROUP_SUB_EXECUTION_ENABLED = ConfigBuilder("spark.ui.groupSQLSubExecutionEnabled")
+    .doc("Whether to group sub executions together in SQL UI when they belong to the same " +
+      "root execution")
+    .version("3.4.0")
+    .booleanConf
+    .createWithDefault(true)
 }

sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala

@@ -30,6 +30,7 @@ import org.apache.spark.util.Utils
 object SQLExecution {
 
   val EXECUTION_ID_KEY = "spark.sql.execution.id"
+  val EXECUTION_ROOT_ID_KEY = "spark.sql.execution.root.id"
 
   private val _nextExecutionId = new AtomicLong(0)
 
@@ -67,6 +68,13 @@ object SQLExecution {
     val oldExecutionId = sc.getLocalProperty(EXECUTION_ID_KEY)
     val executionId = SQLExecution.nextExecutionId
     sc.setLocalProperty(EXECUTION_ID_KEY, executionId.toString)
+    // Track the "root" SQL Execution Id for nested/sub queries. The current execution is the
+    // root execution if the root execution ID is null.
+    // And for the root execution, rootExecutionId == executionId.
+    if (sc.getLocalProperty(EXECUTION_ROOT_ID_KEY) == null) {
+      sc.setLocalProperty(EXECUTION_ROOT_ID_KEY, executionId.toString)
+    }
+    val rootExecutionId = sc.getLocalProperty(EXECUTION_ROOT_ID_KEY).toLong
     executionIdToQueryExecution.put(executionId, queryExecution)
     try {
       // sparkContext.getCallSite() would first try to pick up any call site that was previously
@@ -98,6 +106,7 @@ object SQLExecution {
         try {
           sc.listenerBus.post(SparkListenerSQLExecutionStart(
             executionId = executionId,
+            rootExecutionId = rootExecutionId,
             description = desc,
             details = callSite.longForm,
             physicalPlanDescription = queryExecution.explainString(planDescriptionMode),
@@ -140,6 +149,11 @@ object SQLExecution {
     } finally {
       executionIdToQueryExecution.remove(executionId)
       sc.setLocalProperty(EXECUTION_ID_KEY, oldExecutionId)
+      // Unset the "root" SQL Execution Id once the "root" SQL execution completes.
+      // The current execution is the root execution if rootExecutionId == executionId.
+      if (sc.getLocalProperty(EXECUTION_ROOT_ID_KEY) == executionId.toString) {
+        sc.setLocalProperty(EXECUTION_ROOT_ID_KEY, null)
+      }
     }
   }