(spark) branch master updated: [SPARK-56020][SQL] Improve `GroupPartitions` Spark UI

[pt...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

ptoth pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/master by this push:
     new 2188c9977c58 [SPARK-56020][SQL] Improve `GroupPartitions` Spark UI
2188c9977c58 is described below

commit 2188c9977c581b85e42c285960ed2e416100f65e
Author: Szehon Ho <sz...@gmail.com>
AuthorDate: Thu Mar 19 08:14:52 2026 +0100

    [SPARK-56020][SQL] Improve `GroupPartitions` Spark UI
    
    ### What changes were proposed in this pull request?
    Follow up for https://github.com/apache/spark/pull/54459
    1. It handled explain but not UI
    2. Reducers could be improved a bit
    3. Fix a test typo
    
    ### Why are the changes needed?
    Improve UX for the new SPJ group operator
    
    ### Does this PR introduce _any_ user-facing change?
    Adds to Explain Extended, and changes unreleased Spark UI for the GroupPartitions node a bit
    
    ### How was this patch tested?
    Add unit test
    
    ### Was this patch authored or co-authored using generative AI tooling?
    Yes, cursor
    
    Closes #54842 from szehon-ho/group_partitions_string_follow.
    
    Authored-by: Szehon Ho <sz...@gmail.com>
    Signed-off-by: Peter Toth <pe...@gmail.com>
---
 .../sql/connector/catalog/functions/Reducer.java   |  8 +++++++
 .../datasources/v2/GroupPartitionsExec.scala       | 26 +++++++++++++++-------
 .../connector/KeyGroupedPartitioningSuite.scala    | 22 ++++++++++--------
 .../catalog/functions/transformFunctions.scala     |  1 +
 4 files changed, 40 insertions(+), 17 deletions(-)

diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/functions/Reducer.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/functions/Reducer.java
index 561d66092d64..54d22057b45a 100644
--- a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/functions/Reducer.java
+++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/functions/Reducer.java
@@ -39,4 +39,12 @@ import org.apache.spark.annotation.Evolving;
 @Evolving
 public interface Reducer<I, O> {
   O reduce(I arg);
+
+  /**
+   * Returns a short string representation for plan display (e.g. explain output).
+   * Defaults to the unqualified class name; implementations may override for a richer description.
+   */
+  default String displayName() {
+    return getClass().getSimpleName();
+  }
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala
index d5260cad9c6e..10e59c7b3611 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala
@@ -25,7 +25,7 @@ import org.apache.spark.rdd.{CoalescedRDD, PartitionCoalescer, PartitionGroup, R
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.physical.{KeyedPartitioning, Partitioning}
-import org.apache.spark.sql.catalyst.util.InternalRowComparableWrapper
+import org.apache.spark.sql.catalyst.util.{truncatedString, InternalRowComparableWrapper}
 import org.apache.spark.sql.connector.catalog.functions.Reducer
 import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode}
 import org.apache.spark.sql.types.DataType
@@ -193,13 +193,23 @@ case class GroupPartitionsExec(
   }
 
   override def simpleString(maxFields: Int): String = {
-    val joinKeyPositionsString =
-      joinKeyPositions.map(p => s" JoinKeyPositions: ${p.mkString("[", ",", "]")}").getOrElse("")
-    val expectedPartitionKeysString =
-      expectedPartitionKeys.map(ks => s" ExpectedPartitionKeys: ${ks.size}").getOrElse("")
-    val reducersString = reducers.map(r => s" Reducers: ${r.count(_.isDefined)}").getOrElse("")
-    s"$nodeName$joinKeyPositionsString$expectedPartitionKeysString$reducersString " +
-      s"DistributePartitions: $distributePartitions"
+    s"$nodeName${planSummaryParts(maxFields).map(" " + _).mkString("")}"
+  }
+
+  override protected def stringArgs: Iterator[Any] = planSummaryParts(Int.MaxValue)
+
+  private def planSummaryParts(joinKeyMaxFields: Int): Iterator[String] = {
+    val joinKeyStr = joinKeyPositions.map { p =>
+      s"JoinKeyPositions: ${truncatedString(p, "[", ",", "]", joinKeyMaxFields)}"
+    }.iterator
+    val expectedStr = expectedPartitionKeys.map(ks => s"ExpectedPartitionKeys: ${ks.size}")
+    val reducersStr = reducers.map { seq =>
+      val names = seq.map(_.map(_.toString()).getOrElse("identity"))
+      s"Reducers: ${truncatedString(names, "[", ", ", "]", joinKeyMaxFields)}"
+    }
+    val distributeStr = Iterator(s"DistributePartitions: $distributePartitions")
+    joinKeyStr ++ expectedStr ++ reducersStr ++ distributeStr
+
   }
 }
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
index f2eb02e03846..c29be45b72fa 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
@@ -29,7 +29,7 @@ import org.apache.spark.sql.connector.catalog.functions._
 import org.apache.spark.sql.connector.distributions.Distributions
 import org.apache.spark.sql.connector.expressions._
 import org.apache.spark.sql.connector.expressions.Expressions._
-import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.{ExtendedMode, FormattedMode, SimpleMode, SparkPlan}
 import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation, GroupPartitionsExec}
 import org.apache.spark.sql.execution.exchange.{ShuffleExchangeExec, ShuffleExchangeLike}
 import org.apache.spark.sql.execution.joins.SortMergeJoinExec
