File tree Expand file tree Collapse file tree
sql/core/src/main/scala/org/apache/spark/sql/execution/python Expand file tree Collapse file tree Original file line number Diff line number Diff line change @@ -26,25 +26,30 @@ import org.apache.spark.sql.catalyst.rules.Rule
2626 * Extracts PythonUDFs from operators, rewriting the query plan so that the UDF can be evaluated
2727 * alone in a batch.
2828 *
29+ * Only extracts the PythonUDFs that could be evaluated in Python (the single child is PythonUDFs
30+ * or all the children could be evaluated in JVM).
31+ *
2932 * This has the limitation that the input to the Python UDF is not allowed include attributes from
3033 * multiple child operators.
3134 */
3235private [spark] object ExtractPythonUDFs extends Rule [LogicalPlan ] {
3336
34- private def hasUDF (e : Expression ): Boolean = {
37+ private def hasPythonUDF (e : Expression ): Boolean = {
3538 e.find(_.isInstanceOf [PythonUDF ]).isDefined
3639 }
3740
38- private def canEvaluate (e : PythonUDF ): Boolean = {
41+ private def canEvaluateInPython (e : PythonUDF ): Boolean = {
3942 e.children match {
40- case Seq (u : PythonUDF ) => canEvaluate(u)
41- case children => ! children.exists(hasUDF)
43+ // single PythonUDF child could be chained and evaluated in Python
44+ case Seq (u : PythonUDF ) => canEvaluateInPython(u)
45+ // Python UDF can't be evaluated directly in JVM
46+ case children => ! children.exists(hasPythonUDF)
4247 }
4348 }
4449
4550 private def collectEvaluatableUDF (expr : Expression ): Seq [PythonUDF ] = {
4651 expr.collect {
47- case udf : PythonUDF if canEvaluate (udf) => udf
52+ case udf : PythonUDF if canEvaluateInPython (udf) => udf
4853 }
4954 }
5055
You can’t perform that action at this time.
0 commit comments