diff --git a/backends-velox/src/main/scala/org/apache/gluten/execution/GenerateExecTransformer.scala b/backends-velox/src/main/scala/org/apache/gluten/execution/GenerateExecTransformer.scala index 267ff3897e..9df15769bf 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/execution/GenerateExecTransformer.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/execution/GenerateExecTransformer.scala @@ -186,11 +186,12 @@ object PullOutGenerateProjectHelper extends PullOutProjectHelper { // NOTE: DO NOT use eliminateProjectList to create the project list because // newGeneratorChild can be a duplicated Attribute in generate.child.output. The native // side identifies the last field of projection as generator's input. - val newGeneratorChildren = Seq(expressionMap.values.head) + val aliasExpr = expressionMap.values.head generate.copy( - generator = - generate.generator.withNewChildren(newGeneratorChildren).asInstanceOf[Generator], - child = ProjectExec(generate.child.output ++ newGeneratorChildren, generate.child) + generator = generate.generator + .withNewChildren(Seq(aliasExpr.toAttribute)) + .asInstanceOf[Generator], + child = ProjectExec(generate.child.output ++ Seq(aliasExpr), generate.child) ) } else { // generator.child is Attribute, no need to introduce a Project. diff --git a/backends-velox/src/test/scala/org/apache/gluten/execution/MiscOperatorSuite.scala b/backends-velox/src/test/scala/org/apache/gluten/execution/MiscOperatorSuite.scala index 51cb918aa6..f92f776226 100644 --- a/backends-velox/src/test/scala/org/apache/gluten/execution/MiscOperatorSuite.scala +++ b/backends-velox/src/test/scala/org/apache/gluten/execution/MiscOperatorSuite.scala @@ -24,6 +24,7 @@ import org.apache.spark.shuffle.GlutenShuffleUtils import org.apache.spark.sql.{DataFrame, Row} import org.apache.spark.sql.execution._ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, AQEShuffleReadExec, ShuffleQueryStageExec} +import org.apache.spark.sql.execution.exchange.ReusedExchangeExec import org.apache.spark.sql.execution.joins.BaseJoinExec import org.apache.spark.sql.execution.window.WindowExec import org.apache.spark.sql.functions._ @@ -1146,14 +1147,14 @@ class MiscOperatorSuite extends VeloxWholeStageTransformerSuite with AdaptiveSpa } } - // Fallback for array(struct(...), null) literal. + // array(struct(...), null) literal is not a sub-expression of generator. runQueryAndCompare(s""" |SELECT $func(array( | named_struct('c1', 0, 'c2', 1), | named_struct('c1', 2, 'c2', null), | null)); |""".stripMargin) { - checkSparkPlan[GenerateExec] + checkGlutenPlan[GenerateExecTransformer] } } } @@ -2278,4 +2279,41 @@ class MiscOperatorSuite extends VeloxWholeStageTransformerSuite with AdaptiveSpa assert(GlutenShuffleUtils.getCompressionCodec(conf) === codec) } } + + test("exchange reuse with CTE lateral view explode") { + // This test verifies the fix for Generate pre-project canonicalization. + // When explode(split(col)) is used, a pre-project is inserted to compute the + // split expression as an Alias before feeding it to the generator. Before the + // fix, the generator use an alias expression as its sub-expression, but Spark's + // doCanonicalize does not cope with it. + // This caused structurally identical plans from different query + // parts to have different canonicalized forms, preventing ReuseExchangeExec + // from reusing exchanges. + // With AQE on and shuffle partitions > 1, the self-join on the CTE produces + // two identical shuffle exchanges that should be reused. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + runQueryAndCompare(""" + |WITH exploded AS ( + | SELECT l_orderkey, word + | FROM lineitem + | LATERAL VIEW explode(split(l_comment, ' ')) t AS word + | WHERE l_orderkey < 10 + |) + |SELECT count(*), sum(t1.l_orderkey + t2.l_orderkey) + |FROM exploded t1 + |JOIN exploded t2 ON t1.word = t2.word + |""".stripMargin) { + df => + val plan = df.queryExecution.executedPlan + assert( + collect(plan) { case p: GenerateExecTransformer => p }.nonEmpty, + "Expected GenerateExecTransformer in the plan") + assert( + collect(plan) { case re: ReusedExchangeExec => re }.nonEmpty, + "Expected ReusedExchangeExec in the plan") + } + } + } }