[SPARK-21441][SQL]Incorrect Codegen in SortMergeJoinExec results failures in some cases

[DonnyZone]

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.

[DonnyZone]

Hi, @cloud-fan, @vanzin , could you help to take a look?

[viirya]

Will CodegenFallback be used in wholestage codegen? I think it's not supported.

[DonnyZone]

Yeah, CodegenFallback just provide a fallback mode.
However, in such case, SortMergeJoinExec passes incomplete row as input to hiveUDF that implements CodegenFallback.

[viirya]

No. I meant if there's a CodegenFallback expression, wholestage codegen should not be enabled.

[DonnyZone]

That's interesting, I will take a look at why the codegen is enabled

[DonnyZone]

I notice that the CollapseCodegenStages rule will still enable codegen for SortMergeJoinExec without checking CodegenFallback expressions.

The logic in insertInputAdapter seems to skip validating SortMergeJoinExec.

Actually, I'am not familiar with this part, please correct me if I get something wrong

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)))
......

[viirya]

I think the check for SortMergeJoinExec in insertInputAdapter should be corrected to:

private def insertInputAdapter(plan: SparkPlan): SparkPlan = plan match {
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))
}

Can you try it? Thanks.

[DonnyZone]

Great! I'm also considering to disable codegen for SortMergeJoinExec with CodegenFallback expressions.
Thanks for your advise. I will work on it and validate in our environment.

Moreover, I just wonder whether the current pattern oder in insertInputAdapter is specifically designed to generate code for SortMergeJoinExec in all cases.

Could you give any ideas? @davies

[DonnyZone]

I have validated both cases with and without CodegenFallback expressions for SortMergeJoinExec.
The fix works well.

[viirya]

The previous pattern case already validates j.supportCodegen, we don't need to verify it again.

[DonnyZone]

SortMergeJoinExec.supportCodegen checks whether joinType.isInstanceOf[InnerLike]

[DonnyZone]

Therefore, I think we should still verify it.

[viirya]

supportCodegen will call CodegenSupport.supportCodegen, so SortMergeJoinExec.supportCodegen is called then.

[viirya]

Btw, can you also add a test for this? Thanks.

[viirya]

And please also add SQL tag to the PR title, e.g., [SPARK-21441][SQL]. Thanks.

[DonnyZone]

Thanks for reviewing, I will add a test later.

[viirya]

@cloud-fan Can you help trigger the jenkins test for this? Thanks.

[cloud-fan]

ok to test

[cloud-fan]

LGTM, can you update the PR description?

[SparkQA]

Test build #79738 has finished for PR 18656 at commit 1161ffd.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds no public classes.

[viirya]

LGTM for the code change. But I think we're better to have a test for this.

[SparkQA]

Test build #79749 has finished for PR 18656 at commit e2d8cef.

• This patch fails Scala style tests.
• This patch merges cleanly.
• This patch adds no public classes.

[SparkQA]

Test build #79751 has finished for PR 18656 at commit b53bf1c.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds no public classes.

[cloud-fan]

thanks, merging to master/2.2/2.1!