This is an automated email from the ASF dual-hosted git repository.
rui-mo 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 757ae0759c [GLUTEN-13176][VL] Pass RESPECT NULLS of collect_set to
Velox and re-enable two Spark 4.2 tests (#13441)
757ae0759c is described below
commit 757ae0759c3c5452d9e8469933f15f2d71d5b2f9
Author: Akshay Tayal <[email protected]>
AuthorDate: Tue Oct 27 10:21:20 2026 -0700
[GLUTEN-13176][VL] Pass RESPECT NULLS of collect_set to Velox and re-enable
two Spark 4.2 tests (#13441)
---
.../backendsapi/velox/VeloxSparkPlanExecApi.scala | 4 +-
.../execution/HashAggregateExecTransformer.scala | 17 ++++++++-
.../gluten/extension/CollectRewriteRule.scala | 3 +-
.../gluten/utils/VeloxIntermediateData.scala | 23 +++++++++++-
.../gluten/utils/velox/VeloxTestSettings.scala | 8 +---
.../spark/sql/GlutenDataFrameAggregateSuite.scala | 32 ++++++++++++++++
.../spark/sql/GlutenDataFrameFunctionsSuite.scala | 43 +++++++++++++++++++++-
.../sql/GlutenDataFrameWindowFunctionsSuite.scala | 28 ++++++++++++++
8 files changed, 146 insertions(+), 12 deletions(-)
diff --git
a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala
b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala
index 387215ca66..a62fbf849c 100644
---
a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala
+++
b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala
@@ -28,6 +28,7 @@ import
org.apache.gluten.shuffle.NeedCustomColumnarBatchSerializer
import org.apache.gluten.sql.shims.SparkShimLoader
import org.apache.gluten.substrait.SubstraitContext
import org.apache.gluten.substrait.expression.{ExpressionBuilder,
ExpressionNode, WindowFunctionNode}
+import org.apache.gluten.utils.VeloxIntermediateData
import org.apache.gluten.vectorized.{ColumnarBatchSerializer,
ColumnarBatchSerializeResult}
import org.apache.spark.{ShuffleDependency, SparkEnv}
@@ -1664,7 +1665,8 @@ class VeloxSparkPlanExecApi extends SparkPlanExecApi with
Logging {
}
}
- val childrenNodeList = aggregateFunc.children
+ val childrenNodeList = VeloxIntermediateData
+ .getRawInputs(aggregateFunc)
.map(
ExpressionConverter
.replaceWithExpressionTransformer(_, originalInputAttributes)
diff --git
a/backends-velox/src/main/scala/org/apache/gluten/execution/HashAggregateExecTransformer.scala
b/backends-velox/src/main/scala/org/apache/gluten/execution/HashAggregateExecTransformer.scala
index 8e32c7c9e7..a9d7c42b47 100644
---
a/backends-velox/src/main/scala/org/apache/gluten/execution/HashAggregateExecTransformer.scala
+++
b/backends-velox/src/main/scala/org/apache/gluten/execution/HashAggregateExecTransformer.scala
@@ -295,7 +295,11 @@ abstract class HashAggregateExecTransformer(
val functionInputAttributes = aggFunc.inputAggBufferAttributes
aggFunc match {
case _ if aggregateExpression.mode == Partial ||
aggregateExpression.mode == Complete =>
- val childNodes = aggFunc.children
+ // Literal arguments stay inline in the aggregate (see
getAggRelInternal), because Velox
+ // only treats an argument as constant when it is a literal, not a
projected column.
+ val childNodes = VeloxIntermediateData
+ .getRawInputs(aggFunc)
+ .filterNot(_.isInstanceOf[Literal])
.map(
ExpressionConverter
.replaceWithExpressionTransformer(_, originalInputAttributes)
@@ -504,7 +508,16 @@ abstract class HashAggregateExecTransformer(
val aggregateFunc = aggExpr.aggregateFunction
val childrenNodes = aggExpr.mode match {
case Partial | Complete =>
- aggregateFunc.children.toList.map(toExpressionNode)
+ // rowConstructed is per operator: in mixed-phase aggregates (e.g.
the distinct
+ // stage), Partial arguments are also pre-projected and referenced
by position.
+ // Literals are kept inline so that they reach Velox as constants.
+ VeloxIntermediateData.getRawInputs(aggregateFunc).toList.map {
+ case literal: Literal if rowConstructed =>
+ ExpressionConverter
+ .replaceWithExpressionTransformer(literal,
originalInputAttributes)
+ .doTransform(context)
+ case expr => toExpressionNode(expr)
+ }
case PartialMerge | Final =>
if (rowConstructed) {
// Only occupies one column due to intermediate results are
combined
diff --git
a/backends-velox/src/main/scala/org/apache/gluten/extension/CollectRewriteRule.scala
b/backends-velox/src/main/scala/org/apache/gluten/extension/CollectRewriteRule.scala
index 72e52cf3dd..a7a6a0a906 100644
---
a/backends-velox/src/main/scala/org/apache/gluten/extension/CollectRewriteRule.scala
+++
b/backends-velox/src/main/scala/org/apache/gluten/extension/CollectRewriteRule.scala
@@ -76,7 +76,8 @@ object CollectRewriteRule {
}
}
- private def getIgnoreNulls(expr: Expression): Boolean = {
+ /** Returns the `ignoreNulls` flag of Spark's CollectSet / CollectList (true
before Spark 4.2). */
+ def getIgnoreNulls(expr: Expression): Boolean = {
try {
val method = expr.getClass.getMethod("ignoreNulls")
method.invoke(expr).asInstanceOf[Boolean]
diff --git
a/backends-velox/src/main/scala/org/apache/gluten/utils/VeloxIntermediateData.scala
b/backends-velox/src/main/scala/org/apache/gluten/utils/VeloxIntermediateData.scala
index a226551525..e82efa76f8 100644
---
a/backends-velox/src/main/scala/org/apache/gluten/utils/VeloxIntermediateData.scala
+++
b/backends-velox/src/main/scala/org/apache/gluten/utils/VeloxIntermediateData.scala
@@ -17,8 +17,11 @@
package org.apache.gluten.utils
import org.apache.gluten.expression.ConverterUtils
+import org.apache.gluten.expression.aggregate.VeloxCollectSet
+import org.apache.gluten.extension.CollectRewriteRule
import org.apache.gluten.substrait.`type`.{TypeBuilder, TypeNode}
+import org.apache.spark.sql.catalyst.expressions.{Expression, Literal}
import org.apache.spark.sql.catalyst.expressions.aggregate._
import org.apache.spark.sql.types._
@@ -108,6 +111,24 @@ object VeloxIntermediateData {
}
}
+ /**
+ * Get the arguments passed to a Velox aggregate function when it consumes
raw input (Partial or
+ * Complete mode).
+ *
+ * @param aggregateFunc
+ * The input aggregate function.
+ * @return
+ * The raw input arguments of a Velox aggregate function.
+ */
+ def getRawInputs(aggregateFunc: AggregateFunction): Seq[Expression] =
aggregateFunc match {
+ // Velox's collect_set takes RESPECT NULLS as a constant `ignoreNulls`
argument. Window
+ // functions keep Spark's CollectSet, because CollectRewriteRule only
rewrites aggregates.
+ case c: VeloxCollectSet if !c.ignoreNulls => Seq(c.child,
Literal.FalseLiteral)
+ case c: CollectSet if !CollectRewriteRule.getIgnoreNulls(c) =>
+ Seq(c.child, Literal.FalseLiteral)
+ case _ => aggregateFunc.children
+ }
+
/**
* Get the compatible input types for a Velox aggregate function.
*
@@ -120,7 +141,7 @@ object VeloxIntermediateData {
*/
def getInputTypes(aggregateFunc: AggregateFunction, forMergeCompanion:
Boolean): Seq[DataType] = {
if (!forMergeCompanion) {
- return aggregateFunc.children.map(_.dataType)
+ return getRawInputs(aggregateFunc).map(_.dataType)
}
aggregateFunc match {
case _ @Type(veloxDataTypes: Seq[DataType]) =>
diff --git
a/gluten-ut/spark42/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala
b/gluten-ut/spark42/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala
index adce5b7beb..1eec5810a1 100644
---
a/gluten-ut/spark42/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala
+++
b/gluten-ut/spark42/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala
@@ -1014,9 +1014,6 @@ class VeloxTestSettings extends BackendTestSettings {
.exclude(
"SPARK-57298: collect_set normalizes NaN and -0.0 nested in complex
types"
) // GLUTEN-12569 Spark 4.2 UT enablement (CI-confirmed)
- .exclude(
- "collect_set with RESPECT NULLS preserves null in set"
- ) // GLUTEN-12569 Spark 4.2 UT enablement (CI-confirmed)
enableSuite[GlutenDataFrameAsOfJoinSuite]
enableSuite[GlutenDataFrameComplexTypeSuite]
enableSuite[GlutenDataFrameFunctionsSuite]
@@ -1027,6 +1024,8 @@ class VeloxTestSettings extends BackendTestSettings {
.exclude("transform keys function - primitive data types")
// Overridden.
.exclude("map with arrays")
+ // Overridden: asserts vanilla WholeStageCodegenExec, which Gluten
replaces.
+ .exclude("array_join with nullable nullReplacement under whole-stage
codegen")
.exclude(
"SPARK-54918: array set ops normalize -0.0 and NaN via VALUES inline
table"
) // GLUTEN-12569 Spark 4.2 UT enablement (CI-confirmed)
@@ -1048,9 +1047,6 @@ class VeloxTestSettings extends BackendTestSettings {
.exclude(
"SPARK-54918: array_union normalizes -0.0 to +0.0 - literals"
) // GLUTEN-12569 Spark 4.2 UT enablement (CI-confirmed)
- .exclude(
- "array_join with nullable nullReplacement under whole-stage codegen"
- ) // GLUTEN-12569 Spark 4.2 UT enablement (CI-confirmed)
enableSuite[GlutenDataFrameHintSuite]
enableSuite[GlutenDataFrameImplicitsSuite]
enableSuite[GlutenDataFrameJoinSuite]
diff --git
a/gluten-ut/spark42/src/test/scala/org/apache/spark/sql/GlutenDataFrameAggregateSuite.scala
b/gluten-ut/spark42/src/test/scala/org/apache/spark/sql/GlutenDataFrameAggregateSuite.scala
index a0338f10bc..2326e7ef35 100644
---
a/gluten-ut/spark42/src/test/scala/org/apache/spark/sql/GlutenDataFrameAggregateSuite.scala
+++
b/gluten-ut/spark42/src/test/scala/org/apache/spark/sql/GlutenDataFrameAggregateSuite.scala
@@ -19,6 +19,8 @@ package org.apache.spark.sql
import org.apache.gluten.config.GlutenConfig
import org.apache.gluten.execution.{HashAggregateExecBaseTransformer,
HashAggregateExecTransformer}
+import org.apache.spark.sql.catalyst.expressions.aggregate.{Partial,
PartialMerge}
+import org.apache.spark.sql.catalyst.optimizer.EliminateDistinct
import org.apache.spark.sql.execution.WholeStageCodegenExec
import org.apache.spark.sql.execution.aggregate.{HashAggregateExec,
ObjectHashAggregateExec, SortAggregateExec}
import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
@@ -284,6 +286,36 @@ class GlutenDataFrameAggregateSuite extends
DataFrameAggregateSuite with GlutenS
).foreach(assertNoExceptions)
}
+ testGluten("collect_set RESPECT NULLS in an aggregate that also merges
another aggregate") {
+ // Without EliminateDistinct, the distinct stage runs collect_set
(Partial) and avg
+ // (PartialMerge) in one aggregate, which takes the row-construct path.
+ withSQLConf(SQLConf.OPTIMIZER_EXCLUDED_RULES.key ->
EliminateDistinct.ruleName) {
+ withTempView("t") {
+ Seq((1, Some(2)), (2, None), (3, Some(2)), (4, Some(5))).toDF("a", "b")
+ .createOrReplaceTempView("t")
+ val df = sql("SELECT avg(a), sort_array(collect_set(DISTINCT b)
RESPECT NULLS) FROM t")
+ checkAnswer(df, Row(2.5, Seq(null, 2, 5)))
+ val plan = df.queryExecution.executedPlan
+ assert(
+ collect(plan) {
+ case agg: HashAggregateExecTransformer
+ if agg.aggregateExpressions.exists(_.mode == PartialMerge) &&
+ agg.aggregateExpressions.exists(_.mode == Partial) =>
+ agg
+ }.nonEmpty,
+ plan
+ )
+ assert(
+ collect(plan) {
+ case agg @ (_: HashAggregateExec | _: ObjectHashAggregateExec |
+ _: SortAggregateExec) =>
+ agg
+ }.isEmpty,
+ plan)
+ }
+ }
+ }
+
testGluten("SPARK-22223: ObjectHashAggregate should not introduce
unnecessary shuffle") {
withSQLConf(SQLConf.USE_OBJECT_HASH_AGG.key -> "true") {
val df = Seq(("1", "2", 1), ("1", "2", 2), ("2", "3", 3), ("2", "3", 4))
diff --git
a/gluten-ut/spark42/src/test/scala/org/apache/spark/sql/GlutenDataFrameFunctionsSuite.scala
b/gluten-ut/spark42/src/test/scala/org/apache/spark/sql/GlutenDataFrameFunctionsSuite.scala
index 40ac1edd60..5786921364 100644
---
a/gluten-ut/spark42/src/test/scala/org/apache/spark/sql/GlutenDataFrameFunctionsSuite.scala
+++
b/gluten-ut/spark42/src/test/scala/org/apache/spark/sql/GlutenDataFrameFunctionsSuite.scala
@@ -17,13 +17,19 @@
package org.apache.spark.sql
import org.apache.gluten.exception.GlutenException
+import org.apache.gluten.execution.ProjectExecTransformer
import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.expressions.ArrayJoin
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.functions._
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.{IntegerType, MapType, StringType,
StructField, StructType}
-class GlutenDataFrameFunctionsSuite extends DataFrameFunctionsSuite with
GlutenSQLTestsTrait {
+class GlutenDataFrameFunctionsSuite
+ extends DataFrameFunctionsSuite
+ with GlutenSQLTestsTrait
+ with AdaptiveSparkPlanHelper {
import testImplicits._
testGluten("map with arrays") {
@@ -409,4 +415,39 @@ class GlutenDataFrameFunctionsSuite extends
DataFrameFunctionsSuite with GlutenS
oneRowDF.selectExpr("flatten(null)")
}
}
+
+ testGluten("array_join with nullable nullReplacement under whole-stage
codegen") {
+ withTempView("array_join_codegen") {
+ Seq(
+ (Seq[String]("a", null, "b"), ",", "NR"),
+ (Seq[String]("a", null, "b"), ",", null),
+ (Seq[String]("x", "y"), "-", "NR")
+ ).toDF("arr", "delim_col",
"repl_col").createOrReplaceTempView("array_join_codegen")
+ spark.catalog.cacheTable("array_join_codegen")
+
+ val query =
+ "SELECT array_join(arr, delim_col, repl_col) FROM array_join_codegen "
+
+ "WHERE arr IS NOT NULL AND delim_col IS NOT NULL"
+
+ withSQLConf(SQLConf.CODEGEN_FACTORY_MODE.key -> "CODEGEN_ONLY") {
+ val df = sql(query)
+ // Gluten offloads array_join instead of running it in
WholeStageCodegenExec.
+ assert(
+ find(df.queryExecution.executedPlan) {
+ case p: ProjectExecTransformer =>
+ p.projectList.exists(_.find(_.isInstanceOf[ArrayJoin]).isDefined)
+ case _ => false
+ }.isDefined,
+ "expected array_join to be offloaded"
+ )
+ checkAnswer(df, Seq(Row("a,NR,b"), Row(null), Row("x-y")))
+ }
+
+ withSQLConf(
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false",
+ SQLConf.CODEGEN_FACTORY_MODE.key -> "NO_CODEGEN") {
+ checkAnswer(sql(query), Seq(Row("a,NR,b"), Row(null), Row("x-y")))
+ }
+ }
+ }
}
diff --git
a/gluten-ut/spark42/src/test/scala/org/apache/spark/sql/GlutenDataFrameWindowFunctionsSuite.scala
b/gluten-ut/spark42/src/test/scala/org/apache/spark/sql/GlutenDataFrameWindowFunctionsSuite.scala
index 213c4c1aec..46bb0be76d 100644
---
a/gluten-ut/spark42/src/test/scala/org/apache/spark/sql/GlutenDataFrameWindowFunctionsSuite.scala
+++
b/gluten-ut/spark42/src/test/scala/org/apache/spark/sql/GlutenDataFrameWindowFunctionsSuite.scala
@@ -23,6 +23,7 @@ import
org.apache.spark.sql.catalyst.expressions.{AttributeReference, Expression
import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning
import org.apache.spark.sql.execution.ColumnarShuffleExchangeExec
import org.apache.spark.sql.execution.exchange.ENSURE_REQUIREMENTS
+import org.apache.spark.sql.execution.window.WindowExec
import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions._
import org.apache.spark.sql.internal.SQLConf
@@ -206,4 +207,31 @@ class GlutenDataFrameWindowFunctionsSuite
assert(shuffleByRequirement, "Can't find desired shuffle node from the
query plan")
}
}
+
+ testGluten("collect_set RESPECT NULLS over a window frame") {
+ withTempView("t") {
+ Seq((1, Some(2)), (2, None), (3, Some(2)), (4, Some(5))).toDF("a", "b")
+ .createOrReplaceTempView("t")
+ val df = sql(
+ """
+ |SELECT a,
+ | sort_array(collect_set(b) RESPECT NULLS
+ | OVER (ORDER BY a ROWS BETWEEN 1 PRECEDING AND CURRENT ROW)),
+ | sort_array(collect_set(b) RESPECT NULLS OVER (ORDER BY a)),
+ | sort_array(collect_set(b) OVER (ORDER BY a ROWS BETWEEN 1
PRECEDING AND CURRENT ROW))
+ |FROM t""".stripMargin)
+ checkAnswer(
+ df,
+ Seq(
+ Row(1, Seq(2), Seq(2), Seq(2)),
+ Row(2, Seq(null, 2), Seq(null, 2), Seq(2)),
+ Row(3, Seq(null, 2), Seq(null, 2), Seq(2)),
+ Row(4, Seq(2, 5), Seq(null, 2, 5), Seq(2, 5))
+ )
+ )
+ val plan = df.queryExecution.executedPlan
+ assert(collect(plan) { case w: WindowExecTransformer => w }.nonEmpty,
plan)
+ assert(collect(plan) { case w: WindowExec => w }.isEmpty, plan)
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
