[SPARK-28323][SQL][Python] PythonUDF should be able to use in join condition

[viirya]

What changes were proposed in this pull request?

There is a bug in ExtractPythonUDFs that produces wrong result attributes. It causes a failure when using PythonUDFs among multiple child plans, e.g., join. An example is using PythonUDFs in join condition.

>>> left = spark.createDataFrame([Row(a=1, a1=1, a2=1), Row(a=2, a1=2, a2=2)])                                                                                                                                                                                                                      
>>> right = spark.createDataFrame([Row(b=1, b1=1, b2=1), Row(b=1, b1=3, b2=1)])                                                                       
>>> f = udf(lambda a: a, IntegerType())                                                                                                               
>>> df = left.join(right, [f("a") == f("b"), left.a1 == right.b1])                                                
>>> df.collect()                                                                                                                                      
19/07/10 12:20:49 ERROR Executor: Exception in task 5.0 in stage 0.0 (TID 5)                                                                                                                                                 
java.lang.ArrayIndexOutOfBoundsException: 1                                                                                                           
        at org.apache.spark.sql.catalyst.expressions.GenericInternalRow.genericGet(rows.scala:201)                                                    
        at org.apache.spark.sql.catalyst.expressions.BaseGenericInternalRow.getAs(rows.scala:35)               
        at org.apache.spark.sql.catalyst.expressions.BaseGenericInternalRow.isNullAt(rows.scala:36)               
        at org.apache.spark.sql.catalyst.expressions.BaseGenericInternalRow.isNullAt$(rows.scala:36)                             
        at org.apache.spark.sql.catalyst.expressions.GenericInternalRow.isNullAt(rows.scala:195)
        at org.apache.spark.sql.catalyst.expressions.JoinedRow.isNullAt(JoinedRow.scala:70)
        ...

How was this patch tested?

Added test.

[viirya]

cc @HyukjinKwon @BryanCutler

[HyukjinKwon]

Cool @viirya! I will take a closer look within 2 days and get this in

[viirya]

Thanks! @HyukjinKwon

[SparkQA]

Test build #107433 has finished for PR 25091 at commit 95231a6.

• This patch fails due to an unknown error code, -9.
• This patch merges cleanly.
• This patch adds no public classes.

[viirya]

retest this please.

[SparkQA]

Test build #107445 has finished for PR 25091 at commit 95231a6.

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

[SparkQA]

Test build #107452 has finished for PR 25091 at commit 0c24787.

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

[SparkQA]

Test build #107458 has finished for PR 25091 at commit fbf8be9.

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

[BryanCutler]

LGTM

[BryanCutler]

oof, nice catch!

[BryanCutler]

merged to master, thanks @viirya !

[BryanCutler]

jira seems down, so I wasn't able to resolve the issue. will try later

[HyukjinKwon]

Yea, LGTM too!

[viirya]

Thanks! @HyukjinKwon @BryanCutler