fix: avoid overflow and underflow in vector functions

[namanjain24-sudo]

Which issue does this PR close?

• Closes Vector functions overflow or underflow for finite Float64 inputs #25265.

Rationale for this change

array_distance, cosine_distance and array_normalize square their inputs directly. When a finite value's square falls outside the Float64 range, they return a wrong answer even though the true result is representable:

query	main	this PR
array_distance([1e-200], [0])	0.0	1e-200
array_distance([3e200], [-1e200])	inf	4e200
cosine_distance([3e200, 4e200], [3e200, 4e200])	NaN	-2.220446049250313e-16
cosine_distance([1e-200, 2e-200], [1e-200, 2e-200])	NULL	0.0
array_normalize([3e200, 4e200])	[0.0, 0.0]	[0.6000000000000001, 0.8]
array_normalize([3e-200, 4e-200])	NULL	[0.6, 0.8]

The two results that are off in the last bit are ordinary rounding. The existing formula does the same for ordinary inputs: for identical vectors it already returns a small negative cosine_distance about 11% of the time.

What changes are included in this PR?

Each function still computes the sum of squares as before. It recomputes only when that sum may be wrong: when it is infinite or NaN, or when it is below len * 2^-1012, where len is the number of elements. The recomputation multiplies every value by a power of two that brings the largest magnitude close to 1 (norm_scale in utils.rs):

• array_distance scales the element-wise differences, then divides the result by the same factor.
• cosine_distance scales each vector on its own. The distance does not change when either vector is multiplied by a positive factor.
• array_normalize scales the values. Dividing the scaled values by the scaled magnitude gives the same unit vector.

Rows whose sums are in range run the same code as before, so their results do not change. The threshold comes from the underflow error: a square that underflows is off by at most half the smallest subnormal value, so len of them move the sum by at most len * 2^-1075, which is less than 2^-63 of any sum above the threshold. Because the factor is a power of two, rescaling a value is exact unless the scaled value is subnormal, and such a value is too small to change the result, so a rescaled row gets the same result the unscaled code would have given had it not overflowed or underflowed. Empty, all-zero and infinite inputs are never rescaled, and a NaN input still gives NaN, so their results are also unchanged, including NULL for zero vectors. cosine_distance rescales only the vector whose own sum is out of range, which is enough because the distance does not change when either vector is scaled on its own.

What is the testing strategy for this PR?

• sqllogictests for the overflow and underflow cases in array/array_length.slt, cosine_distance.slt and array_normalize.slt, plus non-finite inputs to pin the unchanged behaviour. .slt rounds floats to 12 decimal places, so the 1e-200 distance is checked as a ratio.
• Unit tests for norm_scale and needs_norm_scale in utils.rs.
• With the change to datafusion/functions-nested/src reverted and the new tests kept, the three new overflow and underflow queries fail (0 Infinity instead of 1 4, NaN NULL NULL instead of 0 0 2, and [0.0, 0.0] NULL instead of [0.6, 0.8] [0.6, -0.8]). The non-finite queries pass either way, as intended.
• To check the "results do not change" claim, I ran copies of the old and new loops on 200,000 random vector pairs (1 to 32 elements, magnitudes between 1e-100 and 1e100) and compared the bits: no differences from main. Around the threshold, on 300,000 pairs (97,001 of them rescaled) and on 228,241 pairs whose elements span 1e-320 to 1e-80 (28,330 with subnormal squares), the new code matched a version that always rescales, bit for bit. A tighter threshold of len * 2^-1021 left some rows 1 to 4 ulps away from the rescaled result. With len * 2^-1012 I found no differences in about 436,000 rows of length 1 to 4096.
• I also ran 29 edge-case queries through datafusion-cli on main and on this branch. Only the six rows in the table above changed.

There is no existing benchmark for these functions, so I measured with a criterion bench that calls the three UDFs through invoke_with_args (8,192 rows at dim 4 and 128, 1,024 rows at dim 1536). It is not part of this PR. Ratios are this PR's time over main (133111f):

input	array_distance	cosine_distance	array_normalize
values in [-1, 1], no row rescaled	0.92–0.95x / 0.99x / 1.02–1.03x	1.02–1.03x / 0.99–1.00x / 1.00–1.02x	0.85–0.98x / 0.63–0.65x / 0.55–0.57x
zero vector (identical vectors for array_distance)	0.94–0.98x / 1.00–1.04x / 1.19–1.20x	1.11x / 1.13x / 1.17x	1.20–1.21x / 1.63–1.65x / 0.99–1.29x

Each cell lists dim 4 / 128 / 1536, as the range over two runs. Rows whose values are around 1e-200 are rescaled and cost more, but those are the rows that returned wrong results before. The array_normalize speedup comes from writing the output with extend.

