fix: normalize scalar float sort and window rank keys

[sunchao]

Why are the changes needed?

A filter such as RANK() <= 1 must keep every row tied for first place. Native execution can currently drop some of those rows when the ordering column contains signed zeros or different NaN representations. This PR fixes that result mismatch for scalar FLOAT and DOUBLE keys, following the native WindowGroupLimitExec support added in #4870.

A rank cutoff can discard a tied row

Consider a table measurements whose DOUBLE column v contains -0.0, +0.0, and 1.0:

SELECT v, rnk
FROM (
  SELECT v, RANK() OVER (ORDER BY v ASC) AS rnk
  FROM measurements
) ranked
WHERE rnk <= 1;

Spark considers the two zeros equal, so both belong at rank 1. The native comparison instead distinguishes their IEEE-754 representations:

Input value	Rank in Spark	Rank with the previous native comparison
-0.0	1	1
+0.0	1	2
1.0	3	3

The query therefore keeps both zero rows in Spark, but only the negative-zero row before this fix. DENSE_RANK has the same problem at the cutoff. Once WindowGroupLimitExec has discarded a qualifying row, the window calculation above it cannot recover that row.

NaNs expose the same mismatch. Spark treats every NaN representation as equal and greater than every non-NaN value. For input containing two differently encoded NaNs and 1.0, ORDER BY v DESC followed by rnk <= 1 should keep both NaNs. Arrow's raw total ordering can separate them; a NaN with its sign bit set can even sort below finite values. The regressions in this PR construct distinct NaN encodings directly so the test actually exercises this distinction.

What changes were proposed in this PR?

The fix gives the native sorting and window operators a common representation for comparing scalar floating values. It reuses Comet's existing NormalizeNaNAndZero expression to map every NaN to one canonical NaN and both zero signs to positive zero in the comparison keys. Arrow can then compare those keys using its existing machinery while agreeing with Spark about which values are tied.

This normalization must happen before sorting as well as when assigning ranks. For example, with ORDER BY v, secondary, the rows (-0.0, 1) and (+0.0, 1) are peers and must precede the rows whose secondary key is 2. A sort that distinguishes the zero signs can instead produce:

(-0.0, 1), (-0.0, 2), (+0.0, 1), (+0.0, 2)

Even a corrected peer comparison would see the first peer group split apart. A streaming rank limit could reach the cutoff at (-0.0, 2) and stop before seeing the other qualifying row. The shared comparison-key construction therefore supplies normalized keys to Sort, Window, and WindowGroupLimit, keeping their ordering and peer decisions consistent.

That consistency includes PARTITION BY, not just ORDER BY. An ordinary query such as RANK() OVER (PARTITION BY p ORDER BY id) must still work when p is a DOUBLE column containing only 1.0 and 2.0. Spark may already have normalized p; wrapping it again in Sort while leaving Window unchanged produces different expressions for the same key. DataFusion uses those expressions to recognize partition ordering, so the mismatch causes execution to fail even though the values compare equally. The planner therefore reuses already-normalized expressions and constructs matching keys for Sort, Window, and WindowGroupLimit. This preserves ordinary windows as well as the rank-limit path.

Native range partitioning uses the same rule for its sampled boundaries. An incoming positive zero and a boundary containing negative zero must compare equal, just as they do during sorting. Normalizing both sides keeps the shuffle's comparisons consistent with the operators that consume its output.

The original input values are preserved: selecting v still returns its original zero sign or NaN payload. Only the temporary keys used for comparison are normalized. The scope remains scalar FLOAT and DOUBLE; the documented limitations for floats nested in arrays or structs and the existing spark.comet.exec.strictFloatingPoint fallback policy are unchanged.

How was this PR tested?

The correctness results in the following paragraphs and table were recorded during the earlier implementation and rebase. They are historical local validation, not a fresh full-suite rerun for this review follow-up. The newly run sort benchmark and checks are described after those results.

The regressions verify both the returned rows and the execution path. The Spark cutoff tests require native Sort and WindowGroupLimit operators, compare results with Spark, and check that the input and output retain their original floating bits. Native tests cover compound keys, peer groups spanning batches, and agreement between range boundaries and incoming keys. The previously ignored signed-zero SQL regression is also enabled.

