[CORE] Only materialize subquery before doing transform (#5862)

We transform subquery(e.g., dpp) during columanr rules which is not actually been executed, so we should not materialize subquery when replacing expression as it is not in concurrent. This pr wraps doTransform with transform to always do materialize subquery before doTransform, so that the subquries can be submitted in concurrent.

backends-clickhouse/src/main/scala/org/apache/gluten/execution/CHHashAggregateExecTransformer.scala

@@ -81,8 +81,8 @@ case class CHHashAggregateExecTransformer(
     }
   }
 
-  override def doTransform(context: SubstraitContext): TransformContext = {
-    val childCtx = child.asInstanceOf[TransformSupport].doTransform(context)
+  override protected def doTransform(context: SubstraitContext): TransformContext = {
+    val childCtx = child.asInstanceOf[TransformSupport].transform(context)
     val operatorId = context.nextOperatorId(this.nodeName)
 
     val aggParams = new AggregationParams

backends-clickhouse/src/test/scala/org/apache/gluten/execution/extension/GlutenCustomAggExpressionSuite.scala

@@ -91,7 +91,7 @@ class GlutenCustomAggExpressionSuite extends GlutenClickHouseTPCHAbstractSuite {
     assert(planExecs(3).isInstanceOf[HashAggregateExec])
 
     val substraitContext = new SubstraitContext
-    planExecs(2).asInstanceOf[CHHashAggregateExecTransformer].doTransform(substraitContext)
+    planExecs(2).asInstanceOf[CHHashAggregateExecTransformer].transform(substraitContext)
 
     // Check the functions
     assert(substraitContext.registeredFunction.containsKey("custom_sum_double:req_fp64"))

backends-clickhouse/src/test/scala/org/apache/spark/sql/execution/benchmarks/CHParquetReadBenchmark.scala

@@ -104,7 +104,7 @@ object CHParquetReadBenchmark extends SqlBasedBenchmark with CHSqlBasedBenchmark
     val scanTime = chFileScan.longMetric("scanTime")
     // Generate Substrait plan
     val substraitContext = new SubstraitContext
-    val transformContext = chFileScan.doTransform(substraitContext)
+    val transformContext = chFileScan.transform(substraitContext)
     val outNames = new java.util.ArrayList[String]()
     for (attr <- outputAttrs) {
       outNames.add(ConverterUtils.genColumnNameWithExprId(attr))

backends-velox/src/main/scala/org/apache/gluten/execution/HashAggregateExecTransformer.scala

@@ -65,8 +65,8 @@ abstract class HashAggregateExecTransformer(
     super.output
   }
 
-  override def doTransform(context: SubstraitContext): TransformContext = {
-    val childCtx = child.asInstanceOf[TransformSupport].doTransform(context)
+  override protected def doTransform(context: SubstraitContext): TransformContext = {
+    val childCtx = child.asInstanceOf[TransformSupport].transform(context)
 
     val aggParams = new AggregationParams
     val operatorId = context.nextOperatorId(this.nodeName)

backends-velox/src/main/scala/org/apache/gluten/execution/TopNTransformer.scala

@@ -67,8 +67,8 @@ case class TopNTransformer(
     doNativeValidation(context, relNode)
   }
 
-  override def doTransform(context: SubstraitContext): TransformContext = {
-    val childCtx = child.asInstanceOf[TransformSupport].doTransform(context)
+  override protected def doTransform(context: SubstraitContext): TransformContext = {
+    val childCtx = child.asInstanceOf[TransformSupport].transform(context)
     val operatorId = context.nextOperatorId(this.nodeName)
     val relNode =
       getRelNode(

gluten-core/src/main/scala/org/apache/gluten/execution/BasicPhysicalOperatorTransformer.scala

@@ -129,8 +129,8 @@ abstract class FilterExecTransformerBase(val cond: Expression, val input: SparkP
     doNativeValidation(substraitContext, relNode)
   }
 
-  override def doTransform(context: SubstraitContext): TransformContext = {
-    val childCtx = child.asInstanceOf[TransformSupport].doTransform(context)
+  override protected def doTransform(context: SubstraitContext): TransformContext = {
+    val childCtx = child.asInstanceOf[TransformSupport].transform(context)
     val remainingCondition = getRemainingCondition
     val operatorId = context.nextOperatorId(this.nodeName)
     if (remainingCondition == null) {
@@ -190,7 +190,7 @@ case class ProjectExecTransformer private (projectList: Seq[NamedExpression], ch
     BackendsApiManager.getMetricsApiInstance.genProjectTransformerMetricsUpdater(metrics)
 
   override def doTransform(context: SubstraitContext): TransformContext = {
-    val childCtx = child.asInstanceOf[TransformSupport].doTransform(context)
+    val childCtx = child.asInstanceOf[TransformSupport].transform(context)
     val operatorId = context.nextOperatorId(this.nodeName)
     if ((projectList == null || projectList.isEmpty) && childCtx != null) {
       // The computing for this project is not needed.

gluten-core/src/main/scala/org/apache/gluten/execution/BasicScanExecTransformer.scala

@@ -80,7 +80,7 @@ trait BasicScanExecTransformer extends LeafTransformSupport with BaseDataSource
     val numOutputVectors = longMetric("outputVectors")
     val scanTime = longMetric("scanTime")
     val substraitContext = new SubstraitContext
-    val transformContext = doTransform(substraitContext)
+    val transformContext = transform(substraitContext)
     val outNames =
       filteRedundantField(outputAttributes()).map(ConverterUtils.genColumnNameWithExprId).asJava
     val planNode =
@@ -117,7 +117,7 @@ trait BasicScanExecTransformer extends LeafTransformSupport with BaseDataSource
     }
 
     val substraitContext = new SubstraitContext
-    val relNode = doTransform(substraitContext).root
+    val relNode = transform(substraitContext).root
 
     doNativeValidation(substraitContext, relNode)
   }
@@ -133,7 +133,7 @@ trait BasicScanExecTransformer extends LeafTransformSupport with BaseDataSource
     }
   }
 
-  override def doTransform(context: SubstraitContext): TransformContext = {
+  override protected def doTransform(context: SubstraitContext): TransformContext = {
     val output = filteRedundantField(outputAttributes())
     val typeNodes = ConverterUtils.collectAttributeTypeNodes(output)
     val nameList = ConverterUtils.collectAttributeNamesWithoutExprId(output)

gluten-core/src/main/scala/org/apache/gluten/execution/BroadcastNestedLoopJoinExecTransformer.scala

@@ -119,12 +119,12 @@ abstract class BroadcastNestedLoopJoinExecTransformer(
       }
   }
 
-  override def doTransform(context: SubstraitContext): TransformContext = {
-    val streamedPlanContext = streamedPlan.asInstanceOf[TransformSupport].doTransform(context)
+  override protected def doTransform(context: SubstraitContext): TransformContext = {
+    val streamedPlanContext = streamedPlan.asInstanceOf[TransformSupport].transform(context)
     val (inputStreamedRelNode, inputStreamedOutput) =
       (streamedPlanContext.root, streamedPlanContext.outputAttributes)
 
-    val buildPlanContext = buildPlan.asInstanceOf[TransformSupport].doTransform(context)
+    val buildPlanContext = buildPlan.asInstanceOf[TransformSupport].transform(context)
     val (inputBuildRelNode, inputBuildOutput) =
       (buildPlanContext.root, buildPlanContext.outputAttributes)
 

gluten-core/src/main/scala/org/apache/gluten/execution/CartesianProductExecTransformer.scala

@@ -77,12 +77,12 @@ case class CartesianProductExecTransformer(
     BackendsApiManager.getMetricsApiInstance.genNestedLoopJoinTransformerMetricsUpdater(metrics)
   }
 
-  override def doTransform(context: SubstraitContext): TransformContext = {
-    val leftPlanContext = left.asInstanceOf[TransformSupport].doTransform(context)
+  override protected def doTransform(context: SubstraitContext): TransformContext = {
+    val leftPlanContext = left.asInstanceOf[TransformSupport].transform(context)
     val (inputLeftRelNode, inputLeftOutput) =
       (leftPlanContext.root, leftPlanContext.outputAttributes)
 
-    val rightPlanContext = right.asInstanceOf[TransformSupport].doTransform(context)
+    val rightPlanContext = right.asInstanceOf[TransformSupport].transform(context)
     val (inputRightRelNode, inputRightOutput) =
       (rightPlanContext.root, rightPlanContext.outputAttributes)
 

gluten-core/src/main/scala/org/apache/gluten/execution/ExpandExecTransformer.scala

@@ -110,8 +110,8 @@ case class ExpandExecTransformer(
     doNativeValidation(substraitContext, relNode)
   }
 
-  override def doTransform(context: SubstraitContext): TransformContext = {
-    val childCtx = child.asInstanceOf[TransformSupport].doTransform(context)
+  override protected def doTransform(context: SubstraitContext): TransformContext = {
+    val childCtx = child.asInstanceOf[TransformSupport].transform(context)
     val operatorId = context.nextOperatorId(this.nodeName)
     if (projections == null || projections.isEmpty) {
       // The computing for this Expand is not needed.

gluten-core/src/main/scala/org/apache/gluten/execution/GenerateExecTransformerBase.scala

@@ -76,8 +76,8 @@ abstract class GenerateExecTransformerBase(
     doNativeValidation(context, relNode)
   }
 
-  override def doTransform(context: SubstraitContext): TransformContext = {
-    val childCtx = child.asInstanceOf[TransformSupport].doTransform(context)
+  override protected def doTransform(context: SubstraitContext): TransformContext = {
+    val childCtx = child.asInstanceOf[TransformSupport].transform(context)
     val relNode = getRelNode(context, childCtx.root, getGeneratorNode(context), validation = false)
     TransformContext(child.output, output, relNode)
   }

gluten-core/src/main/scala/org/apache/gluten/execution/JoinExecTransformer.scala

@@ -227,12 +227,12 @@ trait HashJoinLikeExecTransformer extends BaseJoinExec with TransformSupport {
     doNativeValidation(substraitContext, relNode)
   }
 
-  override def doTransform(context: SubstraitContext): TransformContext = {
-    val streamedPlanContext = streamedPlan.asInstanceOf[TransformSupport].doTransform(context)
+  override protected def doTransform(context: SubstraitContext): TransformContext = {
+    val streamedPlanContext = streamedPlan.asInstanceOf[TransformSupport].transform(context)
     val (inputStreamedRelNode, inputStreamedOutput) =
       (streamedPlanContext.root, streamedPlanContext.outputAttributes)
 
-    val buildPlanContext = buildPlan.asInstanceOf[TransformSupport].doTransform(context)
+    val buildPlanContext = buildPlan.asInstanceOf[TransformSupport].transform(context)
     val (inputBuildRelNode, inputBuildOutput) =
       (buildPlanContext.root, buildPlanContext.outputAttributes)
 

gluten-core/src/main/scala/org/apache/gluten/execution/LimitTransformer.scala

@@ -53,8 +53,8 @@ case class LimitTransformer(child: SparkPlan, offset: Long, count: Long)
     doNativeValidation(context, relNode)
   }
 
-  override def doTransform(context: SubstraitContext): TransformContext = {
-    val childCtx = child.asInstanceOf[TransformSupport].doTransform(context)
+  override protected def doTransform(context: SubstraitContext): TransformContext = {
+    val childCtx = child.asInstanceOf[TransformSupport].transform(context)
     val operatorId = context.nextOperatorId(this.nodeName)
     val relNode = getRelNode(context, operatorId, offset, count, child.output, childCtx.root, false)
     TransformContext(child.output, child.output, relNode)

gluten-core/src/main/scala/org/apache/gluten/execution/SortExecTransformer.scala

@@ -101,8 +101,8 @@ case class SortExecTransformer(
     doNativeValidation(substraitContext, relNode)
   }
 
-  override def doTransform(context: SubstraitContext): TransformContext = {
-    val childCtx = child.asInstanceOf[TransformSupport].doTransform(context)
+  override protected def doTransform(context: SubstraitContext): TransformContext = {
+    val childCtx = child.asInstanceOf[TransformSupport].transform(context)
     val operatorId = context.nextOperatorId(this.nodeName)
     if (sortOrder == null || sortOrder.isEmpty) {
       // The computing for this project is not needed.

gluten-core/src/main/scala/org/apache/gluten/execution/SortMergeJoinExecTransformer.scala

@@ -191,12 +191,12 @@ abstract class SortMergeJoinExecTransformerBase(
     doNativeValidation(substraitContext, relNode)
   }
 
-  override def doTransform(context: SubstraitContext): TransformContext = {
-    val streamedPlanContext = streamedPlan.asInstanceOf[TransformSupport].doTransform(context)
+  override protected def doTransform(context: SubstraitContext): TransformContext = {
+    val streamedPlanContext = streamedPlan.asInstanceOf[TransformSupport].transform(context)
     val (inputStreamedRelNode, inputStreamedOutput) =
       (streamedPlanContext.root, streamedPlanContext.outputAttributes)
 
-    val bufferedPlanContext = bufferedPlan.asInstanceOf[TransformSupport].doTransform(context)
+    val bufferedPlanContext = bufferedPlan.asInstanceOf[TransformSupport].transform(context)
     val (inputBuildRelNode, inputBuildOutput) =
       (bufferedPlanContext.root, bufferedPlanContext.outputAttributes)
 

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

@@ -68,7 +68,22 @@ trait TransformSupport extends GlutenPlan {
    */
   def columnarInputRDDs: Seq[RDD[ColumnarBatch]]
 
-  def doTransform(context: SubstraitContext): TransformContext = {
+  final def transform(context: SubstraitContext): TransformContext = {
+    if (isCanonicalizedPlan) {
+      throw new IllegalStateException(
+        "A canonicalized plan is not supposed to be executed transform.")
+    }
+    if (TransformerState.underValidationState) {
+      doTransform(context)
+    } else {
+      // Materialize subquery first before going to do transform.
+      executeQuery {
+        doTransform(context)
+      }
+    }
+  }
+
+  protected def doTransform(context: SubstraitContext): TransformContext = {
     throw new UnsupportedOperationException(
       s"This operator doesn't support doTransform with SubstraitContext.")
   }
@@ -182,7 +197,7 @@ case class WholeStageTransformer(child: SparkPlan, materializeInput: Boolean = f
     val substraitContext = new SubstraitContext
     val childCtx = child
       .asInstanceOf[TransformSupport]
-      .doTransform(substraitContext)
+      .transform(substraitContext)
     if (childCtx == null) {
       throw new NullPointerException(s"WholeStageTransformer can't do Transform on $child")
     }
@@ -216,8 +231,6 @@ case class WholeStageTransformer(child: SparkPlan, materializeInput: Boolean = f
   }
 
   def doWholeStageTransform(): WholeStageTransformContext = {
-    // invoke SparkPlan.prepare to do subquery preparation etc.
-    super.prepare()
     val context = generateWholeStageTransformContext()
     if (conf.getConf(GlutenConfig.CACHE_WHOLE_STAGE_TRANSFORMER_CONTEXT)) {
       wholeStageTransformerContext = Some(context)
@@ -257,6 +270,12 @@ case class WholeStageTransformer(child: SparkPlan, materializeInput: Boolean = f
 
   override def doExecuteColumnar(): RDD[ColumnarBatch] = {
     val pipelineTime: SQLMetric = longMetric("pipelineTime")
+    // We should do transform first to make sure all subqueries are materialized
+    val wsCtx = GlutenTimeMetric.withMillisTime {
+      doWholeStageTransform()
+    }(
+      t =>
+        logOnLevel(substraitPlanLogLevel, s"$nodeName generating the substrait plan took: $t ms."))
     val inputRDDs = new ColumnarInputRDDsWrapper(columnarInputRDDs)
     // Check if BatchScan exists.
     val basicScanExecTransformers = findAllScanTransformers()
@@ -271,22 +290,11 @@ case class WholeStageTransformer(child: SparkPlan, materializeInput: Boolean = f
       val allScanPartitions = basicScanExecTransformers.map(_.getPartitions)
       val allScanSplitInfos =
         getSplitInfosFromPartitions(basicScanExecTransformers, allScanPartitions)
-
-      val (wsCtx, inputPartitions) = GlutenTimeMetric.withMillisTime {
-        val wsCtx = doWholeStageTransform()
-        val partitions =
-          BackendsApiManager.getIteratorApiInstance.genPartitions(
-            wsCtx,
-            allScanSplitInfos,
-            basicScanExecTransformers)
-
-        (wsCtx, partitions)
-      }(
-        t =>
-          logOnLevel(
-            substraitPlanLogLevel,
-            s"$nodeName generating the substrait plan took: $t ms."))
-
+      val inputPartitions =
+        BackendsApiManager.getIteratorApiInstance.genPartitions(
+          wsCtx,
+          allScanSplitInfos,
+          basicScanExecTransformers)
       val rdd = new GlutenWholeStageColumnarRDD(
         sparkContext,
         inputPartitions,
@@ -321,22 +329,18 @@ case class WholeStageTransformer(child: SparkPlan, materializeInput: Boolean = f
        *      GlutenDataFrameAggregateSuite) in these cases, separate RDDs takes care of SCAN as a
        *      result, genFinalStageIterator rather than genFirstStageIterator will be invoked
        */
-      val resCtx = GlutenTimeMetric.withMillisTime(doWholeStageTransform()) {
-        t =>
-          logOnLevel(substraitPlanLogLevel, s"$nodeName generating the substrait plan took: $t ms.")
-      }
       new WholeStageZippedPartitionsRDD(
         sparkContext,
         inputRDDs,
         numaBindingInfo,
         sparkConf,
-        resCtx,
+        wsCtx,
         pipelineTime,
         BackendsApiManager.getMetricsApiInstance.metricsUpdatingFunction(
           child,
-          resCtx.substraitContext.registeredRelMap,
-          resCtx.substraitContext.registeredJoinParams,
-          resCtx.substraitContext.registeredAggregationParams
+          wsCtx.substraitContext.registeredRelMap,
+          wsCtx.substraitContext.registeredJoinParams,
+          wsCtx.substraitContext.registeredAggregationParams
         ),
         materializeInput
       )

gluten-core/src/main/scala/org/apache/gluten/execution/WindowExecTransformer.scala

@@ -179,8 +179,8 @@ case class WindowExecTransformer(
     doNativeValidation(substraitContext, relNode)
   }
 
-  override def doTransform(context: SubstraitContext): TransformContext = {
-    val childCtx = child.asInstanceOf[TransformSupport].doTransform(context)
+  override protected def doTransform(context: SubstraitContext): TransformContext = {
+    val childCtx = child.asInstanceOf[TransformSupport].transform(context)
     val operatorId = context.nextOperatorId(this.nodeName)
     if (windowExpression == null || windowExpression.isEmpty) {
       // The computing for this operator is not needed.

gluten-core/src/main/scala/org/apache/gluten/execution/WindowGroupLimitExecTransformer.scala

@@ -146,8 +146,8 @@ case class WindowGroupLimitExecTransformer(
     doNativeValidation(substraitContext, relNode)
   }
 
-  override def doTransform(context: SubstraitContext): TransformContext = {
-    val childCtx = child.asInstanceOf[TransformSupport].doTransform(context)
+  override protected def doTransform(context: SubstraitContext): TransformContext = {
+    val childCtx = child.asInstanceOf[TransformSupport].transform(context)
     val operatorId = context.nextOperatorId(this.nodeName)
 
     val currRel =

gluten-core/src/main/scala/org/apache/gluten/execution/WriteFilesExecTransformer.scala

@@ -161,8 +161,8 @@ case class WriteFilesExecTransformer(
     doNativeValidation(substraitContext, relNode)
   }
 
-  override def doTransform(context: SubstraitContext): TransformContext = {
-    val childCtx = child.asInstanceOf[TransformSupport].doTransform(context)
+  override protected def doTransform(context: SubstraitContext): TransformContext = {
+    val childCtx = child.asInstanceOf[TransformSupport].transform(context)
     val operatorId = context.nextOperatorId(this.nodeName)
     val currRel =
       getRelNode(context, getFinalChildOutput(), operatorId, childCtx.root, validation = false)

gluten-core/src/main/scala/org/apache/gluten/expression/ExpressionConverter.scala

@@ -654,7 +654,7 @@ object ExpressionConverter extends SQLConfHelper with Logging {
       // or ColumnarBroadcastExchange was disabled.
       partitionFilters
     } else {
-      val newPartitionFilters = partitionFilters.map {
+      partitionFilters.map {
         case dynamicPruning: DynamicPruningExpression =>
           dynamicPruning.transform {
             // Lookup inside subqueries for duplicate exchanges.
@@ -723,25 +723,6 @@ object ExpressionConverter extends SQLConfHelper with Logging {
           }
         case e: Expression => e
       }
-      updateSubqueryResult(newPartitionFilters)
-      newPartitionFilters
-    }
-  }
-
-  private def updateSubqueryResult(partitionFilters: Seq[Expression]): Unit = {
-    // When it includes some DynamicPruningExpression,
-    // it needs to execute InSubqueryExec first,
-    // because doTransform path can't execute 'doExecuteColumnar' which will
-    // execute prepare subquery first.
-    partitionFilters.foreach {
-      case DynamicPruningExpression(inSubquery: InSubqueryExec) =>
-        if (inSubquery.values().isEmpty) inSubquery.updateResult()
-      case e: Expression =>
-        e.foreach {
-          case s: ScalarSubquery => s.updateResult()
-          case _ =>
-        }
-      case _ =>
     }
   }
 }