Skip to content

Commit c406ee0

Browse files
firestarmanfirestarman
authored
Support collect_set RESPECT NULLS [databricks] (#15372)
Fixes #15219. ### Description - Added Spark 4.2 `collect_set RESPECT NULLS` support on GPU by plumbing `ignoreNulls` through shims, expression metadata, GPU aggregate data types, CPU buffer conversion, and cuDF null policy so GPU results preserve one null like Spark CPU. - Added Spark 4.2 group-by, global reduction, rolling-window, and fully unbounded partitioned-window integration coverage for `collect_set IGNORE NULLS` and `collect_set RESPECT NULLS` because the new SQL clause changes null handling semantics. - Validated the Spark 4.2 jars and focused integration tests with `mvn -s ~/.m2/settings_art.xml -f scala2.13/pom.xml -Dbuildver=400 -Dcuda.version=cuda13 -DskipTests validate`, `mvn -s ~/.m2/settings_art.xml -Dbuildver=330 -Dcuda.version=cuda13 -DskipTests validate`, `mvn -s ~/.m2/settings_art.xml -f scala2.13/pom.xml -pl dist -am -Dbuildver=420 -Dcuda.version=cuda13 -DskipTests -Dmaven.javadoc.skip=true package`, `SPARK_HOME=/path/to/spark-4.2.0-bin-hadoop3 TESTS="hash_aggregate_test.py window_function_test.py" TEST="test_hash_reduction_collect_set_respect_nulls or test_window_aggs_for_rows_collect_set_respect_nulls" TEST_PARALLEL=0 ./run_pyspark_from_build.sh`, and earlier focused coverage for group-by and fully unbounded windows. ### Checklists Documentation - [ ] Updated for new or modified user-facing features or behaviors - [x] No user-facing change Testing - [x] Added or modified tests to cover new code paths - [ ] Covered by existing tests (Please provide the names of the existing tests in the PR description.) - [ ] Not required Performance - [ ] Tests ran and results are added in the PR description - [ ] Issue filed with a link in the PR description - [x] Not required --------- Signed-off-by: Firestarman <firestarmanllc@gmail.com>
1 parent f5f5ccb commit c406ee0

7 files changed

Lines changed: 203 additions & 24 deletions

File tree

‎integration_tests/src/main/python/hash_aggregate_test.py‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -890,6 +890,33 @@ def doit(spark):
890890
conf={'spark.sql.execution.useObjectHashAggregateExec': str(use_obj_hash_agg).lower()})
891891

892892

893+
@pytest.mark.skipif(not is_spark_420_or_later(),
894+
reason='collect_set RESPECT NULLS is introduced in Spark 4.2')
895+
@allow_non_gpu("ProjectExec")
896+
@ignore_order(local=True)
897+
@pytest.mark.parametrize('use_obj_hash_agg', [True, False], ids=idfn)
898+
def test_hash_groupby_collect_set_respect_nulls(use_obj_hash_agg):
899+
def doit(spark):
900+
return spark.sql("""
901+
SELECT a,
902+
sort_array(collect_set(b) IGNORE NULLS) AS ignore_set,
903+
sort_array(collect_set(b) RESPECT NULLS) AS respect_set
904+
FROM VALUES
905+
(1, 1),
906+
(1, NULL),
907+
(1, 1),
908+
(1, NULL),
909+
(2, NULL),
910+
(2, 5)
911+
AS tab(a, b)
912+
GROUP BY a
913+
""")
914+
915+
assert_gpu_and_cpu_are_equal_collect(
916+
doit,
917+
conf={'spark.sql.execution.useObjectHashAggregateExec': str(use_obj_hash_agg).lower()})
918+
919+
893920
@ignore_order(local=True)
894921
@pytest.mark.parametrize('use_obj_hash_agg', [True, False], ids=idfn)
895922
def test_hash_groupby_collect_list_of_maps(use_obj_hash_agg):
@@ -954,6 +981,33 @@ def test_hash_reduction_collect_set(data_gen):
954981
lambda spark: gen_df(spark, data_gen, length=100)
955982
.agg(f.sort_array(f.collect_set('b')), f.count('b')))
956983