Additional regressions protect ordinary windows partitioned by floating columns: RANK, PERCENT_RANK, and NTILE must match Spark for single and compound partitions. A filtered RANK query also checks the Sort → WindowGroupLimit → Window path. A native planner test verifies that both bare and already-normalized partition keys retain recognizable ordering metadata.

With the comparison-key fix removed, the native peer/range regressions and both Spark FLOAT/DOUBLE cutoff regressions fail. In the Spark fixture, unpatched Comet returns 4 rows where Spark returns 8. The fixed runs pass.

On Spark 4.1, the partition-ordering cases were also checked against the library before the key-consistency fix: all six ordinary-window cases failed, while the two WindowGroupLimit controls passed. All eight pass with the fix.

The final review reran all three native planner regressions and the eight Spark 4.1 partition-key regressions; all passed. It also checked the three query shapes behind the Spark SQL CI failures on the earlier revision: a count window partitioned by floating columns, a rank window with computed aggregate/partition keys, and a correlated EXISTS ... LIMIT 1 query that Spark rewrites into a window. Each query matched Spark with native Sort and Window execution asserted. The same probes passed with the pinned base library. These were focused query probes, not reruns of the full CI shards.

Local validation	Result
Native workspace unit tests	909 passed; 4 existing tests ignored
Spark 3.4: floating partition-key regressions	6 passed; 2 WindowGroupLimit cases skipped because Spark 3.4 has no such operator
Spark 3.5: floating partition-key regressions	8 passed
Spark 4.0: floating partition-key regressions	8 passed
Spark 4.1: full window suite, all WindowGroupLimit SQL cases, sort fallback tests, and native range shuffle	87 passed

make core, cargo fmt --all -- --check, and cargo clippy --all-targets --workspace -- -D warnings passed. JVM formatting checks and Spark 4.1 Scala style checks also passed.

The review follow-up adds a reproducible scalar floating-key sort benchmark. It executes DataFusion 54.1.0 SortExec with a bare Column key and with Comet's NormalizeNaNAndZero key in the same optimized binary, representing the before/after key-construction change for previously bare keys. Already-normalized keys are reused by the planner. This measures execution and output draining for an in-memory, single-partition sort; it does not measure Spark planning, JNI, I/O, or a distributed range shuffle. Input generation and physical-plan construction are outside the timed region.

Both variants sort 262,144 rows in 32 batches of 8,192, with one current-thread Tokio runtime, ascending/nulls-last ordering, and no fetch limit. The finite shape has no nulls; the mixed shape has approximately 20% nulls, 10% payload/sign NaNs, 10% signed zeros, and 60% finite values. The release profile uses thin LTO and one codegen unit. Measurements used an AMD EPYC-Milan CPU pinned to one core, Rust 1.98.0, 30 Criterion samples per case, a one-second warmup, and a three-second measurement target. Other workspace builds and tests were paused. A second pass reversed the order of the two variants.

Shape	Bare then normalized: median ms (change)	Normalized then bare: median ms (change)
FLOAT finite	23.617 → 23.711 (+0.4%)	23.314 → 23.897 (+2.5%)
FLOAT mixed	18.873 → 19.763 (+4.7%)	18.811 → 18.487 (−1.7%)
DOUBLE finite	27.064 → 30.155 (+11.4%)	27.058 → 28.843 (+6.6%)
DOUBLE mixed	21.977 → 24.723 (+12.5%)	22.282 → 23.440 (+5.2%)

Each cell reports bare → normalized even when execution order was reversed. FLOAT differences were small and the mixed-shape direction changed between passes. DOUBLE normalization had a measurable cost in both passes: +6.6–11.4% for finite values and +5.2–12.5% for mixed values. These are observed repeat ranges, not confidence intervals or end-to-end Spark performance claims. Normalization deliberately changes the comparison semantics for the mixed input. Folding it into sorting or encoding could avoid temporary arrays, but that optimization is not implemented by this correctness fix.

