[SPARK-40608][SQL] Push necessary predicates from partially translatable V2 filters

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.

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)))

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)