dongjoon-hyun commented on code in PR #59018:
URL: https://github.com/apache/spark/pull/59018#discussion_r4098869004
##########
sql/core/src/test/scala/org/apache/spark/sql/BenchmarkQueryTest.scala:
##########
@@ -98,4 +93,81 @@ abstract class BenchmarkQueryTest extends SharedSparkSession
{
s"and JIT optimization might not work:\n${subtree.treeString}")
}
}
+
Review Comment:
nit. Unnecessary empty line before the closing brace of the class.
##########
sql/core/src/test/scala/org/apache/spark/sql/BenchmarkQueryTest.scala:
##########
@@ -98,4 +93,81 @@ abstract class BenchmarkQueryTest extends SharedSparkSession
{
s"and JIT optimization might not work:\n${subtree.treeString}")
}
}
+
+}
+
+object BenchmarkQueryTest {
+
+ /**
+ * The whole-stage codegen subtrees of `plan` and of its subqueries, in
stage id order, each
+ * with the code it generates for data (see [[withRowBroadcasts]]). The walk
follows an
+ * adaptive plan into its current plan and query stages, as
`debug.codegenStringSeq` does; for
+ * a plan that has not run, that holds no whole-stage codegen subtree yet.
+ */
+ def generatedCode(plan: SparkPlan): Seq[(WholeStageCodegenExec,
CodeAndComment)] = {
+ val subtrees = new mutable.LinkedHashSet[WholeStageCodegenExec]()
+ def findSubtrees(plan: SparkPlan): Unit = plan foreach {
+ case s: WholeStageCodegenExec => subtrees += s
+ case a: AdaptiveSparkPlanExec => findSubtrees(a.executedPlan)
+ case s: QueryStageExec => findSubtrees(s.plan)
+ case s => s.subqueries.foreach(findSubtrees)
+ }
+ findSubtrees(plan)
+ // One replacement per broadcast exchange of the plan, shared by the
subtrees that read it, so
+ // each runs its one-row broadcast once.
+ val replacements = new java.util.IdentityHashMap[BroadcastExchangeExec,
BroadcastExchangeExec]
+ subtrees.toSeq.sortBy(_.codegenStageId).map { s =>
+ s -> withRowBroadcasts(s, replacements).doCodeGen()._2
+ }
+ }
+
+ /**
+ * Returns a copy of a whole-stage codegen subtree in which every broadcast
exchange reads one
+ * row of non-null default values instead of its child, to generate the
subtree's code from.
+ *
+ * The benchmark queries are planned over empty tables. A broadcast join
generates its code
+ * from the value of its broadcast, and for an empty one it generates a stub
in place of the
+ * join, which leaves the rest of the subtree ungenerated too. Over one row
it generates the
+ * code it generates for data. A single row makes every join key unique, so
a hash join
+ * generates its unique-key form, which is what a join on a dimension
table's key generates;
+ * a join whose build side repeats keys also has a loop over the matches,
which is not
+ * measured. Only the copy's code is used: a build side is a subtree of its
own, found and
+ * measured in the original plan.
+ *
+ * Fails if a broadcast join of the copy still reads anything but such a
row, so that no join
+ * can generate its empty-build-side code unnoticed.
+ */
+ def withRowBroadcasts(
+ subtree: WholeStageCodegenExec,
+ replacements: java.util.Map[BroadcastExchangeExec,
BroadcastExchangeExec] =
Review Comment:
nit. This default value looks unused because `generatedCode` is the only
caller and always passes `replacements`. Shall we remove the default and make
this method `private`?
##########
sql/core/src/test/scala/org/apache/spark/sql/BenchmarkQueryTest.scala:
##########
@@ -98,4 +93,81 @@ abstract class BenchmarkQueryTest extends SharedSparkSession
{
s"and JIT optimization might not work:\n${subtree.treeString}")
}
}
+
+}
+
+object BenchmarkQueryTest {
+
+ /**
+ * The whole-stage codegen subtrees of `plan` and of its subqueries, in
stage id order, each
+ * with the code it generates for data (see [[withRowBroadcasts]]). The walk
follows an
+ * adaptive plan into its current plan and query stages, as
`debug.codegenStringSeq` does; for
+ * a plan that has not run, that holds no whole-stage codegen subtree yet.
+ */
+ def generatedCode(plan: SparkPlan): Seq[(WholeStageCodegenExec,
CodeAndComment)] = {
+ val subtrees = new mutable.LinkedHashSet[WholeStageCodegenExec]()
Review Comment:
nit. Since the subtrees are sorted by `codegenStageId` right below, the
insertion order is not used. A plain `HashSet` (as before) would be enough.
Also, this walk duplicates `debug.codegenStringSeq`, so the two may drift
apart in the future.
##########
sql/core/src/test/scala/org/apache/spark/sql/BenchmarkQueryTest.scala:
##########
@@ -98,4 +93,81 @@ abstract class BenchmarkQueryTest extends SharedSparkSession
{
s"and JIT optimization might not work:\n${subtree.treeString}")
}
}
+
+}
+
+object BenchmarkQueryTest {
+
+ /**
+ * The whole-stage codegen subtrees of `plan` and of its subqueries, in
stage id order, each
+ * with the code it generates for data (see [[withRowBroadcasts]]). The walk
follows an
+ * adaptive plan into its current plan and query stages, as
`debug.codegenStringSeq` does; for
+ * a plan that has not run, that holds no whole-stage codegen subtree yet.
+ */
+ def generatedCode(plan: SparkPlan): Seq[(WholeStageCodegenExec,
CodeAndComment)] = {
+ val subtrees = new mutable.LinkedHashSet[WholeStageCodegenExec]()
+ def findSubtrees(plan: SparkPlan): Unit = plan foreach {
+ case s: WholeStageCodegenExec => subtrees += s
+ case a: AdaptiveSparkPlanExec => findSubtrees(a.executedPlan)
+ case s: QueryStageExec => findSubtrees(s.plan)
+ case s => s.subqueries.foreach(findSubtrees)
+ }
+ findSubtrees(plan)
+ // One replacement per broadcast exchange of the plan, shared by the
subtrees that read it, so
+ // each runs its one-row broadcast once.
+ val replacements = new java.util.IdentityHashMap[BroadcastExchangeExec,
BroadcastExchangeExec]
+ subtrees.toSeq.sortBy(_.codegenStageId).map { s =>
+ s -> withRowBroadcasts(s, replacements).doCodeGen()._2
+ }
+ }
+
+ /**
+ * Returns a copy of a whole-stage codegen subtree in which every broadcast
exchange reads one
+ * row of non-null default values instead of its child, to generate the
subtree's code from.
+ *
+ * The benchmark queries are planned over empty tables. A broadcast join
generates its code
+ * from the value of its broadcast, and for an empty one it generates a stub
in place of the
+ * join, which leaves the rest of the subtree ungenerated too. Over one row
it generates the
+ * code it generates for data. A single row makes every join key unique, so
a hash join
+ * generates its unique-key form, which is what a join on a dimension
table's key generates;
+ * a join whose build side repeats keys also has a loop over the matches,
which is not
+ * measured. Only the copy's code is used: a build side is a subtree of its
own, found and
+ * measured in the original plan.
+ *
+ * Fails if a broadcast join of the copy still reads anything but such a
row, so that no join
+ * can generate its empty-build-side code unnoticed.
+ */
+ def withRowBroadcasts(
+ subtree: WholeStageCodegenExec,
+ replacements: java.util.Map[BroadcastExchangeExec,
BroadcastExchangeExec] =
+ new java.util.IdentityHashMap[BroadcastExchangeExec,
BroadcastExchangeExec]
+ ): WholeStageCodegenExec = {
+ def rowBroadcast(b: BroadcastExchangeExec): BroadcastExchangeExec =
+ replacements.computeIfAbsent(b, _ => b.copy(child =
LocalTableScanExec(b.child.output,
+ Seq(InternalRow.fromSeq(b.child.output.map(a =>
Literal.default(a.dataType).value))),
+ None)))
+ // Top down, so an exchange is replaced before its build side would be
copied.
+ val copy = subtree.transformDown {
Review Comment:
nit. `transformDown` also copies stages below `InputAdapter` (e.g. across a
shuffle) and checks their joins. It is harmless because nothing is executed
there, but the assertion message then prints `subtree.treeString`, which is not
the stage that contains the offending join.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
