Fix Spark 4.2 collect_set float buffer conversion for mixed aggs by firestarman · Pull Request #15455 · NVIDIA/cudf-spark

firestarman · 2026-07-30T08:08:55Z

Fixes #15454.

Description

Convert CollectSet float/double aggregation buffers between GPU logical values and Spark 4.2 CPU normalized bit-pattern keys during mixed hashAgg stages, so NaN and signed-zero uniqueness matches pure CPU.
Add collectSetCpuBufferElementType shim (bit-keyed on Spark 4.2, identity earlier) and CollectSet-specific GPU↔CPU buffer converters used by GpuOverrides.
Validate with Spark 4.2.0 / Scala 2.13 and DATAGEN_SEED=1785353212: 352 passed including the previously failing test_hash_groupby_collect_partial_replace_fallback / test_hash_groupby_collect_partial_replace_with_distinct_fallback Float cases; also mvn -f scala2.13/pom.xml -Dbuildver=420 -Dcuda.version=cuda13 -DskipTests -pl sql-plugin verify.

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 · 2026-07-30T08:11:12Z

Greptile Summary

Fixes Spark 4.2 mixed CPU/GPU collect_set aggregation-buffer conversion.

Adds CollectSet-specific converters between logical float/double values and Spark’s normalized integer/long bit-pattern keys.
Adds version-specific CPU buffer element-type shims, preserving identity behavior before Spark 4.2.
Wires the converters into the CollectSet GPU override.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala	Selects the new CollectSet-specific buffer converters at mixed CPU/GPU aggregation boundaries.
sql-plugin/src/main/scala/org/apache/spark/sql/rapids/aggregate/aggregateFunctions.scala	Implements Spark-compatible float/double normalization and bidirectional aggregate-buffer conversion.
sql-plugin/src/main/spark330/scala/com/nvidia/spark/rapids/shims/TypeUtilsShims.scala	Preserves the pre-Spark-4.2 identity buffer representation.
sql-plugin/src/main/spark330db/scala/com/nvidia/spark/rapids/shims/TypeUtilsShims.scala	Preserves identity buffer representation across the listed pre-4.2 and Databricks shims.
sql-plugin/src/main/spark420/scala/com/nvidia/spark/rapids/shims/TypeUtilsShims.scala	Maps top-level float and double CollectSet CPU buffers to Spark 4.2 integer and long bit-key representations.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  CPU["Spark CPU CollectSet buffer<br/>normalized Int/Long bit keys"]
  C2G["CPU-to-GPU converter<br/>bits to Float/Double"]
  GPU["GPU CollectSet buffer<br/>logical Float/Double values"]
  G2C["GPU-to-CPU converter<br/>normalize values to bits"]
  CPU --> C2G --> GPU
  GPU --> G2C --> CPU

_{Reviews (3): Last reviewed commit: "Trigger CI for PR #15455" | Re-trigger Greptile}

Spark 4.2 stores CollectSet float/double buffers as normalized bit patterns, so convert GPU logical values when crossing CPU/GPU stages. Fixes NVIDIA#15454. Signed-off-by: Firestarman <firestarmanllc@gmail.com>

firestarman · 2026-07-30T09:29:29Z

build

ttnghia · 2026-07-30T16:49:56Z

+    elementType match {
+      case FloatType =>
+        while (i < n) {


This loop is running on the CPU and can be a bottleneck, since it iterates over every row and convert one row at a time. The buffer conversion like this is also seen in several other places as well. We should move all these conversion into running on the GPU. I'm not blocking this fix, but we should do the follow up work for migrate all buffer conversion into jni code with native GPU kernels.

Filed #15463.

Signed-off-by: Firestarman <firestarmanllc@gmail.com>

firestarman · 2026-07-31T01:50:37Z

build

firestarman · 2026-07-31T01:51:27Z

~~premerge has already passed, but some required jobs are pending, so submit an empty commit to re-trigger them with skip-ci.~~
It should be [skip ci] not [skip-ci], so premerge has been launched again.

#15546) Contributes to #15463. related to [this comment](#15463 (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): ```sql 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 - [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 - [x] Tests ran and results are added in the PR description - [ ] Issue filed with a link in the PR description - [ ] Not required --------- Signed-off-by: Firestarman <firestarmanllc@gmail.com>

Fix Spark 4.2 collect_set float buffer conversion for mixed aggs

fc8daac

Spark 4.2 stores CollectSet float/double buffers as normalized bit patterns, so convert GPU logical values when crossing CPU/GPU stages. Fixes NVIDIA#15454. Signed-off-by: Firestarman <firestarmanllc@gmail.com>

firestarman force-pushed the fix/collect-set-float-nan-spark420 branch from b3ae0bd to fc8daac Compare July 30, 2026 08:14

firestarman requested a review from a team as a code owner July 30, 2026 08:14

firestarman changed the base branch from main to release/26.08 July 30, 2026 08:14

firestarman requested a review from a team July 30, 2026 08:20

wjxiz1992 approved these changes Jul 30, 2026

View reviewed changes

sameerz added the bug Something isn't working label Jul 30, 2026

ttnghia reviewed Jul 30, 2026

View reviewed changes

ttnghia approved these changes Jul 30, 2026

View reviewed changes

firestarman changed the title ~~Fix Spark 4.2 collect_set float buffer conversion for mixed aggs~~ Fix Spark 4.2 collect_set float buffer conversion for mixed aggs[skip-ci] Jul 31, 2026

Trigger CI for PR NVIDIA#15455

36caa78

Signed-off-by: Firestarman <firestarmanllc@gmail.com>

firestarman changed the title ~~Fix Spark 4.2 collect_set float buffer conversion for mixed aggs[skip-ci]~~ Fix Spark 4.2 collect_set float buffer conversion for mixed aggs Jul 31, 2026

firestarman merged commit 6f42817 into NVIDIA:release/26.08 Jul 31, 2026
57 checks passed

firestarman deleted the fix/collect-set-float-nan-spark420 branch July 31, 2026 05:41

revans2 mentioned this pull request Aug 4, 2026

[FEA] Replace host-side per-row data converters with GPU kernels #15463

Open

firestarman mentioned this pull request Aug 6, 2026

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

Merged

8 tasks

Navigation Menu

Uh oh!

Fix Spark 4.2 collect_set float buffer conversion for mixed aggs - #15455

Fix Spark 4.2 collect_set float buffer conversion for mixed aggs#15455
firestarman merged 2 commits into
NVIDIA:release/26.08from
firestarman:fix/collect-set-float-nan-spark420

firestarman commented Jul 30, 2026

Uh oh!

greptile-apps Bot commented Jul 30, 2026 •

edited

Loading

Uh oh!

firestarman commented Jul 30, 2026

Uh oh!

ttnghia Jul 30, 2026 •

edited

Loading

Uh oh!

ttnghia Jul 30, 2026

Uh oh!

firestarman commented Jul 31, 2026

Uh oh!

firestarman commented Jul 31, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Uh oh!

Conversation

firestarman commented Jul 30, 2026

Description

Checklists

Uh oh!

greptile-apps Bot commented Jul 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Flowchart

Uh oh!

firestarman commented Jul 30, 2026

Uh oh!

ttnghia Jul 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ttnghia Jul 30, 2026

Choose a reason for hiding this comment

Uh oh!

firestarman commented Jul 31, 2026

Uh oh!

firestarman commented Jul 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

greptile-apps Bot commented Jul 30, 2026 •

edited

Loading

ttnghia Jul 30, 2026 •

edited

Loading

firestarman commented Jul 31, 2026 •

edited

Loading