This is an automated email from the ASF dual-hosted git repository.
WangGuangxin pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 8a987f8b81 rewrite multi-children Count in window expressions
8a987f8b81 is described below
commit 8a987f8b81f61b429123431e01d90996c131ad04
Author: wangguangxin.cn <[email protected]>
AuthorDate: Fri Jun 5 17:26:04 2026 +0800
rewrite multi-children Count in window expressions
---
.../gluten/backendsapi/velox/VeloxBackend.scala | 19 +++--
.../functions/WindowFunctionsValidateSuite.scala | 39 +++++++++++
.../rewrite/RewriteMultiChildrenCount.scala | 81 +++++++++++++++++-----
3 files changed, 116 insertions(+), 23 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 01cae32c0c..1415019681 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, HyperLogLogPlusPlus, Percentile}
+import
org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression,
ApproximatePercentile, Count, 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,10 +471,19 @@ object VeloxBackendSettings extends BackendSettingsApi {
case _: NthValue =>
case _: Lag =>
case _: Lead =>
- case aggrExpr: AggregateExpression
- if
!aggrExpr.aggregateFunction.isInstanceOf[ApproximatePercentile]
- && !aggrExpr.aggregateFunction.isInstanceOf[Percentile]
- &&
!aggrExpr.aggregateFunction.isInstanceOf[HyperLogLogPlusPlus] =>
+ 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 _ =>
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 b6ce5ed0cd..06caa16149 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,4 +53,43 @@ 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 1003407191..c4f0a53878 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,17 +19,25 @@ 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.{If, IsNull, Literal, Or}
+import org.apache.spark.sql.catalyst.expressions.{Expression, If, IsNull,
Literal, NamedExpression, 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 the count for aggregates if:
+ * 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:
* - 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)
@@ -43,7 +51,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()
@@ -51,6 +59,7 @@ object RewriteMultiChildrenCount extends RewriteSingleNode
with PullOutProjectHe
override def isRewritable(plan: SparkPlan): Boolean = {
plan match {
case _: BaseAggregateExec => true
+ case _: WindowExec => true
case _ => false
}
}
@@ -69,35 +78,65 @@ 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 =>
- // 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)
- }
+ .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)
+ }
+ .asInstanceOf[NamedExpression]
+ }
+ }
+
override def rewrite(plan: SparkPlan): SparkPlan = {
if (!shouldRewriteCount) {
return plan
@@ -108,6 +147,12 @@ 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]
