holdenk opened a new pull request, #59979:
URL: https://github.com/apache/spark/pull/59979
Backport-of: #59403
### What changes were proposed in this pull request?
`PushPredicateThroughNonJoin` pushes a filter on grouping keys below an
`Aggregate`. This PR makes the pushed filter read the *normalized* float/double
grouping keys wherever the raw value could change its result. It then evaluates
below the `Aggregate` exactly as it would above it.
Float/double grouping keys are normalized during physical planning
(`SparkStrategies`): `-0.0` becomes `0.0` and NaNs are canonicalized. The
normalized key keeps the original `ExprId`, so everything above the `Aggregate`
sees the normalized value, while a filter pushed below it sees the raw value.
The change, in the `Aggregate` case of `PushPredicateThroughNonJoin`:
- A new helper, `pushBelowAggregate`, splits the filter's conjuncts into the
ones pushed below the `Aggregate` and the ones kept above it (the existing
check that a conjunct only reads grouping keys moved into it), expands aliases
as before, and rewrites each pushed conjunct with
`normalizeGroupingExpressions`. It skips the rewrite when there are no
float/double keys.
- `normalizeGroupingExpressions` is a single rewrite. It replaces each
occurrence of a float/double grouping expression with
`NormalizeFloatingNumbers.normalize(groupingExpr)`, the same normalization the
planner applies to the grouping keys, which also covers struct and array keys.
There are two exceptions:
- A float/double attribute that is a direct operand of `EqualTo`,
`EqualNullSafe`, `LessThan`, `LessThanOrEqual`, `GreaterThan`,
`GreaterThanOrEqual`, `In`, `InSet`, `InSubquery`, `IsNull`, `IsNotNull` or
`IsNaN` is left as is. These already treat `-0.0`/`0.0` and all NaNs as equal
(`SQLOrderingUtil.compareDoubles`, `CodegenContext.genEqual`; an IN subquery
becomes a join on `EqualTo` conditions, whose keys are normalized on both sides
when it is an equi-join), so they give the same result on the raw value, and
plain comparisons stay usable by data sources and partition pruning. The
comparisons are listed explicitly rather than matched as `BinaryComparison`, so
a new comparison class is not exempted without review.
- A float attribute widened to double (`cast(f AS DOUBLE)`, as when a
FLOAT key is compared with a DOUBLE or DECIMAL value) is left as is too, since
widening keeps `-0.0` and NaN. Other casts, such as to STRING, are normalized.
- A key that an earlier push already normalized (a filter pushed through
stacked aggregates on the same key) is not normalized again. Any other
normalized subtree is rewritten too, so a lower key inside a normalized upper
key computed from it (e.g. `pow(x, -1) AS k ... GROUP BY k`) is normalized.
For example, `x > 5.0 OR cast(x AS STRING) = '-0.0'` is pushed as `x > 5.0
OR cast(norm(x) AS STRING) = '-0.0'`, and `x = 0.0` is pushed unchanged.
`PushdownPredicatesAndPruneColumnsForCTEDef.removePushedDownFilter` mirrors
this rule's translation through an `Aggregate` so it can find previously pushed
filters. It now uses the same `pushBelowAggregate` split.
The pushed filter reads normalized keys because the `Aggregate` outputs them
normalized. `NormalizeFloatingNumbers` now states when that holds, and the
rules that could break it respect it:
- `outputsNormalizedKeys(aggregate)`: true unless all of the aggregate's
functions are grouped-agg pandas UDFs, which are planned as
`ArrowAggregatePythonExec` on the raw keys.
- `mayChangeKeyForm(aggregate)`: true for an aggregate with any grouped-agg
pandas UDF, whose key form is only settled once `ColumnPruning` has decided
which of its aggregate functions to keep. `pushBelowAggregate` pushes below it
only the conjuncts that normalizing leaves unchanged and keeps the others above
it.
- `readsNormalizedKeys(aggregate, outputs)`: whether a `Project` of those
outputs over the aggregate's child would read a float/double key raw that the
aggregate outputs normalized. `OptimizeOneRowPlan` and
`RemoveRedundantAggregates`' group-only case no longer replace the `Aggregate`
in that case, and `LimitPushDown` keeps it under `LIMIT 1` while still pushing
`LocalLimit(1)` below it. (Normalizing the keys in the `Project` instead is not
possible while keeping the outputs' `ExprId`s: an `Alias` may not reuse the
`ExprId` of an attribute of its child.)
- `RemoveRedundantAggregates` keeps a lower `Aggregate` under a
duplicate-insensitive upper one only if the upper one would read its
float/double key raw: through a lower output computed from the key, inside an
aggregate function, or in a computed grouping expression. A key the upper one
groups by directly is normalized by it, so the lower one is still removed.
- Streaming aggregates are left as before: the rules above still remove
them, and conjuncts are pushed below them unnormalized (a normalized conjunct
would disagree with a streaming aggregate those rules replace with a raw-key
`Project`). Keeping an aggregate, or a filter between two aggregates, would add
a stateful operator, which a query restarted from an existing checkpoint
rejects (`STREAMING_STATEFUL_OPERATOR_NOT_MATCH_IN_STATE_METADATA`).
- The equi-join and single-column null-aware anti join cases of
`NormalizeFloatingNumbers` now normalize the join keys when either side needs
it, not only the left side. An `IN`/`NOT IN` subquery over a pushed, normalized
grouping key becomes a semi/anti join whose left key is already normalized, and
the subquery side was left raw.
Example:
```
-- SELECT * FROM (SELECT DISTINCT x FROM t) WHERE cast(x AS STRING) = '-0.0'
Aggregate [x], [x]
+- Filter (cast(knownfloatingpointnormalized(normalizenanandzero(x)) as
string) = -0.0)
+- t
```
### Why are the changes needed?
The rule assumes a grouping key has the same value above and below the
`Aggregate`, and normalization breaks that assumption for float/double keys. A
pushed filter can then return a row for which the predicate is false:
```python
spark.createDataFrame([(-0.0,)], "x double").createOrReplaceTempView("t")
spark.sql("SELECT x, cast(x AS STRING) = '-0.0' AS p FROM (SELECT DISTINCT x
FROM t)").collect()
# [Row(x=0.0, p=False)]
spark.sql("SELECT * FROM (SELECT DISTINCT x FROM t) WHERE cast(x AS STRING)
= '-0.0'").collect()
# expected: []
# before: [Row(x=0.0)]
```
The same happens with `GROUP BY`, `INTERSECT ALL`, `EXCEPT`/`subtract` and
`dropDuplicates`. Reading the raw key below the aggregate can also split a
group and change its aggregates:
```sql
-- t.x = [-0.0, 0.0]
SELECT x, count(*) FROM t GROUP BY x HAVING cast(x AS STRING) = '0.0'
-- expected: (0.0, 2)
-- before: (0.0, 1)
```
The alternatives were weaker:
- Keeping such conjuncts above the `Aggregate` is correct, but it gives up
early filtering.
- Never pushing filters on float/double keys would also block plain
comparisons like `x > 5` from reaching scans.
This is the same assumption as in SPARK-60093 (`RemoveRedundantAggregates`),
in a different rule. Because a pushed, normalized filter combined with that
rule also returned rows that fail the predicate, this PR keeps the lower
`Aggregate` when the upper one would read its normalized keys raw, which also
fixes the SPARK-60093 repro for batch queries.
### Does this PR introduce _any_ user-facing change?
Yes, it fixes correctness bugs. Queries that filter an aggregate's
float/double grouping key with a predicate that can tell `-0.0` from `0.0`, or
one NaN from another, now return results consistent with the predicate. Such
predicates are still pushed below the aggregate, now reading the normalized
key. Other filters are pushed exactly as before.
A `DISTINCT`/group-only aggregate on a float/double key over at most one row
or under `LIMIT 1`, or a `DISTINCT` under a duplicate-insensitive aggregate
that reads its key raw, is now kept rather than removed, so it returns the
normalized key like any other aggregate (e.g. `0.0` instead of `-0.0`).
`IN`/`NOT IN` subqueries on float/double keys normalize both sides of the join.
Streaming queries are unchanged.
### How was this patch tested?
New tests:
- `FilterPushdownSuite`:
- the ticket's filter is pushed with the normalized key;
- only keys that can change the result are normalized: within an `OR` and
across `AND` conjuncts, a key under a `Cast` is normalized while a direct
comparison operand is not;
- aliases computed from a key are normalized after alias expansion, while
a plain rename is pushed unchanged;
- comparisons on double/float keys are pushed unchanged, including a FLOAT
key compared with DOUBLE values, `InSet` and an IN subquery;
- a grouping expression (`GROUP BY x * 2.0`) and a struct key are
normalized;
- a filter pushed through two aggregates is not normalized twice, and a
lower key inside a computed upper key is normalized;
- a normalization-sensitive conjunct stays above an aggregate with a
grouped-agg pandas UDF (all-pandas and mixed), while a plain comparison is
pushed below it;
- a filter is pushed below a streaming aggregate unnormalized, as before;
- unit tests for `normalizeGroupingExpressions`, including matching
grouping expressions by expression id rather than name.
- `PredicateSuite`: every comparison that `normalizeGroupingExpressions`
leaves unnormalized, including `InSet` and a FLOAT value widened to DOUBLE,
gives the same result for a value and its normalized form, for `-0.0`, `0.0`,
the canonical NaN, a NaN with a different bit pattern, `1.0` and `null`. The
values come from the input row, so generated code sees their raw bits.
- `NormalizeFloatingPointNumbersSuite`: join and null-aware anti join keys
are normalized when only one side is already normalized.
- `OptimizeOneRowPlanSuite`, `RemoveRedundantAggregatesSuite`,
`LimitPushdownSuite`: aggregates that output normalized float/double keys are
kept (with `LocalLimit(1)` still pushed below under `LIMIT 1`), and still
removed when those keys are not read; a lower aggregate is removed when the
upper one groups by its key directly and kept when the upper one reads the key
raw; streaming aggregates are removed as before. Two existing
`RemoveRedundantAggregatesSuite` tests used a double key (`rand(0)`); they now
use an integer key to keep testing the removal, and the double-key upper case
has its own test.
- `PushdownPredicatesForCTEDefStalenessSuite`: a previously pushed,
normalized filter, and one kept above a pandas UDF aggregate, are removed on
rebuild. The first fails without the `removePushedDownFilter` change.
- `DataFrameAggregateSuite` end-to-end tests:
- `distinct`, `groupBy().count()`, `dropDuplicates`, `intersectAll`,
`except`, a float key, a struct key, the ticket's SQL and an aliased key return
the correct rows, and the split-group `count` case returns `(0.0, 2)`;
- `IN` and `NOT IN` on a double key and on a struct key, with `0.0` vs
`-0.0` across the subquery (the struct key is pushed normalized, so it needs
the join-key change on both join paths);
- one-row, `LIMIT 1`, `limit(1)` and `first()` over a group-only
aggregate, and a `DISTINCT` under a duplicate-insensitive aggregate (including
the SPARK-60093 repro, and the removed case with a single remaining aggregate),
filtered with a sign-sensitive predicate;
- a filter pushed through two aggregates with a computed upper key;
- a comparison on a FLOAT key against a DOUBLE value is pushed to the
Parquet scan.
- `PythonUDFSuite`: a sign-sensitive filter above a grouped-agg pandas UDF
aggregate, including a mixed aggregate whose SQL aggregate is pruned and an
all-pandas one whose pandas output is pruned.
- `StreamingAggregationSuite`: a streaming DISTINCT under a
duplicate-insensitive aggregation, with and without a sign-sensitive filter,
keeps one stateful operator. A query restarted from a checkpoint written with
the previous plan was also checked by hand.
Each new end-to-end test fails without its fix.
Existing suites (all pass; Python tests with Python 3.11, pandas and
pyarrow):
- `build/sbt 'catalyst/test'`
- `build/sbt 'sql/testOnly *SQLQueryTestSuite *DataFrameAggregateSuite
*DataFrameSetOperationsSuite *DataFrameSuite *SubquerySuite *JoinSuite
*SQLQuerySuite *DataFrameWindowFunctionsSuite *RemoveRedundantSortsSuite
*ParquetV1FilterSuite *ParquetV2FilterSuite *PythonUDFSuite
*StreamingAggregationSuite *StreamingDeduplicationSuite
*MultiStatefulOperatorsSuite *CTEInlineSuiteAEOff *CTEInlineSuiteAEOn'`
- `build/sbt 'sql/testOnly *PlanStabilitySuite'`. No golden plan changes.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 5.5)
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Closes #59403 from ctagard/SPARK-60150.
Authored-by: Cole Agard <[email protected]>
Signed-off-by: Holden Karau <[email protected]>
(cherry picked from commit 4eecc14b1a6cd872dd819cb1bf5ff73494f1b6aa)
Signed-off-by: Holden Karau <[email protected]>
--
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]
