[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,59 @@ object ExtractPythonUDFFromAggregate extends Rule[LogicalPlan] { | |
| */ | ||
| object ExtractPythonUDFs extends Rule[SparkPlan] with PredicateHelper { | ||
|
|
||
| private def hasScalarPythonUDF(e: Expression): Boolean = { | ||
| e.find(PythonUDF.isScalarPythonUDF).isDefined | ||
| } | ||
|
|
||
| private def canEvaluateInPython(e: PythonUDF, evalType: Int): Boolean = { | ||
| if (e.evalType != evalType) { | ||
|
||
| false | ||
| } else { | ||
| e.children match { | ||
| // single PythonUDF child could be chained and evaluated in Python | ||
| case Seq(u: PythonUDF) => canEvaluateInPython(u, evalType) | ||
| // Python UDF can't be evaluated directly in JVM | ||
| case children => !children.exists(hasScalarPythonUDF) | ||
| } | ||
| } | ||
| } | ||
|
|
||
| private def collectEvaluableUDF(expr: Expression, evalType: Int): Seq[PythonUDF] = expr match { | ||
|
||
| case udf: PythonUDF if PythonUDF.isScalarPythonUDF(udf) && canEvaluateInPython(udf, evalType) => | ||
| Seq(udf) | ||
| case e => e.children.flatMap(collectEvaluableUDF(_, evalType)) | ||
| } | ||
|
|
||
| /** | ||
| * Collect evaluable UDFs from the current node. | ||
| * | ||
| * This function collects Python UDFs or Scalar Python UDFs from expressions of the input node, | ||
| * and returns a list of UDFs of the same eval type. | ||
|
||
| * | ||
| * If expressions contain both UDFs eval types, this function will only return Python UDFs. | ||
| * | ||
| * The caller should call this function multiple times until all evaluable UDFs are collected. | ||
|
||
| */ | ||
| private def collectEvaluableUDFs(plan: SparkPlan): Seq[PythonUDF] = { | ||
| val pythonUDFs = | ||
| plan.expressions.flatMap(collectEvaluableUDF(_, PythonEvalType.SQL_BATCHED_UDF)) | ||
|
|
||
| if (pythonUDFs.isEmpty) { | ||
| plan.expressions.flatMap(collectEvaluableUDF(_, PythonEvalType.SQL_SCALAR_PANDAS_UDF)) | ||
| } else { | ||
| pythonUDFs | ||
|
||
| } | ||
| } | ||
|
|
||
| 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 udfs = collectEvaluableUDFs(plan) | ||
| // 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 +191,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 +230,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) | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,93 @@ | ||
| /* | ||
| * Licensed to the Apache Software Foundation (ASF) under one or more | ||
| * contributor license agreements. See the NOTICE file distributed with | ||
| * this work for additional information regarding copyright ownership. | ||
| * The ASF licenses this file to You under the Apache License, Version 2.0 | ||
| * (the "License"); you may not use this file except in compliance with | ||
| * the License. You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, software | ||
| * distributed under the License is distributed on an "AS IS" BASIS, | ||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| * See the License for the specific language governing permissions and | ||
| * limitations under the License. | ||
| */ | ||
|
|
||
| package org.apache.spark.sql.execution.python | ||
|
|
||
| import org.apache.spark.sql.execution.{SparkPlan, SparkPlanTest} | ||
| import org.apache.spark.sql.functions.col | ||
| import org.apache.spark.sql.test.SharedSQLContext | ||
|
|
||
| class ExtractPythonUDFsSuite extends SparkPlanTest with SharedSQLContext { | ||
| import testImplicits.newProductEncoder | ||
| import testImplicits.localSeqToDatasetHolder | ||
|
|
||
| val batchedPythonUDF = new MyDummyPythonUDF | ||
| val scalarPandasUDF = new MyDummyScalarPandasUDF | ||
|
|
||
| private def collectBatchExec(plan: SparkPlan): Seq[BatchEvalPythonExec] = plan.collect { | ||
| case b: BatchEvalPythonExec => b | ||
| } | ||
|
|
||
| private def collectArrowExec(plan: SparkPlan): Seq[ArrowEvalPythonExec] = plan.collect { | ||
| case b: ArrowEvalPythonExec => b | ||
| } | ||
|
|
||
| test("Chained Batched Python UDFs should be combined to a single physical node") { | ||
| val df = Seq(("Hello", 4)).toDF("a", "b") | ||
| val df2 = df.withColumn("c", batchedPythonUDF(col("a"))) | ||
| .withColumn("d", batchedPythonUDF(col("c"))) | ||
| val pythonEvalNodes = collectBatchExec(df2.queryExecution.executedPlan) | ||
| assert(pythonEvalNodes.size == 1) | ||
| } | ||
|
|
||
| test("Chained Scalar Pandas UDFs should be combined to a single physical node") { | ||
| val df = Seq(("Hello", 4)).toDF("a", "b") | ||
| val df2 = df.withColumn("c", scalarPandasUDF(col("a"))) | ||
| .withColumn("d", scalarPandasUDF(col("c"))) | ||
| val arrowEvalNodes = collectArrowExec(df2.queryExecution.executedPlan) | ||
| assert(arrowEvalNodes.size == 1) | ||
| } | ||
|
|
||
| test("Mixed Batched Python UDFs and Pandas UDF should be separate physical node") { | ||
| val df = Seq(("Hello", 4)).toDF("a", "b") | ||
| val df2 = df.withColumn("c", batchedPythonUDF(col("a"))) | ||
| .withColumn("d", scalarPandasUDF(col("b"))) | ||
|
|
||
| val pythonEvalNodes = collectBatchExec(df2.queryExecution.executedPlan) | ||
| val arrowEvalNodes = collectArrowExec(df2.queryExecution.executedPlan) | ||
| assert(pythonEvalNodes.size == 1) | ||
| assert(arrowEvalNodes.size == 1) | ||
| } | ||
|
|
||
| test("Independent Batched Python UDFs and Scalar Pandas UDFs should be combined separately") { | ||
| val df = Seq(("Hello", 4)).toDF("a", "b") | ||
| val df2 = df.withColumn("c1", batchedPythonUDF(col("a"))) | ||
| .withColumn("c2", batchedPythonUDF(col("c1"))) | ||
| .withColumn("d1", scalarPandasUDF(col("a"))) | ||
| .withColumn("d2", scalarPandasUDF(col("d1"))) | ||
|
|
||
| val pythonEvalNodes = collectBatchExec(df2.queryExecution.executedPlan) | ||
| val arrowEvalNodes = collectArrowExec(df2.queryExecution.executedPlan) | ||
| assert(pythonEvalNodes.size == 1) | ||
| assert(arrowEvalNodes.size == 1) | ||
| } | ||
|
|
||
| test("Dependent Batched Python UDFs and Scalar Pandas UDFs should not be combined") { | ||
| val df = Seq(("Hello", 4)).toDF("a", "b") | ||
| val df2 = df.withColumn("c1", batchedPythonUDF(col("a"))) | ||
| .withColumn("d1", scalarPandasUDF(col("c1"))) | ||
| .withColumn("c2", batchedPythonUDF(col("d1"))) | ||
| .withColumn("d2", scalarPandasUDF(col("c2"))) | ||
|
|
||
| val pythonEvalNodes = collectBatchExec(df2.queryExecution.executedPlan) | ||
| val arrowEvalNodes = collectArrowExec(df2.queryExecution.executedPlan) | ||
| assert(pythonEvalNodes.size == 2) | ||
| assert(arrowEvalNodes.size == 2) | ||
| } | ||
| } | ||
|
|
||
|
|
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?