Are there any user-facing changes?

Only for finite inputs whose squares overflowed or underflowed. Those now return the correct distance or normalized vector instead of 0, inf, NaN, a zero vector or NULL. There are no API changes.

[aoto-tech]

@namanjain24-sudo san
I opened the issue. This is a little different from the fix I had in mind, so I hope you don’t mind me commenting.
1e-180 seems quite high as the rescaling threshold. For example, squaring 1e-100 gives 1e-200, which is still safe in Float64, but this implementation would rescale it.
Since the rescaled path is several times slower, could the threshold be derived from the array length instead?

[namanjain24-sudo]

@aoto-tech thanks, you're right. The threshold only needs to catch sums that an underflowing square could have changed, and 1e-180 is far above that. array_distance([1e-100], [0]) was taking the slow path for no reason.

I've pushed a length-based threshold, as you suggested. A square that underflows is off by at most half the smallest subnormal value, so len of them move the sum by at most len * 2^-1075. Rescaling now happens only when the sum is below len * 2^-1012. Above that bound, the underflow error is less than 2^-63 of the sum, far below its rounding precision.

I first tried the tightest version, len * 2^-1021. On vectors whose elements are around 1e-160, it left some rows 1 to 4 ulps away from the rescaled result. With the 2^10 margin I found no differences in about 436,000 random rows. Those rows had lengths 1 to 4096 and magnitudes from 1e-320 to 1e-90, including rows with subnormal squares.

On 20,000 rows of 1536 elements around 1e-100, nothing is rescaled any more. The ratios to main went from 3.58x / 4.82x / 2.03x (distance / cosine / normalize) to 1.00x / 1.02x / 0.70x. Rows that really underflow, around 1e-200, still take the slow path. I've updated the PR description, and the change is a separate commit so it's easy to review.

Was this close to the fix you had in mind? I'm happy to adjust.

[aoto-tech]

@aoto-tech thanks, you're right. The threshold only needs to catch sums that an underflowing square could have changed, and 1e-180 is far above that. array_distance([1e-100], [0]) was taking the slow path for no reason.

I've pushed a length-based threshold, as you suggested. A square that underflows is off by at most half the smallest subnormal value, so len of them move the sum by at most len * 2^-1075. Rescaling now happens only when the sum is below len * 2^-1012. Above that bound, the underflow error is less than 2^-63 of the sum, far below its rounding precision.

I first tried the tightest version, len * 2^-1021. On vectors whose elements are around 1e-160, it left some rows 1 to 4 ulps away from the rescaled result. With the 2^10 margin I found no differences in about 436,000 random rows. Those rows had lengths 1 to 4096 and magnitudes from 1e-320 to 1e-90, including rows with subnormal squares.

On 20,000 rows of 1536 elements around 1e-100, nothing is rescaled any more. The ratios to main went from 3.58x / 4.82x / 2.03x (distance / cosine / normalize) to 1.00x / 1.02x / 0.70x. Rows that really underflow, around 1e-200, still take the slow path. I've updated the PR description, and the change is a separate commit so it's easy to review.

Was this close to the fix you had in mind? I'm happy to adjust.

Nice, the commit looks good to me. Thanks for making the adjustment!

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 81.92%. Comparing base (133111f) to head (35ccfe7).
⚠️ Report is 5 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main   #25289   +/-   ##
=======================================
  Coverage   81.92%   81.92%           
=======================================
  Files        1134     1134           
  Lines      426000   426085   +85     
  Branches   426000   426085   +85     
=======================================
+ Hits       349000   349075   +75     
- Misses      56304    56308    +4     
- Partials    20696    20702    +6     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[jayzhan211]

I'm worried this check hurts performance 🤔 The numbers drop when I run the benchmark locally.

I think we need a way to handle the edge cases without adding an expensive check — especially for the cases that don't benefit from it.

[jayzhan211]

The rescue path is only reached by rows whose first-pass sum is exactly zero,
infinite, or NaN. For finite data that means a zero vector, or two identical
vectors in array_distance, so ordinary rows never see it and this is not a
blocker. It is still worth making cheap, because a row of exact zeros has to
be scanned to tell it apart from a row whose squares all underflowed.

The per-element return None on non-finite values stops the loop from
vectorizing, which makes that scan about 5x the cost of the sum pass it
duplicates. Checking finiteness once on the max is enough:

    let mut max = 0.0_f64;
    for value in values {
        max = max.max(value.abs());
    }
    if max == 0.0 || !max.is_finite() {
        return None;
    }

