viirya commented on code in PR #59944:
URL: https://github.com/apache/spark/pull/59944#discussion_r4421485633
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala:
##########
@@ -2821,15 +2817,10 @@ object PushPredicateThroughNonJoin extends
Rule[LogicalPlan] with PredicateHelpe
// come from grandchild.
// TODO: non-deterministic predicates could be pushed through some
operators that do not change
// the rows.
- val (candidates, nonDeterministic) =
- splitConjunctivePredicates(filter.condition).partition(_.deterministic)
-
- val (pushDown, rest) = candidates.partition { cond =>
- cond.references.subsetOf(grandchild.outputSet)
+ val (pushDown, stayUp) = splitPushableConjuncts(filter.condition) { cond =>
+ cond.deterministic && cond.references.subsetOf(grandchild.outputSet)
}
Review Comment:
Thanks, confirmed: your `explode(filter(array(id), x => x > 5))` query fails
today. Fixed in 6228e8b3a66. The generic case no longer pushes a throwable
condition through a `Generate` without `outer` or a `ScriptTransformation`,
even when it is written first (`canDropRows`), so a throwable condition stops
the `span` there. An outer `Generate` still gets it, since it outputs a row for
every input row, and so do the other operators of `canPushThrough`, which
output at least one row per input row (`BinBy` included). Added a plan test for
both `outer` values and your query to the `DataFrameSuite` test.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/FilterPushdownSuite.scala:
##########
@@ -750,6 +750,120 @@ class FilterPushdownSuite extends PlanTest {
assertTopFilterConjuncts(Optimize.execute(originalQuery),
Seq(nonDeterministic, unplaced))
}
+ // Raises for the rows with a <= 5, so it must only see the rows the
conditions before it kept.
+ private val throwableOnA = attrA > 5 || IsNull(RaiseError(Literal("boom")))
+
+ test("SPARK-60633: a throwable condition is not pushed through a Python UDF
evaluation " +
+ "ahead of the conditions before it") {
+ Seq(PythonEvalType.SQL_BATCHED_UDF,
PythonEvalType.SQL_SCALAR_PANDAS_UDF).foreach { evalType =>
+ withClue(s"eval type $evalType: ") {
+ val pythonUDF = PythonUDF(
+ "pythonUDF", null, IntegerType, Seq(attrA), evalType,
udfDeterministic = true)
+ val f = AttributeReference("f", IntegerType)()
+ def evalPython(child: LogicalPlan): LogicalPlan =
+ if (evalType == PythonEvalType.SQL_BATCHED_UDF) {
+ BatchEvalPython(Seq(pythonUDF), Seq(f), child)
+ } else {
+ ArrowEvalPython(Seq(pythonUDF), Seq(f), child, evalType)
+ }
+
+ // f > 5 cannot be pushed, so the throwable condition after it stays
above too. b > 1
+ // is still pushed, as that only reduces the rows the throwable
condition sees.
+ val originalQuery = Filter(f > 5 && throwableOnA && attrB > 1,
evalPython(testRelation))
+ val correctAnswer =
+ Filter(f > 5 && throwableOnA, evalPython(Filter(attrB > 1,
testRelation)))
+ comparePlans(Optimize.execute(originalQuery), correctAnswer)
+
+ // With nothing left above before it, a throwable condition is still
pushed.
+ val throwableFirst = Filter(throwableOnA && f > 5,
evalPython(testRelation))
+ val throwableFirstAnswer = Filter(f > 5,
evalPython(Filter(throwableOnA, testRelation)))
+ comparePlans(Optimize.execute(throwableFirst), throwableFirstAnswer)
+
+ // The conditions left above keep their written order.
+ val nonDeterministic = Rand(10) > 0.5
+ val mixed = Filter(nonDeterministic && f > 5 && throwableOnA,
evalPython(testRelation))
+ assertTopFilterConjuncts(
+ Optimize.execute(mixed), Seq(nonDeterministic, f > 5, throwableOnA))
+ }
+ }
+ }
+
+ test("SPARK-60633: a throwable condition is not pushed through a sort ahead
of a " +
+ "non-deterministic condition before it") {
+ val originalQuery = testRelation.orderBy($"b".asc)
+ .where(Rand(10) > 0.5 && throwableOnA && $"b" > 1)
+ .analyze
+ val Seq(nonDeterministic, throwable, _) = splitConjunctivePredicates(
+ originalQuery.asInstanceOf[Filter].condition)
+ val correctAnswer = testRelation.where($"b" > 1).orderBy($"b".asc)
+ .where(Rand(10) > 0.5 && throwableOnA)
+ .analyze
+
+ val optimized = Optimize.execute(originalQuery)
+ comparePlans(optimized, correctAnswer)
+ assertTopFilterConjuncts(optimized, Seq(nonDeterministic, throwable))
+ }
+
+ test("SPARK-60633: a throwable condition is not pushed through an aggregate
ahead of the " +
+ "conditions before it") {
+ val originalQuery = testRelation.groupBy($"a")($"a", count($"b") as "c")
+ .where($"c" > 1 && throwableOnA && $"a" > 1)
+ .analyze
+ val correctAnswer = testRelation.where($"a" > 1).groupBy($"a")($"a",
count($"b") as "c")
+ .where($"c" > 1 && throwableOnA)
+ .analyze
+ comparePlans(Optimize.execute(originalQuery), correctAnswer)
+
+ // With nothing left above before it, a throwable condition is still
pushed.
+ val throwableFirst = testRelation.groupBy($"a")($"a", count($"b") as "c")
+ .where(throwableOnA && $"c" > 1)
+ .analyze
+ val throwableFirstAnswer = testRelation.where(throwableOnA)
+ .groupBy($"a")($"a", count($"b") as "c")
+ .where($"c" > 1)
+ .analyze
+ comparePlans(Optimize.execute(throwableFirst), throwableFirstAnswer)
+ }
+
+ test("SPARK-60633: a throwable condition is not pushed through a window
ahead of the " +
+ "conditions before it") {
+ val winExpr = windowExpr(count($"b"),
+ windowSpec($"a" :: Nil, $"b".asc :: Nil, UnspecifiedFrame))
+ val originalQuery = testRelation
+ .window(winExpr.as("window") :: Nil, $"a" :: Nil, $"b".asc :: Nil)
+ .where($"window" > 1 && throwableOnA && $"a" > 1)
+ .analyze
+ val correctAnswer = testRelation.where($"a" > 1)
+ .window(winExpr.as("window") :: Nil, $"a" :: Nil, $"b".asc :: Nil)
+ .where($"window" > 1 && throwableOnA)
+ .analyze
+ comparePlans(Optimize.execute(originalQuery), correctAnswer)
Review Comment:
Added in 6228e8b3a66 with your snippet. On the merge base, `rest ++
nonDeterministic` makes the top filter `window > 1` then `rand(10) > 0.5`, so
the check would catch a revert.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala:
##########
@@ -2724,15 +2725,10 @@ object PushPredicateThroughNonJoin extends
Rule[LogicalPlan] with PredicateHelpe
if w.partitionSpec.forall(_.isInstanceOf[AttributeReference]) =>
val partitionAttrs = AttributeSet(w.partitionSpec.flatMap(_.references))
- val (candidates, nonDeterministic) =
- splitConjunctivePredicates(condition).partition(_.deterministic)
-
- val (pushDown, rest) = candidates.partition { cond =>
- cond.references.subsetOf(partitionAttrs)
+ val (pushDown, stayUp) = splitPushableConjuncts(condition) { cond =>
+ cond.deterministic && cond.references.subsetOf(partitionAttrs)
}
Review Comment:
Updated in 6228e8b3a66 to the wording you suggested.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala:
##########
@@ -2842,6 +2833,22 @@ object PushPredicateThroughNonJoin extends
Rule[LogicalPlan] with PredicateHelpe
}
}
+ /**
+ * Splits `condition` into the conjuncts to push below an operator and the
conjuncts to keep
+ * above it, where `canPush` tells whether the operator lets a conjunct
through. Every conjunct
+ * written before the first one that cannot be pushed is pushed. After it,
only the conjuncts
+ * that are not throwable (e.g. `raise_error`) are pushed, so a throwable
conjunct is not
+ * evaluated on rows the conjuncts before it would have filtered out. Moving
a non-throwable
+ * conjunct below a throwable one is fine, as it can only reduce the rows
the throwable one
+ * sees. The conjuncts to keep above are returned in the order they were
written.
+ */
+ private def splitPushableConjuncts(condition: Expression)(
+ canPush: Expression => Boolean): (Seq[Expression], Seq[Expression]) = {
Review Comment:
Done in 6228e8b3a66. `splitPushable[T]` takes `canPush` and `throwable`, and
the `Project` case uses it with `replaced.throwable`. `splitPushableConjuncts`
is now a wrapper for the cases that split a plain condition. The `Aggregate`
case also uses `splitPushable` directly now, to check `replaced.throwable`, as
the bot review asked.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala:
##########
@@ -2842,6 +2833,22 @@ object PushPredicateThroughNonJoin extends
Rule[LogicalPlan] with PredicateHelpe
}
}
+ /**
+ * Splits `condition` into the conjuncts to push below an operator and the
conjuncts to keep
+ * above it, where `canPush` tells whether the operator lets a conjunct
through. Every conjunct
+ * written before the first one that cannot be pushed is pushed. After it,
only the conjuncts
+ * that are not throwable (e.g. `raise_error`) are pushed, so a throwable
conjunct is not
+ * evaluated on rows the conjuncts before it would have filtered out. Moving
a non-throwable
+ * conjunct below a throwable one is fine, as it can only reduce the rows
the throwable one
+ * sees. The conjuncts to keep above are returned in the order they were
written.
+ */
+ private def splitPushableConjuncts(condition: Expression)(
+ canPush: Expression => Boolean): (Seq[Expression], Seq[Expression]) = {
+ val (leading, rest) = splitConjunctivePredicates(condition).span(canPush)
+ val (pushedRest, stayUp) = rest.partition(cond => canPush(cond) &&
!cond.throwable)
Review Comment:
Done in 6228e8b3a66, in the shared helper.
--
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]
