diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala index ff3acdfca7e2f..764ad40237942 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala @@ -22,7 +22,7 @@ import scala.collection.mutable import org.apache.spark.SparkException import org.apache.spark.internal.{Logging, LogKeys} import org.apache.spark.sql.AnalysisException -import org.apache.spark.sql.catalyst.expressions.{AttributeReference, AttributeSet, DynamicPruning, DynamicPruningExpression, Expression, ExpressionSet, GetStructField, Literal, NamedExpression, PythonUDF, SchemaPruning, SubqueryExpression, V2ExpressionUtils} +import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference, AttributeSet, DynamicPruning, DynamicPruningExpression, Expression, ExpressionSet, GetStructField, Literal, NamedExpression, Or, PythonUDF, SchemaPruning, SubqueryExpression, V2ExpressionUtils} import org.apache.spark.sql.catalyst.plans.logical.SampleMethod import org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning import org.apache.spark.sql.catalyst.types.DataTypeUtils @@ -69,27 +69,50 @@ object PushDownUtils extends Logging { // pushed down. This map can be used to construct a catalyst filter expression from the // input filter, or a superset(partial push down filter) of the input filter. val translatedFilterToExpr = mutable.HashMap.empty[sources.Filter, Expression] - val translatedFilters = mutable.ArrayBuffer.empty[sources.Filter] + val fullyTranslatedFilters = mutable.ArrayBuffer.empty[sources.Filter] + val extractedFilters = mutable.LinkedHashSet.empty[sources.Filter] // Catalyst filter expression that can't be translated to data source filters. val untranslatableExprs = mutable.ArrayBuffer.empty[Expression] + def translateFilter(expression: Expression): Option[sources.Filter] = { + DataSourceStrategy.translateFilterWithMapping(expression, Some(translatedFilterToExpr), + nestedPredicatePushdownEnabled = true) + } + for (filterExpr <- filters) { - val translated = - DataSourceStrategy.translateFilterWithMapping(filterExpr, Some(translatedFilterToExpr), - nestedPredicatePushdownEnabled = true) - if (translated.isEmpty) { - untranslatableExprs += filterExpr - } else { - translatedFilters += translated.get + translateFilter(filterExpr) match { + case Some(filter) => + fullyTranslatedFilters += filter + case None => + untranslatableExprs += filterExpr + if (filterExpr.deterministic) { + extractPushablePredicate( + filterExpr, + e => DataSourceStrategy.translateFilter( + e, supportNestedPredicatePushdown = true)) + .flatMap(translateFilter) + .foreach(extractedFilters += _) + } } } + // Exclude fully translated inputs from the residuals to omit. + val residualsToOmit = extractedFilters + residualsToOmit --= fullyTranslatedFilters + + // Include extracted filters in the pushdown list. + val filtersToPush = fullyTranslatedFilters + filtersToPush ++= residualsToOmit + // Data source filters that need to be evaluated again after scanning. which means // the data source cannot guarantee the rows returned can pass these filters. // As a result we must return it so Spark can plan an extra filter operator. - val postScanFilters = r.pushFilters(translatedFilters.toArray).map { filter => - DataSourceStrategy.rebuildExpressionFromFilter(filter, translatedFilterToExpr) - } + // Omit returned residuals that came only from predicate extraction. + val postScanFilters = r.pushFilters(filtersToPush.toArray) + .filterNot(residualsToOmit.contains) + .map { filter => + DataSourceStrategy.rebuildExpressionFromFilter(filter, translatedFilterToExpr) + } // Normally translated filters (postScanFilters) are simple filters that can be evaluated // faster, while the untranslated filters are complicated filters that take more time to // evaluate, so we want to evaluate the postScanFilters filters first. @@ -102,21 +125,40 @@ object PushDownUtils extends Logging { // and the data source will return the filters that it cannot guarantee to be true // for all returned rows. val translatedFilterToExpr = mutable.HashMap.empty[Predicate, Expression] - val translatedFilters = mutable.ArrayBuffer.empty[Predicate] + val fullyTranslatedFilters = mutable.ArrayBuffer.empty[Predicate] + val extractedFilters = mutable.LinkedHashSet.empty[Predicate] val untranslatableExprs = mutable.ArrayBuffer.empty[Expression] + def translateFilter(expression: Expression): Option[Predicate] = { + DataSourceV2Strategy.translateFilterV2WithMapping( + expression, Some(translatedFilterToExpr)) + } + for (filterExpr <- deterministicFilters) { - val translated = - DataSourceV2Strategy.translateFilterV2WithMapping( - filterExpr, Some(translatedFilterToExpr)) - if (translated.isEmpty) { - untranslatableExprs += filterExpr - } else { - translatedFilters += translated.get + translateFilter(filterExpr) match { + case Some(filter) => + fullyTranslatedFilters += filter + case None => + untranslatableExprs += filterExpr + extractPushablePredicate( + filterExpr, + DataSourceV2Strategy.translateFilterV2) + .flatMap(translateFilter) + .foreach(extractedFilters += _) } } - val postScanPredicates = r.pushPredicates(translatedFilters.toArray) + // Exclude fully translated inputs from the residuals to omit. + val residualsToOmit = extractedFilters + residualsToOmit --= fullyTranslatedFilters + + // Include extracted filters in the pushdown list. + val filtersToPush = fullyTranslatedFilters + filtersToPush ++= residualsToOmit + + // Omit returned residuals that came only from predicate extraction. + val postScanPredicates = r.pushPredicates(filtersToPush.toArray) + .filterNot(residualsToOmit.contains) val finalPostScanFilters = if (!partitionFields.exists(_.nonEmpty) || !r.supportsIterativePushdown) { @@ -144,6 +186,28 @@ object PushDownUtils extends Logging { } } + // Extract a necessary condition from a deterministic filter that cannot be fully translated. + // The caller must retain the original filter for post-scan evaluation. AND can use either + // child, but OR requires both children. Other expressions, including NOT, must translate whole. + private def extractPushablePredicate[T]( + expression: Expression, + translate: Expression => Option[T]): Option[Expression] = expression match { + case And(left, right) => + val l = extractPushablePredicate(left, translate) + val r = extractPushablePredicate(right, translate) + (l, r) match { + case (Some(a), Some(b)) => Some(And(a, b)) + case _ => l.orElse(r) + } + case Or(left, right) => + for { + l <- extractPushablePredicate(left, translate) + r <- extractPushablePredicate(right, translate) + } yield Or(l, r) + case other => + translate(other).map(_ => other) + } + /** * Rebuilds the Catalyst [[Expression]]s for a sequence of data source [[Predicate]]s, using the * mapping from translated data source predicates to their original Catalyst expressions. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtilsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtilsSuite.scala index f411ca4eef47a..065c54a8e155f 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtilsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtilsSuite.scala @@ -17,16 +17,170 @@ package org.apache.spark.sql.execution.datasources.v2 +import scala.collection.mutable.ArrayBuffer + import org.apache.spark.{SparkFunSuite, SparkNumberFormatException} import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Cast, EqualTo, EvalMode, Literal} +import org.apache.spark.sql.catalyst.expressions.{And, AttributeMap, AttributeReference, BindReferences, Cast, EqualTo, EvalMode, Expression, GetArrayItem, LessThan, Literal, Not, Or, Rand, StringSplit} +import org.apache.spark.sql.catalyst.plans.logical.{ColumnStat, Filter} +import org.apache.spark.sql.catalyst.plans.logical.statsEstimation.FilterEstimation +import org.apache.spark.sql.catalyst.statsEstimation.StatsTestPlan import org.apache.spark.sql.catalyst.types.DataTypeUtils -import org.apache.spark.sql.internal.connector.PartitionPredicateField +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.read.{Scan, ScanBuilder, SupportsPushDownFilters, SupportsPushDownV2Filters} +import org.apache.spark.sql.execution.datasources.DataSourceStrategy +import org.apache.spark.sql.internal.connector.{PartitionPredicateField, PartitionPredicateImpl} +import org.apache.spark.sql.sources import org.apache.spark.sql.types.{IntegerType, StringType, StructField} import org.apache.spark.unsafe.types.UTF8String class PushDownUtilsSuite extends SparkFunSuite { + private val a = AttributeReference("a", IntegerType)() + private val s = AttributeReference("s", StringType)() + private val first = EqualTo(a, Literal(1)) + private val second = EqualTo(a, Literal(2)) + private val unsupported = EqualTo( + GetArrayItem(StringSplit(s, Literal(","), Literal(-1)), Literal(0), false), Literal("fred")) + + private def createScanBuilder( + useV2: Boolean, + returnResidual: Boolean): ScanBuilder = { + if (useV2) { + new SupportsPushDownV2Filters { + private var predicates = Array.empty[Predicate] + override def build(): Scan = throw new UnsupportedOperationException + override def pushPredicates(filters: Array[Predicate]): Array[Predicate] = { + predicates = filters + if (returnResidual) filters else Array.empty + } + override def pushedPredicates(): Array[Predicate] = predicates + } + } else { + new SupportsPushDownFilters { + private var predicates = Array.empty[sources.Filter] + override def build(): Scan = throw new UnsupportedOperationException + override def pushFilters(filters: Array[sources.Filter]): Array[sources.Filter] = { + predicates = filters + if (returnResidual) filters else Array.empty + } + override def pushedFilters(): Array[sources.Filter] = predicates + } + } + } + + private def checkPushdown( + input: Expression, + expected: Option[Expression], + fullyTranslated: Boolean = false): Unit = { + for (useV2 <- Seq(false, true); returnResidual <- Seq(false, true)) { + val builder = createScanBuilder(useV2, returnResidual) + val (pushed, residual) = PushDownUtils.pushFilters(builder, Seq(input), None) + if (useV2) { + assert(pushed.toOption.get == expected.toSeq.map { e => + DataSourceV2Strategy.translateFilterV2(e).get + }) + } else { + assert(pushed.swap.toOption.get == expected.toSeq.map { e => + DataSourceStrategy.translateFilter(e, supportNestedPredicatePushdown = true).get + }) + } + if (fullyTranslated && !returnResidual) { + assert(residual.isEmpty) + } else { + assert(residual == Seq(input)) + } + + if (input.deterministic) { + def passes(e: Expression, row: InternalRow): Boolean = + BindReferences.bindReference(e, Seq(a, s)).eval(row) == true + for (value <- Seq(null, 0, 1, 2, 3); text <- Seq(null, "fred", "other", "fred,other")) { + val row = InternalRow(value, if (text == null) null else UTF8String.fromString(text)) + val actual = expected.forall(passes(_, row)) && residual.forall(passes(_, row)) + assert(actual == passes(input, row), s"$input, value=$value, text=$text") + } + } + } + } + + test("SPARK-40608: push a necessary OR predicate and retain the original filter") { + checkPushdown(Or(first, And(second, unsupported)), Some(Or(first, second))) + checkPushdown(Or(And(unsupported, second), first), Some(Or(second, first))) + checkPushdown(Or(And(first, unsupported), And(second, unsupported)), + Some(Or(first, second))) + checkPushdown(And(unsupported, first), Some(first)) + } + + test("SPARK-40608: do not partially push NOT or an OR with an unsupported branch") { + checkPushdown(Not(Or(first, And(second, unsupported))), None) + checkPushdown(Or(first, unsupported), None) + checkPushdown(And(unsupported, unsupported), None) + checkPushdown(Or(first, And(second, LessThan(Rand(0), Literal(0.5)))), None) + } + + test("SPARK-40608: keep full pushdown of supported predicates") { + val input = Or(first, second) + checkPushdown(input, Some(input), fullyTranslated = true) + } + + test("SPARK-40608: partial pushdown does not discount residual statistics twice") { + val input = Or(first, And(second, unsupported)) + val columnStat = ColumnStat(distinctCount = Some(10), min = Some(1), max = Some(10), + nullCount = Some(0), avgLen = Some(4), maxLen = Some(4)) + val relation = StatsTestPlan(Seq(a, s), 10000, AttributeMap(Seq(a -> columnStat))) + val expected = FilterEstimation(Filter(input, relation)).estimate.get.rowCount + assert(expected.nonEmpty) + + for (useV2 <- Seq(false, true)) { + val builder = createScanBuilder(useV2, returnResidual = true) + val (_, residual) = PushDownUtils.pushFilters(builder, Seq(input), None) + val actual = FilterEstimation(Filter(residual.reduceLeft(And), relation)).estimate.get + assert(actual.rowCount == expected) + } + } + + test("SPARK-40608: keep an original filter that is also extracted from another filter") { + val input = Or(first, And(second, unsupported)) + val extracted = Or(first, second) + for (useV2 <- Seq(false, true); filters <- Seq(Seq(input, extracted), Seq(extracted, input))) { + val builder = createScanBuilder(useV2, returnResidual = true) + val (_, residual) = PushDownUtils.pushFilters(builder, filters, None) + assert(residual == Seq(extracted, input)) + } + } + + test("SPARK-40608: omit extracted residuals before iterative partition pushdown") { + val input = Or(first, And(second, unsupported)) + val extracted = DataSourceV2Strategy.translateFilterV2(Or(first, second)).get + val fields = Seq(a, s).map(attr => PartitionPredicateField(Seq(attr.name), Some(attr))) + for (reportPushed <- Seq(false, true); acceptPartition <- Seq(false, true)) { + val calls = ArrayBuffer.empty[Seq[Predicate]] + val pushed = ArrayBuffer.empty[Predicate] + val builder = new SupportsPushDownV2Filters { + override def build(): Scan = throw new UnsupportedOperationException + override def supportsIterativePushdown(): Boolean = true + override def pushedPredicates(): Array[Predicate] = pushed.toArray + override def pushPredicates(filters: Array[Predicate]): Array[Predicate] = { + calls += filters.toSeq + if (calls.size == 1) { + if (reportPushed) pushed ++= filters + filters + } else if (acceptPartition) { + pushed ++= filters + Array.empty + } else { + filters + } + } + } + val (_, residual) = PushDownUtils.pushFilters(builder, Seq(input), Some(fields)) + assert(calls.size == 2) + assert(calls.head == Seq(extracted)) + assert(calls(1).map(_.asInstanceOf[PartitionPredicateImpl].expression) == Seq(input)) + assert(residual == (if (acceptPartition) Nil else Seq(input))) + } + } + test("SPARK-59572: only a runtime partition predicate keeps a key it cannot evaluate") { val ref = DataTypeUtils.toAttribute(StructField("p", StringType, nullable = true)) val fields = Seq(PartitionPredicateField(Seq("p"), Some(ref))) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCV2Suite.scala b/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCV2Suite.scala index 8f456c3e57ea7..1e6caf0983324 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCV2Suite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCV2Suite.scala @@ -280,6 +280,14 @@ class JDBCV2Suite extends SharedSparkSession with ExplainSuiteHelper { } } + test("SPARK-40608: partially push an OR predicate and evaluate its residual") { + val df = sql("SELECT name, id FROM h2.test.people " + + "WHERE id = 1 OR (id = 2 AND split(name, ',')[0] = 'fred')") + assert(df.queryExecution.optimizedPlan.exists(_.isInstanceOf[Filter])) + checkPushedInfo(df, "PushedFilters: [(ID = 1) OR (ID = 2)]") + checkAnswer(df, Seq(Row("fred", 1))) + } + test("SPARK-57243: IS [NOT] NULL over a composite operand is pushed down and runs on H2") { val df1 = sql("SELECT name FROM h2.test.employee WHERE (salary = 10000) IS NOT NULL") checkFiltersRemoved(df1)