HyukjinKwon commented on code in PR #59402:
URL: https://github.com/apache/spark/pull/59402#discussion_r4285028512
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregates.scala:
##########
@@ -223,27 +230,47 @@ object RewriteDistinctAggregates extends
Rule[LogicalPlan] {
case a: Aggregate if mayNeedtoRewrite(a) => rewrite(a)
}
- def rewrite(a: Aggregate): Aggregate = {
+ /**
+ * Groups distinct aggregate expressions by their unfoldable children.
Aggregates that
+ * share a key can share one Expand projection (or the single-distinct
physical path).
+ */
+ private def groupDistinctAggs(
+ aggExpressions: Seq[AggregateExpression]): Map[ExpressionSet,
Seq[AggregateExpression]] = {
+ aggExpressions.filter(_.isDistinct).groupBy { e =>
+ val unfoldableChildren =
ExpressionSet(e.aggregateFunction.children.filter(!_.foldable))
+ if (unfoldableChildren.nonEmpty) {
+ // Only expand the unfoldable children
+ unfoldableChildren
+ } else {
+ // If aggregateFunction's children are all foldable
+ // we must expand at least one of the children (here we take the first
child),
+ // or If we don't, we will get the wrong result, for example:
+ // count(distinct 1) will be explained to count(1) after the rewrite
function.
+ // Generally, the distinct aggregateFunction should not run
+ // foldable TypeCheck for the first child.
+ ExpressionSet(e.aggregateFunction.children.take(1))
+ }
+ }
+ }
+
+ private def numDistinctAggGroups(a: Aggregate): Int =
+ groupDistinctAggs(collectAggregateExprs(a)).size
+ def rewrite(origAgg: Aggregate): Aggregate = {
+ // Keep the normalized form only when it strictly reduces the number of
distinct
+ // groups. Otherwise shapes like COUNT(DISTINCT IF) + SUM(DISTINCT IF)
(same IF)
+ // or a repeated COUNT(DISTINCT IF) next to COUNT(*) would move from the
+ // single-distinct path onto an Expand with no fan-out win.
+ val normalized = normalizeCountDistinctConditional(origAgg)
+ val a =
+ if (numDistinctAggGroups(normalized) < numDistinctAggGroups(origAgg)) {
Review Comment:
**Non-blocking (P2):** `normalizeCountDistinctConditional` rewrites each
eligible COUNT on its own, and this gate only compares total group counts. So a
COUNT whose `IF` is also the distinct child of an aggregate that stays
un-normalized is still split off, as long as other conditional counts make the
total drop. For example:
```sql
SELECT k,
COUNT(DISTINCT IF(c1, x, NULL)), SUM(DISTINCT IF(c1, x, NULL)),
COUNT(DISTINCT IF(c2, x, NULL)), COUNT(DISTINCT IF(c3, x, NULL))
FROM t GROUP BY k
```
With the config off there are three groups, and COUNT and SUM share `{IF(c1,
x, NULL)}`, so the first aggregate holds d1 + d2 + d3 + 3 keys per `k` (where
di = distinct(x | ci)). With it on, the three COUNTs collapse into a guarded
`{x}` group and the SUM keeps `{IF(c1, x, NULL)}`. That is 2 < 3 groups, so the
gate keeps the rewrite, but c1's keys are now carried twice: distinct(x | c1 OR
c2 OR c3) + 1 in the guarded group plus d1 + 1 in the SUM's group. With a broad
c1 over a near-unique `x`, that roughly doubles the partial-aggregate output
and shuffle, while Expand only goes from 3 to 2 projections. Results are the
same, but this shape no longer meets the 'never worse than the uncollapsed
plan' bound in the PR description.
Could `normalizeCountDistinctConditional` skip a COUNT when its `IF`/`CASE`
child (canonicalized) is also the distinct child of an aggregate it won't
normalize? Then the shared group stays intact, and only counts whose original
groups actually disappear get collapsed.
**Verification:**
- **Regression:** For the four-aggregate shared-IF shape, the optimized plan
has one distinct group whose key is the shared IF(c1, x, NULL) serving both
COUNT and SUM, and a guarded {x} group whose Or covers exactly c2 and c3.
- **Behavior:** That query returns the same results with the config on and
off.
- **Compatibility:** Existing coverage for the guard shape, the
strict-reduction gate shapes, subquery exclusions, and raw-key fallbacks still
passes.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregatesSuite.scala:
##########
@@ -125,4 +126,594 @@ 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 = {
+ // rewrite() adds EqualTo(gid, id) to every distinct aggregate, so
filter.isDefined
+ // alone always passes. Require a user-condition reference (same as
assertCanonicalized).
+ val hasUserConditionFilter = collectAggregateExpressions(optimized)
+ .flatMap(_.filter)
+ .exists(_.references.exists(_.name != "gid"))
+ assert(hasUserConditionFilter,
+ "expected at least one AggregateExpression with a user-condition FILTER
clause")
+ val baseAttr = conditionalTestRelation.output.find(_.name ==
baseColName).get
+ assertGuardedExpandKey(optimized, baseAttr, conds)
+ }
+
+ /**
+ * Asserts that the collapsed distinct group's Expand key is guarded by the
union
+ * of the member filters, i.e. IF(Or(cond1, ..., condN), base, NULL): rows
matching
+ * no predicate project NULL as the distinct key and collapse into a single
key in
+ * the first aggregate, keeping its output cardinality bounded by the
pre-rewrite
+ * plan (distinct(base | Or(conds)) + 1 <= sum_i distinct(base | cond_i) +
N).
+ */
+ private def assertGuardedExpandKey(
+ optimized: LogicalPlan,
+ base: Expression,
+ conds: Seq[Expression]): Unit = {
+ val expand = optimized.collectFirst { case e: Expand => e }.get
+ assert(expand.projections.size == 1,
+ s"expected 1 distinct group but got ${expand.projections.size}")
+ val guard = expand.projections.head.collectFirst {
+ case i @ If(_, trueValue, falseValue)
+ if falseValue == nullInt && trueValue.semanticEquals(base) => i
+ }.getOrElse(fail(s"no union guard over $base in Expand
projections:\n$optimized"))
+ val leaves = orLeaves(guard.predicate)
+ assert(leaves.size == conds.size,
+ s"guard condition should cover all ${conds.size} filters:
${guard.predicate}")
+ conds.foreach { c =>
+ assert(leaves.exists(_.semanticEquals(c)),
+ s"guard condition misses $c: ${guard.predicate}")
+ }
+ }
+
+ private def orLeaves(expr: Expression): Seq[Expression] = expr match {
+ case Or(left, right) => orLeaves(left) ++ orLeaves(right)
+ case other => Seq(other)
+ }
+
+ /** EXISTS / IN / NOT IN nodes in Expand projections. The guard must not
copy these. */
+ private def subqueryExprCount(plan: LogicalPlan): Int = {
+ val exprs = plan.collect { case e: Expand => e
}.flatMap(_.projections.flatten)
+ exprs.map(_.collect { case _: Exists | _: InSubquery => () }.size).sum
+ }
+
+ test("conditional: disabled when config is false") {
+ withSQLConf(SQLConf.REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED.key ->
"false") {
+ val input = conditionalTestRelation
+ .groupBy(Symbol("a"))(
+ countDistinctIf(Symbol("b") > 1, Symbol("c")).as("cnt1"),
+ countDistinctIf(Symbol("b") > 2, Symbol("c")).as("cnt2"))
+ .analyze
+ val optimized = RewriteDistinctAggregates(input)
+ // The general RewriteDistinctAggregates still fires, but conditional
+ // canonicalization is disabled so the two conditional counts stay as 2
+ // distinct groups instead of collapsing to 1.
+ val expands = optimized.collect { case e: Expand => e }
+ assert(expands.head.projections.size == 2,
+ "expected 2 distinct groups when conditional canonicalization is
disabled")
+ }
+ }
+
+ test("conditional: rewrite COUNT(DISTINCT IF(cond, col, NULL)) to
COUNT(DISTINCT col) FILTER") {
+ val bAttr = conditionalTestRelation.output.find(_.name == "b").get
+ val input = conditionalTestRelation
+ .groupBy(Symbol("a"))(
+ countDistinctIf(Symbol("b") > 1, Symbol("c")).as("cnt1"),
+ countDistinctIf(Symbol("b") > 2, Symbol("c")).as("cnt2"))
+ .analyze
+ val optimized = RewriteDistinctAggregates(input)
+ assertSingleDistinctGroupExpand(optimized, "c", Seq(bAttr > 1, bAttr > 2))
+ }
+
+ test("conditional: rewrite COUNT(DISTINCT CASE WHEN cond THEN col END) to " +
+ "COUNT(DISTINCT col) FILTER") {
+ val bAttr = conditionalTestRelation.output.find(_.name == "b").get
+ val input = conditionalTestRelation
+ .groupBy(Symbol("a"))(
+ countDistinctCaseWhen(Symbol("b") > 1, Symbol("c")).as("cnt1"),
+ countDistinctCaseWhen(Symbol("b") > 2, Symbol("c")).as("cnt2"))
+ .analyze
+ val optimized = RewriteDistinctAggregates(input)
+ assertSingleDistinctGroupExpand(optimized, "c", Seq(bAttr > 1, bAttr > 2))
+ }
+
+ test("conditional: rewrite COUNT(DISTINCT CASE WHEN cond THEN col ELSE NULL
END)") {
+ val bAttr = conditionalTestRelation.output.find(_.name == "b").get
+ val input = conditionalTestRelation
+ .groupBy(Symbol("a"))(
+ countDistinctCaseWhenElseNull(Symbol("b") > 1, Symbol("c")).as("cnt1"),
+ countDistinctCaseWhenElseNull(Symbol("b") > 2, Symbol("c")).as("cnt2"))
+ .analyze
+ val optimized = RewriteDistinctAggregates(input)
+ assertSingleDistinctGroupExpand(optimized, "c", Seq(bAttr > 1, bAttr > 2))
+ }
+
+ test("conditional: multiple conditional distinct counts collapse to single
distinct group") {
+ val bAttr = conditionalTestRelation.output.find(_.name == "b").get
+ val input = conditionalTestRelation
+ .groupBy(Symbol("a"))(
+ countDistinctIf(Symbol("b") > 1, Symbol("c")).as("cnt1"),
+ countDistinctIf(Symbol("b") > 2, Symbol("c")).as("cnt2"),
+ countDistinctIf(Symbol("b") > 3, Symbol("c")).as("cnt3"))
+ .analyze
+ val optimized = RewriteDistinctAggregates(input)
+ // All three counts share the same base column c, collapsed to 1 distinct
group.
+ assertSingleDistinctGroupExpand(
+ optimized, "c", Seq(bAttr > 1, bAttr > 2, bAttr > 3))
+ }
+
+ test("conditional: guarded Expand key is IF(union of conditions, base,
NULL)") {
+ val bAttr = conditionalTestRelation.output.find(_.name == "b").get
+ val base = conditionalTestRelation.output.find(_.name == "c").get
+ val input = conditionalTestRelation
+ .groupBy(Symbol("a"))(
+ countDistinctIf(Symbol("b") > 1, Symbol("c")).as("cnt1"),
+ countDistinctIf(Symbol("b") > 2, Symbol("c")).as("cnt2"))
+ .analyze
+ val optimized = RewriteDistinctAggregates(input)
+ assertGuardedExpandKey(optimized, base, Seq(bAttr > 1, bAttr > 2))
+ }
+
+ test("conditional: guarded Expand key covers CASE WHEN conditions too") {
+ val bAttr = conditionalTestRelation.output.find(_.name == "b").get
+ val base = conditionalTestRelation.output.find(_.name == "c").get
+ val input = conditionalTestRelation
+ .groupBy(Symbol("a"))(
+ countDistinctCaseWhen(Symbol("b") > 1, Symbol("c")).as("cnt1"),
+ countDistinctCaseWhenElseNull(Symbol("b") > 2, Symbol("c")).as("cnt2"))
+ .analyze
+ val optimized = RewriteDistinctAggregates(input)
+ assertGuardedExpandKey(optimized, base, Seq(bAttr > 1, bAttr > 2))
+ }
+
+ test("conditional: guarded Expand key covers all conditions of a 3-count
collapse") {
+ val bAttr = conditionalTestRelation.output.find(_.name == "b").get
+ val base = conditionalTestRelation.output.find(_.name == "c").get
+ val input = conditionalTestRelation
+ .groupBy(Symbol("a"))(
+ countDistinctIf(Symbol("b") > 1, Symbol("c")).as("cnt1"),
+ countDistinctIf(Symbol("b") > 2, Symbol("c")).as("cnt2"),
+ countDistinctIf(Symbol("b") > 3, Symbol("c")).as("cnt3"))
+ .analyze
+ val optimized = RewriteDistinctAggregates(input)
+ assertGuardedExpandKey(optimized, base, Seq(bAttr > 1, bAttr > 2, bAttr >
3))
+ }
+
+ test("conditional: user-written FILTER counts sharing a base get the union
guard") {
+ val bAttr = conditionalTestRelation.output.find(_.name == "b").get
+ val base = conditionalTestRelation.output.find(_.name == "c").get
+ val input = conditionalTestRelation
+ .groupBy(Symbol("a"))(
+ Count(base).toAggregateExpression(isDistinct = true, filter =
Some(bAttr > 1)).as("cnt1"),
+ Count(base).toAggregateExpression(isDistinct = true, filter =
Some(bAttr > 2)).as("cnt2"))
+ .analyze
+ val optimized = RewriteDistinctAggregates(input)
+ assertGuardedExpandKey(optimized, base, Seq(bAttr > 1, bAttr > 2))
+ }
+
+ test("conditional: config false disables union guard for user-written
FILTER") {
+ // The config is the kill switch for both IF/CASE normalize and the
union-of-filters
+ // guard. With it off, user-written FILTER distincts still Expand
(mustRewrite) but
+ // keep the raw base as the Expand key.
+ withSQLConf(SQLConf.REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED.key ->
"false") {
+ val bAttr = conditionalTestRelation.output.find(_.name == "b").get
+ val base = conditionalTestRelation.output.find(_.name == "c").get
+ val input = conditionalTestRelation
+ .groupBy(Symbol("a"))(
+ Count(base).toAggregateExpression(isDistinct = true, filter =
Some(bAttr > 1))
+ .as("cnt1"),
+ Count(base).toAggregateExpression(isDistinct = true, filter =
Some(bAttr > 2))
+ .as("cnt2"))
+ .analyze
+ val optimized = RewriteDistinctAggregates(input)
+ val expand = optimized.collectFirst { case e: Expand => e }.get
+ assert(expand.projections.size == 1)
+ assert(expand.projections.head.exists(_.semanticEquals(base)),
+ "raw base expected when rewriteCountDistinctConditional is disabled")
+ assert(
+ !expand.projections.head.exists {
+ case If(_, trueValue, falseValue) =>
+ falseValue == nullInt && trueValue.semanticEquals(base)
+ case _ => false
+ },
+ "union guard should be off when rewriteCountDistinctConditional is
disabled")
+ }
+ }
+
+ test("conditional: unfiltered count in the group keeps the raw base key") {
+ val base = conditionalTestRelation.output.find(_.name == "c").get
+ val input = conditionalTestRelation
+ .groupBy(Symbol("a"))(
+ Count(base).toAggregateExpression(isDistinct = true).as("cnt0"),
+ Count(base).toAggregateExpression(isDistinct = true,
+ filter = Some(Symbol("b") > 1)).as("cnt1"),
+ Count(base).toAggregateExpression(isDistinct = true,
+ filter = Some(Symbol("b") > 2)).as("cnt2"))
+ .analyze
+ val optimized = RewriteDistinctAggregates(input)
+ val expand = optimized.collectFirst { case e: Expand => e }.get
+ // One distinct group {c}, but the unfiltered count needs the raw child
over all
+ // rows, so the Expand key must not be guarded.
+ assert(expand.projections.size == 1)
+ assert(expand.projections.head.exists(_.semanticEquals(base)),
+ "raw base attribute expected when the group contains an unfiltered
count")
+ }
+
+ test("conditional: non-deterministic filter keeps the raw base key") {
+ val base = conditionalTestRelation.output.find(_.name == "c").get
+ val nondet = Rand(42) < Literal(0.5, DoubleType)
+ val input = conditionalTestRelation
+ .groupBy(Symbol("a"))(
+ Count(base).toAggregateExpression(isDistinct = true, filter =
Some(nondet)).as("cnt1"),
+ Count(base).toAggregateExpression(isDistinct = true, filter =
Some(nondet)).as("cnt2"))
+ .analyze
+ val optimized = RewriteDistinctAggregates(input)
+ val expand = optimized.collectFirst { case e: Expand => e }.get
+ // The guard re-evaluates the filters, so non-deterministic ones must not
be
+ // guarded; the Expand key stays the raw attribute.
+ assert(expand.projections.head.exists(_.semanticEquals(base)),
+ "raw base attribute expected with a non-deterministic filter")
+ }
+
+ test("conditional: EXISTS, IN, and NOT IN filters keep the raw key") {
+ // RewritePredicateSubquery has not run. Copying a predicate subquery into
the
+ // Expand key builds a second ExistenceJoin later, so the key stays raw
and each
+ // predicate appears once, as its flag column.
+ val dim = LocalRelation(Symbol("x").int)
+ val existsPred = Exists(dim)
+ val inPred = InSubquery(Seq(Symbol("b")), ListQuery(dim, numCols = 1))
+ val notInPred = Not(InSubquery(Seq(Symbol("c")), ListQuery(dim, numCols =
1)))
+ val b = conditionalTestRelation.output.find(_.name == "b").get
+ val c = conditionalTestRelation.output.find(_.name == "c").get
+ val d = conditionalTestRelation.output.find(_.name == "d").get
+ val input = conditionalTestRelation
+ .groupBy(Symbol("a"))(
+ Count(b).toAggregateExpression(isDistinct = true, filter =
Some(existsPred)).as("cnt1"),
+ Count(c).toAggregateExpression(isDistinct = true, filter =
Some(inPred)).as("cnt2"),
+ Count(d).toAggregateExpression(isDistinct = true, filter =
Some(notInPred)).as("cnt3"))
+ .analyze
+ val optimized = RewriteDistinctAggregates(input)
+ val expand = optimized.collectFirst { case e: Expand => e }.get
+ assert(expand.projections.size == 3)
+ Seq(b, c, d).foreach { base =>
+ val nullOfBase = Literal.create(null, base.dataType)
+ assert(expand.projections.exists(_.exists(_.semanticEquals(base))),
+ s"raw key expected for $base")
+ assert(!expand.projections.exists(_.exists {
+ case If(_, trueValue, falseValue) =>
+ falseValue.semanticEquals(nullOfBase) &&
trueValue.semanticEquals(base)
+ case _ => false
+ }), s"subquery filter must not guard $base")
+ }
+ assert(subqueryExprCount(optimized) == 3,
+ "each predicate subquery should appear once in the Expand")
+ }
+
+ test("conditional: do not canonicalize counts whose condition is a predicate
subquery") {
+ // Normalizing would merge these into one group that cannot be guarded,
and Expand
+ // would project the raw base. Leave each IF/CASE in its own group.
+ val dim = LocalRelation(Symbol("x").int)
+ val x = dim.output.head
+ val input = conditionalTestRelation
+ .groupBy(Symbol("a"))(
+ countDistinctIf(Exists(dim.where(x > 0)), Symbol("c")).as("cnt1"),
+ countDistinctCaseWhen(Exists(dim.where(x > 1)),
Symbol("c")).as("cnt2"),
+ countDistinctCaseWhenElseNull(Exists(dim.where(x > 2)),
Symbol("c")).as("cnt3"))
+ .analyze
+ val optimized = RewriteDistinctAggregates(input)
+ val expand = optimized.collectFirst { case e: Expand => e }.get
+ assert(expand.projections.size == 3,
+ "subquery conditions should stay 3 distinct groups")
+ val hasUserConditionFilter = collectAggregateExpressions(optimized)
+ .flatMap(_.filter)
+ .exists(_.references.exists(_.name != "gid"))
+ assert(!hasUserConditionFilter,
+ "subquery condition should not become a user-condition FILTER clause")
+ assert(subqueryExprCount(optimized) == 3,
+ "each predicate subquery should appear once in the Expand")
+ }
+
+ test("conditional: subquery-free filters stay guarded beside a subquery
condition") {
+ val bAttr = conditionalTestRelation.output.find(_.name == "b").get
+ val base = conditionalTestRelation.output.find(_.name == "c").get
+ val dim = LocalRelation(Symbol("x").int)
+ val input = conditionalTestRelation
+ .groupBy(Symbol("a"))(
+ countDistinctIf(Symbol("b") > 1, Symbol("c")).as("cnt1"),
+ countDistinctIf(Symbol("b") > 2, Symbol("c")).as("cnt2"),
+ countDistinctIf(Exists(dim), Symbol("c")).as("cnt3"))
+ .analyze
+ val optimized = RewriteDistinctAggregates(input)
+ val expand = optimized.collectFirst { case e: Expand => e }.get
+ // The two plain predicates share one guarded group. The EXISTS count
stays separate.
+ assert(expand.projections.size == 2)
+ val guard = expand.projections.flatten.collectFirst {
+ case i @ If(_, trueValue, falseValue)
+ if falseValue == nullInt && trueValue.semanticEquals(base) => i
+ }.getOrElse(fail(s"no union guard over c:\n$optimized"))
+ val leaves = orLeaves(guard.predicate)
+ assert(leaves.size == 2, s"guard should cover only the plain predicates:
${guard.predicate}")
+ Seq(bAttr > 1, bAttr > 2).foreach { cond =>
+ assert(leaves.exists(_.semanticEquals(cond)), s"guard misses $cond:
${guard.predicate}")
+ }
+ assert(subqueryExprCount(optimized) == 1,
+ "the EXISTS condition should appear once, in its own group")
+ }
+
+ test("conditional: throwable distinct child keeps the raw key") {
+ // Division can throw under ANSI on rows the filter rejects. Wrapping it in
+ // the guard would skip those evaluations, so the Expand key stays raw.
+ val input = conditionalTestRelation
+ .groupBy(Symbol("a"))(
+ Count(Symbol("c") / Symbol("b")).toAggregateExpression(
+ isDistinct = true, filter = Some(Symbol("a") > 0)).as("cnt"))
+ .analyze
+ val child = input.collectFirst {
+ case a: Aggregate =>
+ a.aggregateExpressions.flatMap(_.collect {
+ case ae: AggregateExpression if ae.isDistinct =>
+ ae.aggregateFunction.children.headOption
+ }).flatten.head
+ }.get
+ val optimized = RewriteDistinctAggregates(input)
+ val expand = optimized.collectFirst { case e: Expand => e }.get
+ assert(expand.projections.head.exists(_.semanticEquals(child)),
+ s"raw distinct child expected, projections:\n${expand.projections.head}")
+ }
+
+ test("conditional: single conditional distinct count is gated out by
mayNeedtoRewrite") {
+ // A lone COUNT(DISTINCT IF(...)) must NOT be pushed onto the Expand path.
+ // The canonicalization runs inside rewrite(), which is only called when
+ // mayNeedtoRewrite returns true. A single conditional distinct count has
no filter
+ // and forms only one distinct group, so mayNeedtoRewrite returns false,
rewrite()
+ // is never called, and no Expand is produced.
+ val input = conditionalTestRelation
+ .groupBy(Symbol("a"))(
+ countDistinctIf(Symbol("b") > 1, Symbol("c")).as("cnt1"))
+ .analyze
+ val optimized = RewriteDistinctAggregates(input)
+ val expands = optimized.collect { case e: Expand => e }
+ assert(expands.isEmpty, "single conditional distinct count should not
produce an Expand")
+ }
+
+ test("conditional: COUNT+SUM sharing IF keeps single-distinct path") {
+ // Both share one unfiltered distinct child. Normalizing only the COUNT
would
+ // split the group into {c} and {IF(b>1,c,NULL)} and force a 2-projection
Expand.
+ // Keep the original aggregate so AggregationStrategy handles it without
Expand.
+ // Use the same null literal as countDistinctIf so the IF children are one
+ // ExpressionSet key after analysis (as in SQL COUNT/SUM over the same IF).
+ val sharedIf = If(Symbol("b") > 1, Symbol("c"), Literal(null))
+ val input = conditionalTestRelation
+ .groupBy(Symbol("a"))(
+ Count(sharedIf).toAggregateExpression(isDistinct = true).as("cnt"),
+ Sum(sharedIf).toAggregateExpression(isDistinct = true).as("sum"))
+ .analyze
+ val optimized = RewriteDistinctAggregates(input)
+ comparePlans(optimized, input)
+ }
+
+ test("conditional: repeated COUNT(DISTINCT IF) with COUNT(*) keeps
single-distinct path") {
+ // Two identical conditional counts form one unfiltered distinct group.
Normalizing
+ // them to FILTER form would make mustRewrite true and introduce an Expand
+ // (regular branch + distinct branch) with no fan-out reduction.
+ val input = conditionalTestRelation
+ .groupBy(Symbol("a"))(
+ countDistinctIf(Symbol("b") > 1, Symbol("c")).as("cnt"),
+ countDistinctIf(Symbol("b") > 1, Symbol("c")).as("cnt2"),
+ Count(Literal(1)).toAggregateExpression().as("n"))
+ .analyze
+ val optimized = RewriteDistinctAggregates(input)
+ comparePlans(optimized, input)
+ }
+
+ test("conditional: do not rewrite IF with non-null else branch") {
+ val input = conditionalTestRelation
+ .groupBy(Symbol("a"))(
+ Count(If(Symbol("b") > 1, Symbol("c"), Literal(0, IntegerType)))
+ .toAggregateExpression(isDistinct = true)
+ .as("cnt1"),
+ Count(If(Symbol("b") > 2, Symbol("c"), Literal(0, IntegerType)))
+ .toAggregateExpression(isDistinct = true)
+ .as("cnt2"))
+ .analyze
+ val optimized = RewriteDistinctAggregates(input)
+ // Plan is rewritten (2 distinct groups) but not canonicalized
+ checkRewrite(optimized)
+ val expands = optimized.collect { case e: Expand => e }
+ assert(expands.head.projections.size == 2,
+ "non-null else branch should not be collapsed to 1 distinct group")
+ }
+
+ test("conditional: do not rewrite non-distinct COUNT") {
Review Comment:
**Nit (P3):** With one distinct aggregate (`countDistinct(c)`) and no
FILTER, `mayNeedtoRewrite` returns false, so `rewrite()` and
`normalizeCountDistinctConditional` never run, and `comparePlans(optimized,
input)` passes no matter what the pattern does with a non-distinct
`COUNT(IF(...))`. Adding only a second distinct aggregate on another column
wouldn't help either, because the strict-reduction gate discards a
normalization that doesn't reduce the distinct-group count. If the fixture
paired the non-distinct COUNT with two conditional distinct counts on one base,
so that normalization is kept, and asserted that the non-distinct COUNT still
has its `IF` child and no user-condition FILTER, the check could actually fail.
**Verification:**
- **Regression:** The optimized plan shows the conditional distinct counts
collapsed into a guarded group while the non-distinct COUNT keeps its IF child
and has no user-condition FILTER.
- **Inspection:** The fixture reaches rewrite() with normalization kept, so
dropping the isDistinct restriction from the pattern would change the asserted
plan.
--
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]
