Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
xguo27 wants to merge 2 commits into from
Closed

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1065bb2
[SPARK-12981][SQL] Fix Python UDF extraction for aggregation.
xguo27 Jan 26, 2016
e4d2629
Update per Reynold's suggestion
xguo27 Feb 27, 2016
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions python/pyspark/sql/tests.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1283,6 +1283,14 @@ def test_collect_functions(self):
sorted(df.select(functions.collect_list(df.value).alias('r')).collect()[0].r),
["1", "2", "2", "2"])

def test_udf_with_aggregate_function(self):
df = self.sqlCtx.createDataFrame([(1, "1"), (2, "2"), (1, "2"), (1, "2")], ["key", "value"])
from pyspark.sql.functions import udf, col
from pyspark.sql.types import BooleanType

my_filter = udf(lambda a: a == 1, BooleanType())
sel = df.select(col("key")).distinct().filter(my_filter(col("key")))
self.assertEqual(sel.collect(), [Row(key=1)])

if __name__ == "__main__":
from pyspark.sql.tests import *
Expand Down
19 changes: 13 additions & 6 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFs.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
package org.apache.spark.sql.execution.python

import org.apache.spark.sql.catalyst.plans.logical
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan}
import org.apache.spark.sql.catalyst.rules.Rule

/**
Expand Down Expand Up @@ -64,11 +64,18 @@ private[spark] object ExtractPythonUDFs extends Rule[LogicalPlan] {
assert(evaluation != null, "Unable to evaluate PythonUDF. Missing input attributes.")

// Trim away the new UDF value if it was only used for filtering or something.
logical.Project(
plan.output,
plan.transformExpressions {
case p: PythonUDF if p.fastEquals(udf) => evaluation.resultAttribute
}.withNewChildren(newChildren))
val transformed = plan.transformExpressions {
case p: PythonUDF if p.fastEquals(udf) => evaluation.resultAttribute
}.withNewChildren(newChildren)

// If plan is an [[Aggregate]], return evaluated plan as is for
// [[ResolveAggregateFunctions]] rule in a batch for further resolution.
// Otherwise, construct a [[Project]] with evaluated udf.
if (plan.isInstanceOf[Aggregate]) {
transformed
} else {
logical.Project(plan.output, transformed)
}

case None =>
// If there is no Python UDF that is resolved, skip this round.
Expand Down