[SPARK-28585][WebUI] Improve WebUI DAG information: Add extra info to rdd from spark plan

core/src/main/resources/org/apache/spark/ui/static/spark-dag-viz.js

@@ -496,7 +496,7 @@ function connectRDDs(fromRDDId, toRDDId, edgesContainer, svgContainer) {
  * Replace `/n` with `<br/>`
  */
 function replaceLineBreak(str) {
-    return str.replace("\\n", "<br/>");
+    return str.replace(/\\n/g, "<br/>");
 }
 
 /* (Job page only) Helper function to add tooltips for RDDs. */

core/src/main/scala/org/apache/spark/rdd/RDD.scala

@@ -161,6 +161,15 @@ abstract class RDD[T: ClassTag](
     this
   }
 
+  /** Extra information about the RDD */
+  @transient var extraInfo: Option[String] = None
+
+  /** Assign extraInfo to this RDD */
+  def setExtraInfo(_extraInfo: String): this.type = {
+    extraInfo = Option(_extraInfo)
+    this
+  }
+
   /**
    * Mark this RDD for persisting using the specified level.
    *

core/src/main/scala/org/apache/spark/storage/RDDInfo.scala

@@ -31,7 +31,8 @@ class RDDInfo(
     var storageLevel: StorageLevel,
     val parentIds: Seq[Int],
     val callSite: String = "",
-    val scope: Option[RDDOperationScope] = None)
+    val scope: Option[RDDOperationScope] = None,
+    val extraInfo: Option[String] = None)
   extends Ordered[RDDInfo] {
 
   var numCachedPartitions = 0
@@ -68,6 +69,6 @@ private[spark] object RDDInfo {
       rdd.creationSite.shortForm
     }
     new RDDInfo(rdd.id, rddName, rdd.partitions.length,
-      rdd.getStorageLevel, parentIds, callSite, rdd.scope)
+      rdd.getStorageLevel, parentIds, callSite, rdd.scope, rdd.extraInfo)
   }
 }

core/src/main/scala/org/apache/spark/ui/scope/RDDOperationGraph.scala

@@ -42,7 +42,8 @@ private[spark] case class RDDOperationGraph(
     rootCluster: RDDOperationCluster)
 
 /** A node in an RDDOperationGraph. This represents an RDD. */
-private[spark] case class RDDOperationNode(id: Int, name: String, cached: Boolean, callsite: String)
+private[spark] case class RDDOperationNode(id: Int, name: String, cached: Boolean, callsite: String,
+    extraInfo: Option[String] = None)
 
 /**
  * A directed edge connecting two nodes in an RDDOperationGraph.
@@ -143,7 +144,7 @@ private[spark] object RDDOperationGraph extends Logging {
 
       // TODO: differentiate between the intention to cache an RDD and whether it's actually cached
       val node = nodes.getOrElseUpdate(rdd.id, RDDOperationNode(
-        rdd.id, rdd.name, rdd.storageLevel != StorageLevel.NONE, rdd.callSite))
+        rdd.id, rdd.name, rdd.storageLevel != StorageLevel.NONE, rdd.callSite, rdd.extraInfo))
       if (rdd.scope.isEmpty) {
         // This RDD has no encompassing scope, so we put it directly in the root cluster
         // This should happen only if an RDD is instantiated outside of a public RDD API
@@ -227,7 +228,10 @@ private[spark] object RDDOperationGraph extends Logging {
     } else {
       ""
     }
-    val label = s"${node.name} [${node.id}]$isCached\n${node.callsite}"
+    val extraInfo = node.extraInfo.getOrElse("")
+
+    val label = s"${node.name} [${node.id}]$isCached\n${node.callsite}\n${extraInfo}"
+
     s"""${node.id} [label="${StringEscapeUtils.escapeJava(label)}"]"""
   }
 

core/src/main/scala/org/apache/spark/util/JsonProtocol.scala

@@ -462,6 +462,7 @@ private[spark] object JsonProtocol {
     ("Name" -> rddInfo.name) ~
     ("Scope" -> rddInfo.scope.map(_.toJson)) ~
     ("Callsite" -> rddInfo.callSite) ~
+    ("Extra Info" -> rddInfo.extraInfo) ~
     ("Parent IDs" -> parentIds) ~
     ("Storage Level" -> storageLevel) ~
     ("Number of Partitions" -> rddInfo.numPartitions) ~
@@ -1030,6 +1031,7 @@ private[spark] object JsonProtocol {
       .map(_.extract[String])
       .map(RDDOperationScope.fromJson)
     val callsite = jsonOption(json \ "Callsite").map(_.extract[String]).getOrElse("")
+    val extrainfo = jsonOption(json \ "Extra Info").map(_.extract[String])
     val parentIds = jsonOption(json \ "Parent IDs")
       .map { l => l.extract[List[JValue]].map(_.extract[Int]) }
       .getOrElse(Seq.empty)
@@ -1039,7 +1041,8 @@ private[spark] object JsonProtocol {
     val memSize = (json \ "Memory Size").extract[Long]
     val diskSize = (json \ "Disk Size").extract[Long]
 
-    val rddInfo = new RDDInfo(rddId, name, numPartitions, storageLevel, parentIds, callsite, scope)
+    val rddInfo =
+      new RDDInfo(rddId, name, numPartitions, storageLevel, parentIds, callsite, scope, extrainfo)
     rddInfo.numCachedPartitions = numCachedPartitions
     rddInfo.memSize = memSize
     rddInfo.diskSize = diskSize

core/src/test/scala/org/apache/spark/util/JsonProtocolSuite.scala

@@ -141,6 +141,7 @@ class JsonProtocolSuite extends SparkFunSuite {
     val logUrlMap = Map("stderr" -> "mystderr", "stdout" -> "mystdout").toMap
     val attributes = Map("ContainerId" -> "ct1", "User" -> "spark").toMap
     testRDDInfo(makeRddInfo(2, 3, 4, 5L, 6L))
+    testRDDInfo(makeRddInfo(2, 3, 4, 5L, 6L, Some("Extra Info")))
     testStageInfo(makeStageInfo(10, 20, 30, 40L, 50L))
     testTaskInfo(makeTaskInfo(999L, 888, 55, 777L, false))
     testTaskMetrics(makeTaskMetrics(
@@ -840,8 +841,10 @@ private[spark] object JsonProtocolSuite extends Assertions {
     )
   }
 
-  private def makeRddInfo(a: Int, b: Int, c: Int, d: Long, e: Long) = {
-    val r = new RDDInfo(a, "mayor", b, StorageLevel.MEMORY_AND_DISK, Seq(1, 4, 7), a.toString)
+  private def makeRddInfo(a: Int, b: Int, c: Int, d: Long, e: Long,
+      extraInfo: Option[String] = None) = {
+    val r = new RDDInfo(a, "mayor", b, StorageLevel.MEMORY_AND_DISK, Seq(1, 4, 7),
+      a.toString, None, extraInfo)
     r.numCachedPartitions = c
     r.memSize = d
     r.diskSize = e

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

@@ -616,8 +616,13 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan {
     }
   }
 
-  protected override def doExecute(): RDD[InternalRow] =
+  def addExtraInfo: RDD[InternalRow] => RDD[InternalRow] = {
+    _.setExtraInfo(s"Union(${output.mkString(", ")})")
+  }
+
+  protected override def doExecute(): RDD[InternalRow] = addExtraInfo {
     sparkContext.union(children.map(_.execute()))
+  }
 }
 
 /**