[SPARK-18667][PySpark][SQL] Change the way to group row in BatchEvalPythonExec so input_file_name function can work with UDF in pyspark

[viirya]

What changes were proposed in this pull request?

input_file_name doesn't return filename when working with UDF in PySpark. An example shows the problem:

from pyspark.sql.functions import *
from pyspark.sql.types import *

def filename(path):
    return path

sourceFile = udf(filename, StringType())
spark.read.json("tmp.json").select(sourceFile(input_file_name())).show()

+---------------------------+
|filename(input_file_name())|
+---------------------------+
|                           |
+---------------------------+    

The cause of this issue is, we group rows in BatchEvalPythonExec for batching processing of PythonUDF. Currently we group rows first and then evaluate expressions on the rows. If the data is less than the required number of rows for a group, the iterator will be consumed to the end before the evaluation. However, once the iterator reaches the end, we will unset input filename. So the input_file_name expression can't return correct filename.

This patch fixes the approach to group the batch of rows. We evaluate the expression first and then group evaluated results to batch.

How was this patch tested?

Added unit test to PySpark.

Please review http://spark.apache.org/contributing.html before opening a pull request.

[viirya]

cc @davies

[SparkQA]

Test build #69549 has started for PR 16115 at commit 7cd606b.

[viirya]

retest this please.

[SparkQA]

Test build #69553 has finished for PR 16115 at commit 7cd606b.

• This patch fails PySpark unit tests.
• This patch merges cleanly.
• This patch adds no public classes.

[SparkQA]

Test build #69563 has finished for PR 16115 at commit 8b0ef1e.

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

[viirya]

cc @cloud-fan @hvanhovell too

[viirya]

ping @davies again

[cloud-fan]

good catch! Merging to master/2.1, thanks!