Skip to content

[SPARK-48666][SQL] Do not push down filter if it contains PythonUDFs #47313

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
wants to merge 7 commits into from
Closed

[SPARK-48666][SQL] Do not push down filter if it contains PythonUDFs #47313

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
5c327bf
Do not push down filter if it contains Unevaluable expression
weiatwork Jul 8, 2024
7f496d2
Revert "Do not push down filter if it contains Unevaluable expression"
HyukjinKwon Jul 12, 2024
9677b59
Do not push down filter if it contains Unevaluable expression
HyukjinKwon Jul 12, 2024
2a8ec4d
python specific
HyukjinKwon Jul 12, 2024
27bbf3d
Add an assert
HyukjinKwon Jul 12, 2024
c1251c4
fixup
HyukjinKwon Jul 12, 2024
bc89b74
Address comments
HyukjinKwon Jul 12, 2024
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
7 changes: 6 additions & 1 deletion ...src/main/scala/org/apache/spark/sql/execution/datasources/PruneFileSourcePartitions.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,12 @@ private[sql] object PruneFileSourcePartitions extends Rule[LogicalPlan] {
_))
if filters.nonEmpty && fsRelation.partitionSchema.nonEmpty =>
val normalizedFilters = DataSourceStrategy.normalizeExprs(
filters.filter(f => f.deterministic && !SubqueryExpression.hasSubquery(f)),
filters.filter { f =>
f.deterministic &&
!SubqueryExpression.hasSubquery(f) &&
// Python UDFs might exist because this rule is applied before ``ExtractPythonUDFs``.
!f.exists(_.isInstanceOf[PythonUDF])
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

let's add code comments to explain that PythonUDF rewriting happens later so we have to exclude it here.

},
logicalRelation.output)
val (partitionKeyFilters, _) = DataSourceUtils
.getPartitionFiltersAndDataFilters(partitionSchema, normalizedFilters)
Expand Down
7 changes: 5 additions & 2 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/FileScanBuilder.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ package org.apache.spark.sql.execution.datasources.v2
import scala.collection.mutable

import org.apache.spark.sql.{sources, SparkSession}
import org.apache.spark.sql.catalyst.expressions.Expression
import org.apache.spark.sql.catalyst.expressions.{Expression, PythonUDF, SubqueryExpression}
import org.apache.spark.sql.connector.expressions.filter.Predicate
import org.apache.spark.sql.connector.read.{ScanBuilder, SupportsPushDownRequiredColumns}
import org.apache.spark.sql.execution.datasources.{DataSourceStrategy, DataSourceUtils, PartitioningAwareFileIndex, PartitioningUtils}
Expand Down Expand Up @@ -73,7 +73,10 @@ abstract class FileScanBuilder(
val (deterministicFilters, nonDeterminsticFilters) = filters.partition(_.deterministic)
val (partitionFilters, dataFilters) =
DataSourceUtils.getPartitionFiltersAndDataFilters(partitionSchema, deterministicFilters)
this.partitionFilters = partitionFilters
this.partitionFilters = partitionFilters.filter { f =>
// Python UDFs might exist because this rule is applied before ``ExtractPythonUDFs``.
!SubqueryExpression.hasSubquery(f) && !f.exists(_.isInstanceOf[PythonUDF])
}
this.dataFilters = dataFilters
val translatedFilters = mutable.ArrayBuffer.empty[sources.Filter]
for (filterExpr <- dataFilters) {
Expand Down
16 changes: 14 additions & 2 deletions sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonUDFSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,8 +17,8 @@

package org.apache.spark.sql.execution.python

import org.apache.spark.sql.{AnalysisException, IntegratedUDFTestUtils, QueryTest}
import org.apache.spark.sql.functions.{array, count, transform}
import org.apache.spark.sql.{AnalysisException, IntegratedUDFTestUtils, QueryTest, Row}
import org.apache.spark.sql.functions.{array, col, count, transform}
import org.apache.spark.sql.test.SharedSparkSession
import org.apache.spark.sql.types.LongType

Expand Down Expand Up @@ -124,4 +124,16 @@ class PythonUDFSuite extends QueryTest with SharedSparkSession {
context = ExpectedContext(
"transform", s".*${this.getClass.getSimpleName}.*"))
}

test("SPARK-48666: Python UDF execution against partitioned column") {
assume(shouldTestPythonUDFs)
withTable("t") {
spark.range(1).selectExpr("id AS t", "(id + 1) AS p").write.partitionBy("p").saveAsTable("t")
val table = spark.table("t")
val newTable = table.withColumn("new_column", pythonTestUDF(table("p")))
val df = newTable.as("t1").join(
newTable.as("t2"), col("t1.new_column") === col("t2.new_column"))
checkAnswer(df, Row(0, 1, 1, 0, 1, 1))
}
}
}
9 changes: 7 additions & 2 deletions sql/hive/src/main/scala/org/apache/spark/sql/hive/execution/PruneHiveTablePartitions.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ import org.apache.hadoop.hive.common.StatsSetupConst
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.analysis.CastSupport
import org.apache.spark.sql.catalyst.catalog._
import org.apache.spark.sql.catalyst.expressions.{And, AttributeSet, Expression, ExpressionSet, PredicateHelper, SubqueryExpression}
import org.apache.spark.sql.catalyst.expressions.{And, AttributeSet, Expression, ExpressionSet, PredicateHelper, PythonUDF, SubqueryExpression}
import org.apache.spark.sql.catalyst.planning.PhysicalOperation
import org.apache.spark.sql.catalyst.plans.logical.{Filter, LogicalPlan, Project}
import org.apache.spark.sql.catalyst.plans.logical.statsEstimation.FilterEstimation
Expand Down Expand Up @@ -50,7 +50,12 @@ private[sql] class PruneHiveTablePartitions(session: SparkSession)
filters: Seq[Expression],
relation: HiveTableRelation): ExpressionSet = {
val normalizedFilters = DataSourceStrategy.normalizeExprs(
filters.filter(f => f.deterministic && !SubqueryExpression.hasSubquery(f)), relation.output)
filters.filter { f =>
f.deterministic &&
!SubqueryExpression.hasSubquery(f) &&
// Python UDFs might exist because this rule is applied before ``ExtractPythonUDFs``.
!f.exists(_.isInstanceOf[PythonUDF])
}, relation.output)
val partitionColumnSet = AttributeSet(relation.partitionCols)
ExpressionSet(
normalizedFilters.flatMap(extractPredicatesWithinOutputSet(_, partitionColumnSet)))
Expand Down