984+
985+
@pytest.mark.skipif(not is_spark_420_or_later(),
986+
reason='collect_set RESPECT NULLS is introduced in Spark 4.2')
987+
@allow_non_gpu("ProjectExec")
988+
@ignore_order(local=True)
989+
def test_hash_reduction_collect_set_respect_nulls():
990+
def doit(spark):
991+
return spark.sql("""
992+
SELECT
993+
sort_array(collect_set(i) IGNORE NULLS) AS ignore_int,
994+
sort_array(collect_set(i) RESPECT NULLS) AS respect_int,
995+
sort_array(collect_set(d) IGNORE NULLS) AS ignore_double,
996+
sort_array(collect_set(d) RESPECT NULLS) AS respect_double,
997+
sort_array(collect_set(CAST(NULL AS INT)) IGNORE NULLS) AS ignore_all_null,
998+
sort_array(collect_set(CAST(NULL AS INT)) RESPECT NULLS) AS respect_all_null
999+
FROM VALUES
1000+
(CAST(1 AS INT), CAST(1.0 AS DOUBLE)),
1001+
(CAST(NULL AS INT), CAST(NULL AS DOUBLE)),
1002+
(CAST(1 AS INT), CAST('NaN' AS DOUBLE)),
1003+
(CAST(NULL AS INT), CAST('NaN' AS DOUBLE)),
1004+
(CAST(2 AS INT), CAST(2.0 AS DOUBLE))
1005+
AS tab(i, d)
1006+
""")
1007+
1008+
assert_gpu_and_cpu_are_equal_collect(doit)
1009+
1010+
9571011
@ignore_order(local=True)
9581012
@pytest.mark.parametrize('data_gen', _gen_data_for_collect_set_op, ids=idfn)
9591013
@allow_non_gpu(*non_utc_allow)

‎integration_tests/src/main/python/window_function_test.py‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2451,6 +2451,59 @@ def test_window_aggs_for_rows_collect_set():
24512451
'spark.sql.adaptive.enabled': 'false'})
24522452

24532453

2454+
@pytest.mark.skipif(not is_spark_420_or_later(),
2455+
reason='collect_set RESPECT NULLS is introduced in Spark 4.2')
2456+
@allow_non_gpu("ShuffleExchangeExec")
2457+
@ignore_order(local=True)
2458+
@pytest.mark.parametrize('data_type', ['INT', 'FLOAT', 'DOUBLE'], ids=idfn)
2459+
def test_window_aggs_for_rows_collect_set_respect_nulls(data_type):
2460+
def do_it(spark):
2461+
if data_type == 'INT':
2462+
values = """
2463+
(1, 1, '1'),
2464+
(1, 2, NULL),
2465+
(1, 3, '1'),
2466+
(1, 4, NULL),
2467+
(2, 1, NULL),
2468+
(2, 2, '5')
2469+
"""
2470+
else:
2471+
values = """
2472+
(1, 1, '1.0'),
2473+
(1, 2, NULL),
2474+
(1, 3, 'NaN'),
2475+
(1, 4, 'NaN'),
2476+
(2, 1, NULL),
2477+
(2, 2, '5.0')
2478+
"""
2479+
spark.sql(f"""
2480+
SELECT a, b, CAST(c AS {data_type}) AS c
2481+
FROM VALUES
2482+
{values}
2483+
AS tab(a, b, c)
2484+
""").createOrReplaceTempView("window_collect_table")
2485+
return spark.sql("""
2486+
SELECT a, b,
2487+
sort_array(ignore_set) AS ignore_set,
2488+
sort_array(respect_set) AS respect_set
2489+
FROM (
2490+
SELECT a, b,
2491+
collect_set(c) IGNORE NULLS OVER
2492+
(PARTITION BY a ORDER BY b
2493+
ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS ignore_set,
2494+
collect_set(c) RESPECT NULLS OVER
2495+
(PARTITION BY a ORDER BY b
2496+
ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS respect_set
2497+
FROM window_collect_table
2498+
) t
2499+
""")
2500+
2501+
assert_gpu_and_cpu_are_equal_collect(
2502+
do_it,
2503+
conf={'spark.rapids.sql.window.collectSet.enabled': True,
2504+
'spark.sql.adaptive.enabled': 'false'})
2505+
2506+
24542507
@ignore_order(local=True)
24552508
@allow_non_gpu(*non_utc_allow)
24562509
def test_window_aggs_for_fully_unbounded_partitioned_collect_set():
@@ -2520,6 +2573,44 @@ def test_window_aggs_for_fully_unbounded_partitioned_collect_set():
25202573
validate_execs_in_gpu_plan=['GpuUnboundedToUnboundedAggWindowExec'])
25212574

