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

[firestarman]

Contributes to #15463.

related to this comment

Description

• Normalize NaN / -0.0 and store Spark 4.2 CollectSet float/double agg buffers as Int/Long bit keys inside GpuCollectSet, so GPU uniqueness matches CPU without host-side per-row float↔bits converters.
• Reuse the generic Collect buffer converters for mixed CPU/GPU hashAgg stages, because GPU and CPU buffer layouts now match on Spark 4.2+.
• Keep pre-4.2 shims on the previous Float/Double buffer path via TypeUtilsShims.collectSetCpuBufferElementType.
• Opt Spark 4.2 float/double GpuCollectSet out of GpuUnboundedToUnboundedAggWindowExec (bit-key inputProjection is incompatible with that shortcut) and keep the window path on regular GpuWindowExec with GpuNormalizeNaNAndZero.
• Reuse cuDF normalizeNANsAndZeros() instead of a custom NaN/-0.0 normalize helper.
• Validated locally on Spark 4.2.0 / Scala 2.13 / CUDA 13 with DATAGEN_SEED=1785353212: new/updated ITs for mixed-stage Float/Double RESPECT NULLS, deterministic +0/-0/NaN/inf edges, empty typed reduction, and fully-unbounded Float/Double windows (14 passed); also mvn -f scala2.13/pom.xml -Dbuildver=420 -Dcuda.version=cuda13 -DskipTests -pl sql-plugin,dist,integration_tests -am package and mvn -Dbuildver=330 -Dcuda.version=cuda13 -DskipTests -pl sql-plugin -am package.

This follows the direction discussed on #15463 (normalize inside collect_set rather than expanding host-side converters).

Performance

