HyukjinKwon commented on code in PR #59403:
URL: https://github.com/apache/spark/pull/59403#discussion_r4275695866
##########
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:
**Non-blocking (P2):** `outputsNormalizedKeys(aggregate)` is evaluated here,
when the filter is pushed. The Aggregate's function set is not final at that
point. In the operator-optimization rule set `PushDownPredicates`
(Optimizer.scala:108) runs before `ColumnPruning` (Optimizer.scala:115), and
ColumnPruning's `Project(_, a: Aggregate)` case drops aggregate expressions the
parent does not use. `SparkStrategies.Aggregation` plans an all-`PythonUDAF`
Aggregate as `ArrowAggregatePythonExec` on the raw keys and only rejects
mixtures at planning.
Example with rows `(x=-0.0, v=1.0), (x=-0.0, v=2.0)` and a grouped-agg
pandas UDF `pudf`: `df.groupBy("x").agg(pudf(col("v")).as("p"),
count("*").as("c")).where(col("x").cast("string") === "-0.0").select("x",
"p")`. When the filter is pushed the functions are mixed, so it becomes
`cast(knownfloatingpointnormalized(normalizenanandzero(x)) AS STRING) =
'-0.0'`, which is false for both rows. Then `c` is pruned and the Aggregate
runs on the raw key `-0.0`, which does satisfy the predicate. The query returns
no rows, while master returns the `-0.0` group. The new FilterPushdownSuite and
PythonUDFSuite pandas tests only use Aggregates that are all-`PythonUDAF` from
the start. The opposite shape, an all-`PythonUDAF` Aggregate whose UDAF output
is pruned so that it becomes group-only and normalizing, still gets a raw
filter, as on master.
**Recommended change:** For an Aggregate whose aggregate functions include
any PythonUDAF, push below it only the conjuncts that
normalizeGroupingExpressions leaves unchanged, and keep the conjuncts that
would read a float/double key in normalized form in the Filter above the
Aggregate. Update the outputsNormalizedKeys/pushBelowAggregate comments to
describe this. Add optimizer plan tests and an end-to-end grouped-agg pandas
UDF test, including a mixed pandas plus SQL aggregate whose SQL aggregate is
projected away.
**Why this works:** A conjunct evaluated above the Aggregate reads whatever
key form the finally planned operator outputs, so it no longer depends on the
function set at push time. Conjuncts that normalization would not change give
the same result on raw and normalized keys, so they can still be pushed below
either branch.
**Scope:** Make filter push-down below Aggregates that contain grouped-agg
pandas/Arrow UDFs independent of later column pruning.
**Compatibility:** Built-in and PythonAggregate aggregates keep receiving
normalized pushed filters, and plain comparisons keep reaching data sources.
**Risks:** The CTE removal mirror must keep translating pushed conjuncts the
same way and still find conjuncts that stay above the Aggregate.
**Constraints:** Plans for Aggregates without any PythonUDAF must not
change. Plain comparisons on float/double keys must still be pushed below
PythonUDAF aggregates.
**Success:** A mixed pandas UDF plus SQL aggregate whose SQL aggregate is
pruned, filtered with a sign-sensitive predicate on a -0.0 key, returns the
same rows as master. An all-PythonUDAF Aggregate whose pandas output is pruned,
filtered with a sign-sensitive predicate, returns only rows that satisfy the
predicate. Plain comparisons on float/double keys are still pushed below
PythonUDAF aggregates, and filters below built-in aggregates are normalized as
before.
**Verification:**
- **Regression:** A grouped-agg pandas UDF aggregation mixed with a SQL
aggregate that is projected away, over a -0.0 key and filtered with a
sign-sensitive predicate, returns the -0.0 group.
- **Regression:** A grouped-agg pandas UDF aggregation whose pandas output
is projected away, filtered with a sign-sensitive predicate, returns only rows
consistent with the predicate.
- **Behavior:** The optimized plan keeps a normalization-sensitive conjunct
above an Aggregate that contains a PythonUDAF and still pushes plain
comparisons below it.
- **Compatibility:** Pushed-filter plans for Aggregates without PythonUDAF
functions are unchanged.
##########
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:
**Nit (P3):** This sentence doesn't match the code. The guard below matches
with `semanticEquals`, which compares canonicalized trees, and
`HigherOrderFunction.canonicalized` (higherOrderFunctions.scala:245-267)
renumbers lambda variables. So `normalize(g)` for an array or map key is
semantically equal to the earlier normalization even though its
`NamedLambdaVariable`s are fresh, and the case returns `e` unchanged. Array and
map keys are kept like any other key, not normalized again. The results are the
same either way, but the comment, and the matching sentence in the PR
discussion, give the wrong model. No test covers an array or map key on this
path.
**Verification:**
- **Inspection:** The comment matches the short-circuit's semanticEquals
behavior for array and map keys.
##########
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:
**Nit (P3):** With `c = a + rand(0)` as a DOUBLE, the new guard keeps the
upper Aggregate, so this test now only asserts that nothing changes. Nothing in
the suite still checks the original scenario, a group-only upper Aggregate
removed over a lower one kept for its non-deterministic output. A regression
there would pass, and the title still says "Remove". The sibling test above
kept its coverage by using `(rand(0) * 10).cast(IntegerType)`. Doing the same
here, plus a separately named test for the double-key case if you want it
pinned, would restore the signal.
**Verification:**
- **Regression:** The group-only upper Aggregate over a lower Aggregate kept
for its non-deterministic non-grouping output is replaced by a Project for a
non-float key.
- **Inspection:** Each test's title describes its asserted plan.
##########
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:
**Nit (P3):** Nit: this test starts right after the previous test's closing
brace, and three blank lines follow it (345-347). OptimizeOneRowPlanSuite has
the same layout (79-80 and 92-94), and RemoveRedundantAggregatesSuite has two
blank lines after the second new test (281-282). One blank line between tests
would match the rest of these suites.
**Verification:**
- **Inspection:** No new test directly follows a closing brace and no run of
more than one blank line remains around the new tests.
--
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]
