cloud-fan commented on code in PR #55925: URL: https://github.com/apache/spark/pull/55925#discussion_r3596312427
########## sql/core/src/test/scala/org/apache/spark/sql/RewriteDistinctAggregatesConditionalQuerySuite.scala: ########## @@ -0,0 +1,268 @@ +/* + * 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 = { + // Verify the rewrite produces the same result as the explicit FILTER form. + val withRewrite = withSQLConf( + SQLConf.REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED.key -> "true") { + spark.sql(conditionalSql) + } + val withoutRewrite = withSQLConf( + SQLConf.REWRITE_COUNT_DISTINCT_CONDITIONAL_ENABLED.key -> "false") { + spark.sql(conditionalSql) + } + val explicitFilter = spark.sql(filterSql) + + // Rewritten query should match explicit FILTER query. + checkAnswer(withRewrite, explicitFilter) + // Non-rewritten query should also match explicit FILTER query. + checkAnswer(withoutRewrite, explicitFilter) + } + + private def hasFilterClause(plan: LogicalPlan): Boolean = { + val aggregateExpressions = plan.collect { + case a: Aggregate => a.aggregateExpressions + }.flatten + val aggregateExprs = aggregateExpressions.flatMap { + _.collect { case ae: AggregateExpression => ae } + } + aggregateExprs.exists(_.filter.isDefined) + } + + 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") + + val sqlText = "SELECT key, COUNT(DISTINCT IF(col1 > 10, col2, 0)) FROM t GROUP BY key" Review Comment: This negative test uses a **single** conditional distinct count. A lone distinct, filter-less `Count` makes `mayNeedtoRewrite` return `false`, so `rewrite()` is never called and `normalizeCountDistinctConditional` never runs -- the non-null-else guard in `extractCondAndBase` is never reached. Both assertions then pass vacuously: `checkAnswer(withRewrite, withoutRewrite)` compares two non-rewritten (identical) plans, and `assert(!hasFilterClause(...))` would hold even if `isNullExpr` wrongly accepted the non-null `0` else, because the rewrite path is never entered. This is the same single-count vacuity that was fixed for the positive tests in this suite (discussion r3449087013), but that fix listed only the positive tests -- this negative test was left with one count. Give it two conditional distinct counts on the same base so `mayNeedtoRewrite` is true and the rewrite actually fires, mirroring `RewriteDistinctAggregatesSuite`'s "do not rewrite IF with non-null else branch" (which already uses two counts): ```suggestion val sqlText = """SELECT key, | COUNT(DISTINCT IF(col1 > 10, col2, 0)), | COUNT(DISTINCT IF(col1 > 20, col2, 0)) |FROM t GROUP BY key""".stripMargin ``` Non-blocking -- a late catch on my side, not something you introduced. -- 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]