@@ -3217,29 +3217,33 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
     }
   }
 
-  test("SPARK-55992: GroupPartitions textual representatin in plans") {
+  test("SPARK-55992: GroupPartitions string in simple and extended explain") {
     val items_partitions = Array(bucket(4, "id"), years("arrive_time"))
     createTable(items, itemsColumns, items_partitions)
-
     sql(s"INSERT INTO testcat.ns.$items VALUES (1, 'aa', 10.0, cast('2021-01-01' as timestamp))")
-
     val purchases_partitions = Array(bucket(6, "item_id"), years("time"))
     createTable(purchases, purchasesColumns, purchases_partitions)
-
     sql(s"INSERT INTO testcat.ns.$purchases VALUES (2, 10.0, cast('2021-01-01' as timestamp))")
-
     withSQLConf(
       SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false",
       SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true",
       SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
       val df = sql(
         s"""
-           |SELECT *
+           |${selectWithMergeJoinHint("i", "p")}
+           |*
            |FROM testcat.ns.$items i
            |JOIN testcat.ns.$purchases p ON p.item_id = i.id
            |""".stripMargin)
-      checkKeywordsExistsInExplain(df, keywords = "GroupPartitions JoinKeyPositions: [0] " +
-        "ExpectedPartitionKeys: 2 Reducers: 1 DistributePartitions: false")
+      val simpleAndExtendedKeyword =
+        "GroupPartitions JoinKeyPositions: [0] ExpectedPartitionKeys: 2 " +
+        "Reducers: [BucketReducer(2)] DistributePartitions: false"
+      val formattedKeyword =
+        "Arguments: JoinKeyPositions: [0], ExpectedPartitionKeys: 2, " +
+        "Reducers: [BucketReducer(2)], DistributePartitions: false"
+      checkKeywordsExistsInExplain(df, SimpleMode, simpleAndExtendedKeyword)
+      checkKeywordsExistsInExplain(df, ExtendedMode, simpleAndExtendedKeyword)
+      checkKeywordsExistsInExplain(df, FormattedMode, formattedKeyword)
     }
   }
 }
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/catalog/functions/transformFunctions.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/catalog/functions/transformFunctions.scala
index ed2f81d7e8d6..e7d05ab25e2d 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/connector/catalog/functions/transformFunctions.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/catalog/functions/transformFunctions.scala
@@ -114,6 +114,7 @@ object BucketFunction extends ScalarFunction[Int] with ReducibleFunction[Int, In
 
 case class BucketReducer(divisor: Int) extends Reducer[Int, Int] {
   override def reduce(bucket: Int): Int = bucket % divisor
+  override def displayName(): String = toString
 }
 
 object UnboundStringSelfFunction extends UnboundFunction {


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org