ctagard commented on code in PR #59403:
URL: https://github.com/apache/spark/pull/59403#discussion_r4308604904
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala:
##########
@@ -2796,6 +2800,65 @@ object PushPredicateThroughNonJoin extends
Rule[LogicalPlan] with PredicateHelpe
case _ => false
}
+ // Rewrites conjuncts pushed below `aggregate` in terms of the aggregate's
child and ANDs them.
+ // Expressions above the aggregate see its float/double grouping keys
normalized (see
+ // NormalizeFloatingNumbers: -0.0 becomes 0.0 and every NaN becomes one
canonical NaN), so the
+ // pushed conjuncts read them normalized too, wherever the raw value could
change the result.
+ // NormalizeFloatingNumbers.outputsNormalizedKeys says when the keys come
out normalized, whether
+ // the aggregate is planned or a later rule replaces it with a Project.
Aggregates of grouped-agg
+ // pandas UDFs output the raw keys, so conjuncts pushed below them read the
raw keys.
+ private[optimizer] def pushBelowAggregate(
+ conds: Seq[Expression],
+ aggregate: Aggregate): Expression = {
+ val aliasMap = getAliasMap(aggregate)
+ val normalizedGrouping = if
(NormalizeFloatingNumbers.outputsNormalizedKeys(aggregate)) {
Review Comment:
Done in f4d62f0 as suggested: below an Aggregate with any PythonUDAF only
conjuncts that normalizing leaves unchanged are pushed; the rest stay above.
Tests cover the pruned mixed and pruned all-pandas shapes, and the CTE mirror.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala:
##########
@@ -2796,6 +2800,65 @@ object PushPredicateThroughNonJoin extends
Rule[LogicalPlan] with PredicateHelpe
case _ => false
}
+ // Rewrites conjuncts pushed below `aggregate` in terms of the aggregate's
child and ANDs them.
+ // Expressions above the aggregate see its float/double grouping keys
normalized (see
+ // NormalizeFloatingNumbers: -0.0 becomes 0.0 and every NaN becomes one
canonical NaN), so the
+ // pushed conjuncts read them normalized too, wherever the raw value could
change the result.
+ // NormalizeFloatingNumbers.outputsNormalizedKeys says when the keys come
out normalized, whether
+ // the aggregate is planned or a later rule replaces it with a Project.
Aggregates of grouped-agg
+ // pandas UDFs output the raw keys, so conjuncts pushed below them read the
raw keys.
+ private[optimizer] def pushBelowAggregate(
+ conds: Seq[Expression],
+ aggregate: Aggregate): Expression = {
+ val aliasMap = getAliasMap(aggregate)
+ val normalizedGrouping = if
(NormalizeFloatingNumbers.outputsNormalizedKeys(aggregate)) {
+ aggregate.groupingExpressions
+ .filter(e => NormalizeFloatingNumbers.needNormalize(e.dataType))
+ .map(trimAliases)
+ } else {
+ Nil
+ }
+ conds.map { cond =>
+ val replaced = replaceAlias(cond, aliasMap)
+ if (normalizedGrouping.isEmpty) {
+ replaced
+ } else {
+ normalizeGroupingExpressions(replaced, normalizedGrouping)
+ }
+ }.reduce(And)
+ }
+
+ // Rewrites `e` so that it gives the same result on the raw values of
`grouping` as `e` gives on
+ // their normalized values (see NormalizeFloatingNumbers). Each occurrence
of a grouping
+ // expression is wrapped in its normalized form, with one exception: a
float/double attribute
+ // that is a direct operand of one of the comparisons matched below is left
as is, and so is a
+ // float attribute widened to double (as when a FLOAT column is compared
with a DOUBLE or
+ // DECIMAL value), since widening keeps -0.0 and NaN. Those comparisons
already treat -0.0 and
+ // 0.0 as equal and all NaNs as equal (see SQLOrderingUtil.compareDoubles and
+ // CodegenContext.genEqual), and IsNull, IsNotNull and IsNaN cannot tell
normalized values apart
+ // either, so wrapping would not change the result. Leaving them as is keeps
plain comparisons
+ // usable by data sources and partition pruning.
+ private[optimizer] def normalizeGroupingExpressions(
+ e: Expression,
+ grouping: Seq[Expression]): Expression = e match {
+ // A grouping key already normalized when the filter was pushed through an
Aggregate above on
+ // the same key. Any other normalized expression may be computed from a
grouping key (when the
+ // upper Aggregate groups by an expression of this one's key), so the key
inside it still has
+ // to be normalized. `normalize` builds new lambda variables for array and
map keys, so those
Review Comment:
You're right, thanks. Corrected in 44ec4a4, with an array-key unit test.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RemoveRedundantAggregatesSuite.scala:
##########
@@ -180,11 +182,9 @@ class RemoveRedundantAggregatesSuite extends PlanTest {
.groupBy($"a")($"a", ($"a" + rand(0)) as "c")
.groupBy($"a", $"c")($"a", $"c")
.analyze
- val expected = relation
- .groupBy($"a")($"a", ($"a" + rand(0)) as "c")
- .analyze
+ // The upper aggregate is not removed either, since it normalizes the
double key c.
val optimized = Optimize.execute(query)
- comparePlans(optimized, expected)
+ comparePlans(optimized, query)
Review Comment:
Restored in 7483323 with `(rand(0) * 10).cast(IntegerType)`; the double-key
case is its own test.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/LimitPushdownSuite.scala:
##########
@@ -328,6 +328,22 @@ class LimitPushdownSuite extends PlanTest {
Optimize.execute(x.groupBy("x.a".attr)("x.a".attr,
count("x.a".attr)).limit(1).analyze),
x.groupBy("x.a".attr)("x.a".attr, count("x.a".attr)).limit(1).analyze)
}
+ test("SPARK-60150: No limit 1 push down through a group only Aggregate on a
float key") {
Review Comment:
Fixed in 2517918.
--
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]
