sarutak commented on code in PR #59940:
URL: https://github.com/apache/spark/pull/59940#discussion_r4423136202


##########
sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala:
##########
@@ -2949,6 +2949,95 @@ class DataFrameSuite extends SharedSparkSession
     }
   }
 
+  test("SPARK-60730: a filter condition that can fail is not evaluated ahead 
of its order") {
+    val f = udf((i: Long) => i)
+    val g = udf((i: Long) => i)
+    val h = udf((i: Long) => { require(i > 5, s"h called with $i"); i })
+    // Each query fails if the conjunct after the first is evaluated on an id 
<= 5, which the
+    // first conjunct filters out.
+    def divideQuery(): DataFrame = spark.range(0, 10, 1, 1)
+      .select($"id", f($"id") as "f")
+      .where($"f" > 5 && lit(10L) / ($"id" - 3) > 0)
+    def castQuery(): DataFrame = spark.range(0, 10, 1, 1)
+      .select($"id", f($"id") as "f", when($"id" > 5, 
lit("1")).otherwise(lit("x")) as "s")
+      .where($"f" > 5 && $"s".cast("int") > 0)
+    def udfQuery(): DataFrame = spark.range(0, 10, 1, 1)
+      .select($"id", f($"id") as "f")
+      .where($"f" > 5 && h($"id") > 0)
+    // With three expensive aliases, so the projection split is involved too. 
It splits `f` and
+    // `h` off first, as written, and leaves the division over `g` behind them.
+    val k = udf((i: Long) => i)
+    def splitQuery(): DataFrame = spark.range(0, 10, 1, 1)
+      .select($"id", f($"id") as "f", k($"id") as "h", g($"id") as "g")
+      .where($"f" + $"h" > 10 && lit(10L) / ($"g" - 3) > 0)
+    // `q` only works for the ids the first conjunct keeps. The split must put 
that conjunct
+    // lowest, as written, even though the check rejects its addition.
+    val q = udf((i: Long) => { require(i > 4, s"q called with $i"); i })
+    def firstQuery(): DataFrame = spark.range(0, 10, 1, 1)
+      .select($"id", f($"id") as "f", q($"id") as "q")
+      .where($"f" + 1 > 5 && $"q" > 0)
+    // `r` only works for the ids the second conjunct keeps, so the split must 
not put the third
+    // conjunct below the second, which can fail. Without the split, an 
interpreted projection
+    // computes `q` and `r` for every row it gets.
+    val r = udf((i: Long) => { require(i > 5, s"r called with $i"); i })
+    def laterQuery(): DataFrame = spark.range(0, 10, 1, 1)
+      .select($"id", f($"id") as "f", g($"id") as "g", r($"id") as "r")
+      .where($"f" >= 0 && $"g" + 1 > 6 && $"r" > 0)
+    // A non-deterministic conjunct sees the same rows as with the whole 
filter pushed.
+    def randQuery(): DataFrame = spark.range(0, 100, 1, 1)
+      .select($"id", f($"id") as "f")
+      .where($"f" >= 50 && rand(10) > 0.5)
+    val randExpected = withSQLConf(SQLConf.AVOID_DOUBLE_FILTER_EVAL.key -> 
"false") {
+      randQuery().collect().toSeq
+    }
+    for (codegen <- Seq(true, false); split <- Seq(true, false); avoid <- 
Seq(true, false)) {
+      withSQLConf(
+        SQLConf.ANSI_ENABLED.key -> "true",
+        SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> codegen.toString,
+        SQLConf.AVOID_DOUBLE_FILTER_EVAL.key -> avoid.toString,
+        SQLConf.SPLIT_PROJECTION_FOR_EXPENSIVE_FILTERS.key -> split.toString) {
+        checkAnswer(divideQuery(), (6L until 10L).map(i => Row(i, i)))
+        checkAnswer(castQuery(), (6L until 10L).map(i => Row(i, i, "1")))
+        checkAnswer(udfQuery(), (6L until 10L).map(i => Row(i, i)))
+        checkAnswer(splitQuery(), (6L until 10L).map(i => Row(i, i, i, i)))
+        checkAnswer(randQuery(), randExpected)
+        if (codegen || split || !avoid) {
+          checkAnswer(firstQuery(), (5L until 10L).map(i => Row(i, i, i)))
+          checkAnswer(laterQuery(), (6L until 10L).map(i => Row(i, i, i, i)))
+        }
+      }
+    }
+    // The legacy config pushes the division below the projection again.
+    withSQLConf(
+      SQLConf.ANSI_ENABLED.key -> "true",
+      SQLConf.LEGACY_PUSH_FILTER_AHEAD_OF_EXPENSIVE_CONJUNCT.key -> "true") {
+      val e = intercept[SparkArithmeticException](divideQuery().collect())
+      assert(e.getCondition === "DIVIDE_BY_ZERO")
+    }
+  }
+
+  test("SPARK-60730: an expensive filter condition sees the rows a later 
condition drops " +
+    "unless that condition is a plain comparison") {
+    val p = udf((i: Long) => { require(i > 5, s"p called with $i"); i })
+    // `p` only works for the ids the later conjunct keeps, so these queries 
depend on that
+    // conjunct being pushed below the projection and evaluated first.
+    def query(later: Column): DataFrame = spark.range(0, 10, 1, 1)
+      .select($"id", p($"id") as "p")
+      .where($"p" > 0 && later)
+    val expected = (6L until 10L).map(i => Row(i, i))
+    // A plain comparison can not fail, so it is still pushed.
+    checkAnswer(query($"id" > 5), expected)
+    // An addition is not a plain comparison, so it now stays behind `p > 0`, 
as in Spark 4.1,
+    // even with ANSI mode off, where it can not fail.
+    withSQLConf(SQLConf.ANSI_ENABLED.key -> "false") {
+      val e = intercept[SparkException](query($"id" + 1 > 6).collect())
+      assert(e.getCondition === "FAILED_EXECUTE_UDF")

Review Comment:
   Agreed. For the addition, the test now checks the plan instead: it stays in 
the filter above the projection, after `p > 0`. The end-to-end part still runs 
`id > 5`, which is pushed, and the addition under the legacy config.



##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala:
##########
@@ -2651,17 +2688,46 @@ object PushPredicateThroughNonJoin extends 
Rule[LogicalPlan] with PredicateHelpe
         // split orders conditions by cost, and cost is no way to decide how 
many rows a
         // non-deterministic expression sees. Throwable conditions, or 
conditions over throwable
         // aliases, stay there too: ordering by cost could evaluate them (or 
compute the aliases
-        // in a lower projection) ahead of the conditions written before them.
+        // in a lower projection) ahead of the conditions written before them. 
So do the
+        // conditions after the first expensive one that `canPushAhead` 
rejects, e.g. a division
+        // over an expensive alias, or a cheap condition left above by 
`canPushAhead`, which
+        // needs no expensive alias and so would go lowest, and the conditions 
after those.
         val (splitChild, splitAliases, stayUp) =
           if (expensiveWithUsed.nonEmpty && 
SQLConf.get.splitProjectionForExpensiveFilters) {
-            val (splittable, notSplittable) = expensiveWithUsed.partition {
-              case (cond, _, replaced) => cond.deterministic && 
!replaced.throwable
+            // Only the cheap aliases are substituted: `canPushAhead` never 
accepts an expensive
+            // one such as a UDF, so substituting it would reject every 
condition over an
+            // expensive alias and stop the split. Whether the split may 
compute an expensive
+            // alias ahead of the conditions written before it is a separate 
question.
+            val cheapAliasMap = AttributeMap(aliasMap.filter { case (_, a) => 
!a.child.expensive })
+            // A condition after the first expensive one that `canPushAhead` 
rejects stays above
+            // in this pass, and so does every condition written after it, so 
the split does not
+            // put a later condition, or compute the aliases it needs, below 
the rejected one now.
+            // If every condition written before it is placed by a later pass, 
it is then the
+            // first one of the filter left above, and is pushed or split off 
on top of them; the
+            // split orders the conditions after it by cost again. Otherwise 
it stays above,
+            // with the conditions after it. The first expensive condition
+            // is exempt, since every condition written before it is already 
pushed below.
+            // Throwable and non-deterministic conditions stay above as before 
(SPARK-60602),
+            // without holding back the later ones.
+            val firstRejected = expensiveWithUsed.indexWhere({ case (cond, _, 
replaced) =>
+              cond.deterministic && !replaced.throwable &&
+                !canPushAhead(replaceAlias(cond, cheapAliasMap))
+            }, 1) match {
+              case -1 => expensiveWithUsed.length
+              case i => i
             }
+            val (splittableWithIndex, notSplittableWithIndex) =
+              expensiveWithUsed.zipWithIndex.partition { case ((cond, _, 
replaced), i) =>
+                cond.deterministic && !replaced.throwable && i < firstRejected
+              }
+            val splittable = splittableWithIndex.map(_._1)
+            val notSplittable = notSplittableWithIndex.map(_._1)

Review Comment:
   Done along those lines: the index bookkeeping is gone. The conditions are 
now split with `span` at the first one held back and `splitAt(1)`, and the 
split check is a single `exists` over the conditions after it.



##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala:
##########
@@ -2627,16 +2633,47 @@ object PushPredicateThroughNonJoin extends 
Rule[LogicalPlan] with PredicateHelpe
         // eliminate more data quickly and it may also be something which 
could be evaluated at
         // the storage layer. We may wish to improve this heuristic in the 
future.
         def isCheap(used: AttributeMap[Alias]): Boolean = 
!used.exists(_._2.child.expensive)
+        // Whether a component written after an expensive one may be evaluated 
ahead of it, on
+        // rows the components before it would have filtered out. Such a 
component must not be
+        // able to fail there, e.g. on a division by zero or an invalid cast 
in ANSI mode, in a
+        // UDF or in `raise_error`, nor see a different set of rows than as 
written, as a
+        // non-deterministic component would. 
`ExprUtils.canEvaluateUnconditionally` is the
+        // check for that; it only accepts deterministic comparisons, IN, IS 
NULL and the like
+        // over columns and literals, so it also rejects some components that 
can not fail,
+        // such as arithmetic with ANSI mode off. It also accepts some filters 
the optimizer
+        // adds, which have no written position to keep: a runtime bloom 
filter over plain keys
+        // and the `size(...) > 0` that `InferFiltersFromGenerate` infers over 
a plain array or
+        // map, which can not fail, and a dynamic partition pruning filter, 
which the scan
+        // evaluates at run time to prune what it reads, not on the filtered 
rows, so its
+        // position among the components does not matter. These are matched by 
shape, so a
+        // written `size(col) > 0` is accepted too.
+        // Other kinds of added filters stay above, which only costs pushdown. 
The legacy config
+        // only rejects throwable components. Callers pass the component with 
the aliases
+        // substituted that its evaluation would compute, so a part inside 
such an alias counts
+        // too.
+        def canPushAhead(component: Expression): Boolean =
+          if (SQLConf.get.legacyPushFilterAheadOfExpensiveConjunct) {

Review Comment:
   Done: `canPushAhead` is now a private method of the object with its comment 
as Scaladoc.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to