25222575

2576+
@pytest.mark.skipif(not is_spark_420_or_later(),
2577+
reason='collect_set RESPECT NULLS is introduced in Spark 4.2')
2578+
@allow_non_gpu("ShuffleExchangeExec")
2579+
@ignore_order(local=True)
2580+
def test_window_aggs_for_fully_unbounded_partitioned_collect_set_respect_nulls():
2581+
assert_gpu_and_cpu_are_equal_sql(
2582+
lambda spark: spark.sql("""
2583+
SELECT * FROM VALUES
2584+
(1, 1, 1),
2585+
(1, 2, NULL),
2586+
(1, 3, 1),
2587+
(1, 4, NULL),
2588+
(2, 1, NULL),
2589+
(2, 2, 5)
2590+
AS tab(a, b, c)
2591+
"""),
2592+
"window_collect_table",
2593+
"""
2594+
SELECT a, b,
2595+
sort_array(ignore_set) AS ignore_set,
2596+
sort_array(respect_set) AS respect_set
2597+
FROM (
2598+
SELECT a, b,
2599+
collect_set(c) IGNORE NULLS OVER
2600+
(PARTITION BY a ORDER BY b
2601+
ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) AS ignore_set,
2602+
collect_set(c) RESPECT NULLS OVER
2603+
(PARTITION BY a ORDER BY b
2604+
ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) AS respect_set
2605+
FROM window_collect_table
2606+
) t
2607+
""",
2608+
conf={'spark.rapids.sql.window.collectSet.enabled': True,
2609+
'spark.rapids.sql.window.unboundedAgg.enabled': True,
2610+
'spark.sql.adaptive.enabled': 'false'},
2611+
validate_execs_in_gpu_plan=['GpuUnboundedToUnboundedAggWindowExec'])
2612+
2613+
25232614
@ignore_order(local=True)
25242615
@allow_non_gpu(*non_utc_allow)
25252616
def test_window_aggs_for_fully_unbounded_unpartitioned_collect_set():

‎sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3850,15 +3850,17 @@ object GpuOverrides extends Logging {
38503850
}
38513851

38523852
override def convertToGpu(childExprs: Seq[Expression]): GpuExpression =
3853-
GpuCollectSet(childExprs.head, c.mutableAggBufferOffset, c.inputAggBufferOffset)
3853+
GpuCollectSet(childExprs.head, c.mutableAggBufferOffset, c.inputAggBufferOffset,
3854+
TypeUtilsShims.collectSetIgnoreNulls(c))
38543855

38553856
override def aggBufferAttribute: AttributeReference = {
38563857
val aggBuffer = c.aggBufferAttributes.head
38573858
aggBuffer.copy(dataType = c.dataType)(aggBuffer.exprId, aggBuffer.qualifier)
38583859
}
38593860

38603861
override def createCpuToGpuBufferConverter(): CpuToGpuAggregateBufferConverter =
3861-
new CpuToGpuCollectBufferConverter(c.child.dataType)
3862+
new CpuToGpuCollectBufferConverter(c.child.dataType,
3863+
!TypeUtilsShims.collectSetIgnoreNulls(c))
38623864

38633865
override def createGpuToCpuBufferConverter(): GpuToCpuAggregateBufferConverter =
38643866
new GpuToCpuCollectBufferConverter()

‎sql-plugin/src/main/scala/org/apache/spark/sql/rapids/aggregate/aggregateFunctions.scala‎

