[SPARK-24624][SQL][PYTHON] Support mixture of Python UDF and Scalar Pandas UDF #21650
There are no files selected for viewing
| 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,94 @@ object ExtractPythonUDFFromAggregate extends Rule[LogicalPlan] { | |
| */ | ||
| object ExtractPythonUDFs extends Rule[SparkPlan] with PredicateHelper { | ||
|
|
||
| private case class LazyEvalType(var evalType: Int = -1) { | ||
|
Member
There was a problem hiding this comment. hmmmmm looks messier then I thought .. previous one looks a bit better to me .. wdyt @BryanCutler ?
Member
There was a problem hiding this comment. I'm not too fond of the name
Contributor
Author
There was a problem hiding this comment. Yeah the idea of the LazyEvalType is a container object that can be set once. Maybe the name LazyEvalType is confusing. I don't think CurrentEvalType is accurate either because the original idea is that we don't change the value once it's set. Maybe call it |
||
|
|
||
| def isSet: Boolean = evalType >= 0 | ||
|
|
||
| def set(evalType: Int): Unit = { | ||
| if (isSet) { | ||
| throw new IllegalStateException("Eval type has already been set") | ||
| } 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 | ||
| } | ||
|
|
||
| /** | ||
| * Check whether a PythonUDF expression can be evaluated in Python. | ||
| * | ||
| * If the lazy eval type is not set, this method checks for either Batched Python UDF and Scalar | ||
| * Pandas UDF. If the lazy eval type is set, this method checks for the expression of the | ||
| * specified eval type. | ||
| * | ||
| * This method will also set the lazy eval type to be the type of the first evaluable expression, | ||
| * i.e., if lazy eval type is not set and we find a evaluable Python UDF expression, lazy eval | ||
| * type will be set to the eval type of the expression. | ||
| * | ||
| */ | ||
| private def canEvaluateInPython(e: PythonUDF, lazyEvalType: LazyEvalType): Boolean = { | ||
|
Contributor
Author
There was a problem hiding this comment. @BryanCutler I rewrite this function using mutable state based on your suggestion. It's not quite the same as your code so please take a look and let me know if this looks better now. Thanks!
Member
There was a problem hiding this comment. The one method seems overly complicated, so I prefer the code from my suggestion.
Contributor
Author
There was a problem hiding this comment. In your code: I think what's confusing part here is that the value of
Member
There was a problem hiding this comment. I'm not sure I follow how this could get wrong results.
Contributor
Author
There was a problem hiding this comment. Bryan, I tried to apply your implementation and the simple test fails: Do you mind trying this too? Hopefully I didn't do something silly here..
Member
There was a problem hiding this comment. Is the above test part of sql/tests.py?
Contributor
Author
There was a problem hiding this comment. Yes it's in the most recent commit.
Member
There was a problem hiding this comment. Ok, I think I see the problem. Since there was a map over I updated my above code to this, does that look correct now?
Contributor
Author
There was a problem hiding this comment. I applied you new code but the test I mentioned above still fails. I think the issue could be when visiting But similar to #21650 (comment), I think the state machine of the firstEvalType here is fairly complicated (i.e., what is the expected state of the eval type holder before and after The lazyEvalType implementation is better IMHO because the state machine is simpler - lazyEvalType is empty until we find the first evaluable UDF and the value doesn't change once it's set. The first implementation (two pass, immutable state) is probably the simplest in terms of the mental complexity of the algo but is less efficient. I think I am ok with both immutable state or the lazy state. I think @HyukjinKwon prefers the immutable state one. @BryanCutler WDYT? |
||
| if (!lazyEvalType.isSet) { | ||
| e.children match { | ||
| // single PythonUDF child could be chained and evaluated in Python if eval type is the same | ||
| case Seq(u: PythonUDF) => | ||
| // Need to recheck the eval type because lazy eval type will be set if child Python UDF is | ||
| // evaluable | ||
| canEvaluateInPython(u, lazyEvalType) && lazyEvalType.get == e.evalType | ||
| // Python UDF can't be evaluated directly in JVM | ||
| case children => if (!children.exists(hasScalarPythonUDF)) { | ||
| // We found the first evaluable expression, set lazy eval type to its eval type. | ||
| lazyEvalType.set(e.evalType) | ||
| true | ||
| } else { | ||
| false | ||
| } | ||
| } | ||
| } else { | ||
| if (e.evalType != lazyEvalType.get) { | ||
| false | ||
| } else { | ||
| e.children match { | ||
| case Seq(u: PythonUDF) => canEvaluateInPython(u, lazyEvalType) | ||
|
Member
There was a problem hiding this comment. There are 2 paths for recursion here, which is probably not a good idea. This method is much more complicated now and a little difficult to follow. |
||
| case children => !children.exists(hasScalarPythonUDF) | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
| private def collectEvaluableUDFs( | ||
| expr: Expression, | ||
| evalType: LazyEvalType | ||
| ): Seq[PythonUDF] = { | ||
| expr match { | ||
| case udf: PythonUDF if | ||
| PythonUDF.isScalarPythonUDF(udf) && canEvaluateInPython(udf, evalType) => | ||
| Seq(udf) | ||
|
Contributor
Author
There was a problem hiding this comment. @HyukjinKwon In your code this line is |
||
| case e => e.children.flatMap(collectEvaluableUDFs(_, evalType)) | ||
| } | ||
| } | ||
|
|
||
| def apply(plan: SparkPlan): SparkPlan = plan transformUp { | ||
|
Contributor
Author
There was a problem hiding this comment. This is no longer needed because this rule will only extract Python UDF and Scalar Pandas UDF and ignore other types of UDFs |
||
| 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 lazyEvalType = new LazyEvalType | ||
| val udfs = plan.expressions.flatMap(collectEvaluableUDFs(_, lazyEvalType)) | ||
| // 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 +226,8 @@ object ExtractPythonUDFs extends Rule[SparkPlan] with PredicateHelper { | |
| case (vectorizedUdfs, plainUdfs) if vectorizedUdfs.isEmpty => | ||
| BatchEvalPythonExec(plainUdfs, child.output ++ resultAttrs, child) | ||
| case _ => | ||
| throw new AnalysisException( | ||
|
Member
There was a problem hiding this comment. Why change the exception type? Can you make a test that causes this?
Contributor
Author
There was a problem hiding this comment. This is because we shouldn't reach here. (Otherwise it's bug). Don't know what's the best exception type here though. |
||
| "Expected either Scalar Pandas UDFs or Batched UDFs but got both") | ||
| } | ||
|
|
||
| attributeMap ++= validUdfs.zip(resultAttrs) | ||
|
|
@@ -205,7 +265,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) | ||
|
|
||
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?