Normalize float/double inside GpuCollectSet for Spark 4.2 [databricks]

integration_tests/src/main/python/hash_aggregate_test.py

@@ -1156,6 +1156,123 @@ def test_hash_groupby_collect_partial_replace_fallback(data_gen,
         non_exist_classes=','.join(non_exist_clz),
         conf=conf)
 
+
+@pytest.mark.skipif(not is_spark_420_or_later(),
+                    reason='collect_set float/double bit-key buffers and RESPECT NULLS need Spark 4.2+')
+@ignore_order(local=True)
+@allow_non_gpu('ObjectHashAggregateExec', 'SortAggregateExec',
+               'ShuffleExchangeExec', 'HashPartitioning', 'SortExec',
+               'SortArray', 'Alias', 'Literal', 'CollectSet',
+               'AggregateExpression', 'ProjectExec', 'Cast', *non_utc_allow)
+@pytest.mark.parametrize('replace_mode', _replace_modes_non_distinct, ids=idfn)
+@pytest.mark.parametrize('fp_type', ['FLOAT', 'DOUBLE'], ids=idfn)
+def test_hash_groupby_collect_partial_replace_respect_nulls_float_double(replace_mode, fp_type):
+    """Mixed CPU/GPU CollectSet for Float/Double with IGNORE/RESPECT NULLS on Spark 4.2+."""
+    conf = {'spark.rapids.sql.hashAgg.replaceMode': replace_mode,
+            'spark.sql.adaptive.enabled': 'false',
+            'spark.sql.execution.useObjectHashAggregateExec': 'false'}
+
+    cpu_clz, gpu_clz = ['CollectSet'], ['GpuCollectSet']
+    if is_databricks_runtime():
+        if replace_mode == 'partial':
+            exist_clz, non_exist_clz = cpu_clz, gpu_clz
+        else:
+            exist_clz, non_exist_clz = gpu_clz, cpu_clz
+    else:
+        exist_clz = cpu_clz + gpu_clz
+        non_exist_clz = []
+
+    # Mixed-null and all-null groups exercise containsNull on the Spark 4.2 bit-key buffer.
+    sql = f"""
+        SELECT a,
+               sort_array(collect_set(b) IGNORE NULLS) AS ignore_set,
+               sort_array(collect_set(b) RESPECT NULLS) AS respect_set
+        FROM VALUES
+            (1, CAST(1.0 AS {fp_type})),
+            (1, CAST(NULL AS {fp_type})),
+            (1, CAST(1.0 AS {fp_type})),
+            (1, CAST(NULL AS {fp_type})),
+            (2, CAST(NULL AS {fp_type})),
+            (2, CAST(NULL AS {fp_type})),
+            (3, CAST(5.0 AS {fp_type})),
+            (3, CAST(NULL AS {fp_type}))
+        AS tab(a, b)
+        GROUP BY a
+    """
+
+    assert_cpu_and_gpu_are_equal_collect_with_capture(
+        lambda spark: spark.sql(sql),
+        exist_classes=','.join(exist_clz),
+        non_exist_classes=','.join(non_exist_clz),
+        conf=conf)
+
+
+@pytest.mark.skipif(not is_spark_420_or_later(),
+                    reason='collect_set float/double normalized bit-key buffers need Spark 4.2+')
+@ignore_order(local=True)
+@allow_non_gpu('ObjectHashAggregateExec', 'SortAggregateExec',
+               'ShuffleExchangeExec', 'HashPartitioning', 'SortExec',
+               'SortArray', 'Alias', 'Literal', 'CollectSet',
+               'AggregateExpression', 'ProjectExec', 'Cast', *non_utc_allow)
+@pytest.mark.parametrize('replace_mode', _replace_modes_non_distinct, ids=idfn)
+@pytest.mark.parametrize('fp_type', ['FLOAT', 'DOUBLE'], ids=idfn)
+def test_hash_groupby_collect_partial_replace_float_double_edge_cases(replace_mode, fp_type):
+    """Deterministic +0/-0/NaN/inf/null CollectSet round-trip across mixed CPU/GPU stages."""
+    conf = {'spark.rapids.sql.hashAgg.replaceMode': replace_mode,
+            'spark.sql.adaptive.enabled': 'false',
+            'spark.sql.execution.useObjectHashAggregateExec': 'false'}
+
+    cpu_clz, gpu_clz = ['CollectSet'], ['GpuCollectSet']
+    if is_databricks_runtime():
+        if replace_mode == 'partial':
+            exist_clz, non_exist_clz = cpu_clz, gpu_clz
+        else:
+            exist_clz, non_exist_clz = gpu_clz, cpu_clz
+    else:
+        exist_clz = cpu_clz + gpu_clz
+        non_exist_clz = []
+
+    # Put +0, -0, multiple NaN payloads, +/-inf and null in the same group so
+    # normalization/dedup is forced rather than relying on RepeatSeqGen sampling.
+    sql = f"""
+        SELECT a, sort_array(collect_set(b)) AS s
+        FROM VALUES
+            (1, CAST(0.0 AS {fp_type})),
+            (1, CAST(-0.0 AS {fp_type})),
+            (1, CAST('NaN' AS {fp_type})),
+            (1, CAST('NaN' AS {fp_type})),
+            (1, CAST('Infinity' AS {fp_type})),
+            (1, CAST('-Infinity' AS {fp_type})),
+            (1, CAST(NULL AS {fp_type})),
+            (1, CAST(1.5 AS {fp_type})),
+            (2, CAST(NULL AS {fp_type})),
+            (2, CAST(NULL AS {fp_type}))
+        AS tab(a, b)
+        GROUP BY a
+    """
+
+    assert_cpu_and_gpu_are_equal_collect_with_capture(
+        lambda spark: spark.sql(sql),
+        exist_classes=','.join(exist_clz),
+        non_exist_classes=','.join(non_exist_clz),
+        conf=conf)
+
+
+@pytest.mark.skipif(not is_spark_420_or_later(),
+                    reason='collect_set float/double normalized bit-key buffers need Spark 4.2+')
+@ignore_order(local=True)
+@allow_non_gpu('ProjectExec', 'Cast', *non_utc_allow)
+@pytest.mark.parametrize('fp_type', ['FLOAT', 'DOUBLE'], ids=idfn)
+def test_hash_reduction_collect_set_float_double_empty(fp_type):
+    """Empty typed Float/Double collect_set must return an empty array, not null."""
+    assert_gpu_and_cpu_are_equal_collect(
+        lambda spark: spark.sql(f"""
+            SELECT sort_array(collect_set(b)) AS s
+            FROM VALUES (CAST(1.0 AS {fp_type})) AS tab(b)
+            WHERE 1 = 0
+        """))
+
+
 # The special case is to test when the physical plan is being re-written due to the re-optimize
 # of AQE taking effect, which is rare in real world scenarios. So far, this kind of problem only
 # has encountered when there exists a local aggregate ahead of the TypedImperativeAggregate. Then,