Untimed validation checks every row identity, original floating bits and nulls, ordering, and zero spills. Peak DataFusion pool reservations were identical for bare and normalized keys: 16 MiB / 16.0625 MiB for FLOAT finite/mixed and 18 MiB / 18.0625 MiB for DOUBLE finite/mixed. These are reserved-pool peaks, not process RSS or a complete allocation high-water mark: they do not include the temporary normalized values buffers. Each normalization evaluation allocates 4 bytes per row for FLOAT or 8 for DOUBLE (32/64 KiB for an 8,192-row array), sharing the validity bitmap. That is a per-evaluation size, not total allocation traffic; sorting and merging may evaluate the key more than once.

To reproduce from native/:

cargo bench --locked -p datafusion-comet-spark-expr --bench sort_float_keys -- --noplot
COMET_SORT_BENCH_REVERSE_VARIANTS=1 cargo bench --locked -p datafusion-comet-spark-expr --bench sort_float_keys -- --noplot

For this follow-up, the release benchmark build, all eight benchmark smoke cases, release Clippy for this benchmark with -D warnings, cargo fmt --all -- --check, and git diff --check passed. The broader correctness suites above were not rerun for this benchmark/documentation-only commit.

The guide now links #5506 for narrowing strict-mode admission of corrected scalar sort keys and #5507 for nested ORDER BY/rank compatibility. Spark 4.1.3 diagnostic queries reproduced nested array/struct ORDER BY and array RANK mismatches for both FLOAT and DOUBLE; struct RANK already fell back and matched Spark. All eight strict-mode controls matched Spark, and four scalar planning controls confirmed the conservative strict-mode fallback. Those tracker probes reused a prior native build with the unchanged nested path and matching JVM admission sources; they are not a full rebuild or correctness-suite run of this PR head.

Which issue does this PR close?

Closes #5468.

[andygrove]

Note on this review: this was generated by an LLM (Claude Code) at my request while I worked through a review backlog. I have not verified the individual findings myself. Please treat everything below as suggestions to evaluate rather than as authoritative review feedback, and push back on anything that is wrong or already handled.

This is a thorough piece of work. Reusing NormalizeNaNAndZero rather than inventing a new comparator is the right call, and I like that create_normalized_key_expr is shared by sort, window partition keys, window group limit, and range partitioning so the four cannot drift apart. The floating_window_partition_keys_preserve_ordering test that checks get_ordered_partition_by_indices still returns the prefix is a good catch, since silently losing the sorted window mode would have been an easy regression to miss. Flipping the ignore(...) in window_group_limit_rank.sql back to a real assertion is nice to see.

A few things I would like resolved before merge.

No performance data for the sort path

create_sort_expr now wraps every scalar FLOAT and DOUBLE sort key, which means native/core/src/execution/planner.rs:1498 adds a full column copy per batch to every native sort on a floating key, plus the same for window ordering and range partitioning. unary allocates a fresh values buffer each time. That is probably fine relative to the row encoding and the sort itself, but "probably fine" is not a number. Could you run a sort benchmark with a float key before and after and put the result in the description? If it turns out to be measurable, it would be worth discussing whether the normalization can be folded into the row encoder instead of materializing an intermediate array.

strictFloatingPoint now falls back for cases this PR fixes

The docs change says the spark.comet.exec.strictFloatingPoint=true fallback policy is unchanged, "including its conservative fallback for scalar floating-point sort keys." After this PR those scalar keys agree with Spark, so users who turn strict mode on are paying a fallback for something that is now correct. Is narrowing that fallback to the nested array and struct cases in scope here, or should it be a follow-up? If it is a follow-up I would like an issue filed and linked from this PR, because a stale conservative fallback is the kind of thing that quietly stays for a year.

Tracking issue for the nested case

Same question for the nested array and struct divergence the docs now call out. There is no issue link in the new floating-point.md text. A user who reads that paragraph and is affected has nowhere to go. Could you file one and link it from the doc?

Sort-merge join builds an expression it throws away

At planner.rs:2145 the create_sort_expr result is used only for sort_expr.options, and the expression itself is discarded. That call now constructs a NormalizeNaNAndZero node for nothing on float keys. It is not a correctness problem, but it does suggest that site wants a smaller helper that only decodes direction and null ordering. Worth splitting out?

Silent passthrough in normalize_array

