You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/17 12:25:05 UTC
(datafusion) branch gh-readonly-queue/main/pr-25289-902ee8a67e6838163fbe438dbba7d324308c02d5 created (now 7deda30943)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-25289-902ee8a67e6838163fbe438dbba7d324308c02d5
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 7deda30943 fix: avoid overflow and underflow in vector functions (#25289)
This branch includes the following new commits:
new 7deda30943 fix: avoid overflow and underflow in vector functions (#25289)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: fix: avoid overflow and underflow in vector functions (#25289)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-25289-902ee8a67e6838163fbe438dbba7d324308c02d5
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 7deda309436f8b30ebdebb9575f64963f9c365a3
Author: namanjain24-sudo <na...@gmail.com>
AuthorDate: Thu Sep 17 12:24:42 2026 +0000
fix: avoid overflow and underflow in vector functions (#25289)
## Which issue does this PR close?
- Closes #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. So a rescaled row gets the same result the unscaled code
would have given had it not overflowed or underflowed, except that an
`array_normalize` output element that is itself subnormal can differ in
its last bit (5e-324). I checked this by forcing the rescaled path on
200,000 in-range rows with subnormal scaled values: `array_distance` and
`cosine_distance` matched bit for bit, and `array_normalize` differed by
at most 5e-324. 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.
---
datafusion/functions-nested/src/array_normalize.rs | 23 +++++-
datafusion/functions-nested/src/cosine_distance.rs | 57 ++++++++++++---
datafusion/functions-nested/src/distance.rs | 39 +++++++---
datafusion/functions-nested/src/utils.rs | 84 ++++++++++++++++++++++
.../sqllogictest/test_files/array/array_length.slt | 17 +++++
.../sqllogictest/test_files/array_normalize.slt | 15 ++++
.../sqllogictest/test_files/cosine_distance.slt | 27 +++++++
7 files changed, 238 insertions(+), 24 deletions(-)
diff --git a/datafusion/functions-nested/src/array_normalize.rs b/datafusion/functions-nested/src/array_normalize.rs
index 8f4342cd83..988e2679b9 100644
--- a/datafusion/functions-nested/src/array_normalize.rs
+++ b/datafusion/functions-nested/src/array_normalize.rs
@@ -17,7 +17,7 @@
//! [`ScalarUDFImpl`] definitions for array_normalize function.
-use crate::utils::make_scalar_function;
+use crate::utils::{make_scalar_function, needs_norm_scale, norm_scale};
use arrow::array::{
Array, ArrayRef, Float64Array, GenericListArray, NullBufferBuilder, OffsetSizeTrait,
};
@@ -181,6 +181,22 @@ fn general_array_normalize<O: OffsetSizeTrait>(arrays: &[ArrayRef]) -> Result<Ar
sq_sum += vals[i] * vals[i];
}
+ // If a square may have overflowed or underflowed, recompute with scaled
+ // values. Dividing the scaled values by the scaled magnitude gives the
+ // same unit vector.
+ let scale = if needs_norm_scale(sq_sum, len) {
+ norm_scale(vals.iter().copied())
+ } else {
+ None
+ };
+ if let Some(scale) = scale {
+ sq_sum = 0.0;
+ for i in 0..len {
+ let scaled = vals[i] * scale;
+ sq_sum += scaled * scaled;
+ }
+ }
+
// Zero magnitude: undefined normalization. Emit NULL row.
if sq_sum == 0.0 {
nulls.append_null();
@@ -189,8 +205,9 @@ fn general_array_normalize<O: OffsetSizeTrait>(arrays: &[ArrayRef]) -> Result<Ar
}
let mag = sq_sum.sqrt();
- for i in 0..len {
- new_values.push(vals[i] / mag);
+ match scale {
+ Some(scale) => new_values.extend(vals.iter().map(|v| v * scale / mag)),
+ None => new_values.extend(vals.iter().map(|v| v / mag)),
}
nulls.append_non_null();
new_offsets.push(new_offsets[row] + O::usize_as(len));
diff --git a/datafusion/functions-nested/src/cosine_distance.rs b/datafusion/functions-nested/src/cosine_distance.rs
index 56ca071234..1637b2cea1 100644
--- a/datafusion/functions-nested/src/cosine_distance.rs
+++ b/datafusion/functions-nested/src/cosine_distance.rs
@@ -17,7 +17,7 @@
//! [`ScalarUDFImpl`] definitions for cosine_distance function.
-use crate::utils::make_scalar_function;
+use crate::utils::{make_scalar_function, needs_norm_scale, norm_scale};
use arrow::array::{Array, ArrayRef, Float64Array, OffsetSizeTrait};
use arrow::datatypes::{
DataType,
@@ -197,15 +197,30 @@ fn general_cosine_distance<O: OffsetSizeTrait>(arrays: &[ArrayRef]) -> Result<Ar
let vals1 = slice1.values();
let vals2 = slice2.values();
- let mut dot = 0.0;
- let mut sq1 = 0.0;
- let mut sq2 = 0.0;
- for i in 0..len1 {
- let a = vals1[i];
- let b = vals2[i];
- dot += a * b;
- sq1 += a * a;
- sq2 += b * b;
+ let (mut dot, mut sq1, mut sq2) = dot_and_squares(vals1, vals2, 1.0, 1.0);
+ // Cosine distance does not change when either vector is multiplied by a
+ // positive factor, so scale only a vector whose own sum of squares is
+ // out of range, and recompute only if there is something to scale. A
+ // vector whose sum is in range can stay unscaled: its products cannot
+ // overflow, and its underflow error is already negligible.
+ let rescale_both = !dot.is_finite();
+ let scale1 = if rescale_both || needs_norm_scale(sq1, len1) {
+ norm_scale(vals1.iter().copied())
+ } else {
+ None
+ };
+ let scale2 = if rescale_both || 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),
+ );
}
if sq1 == 0.0 || sq2 == 0.0 {
@@ -217,3 +232,25 @@ fn general_cosine_distance<O: OffsetSizeTrait>(arrays: &[ArrayRef]) -> Result<Ar
Ok(Arc::new(builder.finish()) as ArrayRef)
}
+
+/// Returns the dot product of `vals1 * scale1` and `vals2 * scale2`, and the
+/// sum of squares of each.
+#[inline]
+fn dot_and_squares(
+ vals1: &[f64],
+ vals2: &[f64],
+ scale1: f64,
+ scale2: f64,
+) -> (f64, f64, f64) {
+ let mut dot = 0.0;
+ let mut sq1 = 0.0;
+ let mut sq2 = 0.0;
+ for (a, b) in vals1.iter().zip(vals2) {
+ let a = a * scale1;
+ let b = b * scale2;
+ dot += a * b;
+ sq1 += a * a;
+ sq2 += b * b;
+ }
+ (dot, sq1, sq2)
+}
diff --git a/datafusion/functions-nested/src/distance.rs b/datafusion/functions-nested/src/distance.rs
index d9182e871f..36b263870d 100644
--- a/datafusion/functions-nested/src/distance.rs
+++ b/datafusion/functions-nested/src/distance.rs
@@ -17,7 +17,7 @@
//! [ScalarUDFImpl] definitions for array_distance function.
-use crate::utils::make_scalar_function;
+use crate::utils::{make_scalar_function, needs_norm_scale, norm_scale};
use arrow::array::{Array, ArrayRef, Float64Array, OffsetSizeTrait};
use arrow::datatypes::{
DataType,
@@ -189,16 +189,33 @@ fn compute_array_distance(
return exec_err!("Both arrays must have the same length");
}
- let sum_squares: f64 = values1
- .iter()
- .zip(values2.iter())
- .map(|(v1, v2)| {
- let diff = v1.unwrap_or(0.0) - v2.unwrap_or(0.0);
- diff * diff
- })
- .sum();
-
- Ok(Some(sum_squares.sqrt()))
+ let diffs = || {
+ values1
+ .values()
+ .iter()
+ .zip(values2.values().iter())
+ .map(|(v1, v2)| v1 - v2)
+ };
+
+ let sum_squares: f64 = diffs().map(|diff| diff * diff).sum();
+ if !needs_norm_scale(sum_squares, values1.len()) {
+ return Ok(Some(sum_squares.sqrt()));
+ }
+
+ let distance = match norm_scale(diffs()) {
+ Some(scale) => {
+ let scaled_sum_squares: f64 = diffs()
+ .map(|diff| {
+ let scaled = diff * scale;
+ scaled * scaled
+ })
+ .sum();
+ scaled_sum_squares.sqrt() / scale
+ }
+ None => sum_squares.sqrt(),
+ };
+
+ Ok(Some(distance))
}
/// Converts an array of any numeric type to a Float64Array.
diff --git a/datafusion/functions-nested/src/utils.rs b/datafusion/functions-nested/src/utils.rs
index 2906250ec6..a6a93ed5df 100644
--- a/datafusion/functions-nested/src/utils.rs
+++ b/datafusion/functions-nested/src/utils.rs
@@ -463,6 +463,49 @@ where
)?))
}
+/// Returns a power of two that brings the largest magnitude in `values` close
+/// to 1, so that squaring the scaled values neither overflows nor underflows.
+///
+/// Squaring a large finite value overflows (`1e200 * 1e200` is infinity) and
+/// squaring a small one underflows (`1e-200 * 1e-200` is zero), even when the
+/// norm itself is representable. The factor is a power of two, so scaling is
+/// exact whenever the scaled value is normal. A value that becomes subnormal is
+/// rounded, so an `array_normalize` element that is itself subnormal can differ
+/// from the unscaled result in its last bit.
+///
+/// Returns `None` when `values` is empty, all zero, or contains an infinity.
+/// The unscaled computation already gives the expected result for those inputs.
+/// NaN values are ignored, and the scaled computation still produces NaN.
+pub(crate) fn norm_scale(values: impl IntoIterator<Item = f64>) -> Option<f64> {
+ // No early return inside the loop, so that it vectorizes. `f64::max` skips
+ // NaN, so only an infinity can make `max` non-finite.
+ let mut max = 0.0_f64;
+ for value in values {
+ max = max.max(value.abs());
+ }
+ if max == 0.0 || !max.is_finite() {
+ return None;
+ }
+ // Unbiased exponent of `max`. Subnormal values store a biased exponent of 0,
+ // so clamp them to the smallest normal exponent.
+ let exponent = ((max.to_bits() >> 52) as i32 - 1023).max(-1022);
+ Some(2.0_f64.powi(-exponent))
+}
+
+/// Returns whether a sum of `len` squares computed without scaling may be
+/// wrong because a square overflowed or underflowed, in which case it should be
+/// recomputed with the factor from [`norm_scale`].
+///
+/// An overflowing square makes the sum infinite. An underflowing square is off
+/// by at most half the smallest subnormal value, so `len` of them move the sum
+/// by at most `len * 2^-1075`. A sum of at least `len * 2^-1012` is therefore
+/// off by less than `2^-63` of itself, far below its rounding precision.
+pub(crate) fn needs_norm_scale(sum_of_squares: f64, len: usize) -> bool {
+ // 2^-1012 = 2^10 * f64::MIN_POSITIVE
+ let min_unscaled = 1024.0 * len as f64 * f64::MIN_POSITIVE;
+ !(min_unscaled..f64::INFINITY).contains(&sum_of_squares)
+}
+
#[cfg(test)]
mod tests {
use super::*;
@@ -518,4 +561,45 @@ mod tests {
expected_dim
);
}
+
+ #[test]
+ fn norm_scale_brings_largest_magnitude_close_to_one() {
+ assert_eq!(norm_scale([3e200, -4e200]), Some(2.0_f64.powi(-666)));
+ assert_eq!(norm_scale([3.0, 4.0]), Some(0.25));
+ // 2^-1023 and 2^1022, pinned by bit pattern rather than computed
+ 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))
+ );
+ // NaN is ignored; the scaled computation still produces NaN
+ assert_eq!(norm_scale([f64::NAN, 3.0, 4.0]), Some(0.25));
+ }
+
+ #[test]
+ fn norm_scale_skips_inputs_the_unscaled_computation_handles() {
+ assert_eq!(norm_scale([]), None);
+ assert_eq!(norm_scale([0.0, -0.0]), None);
+ assert_eq!(norm_scale([f64::NAN, 0.0]), None);
+ assert_eq!(norm_scale([f64::INFINITY, 1.0]), None);
+ assert_eq!(norm_scale([1.0, f64::NAN, f64::NEG_INFINITY]), None);
+ }
+
+ #[test]
+ fn needs_norm_scale_only_for_sums_that_may_have_overflowed_or_underflowed() {
+ assert!(!needs_norm_scale(1.0, 1));
+ assert!(!needs_norm_scale(f64::MAX, 1));
+ // The square of 1e-100 is 1e-200, which is far from underflowing.
+ assert!(!needs_norm_scale(1e-200, 1));
+ assert!(!needs_norm_scale(1e-200, 1536));
+
+ let min_unscaled = 1024.0 * f64::MIN_POSITIVE;
+ assert!(!needs_norm_scale(min_unscaled, 1));
+ assert!(needs_norm_scale(min_unscaled, 2));
+ assert!(needs_norm_scale(min_unscaled / 2.0, 1));
+
+ assert!(needs_norm_scale(0.0, 1));
+ assert!(needs_norm_scale(f64::INFINITY, 1));
+ assert!(needs_norm_scale(f64::NAN, 1));
+ }
}
diff --git a/datafusion/sqllogictest/test_files/array/array_length.slt b/datafusion/sqllogictest/test_files/array/array_length.slt
index 7741d815bc..bbfda57811 100644
--- a/datafusion/sqllogictest/test_files/array/array_length.slt
+++ b/datafusion/sqllogictest/test_files/array/array_length.slt
@@ -197,6 +197,23 @@ select
----
NULL NULL
+# array_distance scales the differences before squaring them, so finite
+# inputs whose squares overflow or underflow still give the correct distance
+query RR
+select
+ array_distance([CAST(1e-200 AS DOUBLE)], [CAST(0 AS DOUBLE)]) / 1e-200,
+ array_distance([CAST(3e200 AS DOUBLE)], [CAST(-1e200 AS DOUBLE)]) / 1e200;
+----
+1 4
+
+# non-finite inputs propagate as before
+query RR
+select
+ array_distance([CAST('Infinity' AS DOUBLE)], [CAST(0 AS DOUBLE)]),
+ array_distance([CAST('NaN' AS DOUBLE)], [CAST(0 AS DOUBLE)]);
+----
+Infinity NaN
+
# invalid argument count and types
query error DataFusion error: Error during planning: Execution error: Function 'array_distance' user-defined coercion failed with: Execution error: array_distance function requires 2 arguments
select array_distance();
diff --git a/datafusion/sqllogictest/test_files/array_normalize.slt b/datafusion/sqllogictest/test_files/array_normalize.slt
index ba4711d02c..f9fbc14144 100644
--- a/datafusion/sqllogictest/test_files/array_normalize.slt
+++ b/datafusion/sqllogictest/test_files/array_normalize.slt
@@ -144,3 +144,18 @@ select list_normalize(column1) from (values
----
[0.6, 0.8]
NULL
+
+# array_normalize scales the values before squaring them, so finite inputs
+# whose squares overflow or underflow still normalize correctly
+query ??
+select
+ array_normalize([3 * power(2.0, 700), 4 * power(2.0, 700)]),
+ array_normalize([3 * power(2.0, -700), -4 * power(2.0, -700)]);
+----
+[0.6, 0.8] [0.6, -0.8]
+
+# non-finite inputs propagate as before
+query ?
+select array_normalize([CAST('Infinity' AS DOUBLE), 1.0]);
+----
+[NaN, 0.0]
diff --git a/datafusion/sqllogictest/test_files/cosine_distance.slt b/datafusion/sqllogictest/test_files/cosine_distance.slt
index 9142aac8cf..12bc5621ed 100644
--- a/datafusion/sqllogictest/test_files/cosine_distance.slt
+++ b/datafusion/sqllogictest/test_files/cosine_distance.slt
@@ -165,3 +165,30 @@ query RT
select cosine_distance([1.0, 0.0], [0.0, 1.0]), arrow_typeof(cosine_distance([1.0, 0.0], [0.0, 1.0]));
----
1 Float64
+
+# cosine_distance scales each vector before multiplying, so finite inputs
+# whose products overflow or underflow still give the correct distance
+query RRR
+select
+ cosine_distance([CAST(3e200 AS DOUBLE), CAST(4e200 AS DOUBLE)], [CAST(3e200 AS DOUBLE), CAST(4e200 AS DOUBLE)]),
+ cosine_distance([CAST(1e-200 AS DOUBLE), CAST(2e-200 AS DOUBLE)], [CAST(1e-200 AS DOUBLE), CAST(2e-200 AS DOUBLE)]),
+ cosine_distance([CAST(1e200 AS DOUBLE), CAST(0 AS DOUBLE)], [CAST(-1e-200 AS DOUBLE), CAST(0 AS DOUBLE)]);
+----
+0 0 2
+
+# non-finite inputs propagate as before
+query RR
+select
+ cosine_distance([CAST('NaN' AS DOUBLE), 1.0], [1.0, 1.0]),
+ cosine_distance([CAST('Infinity' AS DOUBLE), 1.0], [1.0, 1.0]);
+----
+NaN NaN
+
+# only one vector's sum of squares is out of range, so only that vector is
+# scaled and the dot product mixes a scaled value with an unscaled one
+query RR
+select
+ cosine_distance([CAST(1e200 AS DOUBLE), CAST(0 AS DOUBLE)], [CAST(1 AS DOUBLE), CAST(0 AS DOUBLE)]),
+ cosine_distance([CAST(1e200 AS DOUBLE), CAST(1 AS DOUBLE)], [CAST(1e-200 AS DOUBLE), CAST(1 AS DOUBLE)]);
+----
+0 1
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org