fix transformStageId

Author: marin-ma

backends-velox/src-delta33/main/scala/org/apache/spark/sql/delta/files/GlutenDeltaFileFormatWriter.scala

@@ -260,9 +260,8 @@ object GlutenDeltaFileFormatWriter extends LoggingShims {
             nativeSortPlan
           }
           val newPlan = sortPlan.child match {
-            case WholeStageTransformer(wholeStageChild, materializeInput) =>
-              WholeStageTransformer(addNativeSort(wholeStageChild),
-                materializeInput)(ColumnarCollapseTransformStages.transformStageCounter.incrementAndGet())
+            case wst @ WholeStageTransformer(wholeStageChild, _) =>
+              wst.withNewChildren(Seq(addNativeSort(wholeStageChild)))
             case other =>
               Transitions.toBatchPlan(sortPlan, VeloxBatchType)
           }

gluten-substrait/src/main/scala/org/apache/gluten/execution/TakeOrderedAndProjectExecTransformer.scala

@@ -137,7 +137,7 @@ case class TakeOrderedAndProjectExecTransformer(
           LimitExecTransformer(localSortPlan, limitBeforeShuffleOffset, limit)
       }
       val transformStageCounter: AtomicInteger =
-        ColumnarCollapseTransformStages.transformStageCounter
+        ColumnarCollapseTransformStages.getTransformStageCounter(child)
       val finalLimitPlan = if (hasShuffle) {
         limitBeforeShuffle
       } else {

gluten-substrait/src/main/scala/org/apache/gluten/execution/WholeStageTransformer.scala

@@ -43,6 +43,8 @@ import org.apache.spark.sql.vectorized.ColumnarBatch
 import com.google.common.collect.Lists
 import org.apache.hadoop.fs.viewfs.ViewFileSystemUtils
 
+import java.util.concurrent.atomic.AtomicInteger
+
 import scala.collection.JavaConverters._
 import scala.collection.mutable
 

gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ColumnarCollapseTransformStages.scala

@@ -141,7 +141,7 @@ case class InputIteratorTransformer(child: SparkPlan) extends UnaryTransformSupp
  */
 case class ColumnarCollapseTransformStages(
     glutenConf: GlutenConfig,
-    transformStageCounter: AtomicInteger = ColumnarCollapseTransformStages.transformStageCounter)
+    transformStageCounter: AtomicInteger = new AtomicInteger(0))
   extends Rule[SparkPlan] {
 
   def apply(plan: SparkPlan): SparkPlan = {
@@ -218,4 +218,10 @@ object ColumnarCollapseTransformStages {
   def wrapInputIteratorTransformer(plan: SparkPlan): TransformSupport = {
     InputIteratorTransformer(ColumnarInputAdapter(plan))
   }
+
+  def getTransformStageCounter(plan: SparkPlan): AtomicInteger = {
+    plan
+      .collectFirst { case wst: WholeStageTransformer => new AtomicInteger(wst.transformStageId) }
+      .getOrElse(new AtomicInteger(0))
+  }
 }

gluten-substrait/src/main/scala/org/apache/spark/sql/execution/datasources/GlutenFormatWriterInjectsBase.scala

@@ -23,7 +23,6 @@ import org.apache.gluten.extension.columnar.heuristic.HeuristicTransform
 import org.apache.gluten.extension.columnar.transition.{InsertTransitions, RemoveTransitions}
 
 import org.apache.spark.sql.execution.{ColumnarCollapseTransformStages, SparkPlan}
-import org.apache.spark.sql.execution.ColumnarCollapseTransformStages.transformStageCounter
 
 trait GlutenFormatWriterInjectsBase extends GlutenFormatWriterInjects {
   private lazy val transform = HeuristicTransform.static()
@@ -66,7 +65,7 @@ trait GlutenFormatWriterInjectsBase extends GlutenFormatWriterInjects {
     // and cannot provide const-ness.
     val transformedWithAdapter = injectAdapter(transformed)
     val wst = WholeStageTransformer(transformedWithAdapter, materializeInput = true)(
-      transformStageCounter.incrementAndGet())
+      ColumnarCollapseTransformStages.getTransformStageCounter(transformed).incrementAndGet())
     val wstWithTransitions = BackendsApiManager.getSparkPlanExecApiInstance.genColumnarToCarrierRow(
       InsertTransitions.create(outputsColumnar = true, wst.batchType()).apply(wst))
     wstWithTransitions

gluten-ut/spark40/src/test/scala/org/apache/spark/sql/GlutenSQLQuerySuite.scala

@@ -181,8 +181,8 @@ class GlutenSQLQuerySuite extends SQLQuerySuite with GlutenSQLTestsTrait {
       spark.sql("explain select :first", Map("first" -> 1)),
       """== Physical Plan ==
         |ColumnarToRow
-        |+- ^(2) ProjectExecTransformer [1 AS 1#N]
-        |   +- ^(2) InputIteratorTransformer[]
+        |+- ^(1) ProjectExecTransformer [1 AS 1#N]
+        |   +- ^(1) InputIteratorTransformer[]
         |      +- RowToColumnar
         |         +- *(1) Scan OneRowRelation[]
         |