Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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._
Expand Down Expand Up @@ -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]
}
}
}
Expand Down Expand Up @@ -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")
}
}
}
}
Loading