[SPARK-12981][SQL] Fix Python UDF extraction for aggregation.

[xguo27]

When Aggregate operator being applied ExtractPythonUDFs rule, it becomes a Project. This change fixes that and maintain Aggregate operator to the original type.

[AmplabJenkins]

Can one of the admins verify this patch?

[xguo27]

@rxin Does this fix look good to you?

[rxin]

cc @davies

[rxin]

a style nit: put else on the same line as the previous }

also can you add some comment explaining what's happening

[xguo27]

Using these two functionally equavalent code snippets:

Scala

val data = Seq((1, "1"), (2, "2"), (3, "2"), (1, "3")).toDF("a","b")
val my_filter = sqlContext.udf.register("my_filter", (a:Int) => a==1)
data.select(col("a")).distinct().filter(my_filter(col("a")))

Python

data = sqlContext.createDataFrame([(1, "1"), (2, "2"), (3, "2"), (1, "3")], ["a", "b"])
my_filter = udf(lambda a: a == 1, BooleanType())
data.select(col("a")).distinct().filter(my_filter(col("a")))

The logical plan comes out execute(aggregateCondition) in here is as below:

spark/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala

Line 801 in 916fc34

val resolvedOperator = execute(aggregatedCondition)

Scala

Aggregate [a#8], [UDF(a#8) AS havingCondition#11]
+- Project [a#8]
   +- Project [_1#6 AS a#8,_2#7 AS b#9]
      +- LocalRelation [_1#6,_2#7], [[1,1],[2,2],[3,2],[1,3]]

Python

Project [havingCondition#2]
+- Aggregate [a#0L], [pythonUDF#3 AS havingCondition#2]
   +- EvaluatePython PythonUDF#<lambda>(a#0L), pythonUDF#3: boolean
      +- Project [a#0L]
         +- LogicalRDD [a#0L,b#1], MapPartitionsRDD[4] at applySchemaToPythonRDD at NativeMethodAccessorImpl.java:-2

We can see in Python's case, we inject an extra Project when execute(aggregateCondition)going through ExtractPythonUDFs, but ResolveAggregateFunctions expects an Aggregate here:

spark/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala

Lines 801 to 805 in 916fc34

	val resolvedOperator = execute(aggregatedCondition)
	def resolvedAggregateFilter =
	resolvedOperator
	.asInstanceOf[Aggregate]
	.aggregateExpressions.head

With this fix, the logical plan generated for Python UDFs does not construct a Project if it is an Aggregate, making it consistent with its Scala counterpart, which gives correct results for ResolveAggregateFunctions to consume:

After fix, Python:

Aggregate [a#0L], [pythonUDF#3 AS havingCondition#2]
+- EvaluatePython PythonUDF#<lambda>(a#0L), pythonUDF#3: boolean
   +- Project [a#0L]
      +- LogicalRDD [a#0L,b#1], MapPartitionsRDD[4] at applySchemaToPythonRDD at NativeMethodAccessorImpl.java:-2

[davies]

@xguo27 Thanks for working on this. I think the root cause here is that we extract Python UDFs too early (in analyzer), EvaluatePython is an special logical plan, many rules have no knowledge of it, which will break many things. We should extract Python UDFs later, in end of optimizer, or physical plan, I will send an PR to fix that.

[xguo27]

Sure @davies . I will close this PR.