From d8b8f346318918aaa0a88ec93673b5b83eedc3d2 Mon Sep 17 00:00:00 2001
From: Hongze Zhang <hongze.zhang@intel.com>
Date: Thu, 2 Nov 2023 16:10:45 +0800
Subject: [PATCH] [VL] Fix for a corner case may break bhj + exchange
 validation

---
 .../extension/columnar/TransformHintRule.scala         | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/gluten-core/src/main/scala/io/glutenproject/extension/columnar/TransformHintRule.scala b/gluten-core/src/main/scala/io/glutenproject/extension/columnar/TransformHintRule.scala
index 7bede35f7abc..8aae39e8e994 100644
--- a/gluten-core/src/main/scala/io/glutenproject/extension/columnar/TransformHintRule.scala
+++ b/gluten-core/src/main/scala/io/glutenproject/extension/columnar/TransformHintRule.scala
@@ -313,8 +313,10 @@ case class AddTransformHintRule() extends Rule[SparkPlan] {
 
   /** Inserts a transformable tag on top of those that are not supported. */
   private def addTransformableTags(plan: SparkPlan): SparkPlan = {
-    addTransformableTag(plan)
-    plan.withNewChildren(plan.children.map(addTransformableTags))
+    // Walk the tree with post-order
+    val out = plan.withNewChildren(plan.children.map(addTransformableTags))
+    addTransformableTag(out)
+    out
   }
 
   private def addTransformableTag(plan: SparkPlan): Unit = {
@@ -551,7 +553,9 @@ case class AddTransformHintRule() extends Rule[SparkPlan] {
             maybeExchange match {
               case Some(exchange @ BroadcastExchangeExec(mode, child)) =>
                 TransformHints.tag(bhj, isBhjTransformable.toTransformHint)
-                TransformHints.tagNotTransformable(exchange, isBhjTransformable)
+                if (!isBhjTransformable.isValid) {
+                  TransformHints.tagNotTransformable(exchange, isBhjTransformable)
+                }
               case None =>
                 // we are in AQE, find the hidden exchange
                 // FIXME did we consider the case that AQE: OFF && Reuse: ON ?