fix ut

Author: liujiayi771

backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHSparkPlanExecApi.scala

@@ -28,7 +28,7 @@ import org.apache.gluten.shuffle.NeedCustomColumnarBatchSerializer
 import org.apache.gluten.sql.shims.SparkShimLoader
 import org.apache.gluten.substrait.SubstraitContext
 import org.apache.gluten.substrait.expression.{ExpressionBuilder, ExpressionNode, WindowFunctionNode}
-import org.apache.gluten.utils.{CHJoinValidateUtil, UnknownJoinStrategy}
+import org.apache.gluten.utils.{CHJoinValidateUtil, PullOutProjectHelper, UnknownJoinStrategy}
 import org.apache.gluten.vectorized.{BlockOutputStream, CHColumnarBatchSerializer, CHNativeBlock, CHStreamReader}
 
 import org.apache.spark.{ShuffleDependency, SparkEnv}
@@ -66,7 +66,7 @@ import java.util.{ArrayList => JArrayList, List => JList}
 
 import scala.collection.JavaConverters._
 
-class CHSparkPlanExecApi extends SparkPlanExecApi with Logging {
+class CHSparkPlanExecApi extends SparkPlanExecApi with Logging with PullOutProjectHelper {
 
   /** Transform GetArrayItem to Substrait. */
   override def genGetArrayItemTransformer(
@@ -488,22 +488,33 @@ class CHSparkPlanExecApi extends SparkPlanExecApi with Logging {
         (Seq.empty, false)
     }
 
+    val (newChild, newOutput, newBuildKeys) = if (buildKeys.exists(isNotAttribute)) {
+      val (newChild, newBuildKeys, _) = pullOutPreProjectForJoin(child, buildKeys)
+      (newChild, newChild.output, newBuildKeys)
+    } else {
+      (child, child.output, Seq.empty)
+    }
+
     // find the key index in the output
-    // Note: PullOutPreProject rule ensures join keys are AttributeReferences before this point,
-    // so no pre-projection is needed here.
     val keyColumnIndex = if (isNullAware) {
-      buildKeys(0) match {
-        case b: BoundReference => b.ordinal
-        case n: NamedExpression =>
-          child.output.indexWhere(o => o.name.equals(n.name) && o.exprId == n.exprId)
-        case key =>
-          throw new GlutenException(s"Cannot find $key in the child's output: ${child.output}")
+      def findKeyOrdinal(key: Expression, output: Seq[Attribute]): Int = {
+        key match {
+          case b: BoundReference => b.ordinal
+          case n: NamedExpression =>
+            output.indexWhere(o => (o.name.equals(n.name) && o.exprId == n.exprId))
+          case _ => throw new GlutenException(s"Cannot find $key in the child's output: $output")
+        }
+      }
+      if (newBuildKeys.isEmpty) {
+        findKeyOrdinal(buildKeys(0), newOutput)
+      } else {
+        findKeyOrdinal(newBuildKeys(0), newOutput)
       }
     } else {
       0
     }
     val countsAndBytes =
-      CHExecUtil.buildSideRDD(dataSize, child, isNullAware, keyColumnIndex).collect
+      CHExecUtil.buildSideRDD(dataSize, newChild, isNullAware, keyColumnIndex).collect
 
     val batches = countsAndBytes.map(_._2)
     val totalBatchesSize = batches.map(_.length).sum
@@ -524,10 +535,10 @@ class CHSparkPlanExecApi extends SparkPlanExecApi with Logging {
     numOutputRows += rowCount
     ClickHouseBuildSideRelation(
       mode,
-      child.output,
+      newOutput,
       batches.flatten,
       rowCount,
-      Seq.empty,
+      newBuildKeys,
       hasNullKeyValues)
   }
 

gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/rewrite/PullOutPreProject.scala

@@ -17,7 +17,6 @@
 package org.apache.gluten.extension.columnar.rewrite
 
 import org.apache.gluten.backendsapi.BackendsApiManager
-import org.apache.gluten.extension.columnar.heuristic.RewrittenNodeWall
 import org.apache.gluten.sql.shims.SparkShimLoader
 import org.apache.gluten.utils.PullOutProjectHelper
 
@@ -303,28 +302,6 @@ object PullOutPreProject extends RewriteSingleNode with PullOutProjectHelper {
         case _ =>
           (join.leftKeys, join.rightKeys)
       }
-
-      def pullOutPreProjectForJoin(joinChild: SparkPlan, joinKeys: Seq[Expression])
-          : (SparkPlan, Seq[Expression], mutable.HashMap[Expression, NamedExpression]) = {
-        val expressionMap = new mutable.HashMap[Expression, NamedExpression]()
-        if (joinKeys.exists(isNotAttribute)) {
-          val newJoinKeys =
-            joinKeys.toIndexedSeq.map(replaceExpressionWithAttribute(_, expressionMap))
-          val preProject = ProjectExec(
-            eliminateProjectList(joinChild.outputSet, expressionMap.values.toSeq),
-            joinChild)
-          joinChild match {
-            case r: RewrittenNodeWall =>
-              r.originalChild.logicalLink.foreach(preProject.setLogicalLink)
-            case _ =>
-              joinChild.logicalLink.foreach(preProject.setLogicalLink)
-          }
-          (preProject, newJoinKeys, expressionMap)
-        } else {
-          (joinChild, joinKeys, expressionMap)
-        }
-      }
-
       val (newLeft, newLeftKeys, leftMap) = pullOutPreProjectForJoin(join.left, leftKeys)
       val (newRight, newRightKeys, rightMap) = pullOutPreProjectForJoin(join.right, rightKeys)
       val newCondition = if (leftMap.nonEmpty || rightMap.nonEmpty) {

gluten-substrait/src/main/scala/org/apache/gluten/utils/PullOutProjectHelper.scala

@@ -18,14 +18,15 @@ package org.apache.gluten.utils
 
 import org.apache.gluten.backendsapi.BackendsApiManager
 import org.apache.gluten.exception.{GlutenException, GlutenNotSupportException}
+import org.apache.gluten.extension.columnar.heuristic.RewrittenNodeWall
 
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateFunction, Complete, Partial}
-import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.{ProjectExec, SparkPlan}
 import org.apache.spark.sql.execution.aggregate._
-import org.apache.spark.sql.execution.joins.{BaseJoinExec, BroadcastHashJoinExec, BroadcastNestedLoopJoinExec, CartesianProductExec, ShuffledHashJoinExec, SortMergeJoinExec}
+import org.apache.spark.sql.execution.joins._
 import org.apache.spark.sql.execution.window.WindowExec
-import org.apache.spark.sql.types.{ByteType, DateType, IntegerType, LongType, ShortType}
+import org.apache.spark.sql.types._
 
 import java.sql.Date
 import java.util.concurrent.atomic.AtomicInteger
@@ -332,4 +333,25 @@ trait PullOutProjectHelper {
     newWe.copyTagsFrom(we)
     newWe
   }
+
+  protected def pullOutPreProjectForJoin(joinChild: SparkPlan, joinKeys: Seq[Expression])
+      : (SparkPlan, Seq[Expression], mutable.HashMap[Expression, NamedExpression]) = {
+    val expressionMap = new mutable.HashMap[Expression, NamedExpression]()
+    if (joinKeys.exists(isNotAttribute)) {
+      val newJoinKeys =
+        joinKeys.toIndexedSeq.map(replaceExpressionWithAttribute(_, expressionMap))
+      val preProject = ProjectExec(
+        eliminateProjectList(joinChild.outputSet, expressionMap.values.toSeq),
+        joinChild)
+      joinChild match {
+        case r: RewrittenNodeWall =>
+          r.originalChild.logicalLink.foreach(preProject.setLogicalLink)
+        case _ =>
+          joinChild.logicalLink.foreach(preProject.setLogicalLink)
+      }
+      (preProject, newJoinKeys, expressionMap)
+    } else {
+      (joinChild, joinKeys, expressionMap)
+    }
+  }
 }