integration_tests/src/main/python/window_function_test.py

@@ -2536,7 +2536,15 @@ def test_window_aggs_for_fully_unbounded_partitioned_collect_set():
     runs through the `GpuUnboundedToUnboundedAggWindowExec` (which optimizes it to run via sort-based group-by
     aggregations).
     Note: This optimization only holds for the partitioned case.  Unpartitioned windows are not supported yet.
+
+    On Spark 4.2+, Float/Double CollectSet uses a bit-key hash-agg projection that is incompatible with
+    GpuUnboundedToUnboundedAggWindowExec. Mixed-type unbounded windows that include Float/Double therefore
+    fall back to regular GpuWindowExec for the whole Window node (allBatched=false). Float/Double-only
+    unbounded coverage is in test_window_aggs_for_fully_unbounded_partitioned_collect_set_float_double_spark420.
     """
+    # On Spark 4.2+ float/double force the mixed WindowExec onto GpuWindowExec.
+    expected_exec = (['GpuWindowExec'] if is_spark_420_or_later()
+                     else ['GpuUnboundedToUnboundedAggWindowExec'])
     assert_gpu_and_cpu_are_equal_sql(
         lambda spark: gen_df(spark, _gen_data_for_collect_set, length=2048),
         "window_collect_table",
@@ -2594,7 +2602,53 @@ def test_window_aggs_for_fully_unbounded_partitioned_collect_set():
               'spark.rapids.sql.window.unboundedAgg.enabled': True,
               'spark.sql.parquet.int96RebaseModeInWrite': 'LEGACY',
               'spark.sql.adaptive.enabled': 'false'},
-        validate_execs_in_gpu_plan=['GpuUnboundedToUnboundedAggWindowExec'])
+        validate_execs_in_gpu_plan=expected_exec)
+
+
+@pytest.mark.skipif(not is_spark_420_or_later(),
+                    reason='Spark 4.2 float/double CollectSet uses bit-key hash path incompatible '
+                           'with GpuUnboundedToUnboundedAggWindowExec')
+@ignore_order(local=True)
+@allow_non_gpu('ShuffleExchangeExec')
+@pytest.mark.parametrize('fp_type', ['FLOAT', 'DOUBLE'], ids=idfn)
+def test_window_aggs_for_fully_unbounded_partitioned_collect_set_float_double_spark420(fp_type):
+    """
+    Spark 4.2+ Float/Double collect_set over fully unbounded frames must fall back to
+    GpuWindowExec (normalize in-place) rather than the unbounded group-by shortcut.
+    """
+    assert_gpu_and_cpu_are_equal_sql(
+        lambda spark: spark.sql(f"""
+            SELECT * FROM VALUES
+                (1, 1, CAST(0.0 AS {fp_type})),
+                (1, 2, CAST(-0.0 AS {fp_type})),
+                (1, 3, CAST('NaN' AS {fp_type})),
+                (1, 4, CAST('NaN' AS {fp_type})),
+                (1, 5, CAST(NULL AS {fp_type})),
+                (1, 6, CAST('Infinity' AS {fp_type})),
+                (2, 1, CAST(1.5 AS {fp_type})),
+                (2, 2, CAST(NULL AS {fp_type}))
+            AS tab(a, b, c)
+        """),
+        "window_collect_table",
+        """
+        SELECT a, b,
+               sort_array(ignore_set) AS ignore_set,
+               sort_array(respect_set) AS respect_set
+        FROM (
+            SELECT a, b,
+                   collect_set(c) IGNORE NULLS OVER
+                     (PARTITION BY a ORDER BY b
+                      ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) AS ignore_set,
+                   collect_set(c) RESPECT NULLS OVER
+                     (PARTITION BY a ORDER BY b
+                      ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) AS respect_set
+            FROM window_collect_table
+        ) t
+        """,
+        conf={'spark.rapids.sql.window.collectSet.enabled': True,
+              'spark.rapids.sql.window.unboundedAgg.enabled': True,
+              'spark.sql.adaptive.enabled': 'false'},
+        validate_execs_in_gpu_plan=['GpuWindowExec'])
 
 
 @pytest.mark.skipif(not is_spark_420_or_later(),

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

@@ -3862,16 +3862,23 @@ object GpuOverrides extends Logging {
 
         override def aggBufferAttribute: AttributeReference = {
           val aggBuffer = c.aggBufferAttributes.head
-          aggBuffer.copy(dataType = c.dataType)(aggBuffer.exprId, aggBuffer.qualifier)
+          // Match Spark 4.2+ CollectSet buffer layout for float/double (normalized bit keys).
+          val ignoreNulls = TypeUtilsShims.collectSetIgnoreNulls(c)
+          val bufferElementType =
+            TypeUtilsShims.collectSetCpuBufferElementType(c.child.dataType)
+          aggBuffer.copy(dataType = ArrayType(bufferElementType, !ignoreNulls))(
+            aggBuffer.exprId, aggBuffer.qualifier)
         }
 
-        override def createCpuToGpuBufferConverter(): CpuToGpuAggregateBufferConverter =
-          new CpuToGpuCollectSetBufferConverter(c.child.dataType,
-            !TypeUtilsShims.collectSetIgnoreNulls(c))
+        override def createCpuToGpuBufferConverter(): CpuToGpuAggregateBufferConverter = {
+          val ignoreNulls = TypeUtilsShims.collectSetIgnoreNulls(c)
+          new CpuToGpuCollectBufferConverter(
+            TypeUtilsShims.collectSetCpuBufferElementType(c.child.dataType),
+            !ignoreNulls)
+        }
 
         override def createGpuToCpuBufferConverter(): GpuToCpuAggregateBufferConverter =
-          new GpuToCpuCollectSetBufferConverter(c.child.dataType,
-            !TypeUtilsShims.collectSetIgnoreNulls(c))
+          new GpuToCpuCollectBufferConverter()
 
         override val supportBufferConversion: Boolean = true
 

sql-plugin/src/main/scala/com/nvidia/spark/rapids/window/GpuWindowExecMeta.scala

@@ -557,8 +557,10 @@ object GpuWindowExecMeta {
         false // Must be both unbounded, and have a group by specification.
       } else {
         func match {
-          case _: GpuUnboundedToUnboundedWindowAgg => true
-          case GpuAggregateExpression(_: GpuUnboundedToUnboundedWindowAgg, _, _, _, _) => true
+          case a: GpuUnboundedToUnboundedWindowAgg =>
+            a.supportsUnboundedToUnboundedWindowExec
+          case GpuAggregateExpression(a: GpuUnboundedToUnboundedWindowAgg, _, _, _, _) =>
+            a.supportsUnboundedToUnboundedWindowExec
           case _ => false
         }
       }

sql-plugin/src/main/scala/com/nvidia/spark/rapids/window/GpuWindowExpression.scala

@@ -985,8 +985,13 @@ trait GpuUnboundToUnboundWindowWithFixer {
 /**
  * This is used to tag a GpuAggregateFunction that it has been tested to work properly
  * with `GpuUnboundedToUnboundedAggWindowExec`.
+ *
+ * Implementations may opt out when their hash-agg [[GpuAggregateFunction.inputProjection]]
+ * is incompatible with the unbounded-agg shortcut (which only accepts BoundReference/Literal).
  */
-trait GpuUnboundedToUnboundedWindowAgg extends GpuAggregateFunction
+trait GpuUnboundedToUnboundedWindowAgg extends GpuAggregateFunction {
+  def supportsUnboundedToUnboundedWindowExec: Boolean = true
+}
 
 /**
  * Fixes up a count operation for unbounded preceding to unbounded following