Lines changed: 45 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -115,28 +115,47 @@ class CudfMergeLists(override val dataType: DataType) extends CudfAggregate {
115115
* for the non-nested NaN equality in CudfCollectSet and CudfMergeSets.
116116
* Note that dataType is ArrayType(child.dataType) here.
117117
*/
118-
class CudfCollectSet(override val dataType: DataType) extends CudfAggregate {
119-
override lazy val reductionAggregate: cudf.ColumnVector => cudf.Scalar =
120-
(col: cudf.ColumnVector) => {
121-
val collectSet = dataType match {
122-
case ArrayType(FloatType | DoubleType, _) =>
123-
ReductionAggregation.collectSet(
124-
NullPolicy.EXCLUDE, NullEquality.EQUAL, TypeUtilsShims.collectSetFloatNanEquality)
125-
case _: DataType =>
126-
ReductionAggregation.collectSet(
127-
NullPolicy.EXCLUDE, NullEquality.EQUAL, NaNEquality.ALL_EQUAL)
128-
}
129-
col.reduce(collectSet, DType.LIST)
130-
}
118+
class CudfCollectSet(
119+
override val dataType: DataType,
120+
nullPolicy: NullPolicy) extends CudfAggregate {
121+
private lazy val reductionCollectSet: ReductionAggregation = dataType match {
122+
case ArrayType(FloatType | DoubleType, _) =>
123+
ReductionAggregation.collectSet(
124+
nullPolicy, NullEquality.EQUAL, TypeUtilsShims.collectSetFloatNanEquality)
125+
case _: DataType =>
126+
ReductionAggregation.collectSet(
127+
nullPolicy, NullEquality.EQUAL, NaNEquality.ALL_EQUAL)
128+
}
129+
131130
override lazy val groupByAggregate: GroupByAggregation = dataType match {
132131
case ArrayType(FloatType | DoubleType, _) =>
133132
GroupByAggregation.collectSet(
134-
NullPolicy.EXCLUDE, NullEquality.EQUAL, TypeUtilsShims.collectSetFloatNanEquality)
133+
nullPolicy, NullEquality.EQUAL, TypeUtilsShims.collectSetFloatNanEquality)
135134
case _: DataType =>
136135
GroupByAggregation.collectSet(
137-
NullPolicy.EXCLUDE, NullEquality.EQUAL, NaNEquality.ALL_EQUAL)
136+
nullPolicy, NullEquality.EQUAL, NaNEquality.ALL_EQUAL)
138137
}
139138
override val name: String = "CudfCollectSet"
139+
140+
override lazy val reductionAggregate: cudf.ColumnVector => cudf.Scalar =
141+
(col: cudf.ColumnVector) => {
142+
val rowCount = Math.toIntExact(col.getRowCount)
143+
if (nullPolicy == NullPolicy.INCLUDE && rowCount > 0) {
144+
// cuDF reduction collectSet currently drops nulls for some INCLUDE cases. Use the
145+
// group-by implementation for single-group reductions to preserve Spark's null semantics.
146+
withResource(Scalar.fromInt(0)) { keyScalar =>
147+
withResource(ColumnVector.fromScalar(keyScalar, rowCount)) { keys =>
148+
withResource(new cudf.Table(keys, col)) { table =>
149+
withResource(table.groupBy(0).aggregate(groupByAggregate.onColumn(1))) { result =>
150+
result.getColumn(1).getScalarElement(0)
151+
}
152+
}
153+
}
154+
}
155+
} else {
156+
col.reduce(reductionCollectSet, DType.LIST)
157+
}
158+
}
140159
}
141160