f64::max skips NaN, so a vector containing NaN may now return Some. That
is fine: the scaled recompute still produces NaN. The doc comment and the
[1.0, f64::NAN] assertion need updating to match. While there, the doc's
"scaling is exact" only holds when the scaled value is normal; a scaled
subnormal is rounded but is then too small to matter.

This only affects a rare input: one vector all zeros, where the result is
NULL either way. Normal rows never enter this branch, and the fast path is
unchanged or faster than main, so this is a minor cleanup rather than a
correctness or hot-path concern.

For that input the code scans both vectors, finds a scale for the normal
one, and re-runs dot_and_squares for a result that cannot change. In a
microbenchmark at dim 1536 that makes such rows about 2.5x slower than main.
Scan a vector only when its own sum is out of range, and recompute only if a
scan found something to scale:

    let rescale = !dot.is_finite();
    let scale1 = if rescale || needs_norm_scale(sq1, len1) {
        norm_scale(vals1.iter().copied())
    } else {
        None
    };
    let scale2 = if rescale || needs_norm_scale(sq2, len1) {
        norm_scale(vals2.iter().copied())
    } else {
        None
    };
    if scale1.is_some() || scale2.is_some() {
        (dot, sq1, sq2) = dot_and_squares(
            vals1, vals2, scale1.unwrap_or(1.0), scale2.unwrap_or(1.0),
        );
    }

Scaling one side only is valid: cosine distance is invariant to scaling each
vector independently, the unscaled vector's products cannot overflow because
its own sum of squares is finite, and the underflow error bound in the PR
description still holds. With this and the norm_scale change, zero-vector
rows cost about 20% over main at dim 1536, which is acceptable for an input
that produces NULL.

Nit: the norm_scale tests for f64::MAX and f64::MIN_POSITIVE / 4.0
compare against powi itself, so they would pass even if powi were
inexact at those exponents. Pinning the bit patterns makes them independent:

    assert_eq!(norm_scale([f64::MAX]), Some(f64::from_bits(1 << 51)));
    assert_eq!(
        norm_scale([f64::MIN_POSITIVE / 4.0]),
        Some(f64::from_bits(2045 << 52))
    );

[namanjain24-sudo]

@jayzhan211 thanks, you're right. My earlier timings were on copies of the loops outside DataFusion, so this time I measured inside it: a criterion bench that calls the three UDFs, comparing main (133111f), this PR before the change, and this PR with your suggestions. I ran it twice.

Rows with values in [-1, 1] never reach the rescue path. They run at 0.92x to 1.03x of main for array_distance and cosine_distance, and array_normalize is still 0.55x to 0.65x at dim 128 and 1536. The regression was in rows that reach the rescue path without benefiting from it (time over main, before → after your changes):

input	dim 4	dim 128	dim 1536
cosine_distance, one vector all zero	1.35x → 1.11x	3.45x → 1.13x	4.67x → 1.17x
array_distance, identical vectors	0.94x → 0.94–0.98x	1.51x → 1.00–1.04x	2.35x → 1.19–1.20x
array_normalize, zero vector	1.18x → 1.20–1.21x	2.39x → 1.63–1.65x	1.98x → 0.99–1.29x

A range means the two runs differed. For array_normalize at dim 1536, main itself moved between runs.

I've pushed both suggestions as b3fcbe1:

• norm_scale checks finiteness once on the max, so the loop vectorizes.
• cosine_distance scans a vector only when its own sum is out of range, and recomputes only if there is something to scale.

The commit also includes your doc fix for scaled subnormals and the bit-pattern assertions.

To check that no result moved, I ran the previous and new kernels on about 3.7 million vector pairs. They covered lengths from 1 to 1536, magnitudes from subnormal up to 1e300, and zero, NaN, infinite and one-side-only inputs. All three functions are bit-identical to the previous commit, and cosine_distance also matches a version that always rescales both vectors.

The length-based threshold from @aoto-tech's review is unchanged, so rows around 1e-100 still never reach the rescue path (0 scans in about 87,000 random rows). On the inputs they flagged in the issue, NaN, infinite, zero and identical vectors give the same bits as main. A dot product that overflows and cancels ([1e200, 1e200] with [1e200, -1e200]) gives 1 like the always-rescaled version, where main returns NaN.

Two small additions to your comment:

• Rows whose sum is not zero but below len * 2^-1012 also take the rescue path. Those are the rows that actually underflowed, so they are the ones it is for.
• A vector left unscaled can also be one containing an infinity, where norm_scale returns None, so its sum of squares is not always finite. The result is still NaN exactly as on main, so I worded the code comment around the sum being in range.

The bench isn't part of this PR. I can add it as benches/vector_functions.rs if that would help.