This is an automated email from the ASF dual-hosted git repository. WangGuangxin pushed a commit to branch revert-12256-window_count in repository https://gitbox.apache.org/repos/asf/gluten.git
commit 278faa4b9497a5d406c66a2620c42d18a49f04f6 Author: WangGuangxin <[email protected]> AuthorDate: Mon Jul 6 22:57:38 2026 +0800 Revert "rewrite multi-children Count in window expressions" This reverts commit 8a987f8b81f61b429123431e01d90996c131ad04. --- .../gluten/backendsapi/velox/VeloxBackend.scala | 19 ++--- .../functions/WindowFunctionsValidateSuite.scala | 39 ----------- .../rewrite/RewriteMultiChildrenCount.scala | 81 +++++----------------- 3 files changed, 23 insertions(+), 116 deletions(-) diff --git a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxBackend.scala b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxBackend.scala index 1415019681..01cae32c0c 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxBackend.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxBackend.scala @@ -33,7 +33,7 @@ import org.apache.gluten.utils._ import org.apache.spark.sql.catalyst.catalog.BucketSpec import org.apache.spark.sql.catalyst.expressions.{Alias, CumeDist, DenseRank, Descending, Expression, Lag, Lead, NamedExpression, NthValue, NTile, PercentRank, RangeFrame, Rank, RowNumber, SortOrder, SpecialFrameBoundary, SpecifiedWindowFrame} -import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, ApproximatePercentile, Count, HyperLogLogPlusPlus, Percentile} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, ApproximatePercentile, HyperLogLogPlusPlus, Percentile} import org.apache.spark.sql.catalyst.plans.{JoinType, LeftOuter, RightOuter} import org.apache.spark.sql.catalyst.util.{CaseInsensitiveMap, CharVarcharUtils} import org.apache.spark.sql.connector.read.Scan @@ -471,19 +471,10 @@ object VeloxBackendSettings extends BackendSettingsApi { case _: NthValue => case _: Lag => case _: Lead => - case ae: AggregateExpression => - // Velox only supports count() and count(T) signatures for the window count - // function. Spark's count(c1, c2, ...) (counts rows where ALL arguments are - // non-null) is normally rewritten to single-arg form by RewriteMultiChildrenCount. - // If the rewrite did not run (e.g., the rule is disabled), keep us safe by - // falling back to vanilla Spark instead of crashing inside Velox. - ae.aggregateFunction match { - case c: Count if c.children.size > 1 => - allSupported = false - case _: ApproximatePercentile | _: Percentile | _: HyperLogLogPlusPlus => - allSupported = false - case _ => - } + case aggrExpr: AggregateExpression + if !aggrExpr.aggregateFunction.isInstanceOf[ApproximatePercentile] + && !aggrExpr.aggregateFunction.isInstanceOf[Percentile] + && !aggrExpr.aggregateFunction.isInstanceOf[HyperLogLogPlusPlus] => case _ => allSupported = false } diff --git a/backends-velox/src/test/scala/org/apache/gluten/functions/WindowFunctionsValidateSuite.scala b/backends-velox/src/test/scala/org/apache/gluten/functions/WindowFunctionsValidateSuite.scala index 06caa16149..b6ce5ed0cd 100644 --- a/backends-velox/src/test/scala/org/apache/gluten/functions/WindowFunctionsValidateSuite.scala +++ b/backends-velox/src/test/scala/org/apache/gluten/functions/WindowFunctionsValidateSuite.scala @@ -53,43 +53,4 @@ class WindowFunctionsValidateSuite extends FunctionsValidateSuite { checkGlutenPlan[WindowExecTransformer] } } - - test("count window function with multiple arguments is rewritten and offloaded") { - // Velox only supports count() / count(T) for window functions. Spark's - // count(c1, c2, ...) variant must be rewritten into count(if(or(isnull(c1), - // isnull(c2), ...), null, 1)) so the WindowExec can still be offloaded. - runQueryAndCompare( - "select l_orderkey, " + - "count(l_partkey, l_suppkey, l_linenumber) " + - "over (partition by l_orderkey) as cnt " + - "from lineitem") { - checkGlutenPlan[WindowExecTransformer] - } - } - - test("count window function with multiple arguments returns vanilla Spark result") { - // Validate semantic equivalence with vanilla Spark: counts rows where ALL - // arguments are non-null in the window partition. - withTable("nullable_window_data") { - spark - .sql(""" - |select * from values - | (1, 'a', 10), - | (1, 'a', null), - | (1, null, 20), - | (2, 'b', 30), - | (2, 'b', 40), - | (2, 'c', null) - |as t(k, s, v) - """.stripMargin) - .write - .saveAsTable("nullable_window_data") - - runQueryAndCompare( - "select k, count(s, v) over (partition by k) as cnt " + - "from nullable_window_data") { - checkGlutenPlan[WindowExecTransformer] - } - } - } } diff --git a/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/rewrite/RewriteMultiChildrenCount.scala b/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/rewrite/RewriteMultiChildrenCount.scala index c4f0a53878..1003407191 100644 --- a/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/rewrite/RewriteMultiChildrenCount.scala +++ b/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/rewrite/RewriteMultiChildrenCount.scala @@ -19,25 +19,17 @@ package org.apache.gluten.extension.columnar.rewrite import org.apache.gluten.backendsapi.BackendsApiManager import org.apache.gluten.utils.PullOutProjectHelper -import org.apache.spark.sql.catalyst.expressions.{Expression, If, IsNull, Literal, NamedExpression, Or} +import org.apache.spark.sql.catalyst.expressions.{If, IsNull, Literal, Or} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Count, Partial} import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.aggregate.BaseAggregateExec -import org.apache.spark.sql.execution.window.WindowExec import org.apache.spark.sql.types.IntegerType /** - * This Rule is used to rewrite multi-children count for both partial aggregates and window - * functions, since Velox only supports `count()` and `count(T)` signatures. - * - * For partial aggregates the rewrite triggers when: + * This Rule is used to rewrite the count for aggregates if: * - the count is partial * - the count has more than one child * - * For window functions the rewrite triggers whenever the inner aggregate function of the window - * expression is a `Count` with more than one child. The aggregate mode for window functions is - * always `Complete`, so we don't gate on mode there. - * * A rewrite count multi-children example: * {{{ * Count(c1, c2) @@ -51,7 +43,7 @@ import org.apache.spark.sql.types.IntegerType * ) * }}} * - * TODO: Remove this rule when Velox support multi-children Count. + * TODO: Remove this rule when Velox support multi-children Count */ object RewriteMultiChildrenCount extends RewriteSingleNode with PullOutProjectHelper { private lazy val shouldRewriteCount = BackendsApiManager.getSettings.shouldRewriteCount() @@ -59,7 +51,6 @@ object RewriteMultiChildrenCount extends RewriteSingleNode with PullOutProjectHe override def isRewritable(plan: SparkPlan): Boolean = { plan match { case _: BaseAggregateExec => true - case _: WindowExec => true case _ => false } } @@ -78,62 +69,32 @@ object RewriteMultiChildrenCount extends RewriteSingleNode with PullOutProjectHe } } - /** Window functions don't have Partial/Final modes; rewrite any multi-child Count. */ - private def extractWindowCountForRewrite(aggExpr: AggregateExpression): Option[Count] = { - aggExpr.aggregateFunction match { - case c: Count if c.children.size > 1 => Option(c) - case _ => None - } - } - private def shouldRewrite(aggregateExpressions: Seq[AggregateExpression]): Boolean = { aggregateExpressions.exists(extractCountForRewrite(_).isDefined) } - private def shouldRewriteWindow(windowExpressions: Seq[NamedExpression]): Boolean = { - windowExpressions.exists(_.find { - case ae: AggregateExpression => extractWindowCountForRewrite(ae).isDefined - case _ => false - }.isDefined) - } - - private def buildSingleChildCount(count: Count): Count = { - // Follow vanilla Spark Count semantics: count rows where every argument is non-null. - val nullableChildren = count.children.filter(_.nullable) - val newChild: Expression = if (nullableChildren.isEmpty) { - Literal.create(1, IntegerType) - } else { - If( - nullableChildren.map(IsNull).reduce(Or), - Literal.create(null, IntegerType), - Literal.create(1, IntegerType)) - } - count.copy(children = newChild :: Nil) - } - private def rewriteCount( aggregateExpressions: Seq[AggregateExpression]): Seq[AggregateExpression] = { aggregateExpressions.map { aggExpr => val countOpt = extractCountForRewrite(aggExpr) countOpt - .map(count => aggExpr.copy(aggregateFunction = buildSingleChildCount(count))) - .getOrElse(aggExpr) - } - } - - private def rewriteWindowExpressions( - windowExpressions: Seq[NamedExpression]): Seq[NamedExpression] = { - windowExpressions.map { - expr => - expr - .transformDown { - case ae: AggregateExpression => - extractWindowCountForRewrite(ae) - .map(count => ae.copy(aggregateFunction = buildSingleChildCount(count))) - .getOrElse(ae) + .map { + count => + // Follow vanillas Spark Count function + val nullableChildren = count.children.filter(_.nullable) + val newChild = if (nullableChildren.isEmpty) { + Literal.create(1, IntegerType) + } else { + If( + nullableChildren.map(IsNull).reduce(Or), + Literal.create(null, IntegerType), + Literal.create(1, IntegerType)) + } + val newCount = count.copy(children = newChild :: Nil) + aggExpr.copy(aggregateFunction = newCount) } - .asInstanceOf[NamedExpression] + .getOrElse(aggExpr) } } @@ -147,12 +108,6 @@ object RewriteMultiChildrenCount extends RewriteSingleNode with PullOutProjectHe val newAggExprs = rewriteCount(agg.aggregateExpressions) copyBaseAggregateExec(agg)(newAggregateExpressions = newAggExprs) - case window: WindowExec if shouldRewriteWindow(window.windowExpression) => - val newWindow = - window.copy(windowExpression = rewriteWindowExpressions(window.windowExpression)) - newWindow.copyTagsFrom(window) - newWindow - case _ => plan } } --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