142161
class CudfMergeSets(override val dataType: DataType) extends CudfAggregate {
@@ -1970,10 +1989,17 @@ case class GpuCollectList(
19701989
case class GpuCollectSet(
19711990
child: Expression,
19721991
mutableAggBufferOffset: Int = 0,
1973-
inputAggBufferOffset: Int = 0)
1992+
inputAggBufferOffset: Int = 0,
1993+
ignoreNulls: Boolean = true)
19741994
extends GpuCollectBase with GpuUnboundedToUnboundedWindowAgg {
19751995

1976-
override lazy val updateAggregates: Seq[CudfAggregate] = Seq(new CudfCollectSet(dataType))
1996+
private lazy val nullPolicy: NullPolicy =
1997+
if (ignoreNulls) NullPolicy.EXCLUDE else NullPolicy.INCLUDE
1998+
1999+
override protected def arrayContainsNull: Boolean = !ignoreNulls
2000+
2001+
override lazy val updateAggregates: Seq[CudfAggregate] =
2002+
Seq(new CudfCollectSet(dataType, nullPolicy))
19772003
override lazy val mergeAggregates: Seq[CudfAggregate] = Seq(new CudfMergeSets(dataType))
19782004
override lazy val evaluateExpression: Expression = outputBuf
19792005
override def aggBufferAttributes: Seq[AttributeReference] = outputBuf :: Nil
@@ -1987,10 +2013,10 @@ case class GpuCollectSet(
19872013
override def windowAggregation(
19882014
inputs: Seq[(ColumnVector, Int)]): RollingAggregationOnColumn = child.dataType match {
19892015
case FloatType | DoubleType =>
1990-
RollingAggregation.collectSet(NullPolicy.EXCLUDE, NullEquality.EQUAL,
2016+
RollingAggregation.collectSet(nullPolicy, NullEquality.EQUAL,
19912017
TypeUtilsShims.collectSetFloatNanEquality).onColumn(inputs.head._2)
19922018
case _ =>
1993-
RollingAggregation.collectSet(NullPolicy.EXCLUDE, NullEquality.EQUAL,
2019+
RollingAggregation.collectSet(nullPolicy, NullEquality.EQUAL,
19942020
NaNEquality.ALL_EQUAL).onColumn(inputs.head._2)
19952021
}
19962022

‎sql-plugin/src/main/spark330/scala/com/nvidia/spark/rapids/shims/TypeUtilsShims.scala‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ package com.nvidia.spark.rapids.shims
2525

2626
import ai.rapids.cudf.NaNEquality
2727

28-
import org.apache.spark.sql.catalyst.expressions.aggregate.CollectList
28+
import org.apache.spark.sql.catalyst.expressions.aggregate.{CollectList, CollectSet}
2929
import org.apache.spark.sql.catalyst.util.TypeUtils
3030

3131
/**
@@ -39,6 +39,8 @@ object TypeUtilsShims {
3939

4040
def collectListIgnoreNulls(_collectList: CollectList): Boolean = true
4141

42+
def collectSetIgnoreNulls(_collectSet: CollectSet): Boolean = true
43+
4244
val useImprovedAsinhByDefault: Boolean = false
4345

4446
def isUnsupportedArrowAggregatePythonEvalType(evalType: Int): Boolean = false

‎sql-plugin/src/main/spark330db/scala/com/nvidia/spark/rapids/shims/TypeUtilsShims.scala‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ package com.nvidia.spark.rapids.shims
4949
import ai.rapids.cudf.NaNEquality
5050

5151
import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
52-
import org.apache.spark.sql.catalyst.expressions.aggregate.CollectList
52+
import org.apache.spark.sql.catalyst.expressions.aggregate.{CollectList, CollectSet}
5353
import org.apache.spark.sql.types.{DataType, NullType, NumericType}
5454

5555
/**
@@ -69,6 +69,8 @@ object TypeUtilsShims {
6969

7070
def collectListIgnoreNulls(_collectList: CollectList): Boolean = true
7171

72+
def collectSetIgnoreNulls(_collectSet: CollectSet): Boolean = true
73+
7274
val useImprovedAsinhByDefault: Boolean = false
7375

7476
def isUnsupportedArrowAggregatePythonEvalType(evalType: Int): Boolean = false

‎sql-plugin/src/main/spark420/scala/com/nvidia/spark/rapids/shims/TypeUtilsShims.scala‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ package com.nvidia.spark.rapids.shims
2121
import ai.rapids.cudf.NaNEquality
2222

2323
import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
24-
import org.apache.spark.sql.catalyst.expressions.aggregate.CollectList
24+
import org.apache.spark.sql.catalyst.expressions.aggregate.{CollectList, CollectSet}
2525
import org.apache.spark.sql.types.{DataType, NullType, NumericType}
2626

2727
/**
@@ -42,6 +42,8 @@ object TypeUtilsShims {
4242

4343
def collectListIgnoreNulls(collectList: CollectList): Boolean = collectList.ignoreNulls
4444

45+
def collectSetIgnoreNulls(collectSet: CollectSet): Boolean = collectSet.ignoreNulls
46+
4547
// Spark 4.2 uses fdlibm asinh, so the default CPU behavior matches the stable cuDF kernel.
4648
val useImprovedAsinhByDefault: Boolean = true
4749

0 commit comments

Comments
 (0)