Operator-level microbench (not NDS) comparing current main (#15455 host CollectSet float↔bits converters) vs this PR.

• Hardware: NVIDIA RTX 5880 Ada Generation
• Spark 4.2.0 / Scala 2.13 / CUDA 13 / spark.rapids.memory.gpu.allocSize=8192m
• Data includes ~2% NaN and ~2% -0.0 in float/double columns
• Method: 1 warmup + 3 iters, report median wall time
• Timed SQL (equivalent to the DataFrame microbench):

SELECT SUM(sf) AS sum_f, SUM(sd) AS sum_d
FROM (
  SELECT
    k,
    SIZE(COLLECT_SET(f)) AS sf,
    SIZE(COLLECT_SET(d)) AS sd
  FROM (
    SELECT
      CAST(id % ${num_groups} AS INT) AS k,
      CASE
        WHEN (id % 50) = 0 THEN CAST('NaN' AS FLOAT)
        WHEN (id % 50) = 1 THEN CAST(-0.0 AS FLOAT)
        ELSE CAST(CAST((id % 997) AS FLOAT) / 10.0 AS FLOAT)
      END AS f,
      CASE
        WHEN (id % 50) = 2 THEN CAST('NaN' AS DOUBLE)
        WHEN (id % 50) = 3 THEN CAST(-0.0 AS DOUBLE)
        ELSE CAST(CAST((id % 1009) AS DOUBLE) / 10.0 AS DOUBLE)
      END AS d
    FROM range(0, ${num_rows})
  )
  GROUP BY k
)

Suite C: 200,000,000 rows / 500,000 groups (mixed modes ~12–13s)

Case	main median (s)	PR median (s)	Speedup (main/PR)	Notes
pure_gpu	3.160	3.235	0.98x	pure GPU (short at this scale)
mixed_partial_gpu	12.443	11.553	1.08x	GPU partial + CPU final
mixed_final_gpu	13.133	12.935	1.02x	CPU partial + GPU final

Checksums matched (sum_f=sum_d=192020000).

Suite D: 1,000,000,000 rows / 1,000,000 groups (pure GPU ~15s)

Case	main median (s)	PR median (s)	Speedup (main/PR)	Notes
pure_gpu	15.725	15.436	1.02x	pure GPU

Checksums matched (sum_f=957160000, sum_d=960040000).

Takeaway: no meaningful wall-time regression vs main at multi-second scale. Mixed partial lean (~1.08x) is consistent with avoiding host float↔bits conversion on GPU buffer boundaries; pure GPU is within noise (~±2%).

Checklists

Documentation

• Updated for new or modified user-facing features or behaviors
• No user-facing change

Testing

• 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
• Not required

[greptile-apps]

Greptile Summary

The PR changes Spark 4.2 float/double GpuCollectSet to normalize values and retain INT/LONG bit keys throughout aggregate buffers, allowing mixed CPU/GPU stages to share Spark’s buffer layout.

• Adds normalized float/double bit-key projection and final value reconstruction.
• Reuses generic aggregate-buffer converters and preserves pre-4.2 behavior through shims.
• Routes incompatible fully unbounded float/double windows through regular GpuWindowExec.
• Adds mixed-stage, null-handling, floating-point edge-case, empty-input, and window coverage.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
sql-plugin/src/main/scala/org/apache/spark/sql/rapids/aggregate/aggregateFunctions.scala	Implements normalized bit-key aggregate buffers, value reconstruction, and window-specific normalization with consistent resource ownership.
sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala	Aligns declared CollectSet buffer attributes and mixed-stage converters with the Spark-version-specific CPU buffer element type.
sql-plugin/src/main/scala/com/nvidia/spark/rapids/window/GpuWindowExpression.scala	Adds an aggregate capability hook allowing implementations to opt out of the unbounded aggregate-window shortcut.
sql-plugin/src/main/scala/com/nvidia/spark/rapids/window/GpuWindowExecMeta.scala	Consults the new capability hook when selecting the optimized fully unbounded window path.
sql-plugin/src/main/spark420/scala/com/nvidia/spark/rapids/shims/TypeUtilsShims.scala	Documents and exposes Spark 4.2’s INT/LONG CollectSet buffer representation while retaining version isolation.
integration_tests/src/main/python/hash_aggregate_test.py	Adds CPU/GPU mixed-stage coverage for null modes, floating-point edge cases, and empty aggregate input.
integration_tests/src/main/python/window_function_test.py	Verifies Spark 4.2 float/double fully unbounded CollectSet uses regular GPU window execution with matching results.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  Input[Float or Double input] --> Normalize[Normalize NaN and signed zero]
  Normalize --> Bits[Bit-cast to INT or LONG keys]
  Bits --> Aggregate[GPU collect_set update or merge]
  CPU[CPU partial buffer] <--> Convert[Generic buffer converters]
  Convert <--> Aggregate
  Aggregate --> Buffer[INT or LONG array buffer]
  Buffer --> Values[Bit-cast keys back to Float or Double]
  Values --> Result[CollectSet result array]
  Normalize --> Window[Regular GpuWindowExec rolling collect_set]

Loading

_{Reviews (3): Last reviewed commit: "Address Spark 4.2 CollectSet review: win..." | Re-trigger Greptile}

[firestarman]

build

[jihoonson]

This is a performance improvement. Please evaluate performance and report the result.

[firestarman]

@jihoonson Thanks for the review. I ran a local micro-benchmark (operator-level) comparing current main (#15455 host float↔bits converters) vs this PR on Spark 4.2.0 / Scala 2.13 / RTX 5880 Ada.

Summary (median, 1 warmup + 3 iters):

• Suite C (200M rows / 500k groups): pure_gpu 0.98x, mixed_partial_gpu 1.08x, mixed_final_gpu 1.02x (speedup = main/PR)
• Suite D (1B rows / 1M groups, pure GPU): 1.02x

Details and tables are in the updated PR description. Please take another look.

[wjxiz1992]

Requesting changes for the Spark 4.2 fully-unbounded window correctness issue described inline. The other inline comments cover nullable mixed-stage coverage, deterministic floating-point edge coverage, and reuse of fused normalization.

PR description follow-up: link the tracking issue with Contributes to #15463 or Closes #15463.

The existing performance request appears addressed by the benchmark section. I did not run a local GPU runtime test; the blocking finding is based on exact-head source tracing through the existing integration-test path.

[wjxiz1992]

Could you add a Spark 4.2 mixed CPU/GPU test for Float/Double collect_set(... RESPECT NULLS) here? This buffer schema derives containsNull from ignoreNulls, but the mixed-stage tests only exercise default IGNORE NULLS; the existing RESPECT NULLS case uses integral input and does not cross a CPU/GPU aggregate boundary. Please cover both CPU-to-GPU and GPU-to-CPU conversion with mixed-null and all-null groups.

[wjxiz1992]

Blocking: for Spark 4.2 Float/Double, this becomes GpuCollectSetNormalizedBitKey(child), while GpuCollectSet still opts into GpuUnboundedToUnboundedAggWindowExec. That executor selects every marker implementation, but its stage builder accepts only GpuBoundReference/GpuLiteral, so the existing fully-unbounded Float/Double window IT reaches IllegalStateException("Unexpected expression"). It also exposes aggregate result types directly and never applies evaluateExpression, so simply accepting this unary projection would still omit GpuCollectSetBitKeysToValues. Please either guard this path back to regular GpuWindowExec or add both projection and finalizer support, with a Spark 4.2 regression.

[wjxiz1992]

Could you add deterministic Spark 4.2 Float/Double regression data for this normalization and round trip? The current RepeatSeqGen(FloatGen/DoubleGen) data samples before repeating, so it does not guarantee that +0/-0 and distinct NaN payloads land in the same group. Please include +0, -0, multiple NaN payloads, infinities, null, both mixed replacement directions, and an empty typed Float/Double aggregation.

[wjxiz1992]

Could this reuse the existing GpuNormalizeNaNAndZero / cuDF normalizeNANsAndZeros() path? The current sequence materializes isNan + ifElse, followed by another equality mask + ifElse in HashUtils.normalizeInput, for every hash/window input batch. The existing native operation performs the same NaN/signed-zero normalization in one transform and avoids duplicate normalization logic and temporaries.

[firestarman]

Contributes to #15463 or Closes #15463

We can not close that issue because it covers more things than what this PR can fix.

[firestarman]

@wjxiz1992 Thanks for the detailed review. Addressed in ef366d0:

1. Blocking unbounded window: Spark 4.2 float/double GpuCollectSet now opts out of GpuUnboundedToUnboundedAggWindowExec via supportsUnboundedToUnboundedWindowExec (bit-key inputProjection is incompatible with that shortcut). The window path stays on regular GpuWindowExec with windowInputProjection + rolling aggregation. Added test_window_aggs_for_fully_unbounded_partitioned_collect_set_float_double_spark420, and updated the mixed-type unbounded IT to expect GpuWindowExec on Spark 4.2+.

2. Normalization reuse: Removed CollectSetFloatingNormalize / GpuNormalizeFloatingPointForCollectSet; hash/window paths now use cuDF normalizeNANsAndZeros() / GpuNormalizeNaNAndZero.

3. RESPECT NULLS mixed-stage Float/Double: Added test_hash_groupby_collect_partial_replace_respect_nulls_float_double covering both replace modes with mixed-null and all-null groups.

4. Deterministic edge coverage: Added test_hash_groupby_collect_partial_replace_float_double_edge_cases (+0/-0/NaN/inf/null) and test_hash_reduction_collect_set_float_double_empty.

PR description updated with Contributes to #15463. Local IT result: 14 passed for the new/updated cases above. Please take another look.

[firestarman]

build

[jihoonson]

@wjxiz1992 Thanks for the detailed review. Addressed in ef366d0:

1. Blocking unbounded window: Spark 4.2 float/double GpuCollectSet now opts out of GpuUnboundedToUnboundedAggWindowExec via supportsUnboundedToUnboundedWindowExec (bit-key inputProjection is incompatible with that shortcut). The window path stays on regular GpuWindowExec with windowInputProjection + rolling aggregation. Added test_window_aggs_for_fully_unbounded_partitioned_collect_set_float_double_spark420, and updated the mixed-type unbounded IT to expect GpuWindowExec on Spark 4.2+.
2. Normalization reuse: Removed CollectSetFloatingNormalize / GpuNormalizeFloatingPointForCollectSet; hash/window paths now use cuDF normalizeNANsAndZeros() / GpuNormalizeNaNAndZero.
3. RESPECT NULLS mixed-stage Float/Double: Added test_hash_groupby_collect_partial_replace_respect_nulls_float_double covering both replace modes with mixed-null and all-null groups.
4. Deterministic edge coverage: Added test_hash_groupby_collect_partial_replace_float_double_edge_cases (+0/-0/NaN/inf/null) and test_hash_reduction_collect_set_float_double_empty.

PR description updated with Contributes to #15463. Local IT result: 14 passed for the new/updated cases above. Please take another look.

Please take the perspective of readers/reviewers into account. Seeing @wjxiz1992's comments (such as #15546 (comment)) unanswered and unresolved can make readers and reviewers easily think that those comments are left unaddressed. Now I know that those comments are addressed, but need to scroll up and down to map your answers and his comments. Please leave your answer directly to his comment.