feat: Add config option for showing all Comet plan transformations (#1780)

Author: andygrove

common/src/main/scala/org/apache/comet/CometConf.scala

@@ -454,6 +454,14 @@ object CometConf extends ShimCometConf {
       .booleanConf
       .createWithDefault(false)
 
+  val COMET_EXPLAIN_TRANSFORMATIONS: ConfigEntry[Boolean] =
+    conf("spark.comet.explain.rules")
+      .doc("When this setting is enabled, Comet will log all plan transformations performed " +
+        "in physical optimizer rules. Default: false")
+      .internal()
+      .booleanConf
+      .createWithDefault(false)
+
   val COMET_EXPLAIN_FALLBACK_ENABLED: ConfigEntry[Boolean] =
     conf("spark.comet.explainFallback.enabled")
       .doc(

spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala

@@ -44,6 +44,9 @@ import org.apache.comet.serde.QueryPlanSerde
  * Spark physical optimizer rule for replacing Spark operators with Comet operators.
  */
 case class CometExecRule(session: SparkSession) extends Rule[SparkPlan] {
+
+  private lazy val showTransformations = CometConf.COMET_EXPLAIN_TRANSFORMATIONS.get()
+
   private def applyCometShuffle(plan: SparkPlan): SparkPlan = {
     plan.transformUp {
       case s: ShuffleExchangeExec
@@ -619,6 +622,14 @@ case class CometExecRule(session: SparkSession) extends Rule[SparkPlan] {
   }
 
   override def apply(plan: SparkPlan): SparkPlan = {
+    val newPlan = _apply(plan)
+    if (showTransformations) {
+      logInfo(s"\nINPUT: $plan\nOUTPUT: $newPlan")
+    }
+    newPlan
+  }
+
+  private def _apply(plan: SparkPlan): SparkPlan = {
     // DataFusion doesn't have ANSI mode. For now we just disable CometExec if ANSI mode is
     // enabled.
     if (isANSIEnabled(conf)) {

spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala

@@ -42,7 +42,18 @@ import org.apache.comet.parquet.{CometParquetScan, SupportsComet}
  * Spark physical optimizer rule for replacing Spark scans with Comet scans.
  */
 case class CometScanRule(session: SparkSession) extends Rule[SparkPlan] {
+
+  private lazy val showTransformations = CometConf.COMET_EXPLAIN_TRANSFORMATIONS.get()
+
   override def apply(plan: SparkPlan): SparkPlan = {
+    val newPlan = _apply(plan)
+    if (showTransformations) {
+      logInfo(s"\nINPUT: $plan\nOUTPUT: $newPlan")
+    }
+    newPlan
+  }
+
+  private def _apply(plan: SparkPlan): SparkPlan = {
     if (!isCometLoaded(conf) || !isCometScanEnabled(conf)) {
       if (!isCometLoaded(conf)) {
         withInfo(plan, "Comet is not enabled")

spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala

@@ -26,6 +26,8 @@ import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, Comet
 import org.apache.spark.sql.execution.{ColumnarToRowExec, RowToColumnarExec, SparkPlan}
 import org.apache.spark.sql.execution.adaptive.QueryStageExec
 
+import org.apache.comet.CometConf
+
 // This rule is responsible for eliminating redundant transitions between row-based and
 // columnar-based operators for Comet. Currently, three potential redundant transitions are:
 // 1. `ColumnarToRowExec` on top of an ending `CometCollectLimitExec` operator, which is
@@ -48,7 +50,18 @@ import org.apache.spark.sql.execution.adaptive.QueryStageExec
 // another `ColumnarToRowExec` on top of `CometSparkToColumnarExec`. In this case, the pair could
 // be removed.
 case class EliminateRedundantTransitions(session: SparkSession) extends Rule[SparkPlan] {
+
+  private lazy val showTransformations = CometConf.COMET_EXPLAIN_TRANSFORMATIONS.get()
+
   override def apply(plan: SparkPlan): SparkPlan = {
+    val newPlan = _apply(plan)
+    if (showTransformations) {
+      logInfo(s"\nINPUT: $plan\nOUTPUT: $newPlan")
+    }
+    newPlan
+  }
+
+  private def _apply(plan: SparkPlan): SparkPlan = {
     val eliminatedPlan = plan transformUp {
       case ColumnarToRowExec(sparkToColumnar: CometSparkToColumnarExec) =>
         if (sparkToColumnar.child.supportsColumnar) {