This is an automated email from the ASF dual-hosted git repository. cloud-fan pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/spark.git
commit e3e7720cf68db7a8837ba8f0b922fb9dcdbdb5e7 Author: Wenchen Fan <[email protected]> AuthorDate: Tue Sep 29 07:11:37 2026 +0000 Revert "[SPARK-56898][SQL] Rewrite COUNT(DISTINCT IF) to COUNT(DISTINCT) FILTER for Expand reduction" This reverts commit 9deb4037705a6a6e8d548ee803d323ddf75c693b. --- .../optimizer/RewriteDistinctAggregates.scala | 63 +---- .../org/apache/spark/sql/internal/SQLConf.scala | 18 -- .../optimizer/RewriteDistinctAggregatesSuite.scala | 308 +-------------------- ...teDistinctAggregatesConditionalQuerySuite.scala | 307 -------------------- 4 files changed, 5 insertions(+), 691 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregates.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregates.scala index b33098892746..5aef82b64ed3 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregates.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregates.scala @@ -21,8 +21,7 @@ import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.expressions.aggregate._ import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Expand, LogicalPlan} import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.catalyst.trees.TreePattern.{AGGREGATE, CASE_WHEN, IF} -import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.catalyst.trees.TreePattern.AGGREGATE import org.apache.spark.sql.types.IntegerType import org.apache.spark.util.collection.Utils @@ -224,9 +223,8 @@ object RewriteDistinctAggregates extends Rule[LogicalPlan] { case a: Aggregate if mayNeedtoRewrite(a) => rewrite(a) } - def rewrite(origAgg: Aggregate): Aggregate = { + def rewrite(a: Aggregate): Aggregate = { - val a = normalizeCountDistinctConditional(origAgg) val aggExpressions = collectAggregateExprs(a) val distinctAggs = aggExpressions.filter(_.isDistinct) @@ -421,63 +419,6 @@ 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 - * originally it was only evaluated on rows where `cond` holds. To preserve the - * short-circuit semantics of IF/CASE WHEN, the rewrite is restricted to base - * expressions that can be evaluated unconditionally, i.e. cannot raise errors or change - * results when evaluated on extra rows (see [[ExprUtils.canEvaluateUnconditionally]]). - */ - private def normalizeCountDistinctConditional(a: Aggregate): Aggregate = { - if (!SQLConf.get.rewriteCountDistinctConditionalEnabled) return a - a.transformExpressionsUpWithPruning( - _.containsAnyPattern(IF, CASE_WHEN)) { - case ae @ AggregateExpression(count: Count, _, true, None, _) - if count.children.size == 1 => - extractCondAndBase(count.children.head) match { - case Some((cond, base)) => - ae.copy( - aggregateFunction = Count(base), - filter = Some(cond)) - case None => ae - } - } - } - - /** - * Matches IF(cond, base, null), CASE WHEN cond THEN base END, and - * CASE WHEN cond THEN base ELSE NULL END (including null wrapped in Cast). - * Multi-branch CaseWhen is intentionally not rewritten -- Or-flattening is out of scope. - * The base must be safe to evaluate unconditionally (the rewrite evaluates it on rows - * where the original branch would not have been taken, see - * [[ExprUtils.canEvaluateUnconditionally]]); the condition needs no check because it is - * evaluated unconditionally as the IF/CASE WHEN predicate anyway. - * Returns None for anything else. - */ - private def extractCondAndBase(expr: Expression): Option[(Expression, Expression)] = - expr match { - case If(cond, base, e) if isNullExpr(e) && ExprUtils.canEvaluateUnconditionally(base) => - Some((cond, base)) - case CaseWhen(Seq((cond, base)), None) if ExprUtils.canEvaluateUnconditionally(base) => - Some((cond, base)) - case CaseWhen(Seq((cond, base)), Some(e)) - if isNullExpr(e) && ExprUtils.canEvaluateUnconditionally(base) => - Some((cond, base)) - case _ => None - } - - private def isNullExpr(e: Expression): Boolean = e match { - case Literal(null, _) => true - case Cast(child, _, _, _) => isNullExpr(child) - case _ => false - } - private def collectAggregateExprs(a: Aggregate): Seq[AggregateExpression] = { // Collect all aggregate expressions. a.aggregateExpressions.flatMap { _.collect { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index e62971429dfd..dbec87f3cba2 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -1729,21 +1729,6 @@ object SQLConf { .booleanConf .createWithDefault(true) - val REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED = - buildConf("spark.sql.optimizer.rewriteCountDistinctConditional.enabled") - .internal() - .doc("When true, rewrites COUNT(DISTINCT IF(cond, base, NULL)) and " + - "COUNT(DISTINCT CASE WHEN cond THEN base END) into " + - "COUNT(DISTINCT base) FILTER (WHERE cond). This reduces the Expand factor " + - "in RewriteDistinctAggregates from Nx to 1x when multiple conditional distinct " + - "counts share the same base column. The rewrite is only applied to base " + - "expressions that are safe to evaluate unconditionally (e.g. plain columns), " + - "so the short-circuit semantics of IF/CASE WHEN are preserved.") - .version("4.3.0") - .withBindingPolicy(ConfigBindingPolicy.SESSION) - .booleanConf - .createWithDefault(true) - val ESCAPED_STRING_LITERALS = buildConf("spark.sql.parser.escapedStringLiterals") .internal() .doc("When true, string literals (including regex patterns) remain escaped in our SQL " + @@ -9936,9 +9921,6 @@ class SQLConf extends Serializable with Logging with SqlApiConf { def decorrelateInnerQueryEnabledForExistsIn: Boolean = !getConf(SQLConf.DECORRELATE_EXISTS_IN_SUBQUERY_LEGACY_INCORRECT_COUNT_HANDLING_ENABLED) - def rewriteCountDistinctConditionalEnabled: Boolean = - getConf(SQLConf.REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED) - def maxConcurrentOutputFileWriters: Int = getConf(SQLConf.MAX_CONCURRENT_OUTPUT_FILE_WRITERS) def plannedWriteEnabled: Boolean = getConf(SQLConf.PLANNED_WRITE_ENABLED) diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregatesSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregatesSuite.scala index b67eb88f0d68..08dd4011f04d 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregatesSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteDistinctAggregatesSuite.scala @@ -18,12 +18,11 @@ package org.apache.spark.sql.catalyst.optimizer import org.apache.spark.sql.catalyst.dsl.expressions._ import org.apache.spark.sql.catalyst.dsl.plans._ -import org.apache.spark.sql.catalyst.expressions.{AttributeReference, CaseWhen, Cast, Coalesce, Expression, GetMapValue, GetStructField, If, Literal, Lower, Round} -import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, CollectSet, Count, Sum} +import org.apache.spark.sql.catalyst.expressions.{Literal, Round} +import org.apache.spark.sql.catalyst.expressions.aggregate.CollectSet import org.apache.spark.sql.catalyst.plans.PlanTest import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Expand, LocalRelation, LogicalPlan} -import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.{IntegerType, MapType, StringType, StructField, StructType} +import org.apache.spark.sql.types.{IntegerType, StringType} class RewriteDistinctAggregatesSuite extends PlanTest { val nullInt = Literal(null, IntegerType) @@ -126,305 +125,4 @@ 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 projection contains `baseColName` as a plain attribute, that it - * contains no expression of `removedWrapperType` (the IF/CaseWhen that was stripped), - * and that the outer aggregate has moved the condition into a FILTER clause. - */ - private def assertSingleDistinctGroupExpand( - optimized: LogicalPlan, - baseColName: String, - removedWrapperType: Class[_]): 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 baseAttr = conditionalTestRelation.output.find(_.name == baseColName).get - assert(expand.projections.head.exists(_.semanticEquals(baseAttr)), - s"expected base column $baseColName in Expand projection") - assert(!expand.projections.head.exists(e => removedWrapperType.isInstance(e)), - s"${removedWrapperType.getSimpleName} wrapper should have been removed " + - "from Expand projection") - assert(collectAggregateExpressions(optimized).exists(_.filter.isDefined), - "expected at least one AggregateExpression with a FILTER clause") - } - - 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 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", classOf[If]) - } - - test("conditional: rewrite COUNT(DISTINCT CASE WHEN cond THEN col END) to " + - "COUNT(DISTINCT col) FILTER") { - 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", classOf[CaseWhen]) - } - - test("conditional: rewrite COUNT(DISTINCT CASE WHEN cond THEN col ELSE NULL END)") { - 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", classOf[CaseWhen]) - } - - test("conditional: multiple conditional distinct counts collapse to single distinct group") { - 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", classOf[If]) - } - - 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: 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") { - val input = conditionalTestRelation - .groupBy(Symbol("a"))( - Count(If(Symbol("b") > 1, Symbol("c"), Literal(null, IntegerType))) - .toAggregateExpression(isDistinct = false) - .as("cnt1"), - countDistinct(Symbol("c")).as("cnt2")) - .analyze - val optimized = RewriteDistinctAggregates(input) - // Still a single distinct group (cnt2), no canonicalization of the non-distinct agg - comparePlans(optimized, input) - } - - test("conditional: do not rewrite when FILTER already exists") { - val input = conditionalTestRelation - .groupBy(Symbol("a"))( - Count(If(Symbol("b") > 1, Symbol("c"), Literal(null, IntegerType))) - .toAggregateExpression(isDistinct = true, filter = Some(Symbol("d") === "x")) - .as("cnt1"), - Count(If(Symbol("b") > 2, Symbol("c"), Literal(null, IntegerType))) - .toAggregateExpression(isDistinct = true, filter = Some(Symbol("d") === "y")) - .as("cnt2")) - .analyze - val optimized = RewriteDistinctAggregates(input) - checkRewrite(optimized) - val expands = optimized.collect { case e: Expand => e } - // 2 groups because existing FILTERs prevent canonicalization - assert(expands.head.projections.size == 2) - } - - test("conditional: do not rewrite multi-branch CASE WHEN") { - val caseWhen1 = new CaseWhen( - Seq( - (Symbol("b") > Literal(1), Symbol("c")), - (Symbol("b") > Literal(2), Symbol("a"))), - Some(Literal(null))) - val caseWhen2 = new CaseWhen( - Seq( - (Symbol("b") > Literal(3), Symbol("c")), - (Symbol("b") > Literal(4), Symbol("a"))), - Some(Literal(null))) - val input = conditionalTestRelation - .groupBy(Symbol("a"))( - Count(caseWhen1).toAggregateExpression(isDistinct = true).as("cnt1"), - Count(caseWhen2).toAggregateExpression(isDistinct = true).as("cnt2")) - .analyze - val optimized = RewriteDistinctAggregates(input) - checkRewrite(optimized) - val expands = optimized.collect { case e: Expand => e } - // 2 groups - multi-branch CASE was not canonicalized - assert(expands.head.projections.size == 2) - } - - test("conditional: do not rewrite SUM(DISTINCT IF(...))") { - val input = conditionalTestRelation - .groupBy(Symbol("a"))( - Sum(If(Symbol("b") > 1, Symbol("c"), Literal(null, IntegerType))) - .toAggregateExpression(isDistinct = true) - .as("sum1"), - Sum(If(Symbol("b") > 2, Symbol("c"), Literal(null, IntegerType))) - .toAggregateExpression(isDistinct = true) - .as("sum2")) - .analyze - val optimized = RewriteDistinctAggregates(input) - checkRewrite(optimized) - val expands = optimized.collect { case e: Expand => e } - // 2 groups - SUM(DISTINCT IF) is not canonicalized - assert(expands.head.projections.size == 2) - } - - /** - * Asserts that two conditional distinct counts over `base` collapse into a single - * distinct group with a FILTER clause, i.e. the conditional canonicalization fired. - */ - private def assertCanonicalized(relation: LocalRelation, base: Expression): Unit = { - val input = relation - .groupBy(Symbol("a"))( - countDistinctIf(Symbol("b") > 1, base).as("cnt1"), - countDistinctIf(Symbol("b") > 2, base).as("cnt2")) - .analyze - val optimized = RewriteDistinctAggregates(input) - val expand = optimized.collectFirst { case e: Expand => e }.get - assert(expand.projections.size == 1, - s"branch-safe base should collapse to 1 distinct group: $base") - // The Expand rewrite adds gid-routing filters to every distinct aggregate; the - // canonicalization additionally introduces filters referencing the user condition. - val hasUserConditionFilter = collectAggregateExpressions(optimized) - .flatMap(_.filter) - .exists(_.references.exists(_.name != "gid")) - assert(hasUserConditionFilter, - s"expected a user-condition FILTER clause after canonicalizing base: $base") - } - - /** - * Asserts that two conditional distinct counts over `base` are rewritten by the - * general Expand path but are NOT canonicalized, i.e. they stay 2 distinct groups - * and no FILTER clause is introduced. - */ - private def assertNotCanonicalized(relation: LocalRelation, base: Expression): Unit = { - val input = relation - .groupBy(Symbol("a"))( - countDistinctIf(Symbol("b") > 1, base).as("cnt1"), - countDistinctIf(Symbol("b") > 2, base).as("cnt2")) - .analyze - val optimized = RewriteDistinctAggregates(input) - checkRewrite(optimized) - val expand = optimized.collectFirst { case e: Expand => e }.get - assert(expand.projections.size == 2, - s"non-branch-safe base should stay 2 distinct groups: $base") - // The Expand rewrite adds gid-routing filters to every distinct aggregate; a - // user-condition filter introduced by the canonicalization would reference the - // condition's columns instead of only gid. - val hasUserConditionFilter = collectAggregateExpressions(optimized) - .flatMap(_.filter) - .exists(_.references.exists(_.name != "gid")) - assert(!hasUserConditionFilter, - s"non-branch-safe base should not produce a user-condition FILTER clause: $base") - } - - test("conditional: do not rewrite when base is not branch-safe") { - val unsafeBases = Seq( - // Division may throw (e.g. divide-by-zero under ANSI mode) on rows where the - // original IF branch is not taken. - Symbol("c") / Symbol("b"), - // Cast may throw under ANSI mode. - Cast(Symbol("c"), StringType), - // String functions are not in the branch-safe whitelist. - Lower(Symbol("d"))) - unsafeBases.foreach { base => - assertNotCanonicalized(conditionalTestRelation, base) - } - } - - test("conditional: rewrite branch-safe expression bases") { - // Predicate and null-handling bases on the flat test relation. - Seq( - Symbol("b") > Symbol("c"), - Coalesce(Seq(Symbol("b"), Literal(0, IntegerType)))) - .foreach { base => - assertCanonicalized(conditionalTestRelation, base) - } - - // Total accessor bases need struct/map attributes. - val structAttr = AttributeReference("s", StructType(Seq(StructField("f", IntegerType))))() - val mapAttr = AttributeReference("m", MapType(IntegerType, StringType))() - val richRelation = LocalRelation(Symbol("a").int, Symbol("b").int, structAttr, mapAttr) - Seq( - GetStructField(structAttr, 0), - GetMapValue(mapAttr, Literal(1, IntegerType))) - .foreach { base => - assertCanonicalized(richRelation, base) - } - } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/RewriteDistinctAggregatesConditionalQuerySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/RewriteDistinctAggregatesConditionalQuerySuite.scala deleted file mode 100644 index f99429cd1649..000000000000 --- a/sql/core/src/test/scala/org/apache/spark/sql/RewriteDistinctAggregatesConditionalQuerySuite.scala +++ /dev/null @@ -1,307 +0,0 @@ -/* - * 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("rewrite is present in optimized plan") { - withTempView("t") { - spark.range(2) - .selectExpr( - "cast(id + 1 as int) as key", - "cast(id * 10 as int) as col1", - "cast(id * 100 as int) as col2") - .createOrReplaceTempView("t") - - // Two conditional distinct counts on the same base column trigger canonicalization, - // so the optimized plan must contain user-condition FILTER clauses. - val hasFilter = withSQLConf( - SQLConf.REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED.key -> "true") { - val df = spark.sql( - """SELECT key, - | COUNT(DISTINCT IF(col1 > 10, col2, NULL)) as cnt1, - | COUNT(DISTINCT IF(col1 > 5, col2, NULL)) as cnt2 - |FROM t GROUP BY key""".stripMargin) - hasUserConditionFilter(df.queryExecution.optimizedPlan) - } - - assert(hasFilter, "Optimized plan should contain user-condition FILTER clause") - } - } -} --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
