Use args to display lore info.

nvliyuan · Jun 26, 2024 · 7730d32 · 7730d32
1 parent b4610cd
commit 7730d32
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 7 deletions.
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAggregateExec.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAggregateExec.scala
@@ -1758,6 +1758,7 @@ case class GpuHashAggregateExec(
        |${ExplainUtils.generateFieldString("Functions", aggregateExpressions)}
        |${ExplainUtils.generateFieldString("Aggregate Attributes", aggregateAttributes)}
        |${ExplainUtils.generateFieldString("Results", resultExpressions)}
+       |Lore: ${loreArgs.mkString(", ")}
        |""".stripMargin
   }
 

diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuExec.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuExec.scala
@@ -20,7 +20,7 @@ import ai.rapids.cudf.NvtxColor
 import com.nvidia.spark.rapids.Arm.withResource
 import com.nvidia.spark.rapids.filecache.FileCacheConf
 import com.nvidia.spark.rapids.lore.{GpuLore, GpuLoreDumpRDD}
-import com.nvidia.spark.rapids.lore.GpuLore.{lordIdOf, LORE_DUMP_PATH_TAG, LORE_DUMP_RDD_TAG}
+import com.nvidia.spark.rapids.lore.GpuLore.{loreIdOf, LORE_DUMP_PATH_TAG, LORE_DUMP_RDD_TAG}
 import com.nvidia.spark.rapids.shims.SparkShimImpl
 import org.apache.hadoop.fs.Path
 
@@ -378,11 +378,14 @@ trait GpuExec extends SparkPlan {
     }.getOrElse(orig)
   }
 
-  override def nodeName: String = {
-    val loreId = lordIdOf(this)
-    val lorePath = getTagValue(LORE_DUMP_PATH_TAG)
-    val loreRDDInfo = getTagValue(LORE_DUMP_RDD_TAG)
-    s"${super.nodeName} [loreId=$loreId] [lorePath=$lorePath] [loreRDDInfo=$loreRDDInfo]"
+  override def stringArgs: Iterator[Any] = super.stringArgs ++ loreArgs
+
+  protected def loreArgs: Iterator[String] = {
+    val loreIdStr = loreIdOf(this).map(id => s"[loreId=$id]")
+    val lorePathStr = getTagValue(LORE_DUMP_PATH_TAG).map(path => s"[lorePath=$path]")
+    val loreRDDInfoStr = getTagValue(LORE_DUMP_RDD_TAG).map(info => s"[loreRDDInfo=$info]")
+
+    List(loreIdStr, lorePathStr, loreRDDInfoStr).flatten.iterator
   }
 
   private def dumpLoreMetaInfo(): Unit = {

diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/lore/GpuLore.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/lore/GpuLore.scala
@@ -226,7 +226,7 @@ object GpuLore {
     newPlan
   }
 
-  def lordIdOf(node: SparkPlan): Option[String] = {
+  def loreIdOf(node: SparkPlan): Option[String] = {
     node.getTagValue(LORE_ID_TAG)
   }