NormalizeNaNAndZero::normalize_array returns the input untouched for any type that is not Float32 or Float64, including Dictionary(_, Float64), List<Float64>, and structs with float fields. In the range-partitioning call site that is consistent with create_normalized_key_expr, which also only handles the two scalar types, so I do not think there is a bug here today. But the function is pub and the passthrough is silent. Could the doc comment say explicitly that nested and dictionary-encoded floats are deliberately not normalized, and that callers must match create_normalized_key_expr's coverage? Otherwise the next caller gets a silently wrong answer with no compile error.

One test question

In CometWindowExecSuite, the new tests set COMET_EXEC_STRICT_FLOATING_POINT to false explicitly, which is already the default. Is that there to document intent, or to guard against the default flipping? If the former, a short comment saying so would help, since otherwise it reads like the tests only pass in non-strict mode.

[comphead]

Thanks @sunchao this PR looks good to me from manual review, let me do 1 pass of LLM review as well

[sunchao]

Updated in 38825e26a.

The follow-up adds the sort benchmark and links the two remaining compatibility issues from the guide. I checked the six points separately:

1. Sort cost: the new benchmark runs DataFusion 54.1.0 SortExec with bare versus normalized keys in the same release binary, using 262,144 rows, 32 batches, one partition/thread, and both finite/non-null and null/NaN/signed-zero inputs. Two passes reverse the variant order. FLOAT median changes were small (−1.7% to +4.7%, with mixed-input direction varying); DOUBLE cost was measurable in both passes: +6.6–11.4% finite and +5.2–12.5% mixed. The description contains both passes and reproduction commands. This is a SortExec benchmark, not an end-to-end Spark result. Peak pool reservations were unchanged, but they exclude temporary normalized arrays: each evaluation adds a 4/8-byte-per-row values buffer and shares validity. Already-normalized keys are reused, so this measures previously bare keys rather than another copy at every call. Folding normalization into sorting/encoding is a worthwhile optimization to investigate separately; this change does not claim the cost is free.

2. Strict-mode admission: tracked in #5506, now linked from the guide. The follow-up is scoped to scalar Sort admission with strict-mode execution tests; this PR retains the existing guards for other consumers and nested keys.

3. Nested ordering/rank: tracked in #5507, also linked from the guide. Diagnostic SQL reproduced array/struct ORDER BY and array RANK differences for FLOAT and DOUBLE with native operators asserted. Struct RANK already falls back in those probes; it is not a reproduced native wrong result. Strict-mode controls matched Spark. The issue records the reused-native-build limitation, so these are evidence for the remaining limitation, not a fresh full-head validation claim.

4. Sort-merge-join options: the observation is correct, but that site constructs and discards a planning-time expression. It does not evaluate the normalizer or copy a column there; it already built an expression just to obtain options before this PR. I left the optional helper cleanup out of this correctness fix.

5. normalize_array coverage: the existing doc comment specifies scalar floating comparison keys and says other types are unchanged. Its Float32/Float64 match and create_normalized_key_expr use the same gate: dictionary, list, and struct arrays deliberately pass through. There is no recursive or dictionary-normalization promise in that API; the nested limitation now has the explicit tracker above.

6. Explicit strictFloatingPoint=false in tests: it pins the native execution path under test, independently of externally supplied configuration or a future default change. The tests assert native Sort/WindowGroupLimit operators; true deliberately still invokes the conservative fallback tracked by Narrow strict floating-point admission for corrected scalar sort keys #5506. I kept the setting for that reason.

The new release benchmark built, all eight smoke cases passed, and benchmark Clippy (-D warnings), formatting, and diff checks passed. I also corrected the stale validation paragraph: the earlier full-suite results are historical, not rerun for this benchmark/documentation-only follow-up.

[comphead]

One thing to consider is https://issues.apache.org/jira/browse/SPARK-54918 which fixes +0, -0 in Spark 4.2.0

[sunchao]

@comphead Thanks for the pointer. I checked the SPARK-54918 patch. It extends normalization to array set operations such as array_distinct and array_union.

For the scalar sorting and ranking cases here, Spark already treats -0.0 and +0.0 as equal before 4.2. The problem is that Comet’s native comparison can distinguish them, splitting rows that should share a rank and potentially dropping qualifying rows at a rank cutoff.

This PR normalizes the comparison keys consistently across sorting and ranking, so the fix is still needed with Spark 4.2.

[comphead]

Thanks @sunchao LGTM