cloud-fan commented on code in PR #59402:
URL: https://github.com/apache/spark/pull/59402#discussion_r4256199476
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregates.scala:
##########
@@ -223,8 +227,9 @@ object RewriteDistinctAggregates extends Rule[LogicalPlan] {
case a: Aggregate if mayNeedtoRewrite(a) => rewrite(a)
}
- def rewrite(a: Aggregate): Aggregate = {
+ def rewrite(origAgg: Aggregate): Aggregate = {
+ val a = normalizeCountDistinctConditional(origAgg)
Review Comment:
**Non-blocking (P2):** `normalizeCountDistinctConditional` runs on every
aggregate that reaches `rewrite()` and rewrites each matching COUNT on its own.
Nothing checks that the rewrite reduces the number of distinct groups, and two
shapes get worse than before:
- `COUNT(DISTINCT IF(c, x, NULL))` next to `SUM(DISTINCT IF(c, x, NULL))`.
Before, both shared one unfiltered group, and `rewrite()` returned the
aggregate unchanged (single-distinct planning, no Expand). Now only the COUNT
is normalized, so the groups become `{x}` and `{IF(c, x, NULL)}` and the plan
gets a 2-projection Expand.
- A repeated `COUNT(DISTINCT IF(c, x, NULL))`, e.g. once as a value and once
in a ratio over `COUNT(*)`. Before, this was one unfiltered group with no
Expand. Now the FILTER makes `mustRewrite` true, so the plan gets an Expand
with a regular branch and a distinct branch.
Results are the same, but every input row is now duplicated where it wasn't
before. Counts over different bases also pick up flag columns and `max` buffers
with no fan-out reduction. Could we keep the normalized aggregate only when it
has strictly fewer distinct groups than `origAgg` (using the same
`ExpressionSet` keying), and otherwise continue with `origAgg`?
This extends an [existing
discussion](https://github.com/apache/spark/pull/59402#issuecomment-6124805791).
**Verification:**
- **Regression:** COUNT(DISTINCT IF(c, x, NULL)) together with SUM(DISTINCT
IF(c, x, NULL)), and a repeated COUNT(DISTINCT IF(c, x, NULL)) next to
COUNT(*), produce no Expand, matching the merge-base plan shape.
- **Behavior:** Two or more conditional distinct counts over one base still
produce a single distinct Expand projection with the guarded key, and the
results match the config-off run.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregates.scala:
##########
@@ -419,6 +446,63 @@ object RewriteDistinctAggregates extends Rule[LogicalPlan]
{
}
}
+ /**
+ * Canonicalizes COUNT(DISTINCT IF(cond, base, NULL)) and
+ * COUNT(DISTINCT CASE WHEN cond THEN base END) to COUNT(DISTINCT base)
FILTER (WHERE cond).
+ * This reduces the number of distinct groups: multiple conditional counts
on the same base
+ * column collapse into one group, shrinking the Expand fan-out from Nx to
1x.
+ *
+ * Note that the rewrite moves `base` out of the protective conditional
branch: after the
+ * rewrite the Expand operator evaluates the distinct child for every input
row, while
Review Comment:
**Nit (P3):** This describes the pre-guard plan. With the union-of-filters
guard, an all-filtered group projects `IF(Or(conds), base, NULL)`. The base is
then evaluated only on rows matching some member's condition, which are the
same rows the per-condition `IF`s evaluated it on before. Only a group that
also has an unfiltered aggregate on the same base projects the raw base for
every row, and the uncollapsed plan already did that too. Could you reword this
to describe the evaluation domain the guard actually produces, and say why
`canEvaluateUnconditionally` is still required under it? As written, it gives
the wrong reason for the restriction.
**Verification:**
- **Inspection:** The scaladoc's account of base evaluation matches the
guard's applicability conditions and the normalization's admission rules.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregatesSuite.scala:
##########
@@ -125,4 +126,418 @@ class RewriteDistinctAggregatesSuite extends PlanTest {
fail(s"Plan is not as expected:\n$rewrite")
}
}
+
+ //
---------------------------------------------------------------------------
+ // COUNT(DISTINCT IF/CASE) canonicalization (SPARK-56898)
+ //
---------------------------------------------------------------------------
+
+ val conditionalTestRelation = LocalRelation(
+ Symbol("a").int, Symbol("b").int, Symbol("c").int, Symbol("d").string)
+
+ private def countDistinctIf(cond: Expression, base: Expression): Expression
= {
+ Count(If(cond, base, Literal(null))).toAggregateExpression(isDistinct =
true)
+ }
+
+ private def countDistinctCaseWhen(cond: Expression, base: Expression):
Expression = {
+ val caseWhen = CaseWhen(
+ Seq((cond, base)),
+ None)
+ Count(caseWhen).toAggregateExpression(isDistinct = true)
+ }
+
+ private def countDistinctCaseWhenElseNull(cond: Expression, base:
Expression): Expression = {
+ val caseWhen = CaseWhen(
+ Seq((cond, base)),
+ Some(Literal(null)))
+ Count(caseWhen).toAggregateExpression(isDistinct = true)
+ }
+
+ private def collectAggregateExpressions(plan: LogicalPlan):
Seq[AggregateExpression] = {
+ plan.collect { case a: Aggregate => a.aggregateExpressions }
+ .flatten
+ .flatMap(_.collect { case ae: AggregateExpression => ae })
+ }
+
+ /**
+ * Asserts that the optimized plan has exactly one Expand node with one
projection,
+ * that the shared distinct key is guarded by the union of the member
filters,
+ * and that the outer aggregate has moved the conditions into FILTER clauses.
+ */
+ private def assertSingleDistinctGroupExpand(
+ optimized: LogicalPlan,
+ baseColName: String,
+ conds: Seq[Expression]): Unit = {
+ assert(collectAggregateExpressions(optimized).exists(_.filter.isDefined),
Review Comment:
**Nit (P3):** `rewrite()` gives every distinct aggregate an `EqualTo(gid,
id)` filter, so `exists(_.filter.isDefined)` is true for any rewritten plan.
The doc says this checks that the conditions were moved into FILTER clauses.
But if each aggregate's own `max_cond` were dropped from its final FILTER, all
four tests that use this helper would still pass. `assertCanonicalized` below
already excludes gid-only filters (`_.references.exists(_.name != "gid")`).
Doing the same here, or matching each condition's flag attribute, would let
this check actually fail.
**Verification:**
- **Regression:** Collapsed conditional counts each keep a final FILTER that
combines gid routing with their own condition's max flag, and a plan with only
gid routing filters fails the helper.
##########
sql/core/src/test/scala/org/apache/spark/sql/RewriteDistinctAggregatesConditionalQuerySuite.scala:
##########
@@ -0,0 +1,412 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql
+
+import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+
+class RewriteDistinctAggregatesConditionalQuerySuite extends QueryTest with
SharedSparkSession {
+
+ private def checkRewriteAndResult(
+ conditionalSql: String,
+ filterSql: String): Unit = {
+ val expectedRows = spark.sql(filterSql).collect()
+
+ // Verify the rewrite produces the same result as the explicit FILTER form.
+ withSQLConf(
+ SQLConf.REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED.key -> "true") {
+ checkAnswer(spark.sql(conditionalSql), expectedRows)
+ }
+
+ // Verify the non-rewritten form also matches the explicit FILTER form.
+ withSQLConf(
+ SQLConf.REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED.key -> "false") {
+ checkAnswer(spark.sql(conditionalSql), expectedRows)
+ }
+ }
+
+ private def hasUserConditionFilter(plan: LogicalPlan): Boolean = {
+ val aggregateExpressions = plan.collect {
+ case a: Aggregate => a.aggregateExpressions
+ }.flatten
+ val aggregateExprs = aggregateExpressions.flatMap {
+ _.collect { case ae: AggregateExpression => ae }
+ }
+ // The Expand rewrite adds gid-routing filters to every distinct
aggregate; we only
+ // care about user-condition filters introduced by the IF/CASE
canonicalization.
+ aggregateExprs.flatMap(_.filter).exists(_.references.exists(_.name !=
"gid"))
+ }
+
+ test("rewrite COUNT(DISTINCT IF(cond, col, NULL)) correctness") {
+ withTempView("t") {
+ spark.range(7)
+ .selectExpr(
+ "cast(id % 3 + 1 as int) as key",
+ "cast(id * 10 as int) as col1",
+ "case when id % 4 = 0 then null else cast(id * 100 as int) end as
col2")
+ .createOrReplaceTempView("t")
+
+ // Two conditional distinct counts on the same base so the rewrite fires
(size > 1).
+ checkRewriteAndResult(
+ """SELECT key,
+ | COUNT(DISTINCT IF(col1 > 10, col2, NULL)),
+ | COUNT(DISTINCT IF(col1 > 20, col2, NULL))
+ |FROM t GROUP BY key""".stripMargin,
+ """SELECT key,
+ | COUNT(DISTINCT col2) FILTER (WHERE col1 > 10),
+ | COUNT(DISTINCT col2) FILTER (WHERE col1 > 20)
+ |FROM t GROUP BY key""".stripMargin)
+ }
+ }
+
+ test("rewrite COUNT(DISTINCT CASE WHEN cond THEN col END) correctness") {
+ withTempView("t") {
+ spark.range(7)
+ .selectExpr(
+ "cast(id % 3 + 1 as int) as key",
+ "cast(id * 10 as int) as col1",
+ "case when id % 4 = 0 then null else cast(id * 100 as string) end as
col2")
+ .createOrReplaceTempView("t")
+
+ // Two CASE WHEN conditional distinct counts on the same base so the
rewrite fires.
+ checkRewriteAndResult(
+ """SELECT key,
+ | COUNT(DISTINCT CASE WHEN col1 > 10 THEN col2 END),
+ | COUNT(DISTINCT CASE WHEN col1 > 20 THEN col2 END)
+ |FROM t GROUP BY key""".stripMargin,
+ """SELECT key,
+ | COUNT(DISTINCT col2) FILTER (WHERE col1 > 10),
+ | COUNT(DISTINCT col2) FILTER (WHERE col1 > 20)
+ |FROM t GROUP BY key""".stripMargin)
+ }
+ }
+
+ test("rewrite COUNT(DISTINCT CASE WHEN cond THEN col ELSE NULL END)
correctness") {
+ withTempView("t") {
+ spark.range(6)
+ .selectExpr(
+ "cast(id % 2 + 1 as int) as key",
+ "cast(id * 10 as int) as col1",
+ "case when id % 4 = 0 then null else cast(id * 1.0 as double) end as
col2")
+ .createOrReplaceTempView("t")
+
+ // Two CASE WHEN ... ELSE NULL counts on the same base so the rewrite
fires.
+ checkRewriteAndResult(
+ """SELECT key,
+ | COUNT(DISTINCT CASE WHEN col1 > 10 THEN col2 ELSE NULL END),
+ | COUNT(DISTINCT CASE WHEN col1 > 20 THEN col2 ELSE NULL END)
+ |FROM t GROUP BY key""".stripMargin,
+ """SELECT key,
+ | COUNT(DISTINCT col2) FILTER (WHERE col1 > 10),
+ | COUNT(DISTINCT col2) FILTER (WHERE col1 > 20)
+ |FROM t GROUP BY key""".stripMargin)
+ }
+ }
+
+ test("rewrite with no GROUP BY") {
+ withTempView("t") {
+ spark.range(5)
+ .selectExpr(
+ "cast(id * 10 as int) as col1",
+ "case when id % 3 = 0 then null else cast(id * 100 as int) end as
col2")
+ .createOrReplaceTempView("t")
+
+ // Two counts so the rewrite fires without GROUP BY.
+ checkRewriteAndResult(
+ """SELECT
+ | COUNT(DISTINCT IF(col1 > 10, col2, NULL)),
+ | COUNT(DISTINCT IF(col1 > 20, col2, NULL))
+ |FROM t""".stripMargin,
+ """SELECT
+ | COUNT(DISTINCT col2) FILTER (WHERE col1 > 10),
+ | COUNT(DISTINCT col2) FILTER (WHERE col1 > 20)
+ |FROM t""".stripMargin)
+ }
+ }
+
+ test("rewrite with all NULLs in conditional branch") {
+ withTempView("t") {
+ spark.range(3)
+ .selectExpr(
+ "cast(id % 2 + 1 as int) as key",
+ "cast(id * 5 as int) as col1",
+ "cast(id * 100 as int) as col2")
+ .createOrReplaceTempView("t")
+
+ // col1 values are 0, 5, 10 - both thresholds (> 10 and > 20) yield zero
matches,
+ // so both counts return 0. Two counts so the rewrite fires.
+ checkRewriteAndResult(
+ """SELECT key,
+ | COUNT(DISTINCT IF(col1 > 10, col2, NULL)),
+ | COUNT(DISTINCT IF(col1 > 20, col2, NULL))
+ |FROM t GROUP BY key""".stripMargin,
+ """SELECT key,
+ | COUNT(DISTINCT col2) FILTER (WHERE col1 > 10),
+ | COUNT(DISTINCT col2) FILTER (WHERE col1 > 20)
+ |FROM t GROUP BY key""".stripMargin)
+ }
+ }
+
+ test("rewrite with duplicates in base column") {
+ withTempView("t") {
+ spark.range(6)
+ .selectExpr(
+ "cast(id % 2 + 1 as int) as key",
+ "cast(id * 10 as int) as col1",
+ "case when id % 3 = 0 then 100 when id % 3 = 1 then 100 else 200 end
as col2")
+ .createOrReplaceTempView("t")
+
+ // Two counts on the same base (with duplicates) so the rewrite fires.
+ checkRewriteAndResult(
+ """SELECT key,
+ | COUNT(DISTINCT IF(col1 > 10, col2, NULL)),
+ | COUNT(DISTINCT IF(col1 > 20, col2, NULL))
+ |FROM t GROUP BY key""".stripMargin,
+ """SELECT key,
+ | COUNT(DISTINCT col2) FILTER (WHERE col1 > 10),
+ | COUNT(DISTINCT col2) FILTER (WHERE col1 > 20)
+ |FROM t GROUP BY key""".stripMargin)
+ }
+ }
+
+ test("multiple conditional distinct counts collapse and produce correct
results") {
+ withTempView("t") {
+ spark.range(5)
+ .selectExpr(
+ "cast(id % 2 + 1 as int) as key",
+ "cast(id * 10 as int) as col1",
+ "case when id % 3 = 0 then null else cast(id * 100 as int) end as
col2",
+ "case when id % 4 = 0 then null else cast(id * 10 as string) end as
col3")
+ .createOrReplaceTempView("t")
+
+ val conditionalSql =
+ """SELECT key,
+ | COUNT(DISTINCT IF(col1 > 10, col2, NULL)) as cnt1,
+ | COUNT(DISTINCT IF(col1 > 5, col3, NULL)) as cnt2
+ |FROM t GROUP BY key""".stripMargin
+
+ val filterSql =
+ """SELECT key,
+ | COUNT(DISTINCT col2) FILTER (WHERE col1 > 10) as cnt1,
+ | COUNT(DISTINCT col3) FILTER (WHERE col1 > 5) as cnt2
+ |FROM t GROUP BY key""".stripMargin
+
+ checkRewriteAndResult(conditionalSql, filterSql)
+ }
+ }
+
+ test("rewrite does not affect COUNT(DISTINCT IF(cond, col, non_null))") {
+ withTempView("t") {
+ spark.range(3)
+ .selectExpr(
+ "cast(id % 2 + 1 as int) as key",
+ "cast(id * 10 as int) as col1",
+ "cast(id * 100 as int) as col2")
+ .createOrReplaceTempView("t")
+
+ // Two counts so mayNeedtoRewrite fires and the non-null-else guard is
exercised.
+ val sqlText =
+ """SELECT key,
+ | COUNT(DISTINCT IF(col1 > 10, col2, 0)),
+ | COUNT(DISTINCT IF(col1 > 20, col2, 0))
+ |FROM t GROUP BY key""".stripMargin
+
+ // Collect the conf-off result as the semantic baseline; the conf-on
query must match it.
+ val withoutRewriteRows = withSQLConf(
+ SQLConf.REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED.key -> "false") {
+ spark.sql(sqlText).collect()
+ }
+ withSQLConf(
+ SQLConf.REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED.key -> "true") {
+ checkAnswer(spark.sql(sqlText), withoutRewriteRows)
+ }
+
+ // When the switch is enabled, the non-null ELSE branch must not be
canonicalized
+ // into a user-condition FILTER; only the internal gid-routing filters
introduced
+ // by the Expand rewrite may be present.
+ val hasFilter = withSQLConf(
+ SQLConf.REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED.key -> "true") {
+ hasUserConditionFilter(spark.sql(sqlText).queryExecution.optimizedPlan)
+ }
+ assert(!hasFilter,
+ "Non-null ELSE branch should not produce a user-condition FILTER
clause")
+ }
+ }
+
+ test("rewrite is not applied to error-prone base expressions (ANSI divide by
zero)") {
+ withTempView("t") {
+ // Rows with x = 0 make both conditions false, so the IF branches
short-circuit
+ // and 10 / x is never evaluated for them.
+ spark.range(5)
+ .selectExpr(
+ "cast(id % 2 + 1 as int) as key",
+ "cast(id % 2 as int) as x")
+ .createOrReplaceTempView("t")
+
+ // Two conditional distinct counts on the same error-prone base so that
rewrite()
+ // is reached; the canonicalization must NOT fire for such a base.
+ val sqlText =
+ """SELECT key,
+ | COUNT(DISTINCT IF(x <> 0, 10 / x, NULL)),
+ | COUNT(DISTINCT IF(x > 0, 10 / x, NULL))
+ |FROM t GROUP BY key""".stripMargin
+
+ withSQLConf(
+ SQLConf.ANSI_ENABLED.key -> "true",
+ SQLConf.REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED.key -> "true") {
+ // Must not raise DIVIDE_BY_ZERO: the canonicalization must leave the
IFs intact
+ // so that their branch short-circuiting is preserved.
+ checkAnswer(spark.sql(sqlText), Seq(Row(1, 0L, 0L), Row(2, 1L, 1L)))
+
assert(!hasUserConditionFilter(spark.sql(sqlText).queryExecution.optimizedPlan),
+ "error-prone base should not produce a user-condition FILTER clause")
+ }
+ }
+ }
+
+ test("guard: high-cardinality base with low-selectivity disjoint
predicates") {
+ withTempView("t") {
+ // ord is unique (high-cardinality base); each predicate matches ~1% of
rows and
+ // they are disjoint. This is the workload shape that regressed before
the
+ // union-of-filters guard: the guarded key must stay bounded by the
pre-rewrite
+ // plan's cardinality while the Expand stays collapsed.
+ spark.range(1000)
+ .selectExpr(
+ "cast(id % 7 as int) as key",
+ "cast(id as bigint) as ord",
+ "cast(id % 100 as int) as bucket")
+ .createOrReplaceTempView("t")
+
+ checkRewriteAndResult(
+ """SELECT key,
+ | COUNT(DISTINCT IF(bucket = 0, ord, NULL)),
+ | COUNT(DISTINCT IF(bucket = 3, ord, NULL))
+ |FROM t GROUP BY key""".stripMargin,
+ """SELECT key,
+ | COUNT(DISTINCT ord) FILTER (WHERE bucket = 0),
+ | COUNT(DISTINCT ord) FILTER (WHERE bucket = 3)
+ |FROM t GROUP BY key""".stripMargin)
+ }
+ }
+
+ test("guard: mixed overlapping windows and disjoint flags") {
+ withTempView("t") {
+ spark.range(1000)
+ .selectExpr(
+ "cast(id % 7 as int) as key",
+ "cast(id as bigint) as ord",
+ "cast(id % 30 as int) as day",
+ "cast(id % 100 as int) as bucket")
+ .createOrReplaceTempView("t")
+
+ // Overlapping windows (day < 10 is a subset of day < 30) plus disjoint
selective flags in
+ // the same query: the union guard must keep every count correct.
+ checkRewriteAndResult(
+ """SELECT key,
+ | COUNT(DISTINCT IF(day < 10, ord, NULL)),
+ | COUNT(DISTINCT IF(day < 30, ord, NULL)),
+ | COUNT(DISTINCT IF(bucket = 0, ord, NULL)),
+ | COUNT(DISTINCT IF(bucket = 7, ord, NULL))
+ |FROM t GROUP BY key""".stripMargin,
+ """SELECT key,
+ | COUNT(DISTINCT ord) FILTER (WHERE day < 10),
+ | COUNT(DISTINCT ord) FILTER (WHERE day < 30),
+ | COUNT(DISTINCT ord) FILTER (WHERE bucket = 0),
+ | COUNT(DISTINCT ord) FILTER (WHERE bucket = 7)
+ |FROM t GROUP BY key""".stripMargin)
+ }
+ }
+
+ test("guard: NULL condition keeps the count empty") {
+ withTempView("t") {
+ spark.range(50)
+ .selectExpr(
+ "cast(id % 7 as int) as key",
+ "cast(id as bigint) as ord",
+ "cast(id % 100 as int) as bucket")
+ .createOrReplaceTempView("t")
+
+ // A NULL condition is never true; with the union guard its rows may
share the
+ // Expand key of the other count, but must never be counted. The FILTER
syntax
+ // cannot express a NULL condition (non-boolean type is rejected), so
compare
+ // the rewritten plan against the conf-off result as the semantic
baseline.
+ val sqlText =
+ """SELECT key,
+ | COUNT(DISTINCT IF(NULL, ord, NULL)),
Review Comment:
**Nit (P3):** `IF(NULL, ord, NULL)` is gone before this rule runs.
`SimplifyConditionals` folds `If(null, _, f)` to `f` during operator
optimization, which comes before the Distinct Aggregate Rewrite batch. When
`RewriteDistinctAggregates` runs, this aggregate is already `COUNT(DISTINCT
CAST(NULL AS BIGINT))`, a separate foldable group, so the guarded group only
contains `bucket = 0`. The test passes however the guard handles NULL
predicates, and every condition column in the suite is non-null. A predicate
over a nullable column that is NULL for some rows, e.g. `IF(nullable_col > 0,
ord, NULL)`, would actually reach `Or(...)` and `max(flag)`. Asserting the
guarded key in the optimized plan would also stop the case from being folded
away again.
**Verification:**
- **Regression:** A conditional distinct count whose condition is NULL for
some rows, collapsed with another conditional count on the same base, returns
the config-off result, and the optimized plan shows the guarded key over both
conditions.
--
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]
