[SPARK-21441][SQL] Incorrect Codegen in SortMergeJoinExec results failures in some cases
## What changes were proposed in this pull request? https://issues.apache.org/jira/projects/SPARK/issues/SPARK-21441 This issue can be reproduced by the following example: ``` val spark = SparkSession .builder() .appName("smj-codegen") .master("local") .config("spark.sql.autoBroadcastJoinThreshold", "1") .getOrCreate() val df1 = spark.createDataFrame(Seq((1, 1), (2, 2), (3, 3))).toDF("key", "int") val df2 = spark.createDataFrame(Seq((1, "1"), (2, "2"), (3, "3"))).toDF("key", "str") val df = df1.join(df2, df1("key") === df2("key")) .filter("int = 2 or reflect('java.lang.Integer', 'valueOf', str) = 1") .select("int") df.show() ``` To conclude, the issue happens when: (1) SortMergeJoin condition contains CodegenFallback expressions. (2) In PhysicalPlan tree, SortMergeJoin node is the child of root node, e.g., the Project in above example. This patch fixes the logic in `CollapseCodegenStages` rule. ## How was this patch tested? Unit test and manual verification in our cluster. Author: donnyzone <wellfengzhu@gmail.com> Closes #18656 from DonnyZone/Fix_SortMergeJoinExec.
This commit is contained in:
parent
4eb081cc87
commit
6b6dd682e8
|
@ -489,13 +489,13 @@ case class CollapseCodegenStages(conf: SQLConf) extends Rule[SparkPlan] {
|
|||
* Inserts an InputAdapter on top of those that do not support codegen.
|
||||
*/
|
||||
private def insertInputAdapter(plan: SparkPlan): SparkPlan = plan match {
|
||||
case j @ SortMergeJoinExec(_, _, _, _, left, right) if j.supportCodegen =>
|
||||
// The children of SortMergeJoin should do codegen separately.
|
||||
j.copy(left = InputAdapter(insertWholeStageCodegen(left)),
|
||||
right = InputAdapter(insertWholeStageCodegen(right)))
|
||||
case p if !supportCodegen(p) =>
|
||||
// collapse them recursively
|
||||
InputAdapter(insertWholeStageCodegen(p))
|
||||
case j @ SortMergeJoinExec(_, _, _, _, left, right) =>
|
||||
// The children of SortMergeJoin should do codegen separately.
|
||||
j.copy(left = InputAdapter(insertWholeStageCodegen(left)),
|
||||
right = InputAdapter(insertWholeStageCodegen(right)))
|
||||
case p =>
|
||||
p.withNewChildren(p.children.map(insertInputAdapter))
|
||||
}
|
||||
|
|
|
@ -22,8 +22,10 @@ import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute
|
|||
import org.apache.spark.sql.catalyst.expressions.{Add, Literal, Stack}
|
||||
import org.apache.spark.sql.execution.aggregate.HashAggregateExec
|
||||
import org.apache.spark.sql.execution.joins.BroadcastHashJoinExec
|
||||
import org.apache.spark.sql.execution.joins.SortMergeJoinExec
|
||||
import org.apache.spark.sql.expressions.scalalang.typed
|
||||
import org.apache.spark.sql.functions.{avg, broadcast, col, max}
|
||||
import org.apache.spark.sql.internal.SQLConf
|
||||
import org.apache.spark.sql.test.SharedSQLContext
|
||||
import org.apache.spark.sql.types.{IntegerType, StringType, StructType}
|
||||
|
||||
|
@ -127,4 +129,24 @@ class WholeStageCodegenSuite extends SparkPlanTest with SharedSQLContext {
|
|||
"named_struct('a',id+2, 'b',id+2) as col2")
|
||||
.filter("col1 = col2").count()
|
||||
}
|
||||
|
||||
test("SPARK-21441 SortMergeJoin codegen with CodegenFallback expressions should be disabled") {
|
||||
withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1") {
|
||||
import testImplicits._
|
||||
|
||||
val df1 = Seq((1, 1), (2, 2), (3, 3)).toDF("key", "int")
|
||||
val df2 = Seq((1, "1"), (2, "2"), (3, "3")).toDF("key", "str")
|
||||
|
||||
val df = df1.join(df2, df1("key") === df2("key"))
|
||||
.filter("int = 2 or reflect('java.lang.Integer', 'valueOf', str) = 1")
|
||||
.select("int")
|
||||
|
||||
val plan = df.queryExecution.executedPlan
|
||||
assert(!plan.find(p =>
|
||||
p.isInstanceOf[WholeStageCodegenExec] &&
|
||||
p.asInstanceOf[WholeStageCodegenExec].child.children(0)
|
||||
.isInstanceOf[SortMergeJoinExec]).isDefined)
|
||||
assert(df.collect() === Array(Row(1), Row(2)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
Loading…
Reference in a new issue