[SPARK-24624][SQL][PYTHON] Support mixture of Python UDF and Scalar Pandas UDF #21650
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -21,6 +21,7 @@ import scala.collection.mutable | |
| import scala.collection.mutable.ArrayBuffer | ||
|
|
||
| import org.apache.spark.api.python.PythonEvalType | ||
| import org.apache.spark.sql.AnalysisException | ||
| import org.apache.spark.sql.catalyst.expressions._ | ||
| import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression | ||
| import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan, Project} | ||
|
|
@@ -94,36 +95,60 @@ object ExtractPythonUDFFromAggregate extends Rule[LogicalPlan] { | |
| */ | ||
| object ExtractPythonUDFs extends Rule[SparkPlan] with PredicateHelper { | ||
|
|
||
| private case class EvalTypeHolder(private var evalType: Int = -1) { | ||
|
||
| def isSet: Boolean = evalType >= 0 | ||
|
|
||
| def set(evalType: Int): Unit = { | ||
| if (isSet && evalType != this.evalType) { | ||
| throw new IllegalStateException("Cannot reset eval type to a different value") | ||
| } else { | ||
| this.evalType = evalType | ||
| } | ||
| } | ||
|
|
||
| def get(): Int = { | ||
| if (!isSet) { | ||
| throw new IllegalStateException("Eval type is not set") | ||
| } else { | ||
| evalType | ||
| } | ||
| } | ||
| } | ||
|
|
||
| private def hasScalarPythonUDF(e: Expression): Boolean = { | ||
| e.find(PythonUDF.isScalarPythonUDF).isDefined | ||
| } | ||
|
|
||
| private def canEvaluateInPython(e: PythonUDF): Boolean = { | ||
| e.children match { | ||
| // single PythonUDF child could be chained and evaluated in Python | ||
| case Seq(u: PythonUDF) => e.evalType == u.evalType && canEvaluateInPython(u) | ||
| // Python UDF can't be evaluated directly in JVM | ||
| case children => !children.exists(hasScalarPythonUDF) | ||
| } | ||
| } | ||
|
|
||
| private def collectEvaluableUDFs( | ||
| expr: Expression, | ||
| firstEvalType: EvalTypeHolder): Seq[PythonUDF] = expr match { | ||
| case udf: PythonUDF if PythonUDF.isScalarPythonUDF(udf) | ||
| && (!firstEvalType.isSet || firstEvalType.get == udf.evalType) | ||
| && canEvaluateInPython(udf) => | ||
| firstEvalType.set(udf.evalType) | ||
| Seq(udf) | ||
| case e => e.children.flatMap(collectEvaluableUDFs(_, firstEvalType)) | ||
| } | ||
|
|
||
| def apply(plan: SparkPlan): SparkPlan = plan transformUp { | ||
|
||
| case plan: SparkPlan => extract(plan) | ||
| } | ||
|
|
||
| /** | ||
| * Extract all the PythonUDFs from the current operator and evaluate them before the operator. | ||
| */ | ||
| private def extract(plan: SparkPlan): SparkPlan = { | ||
| val firstEvalType = new EvalTypeHolder | ||
| val udfs = plan.expressions.flatMap(collectEvaluableUDFs(_, firstEvalType)) | ||
| // ignore the PythonUDF that come from second/third aggregate, which is not used | ||
| .filter(udf => udf.references.subsetOf(plan.inputSet)) | ||
| if (udfs.isEmpty) { | ||
|
|
@@ -167,7 +192,8 @@ object ExtractPythonUDFs extends Rule[SparkPlan] with PredicateHelper { | |
| case (vectorizedUdfs, plainUdfs) if vectorizedUdfs.isEmpty => | ||
| BatchEvalPythonExec(plainUdfs, child.output ++ resultAttrs, child) | ||
| case _ => | ||
| throw new AnalysisException( | ||
|
||
| "Expected either Scalar Pandas UDFs or Batched UDFs but got both") | ||
| } | ||
|
|
||
| attributeMap ++= validUdfs.zip(resultAttrs) | ||
|
|
@@ -205,7 +231,7 @@ object ExtractPythonUDFs extends Rule[SparkPlan] with PredicateHelper { | |
| case filter: FilterExec => | ||
| val (candidates, nonDeterministic) = | ||
| splitConjunctivePredicates(filter.condition).partition(_.deterministic) | ||
| val (pushDown, rest) = candidates.partition(!hasScalarPythonUDF(_)) | ||
| if (pushDown.nonEmpty) { | ||
| val newChild = FilterExec(pushDown.reduceLeft(And), filter.child) | ||
| FilterExec((rest ++ nonDeterministic).reduceLeft(And), newChild) | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This looks testing udf + udf
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yeah, the way the test is written is that I am trying to test many combinations so some combinations might not be mixed UDF. Do you prefer that I remove these cases?