From b00beb3715b44bf0eacb0db855a9ec815bf6993c Mon Sep 17 00:00:00 2001 From: naman Date: Mon, 14 Sep 2026 14:42:06 +0530 Subject: [PATCH 1/5] fix: avoid overflow and underflow in vector functions array_distance, cosine_distance and array_normalize squared their inputs directly, so finite values whose squares leave the Float64 range returned 0, inf, NaN, a zero vector or NULL instead of the representable result. When a sum of squares may have overflowed or underflowed, recompute it with the values multiplied by a power of two that brings the largest magnitude close to 1. Other rows run the same code as before, and empty, all-zero and non-finite inputs are never rescaled. --- .../functions-nested/src/array_normalize.rs | 23 ++++++- .../functions-nested/src/cosine_distance.rs | 41 ++++++++--- datafusion/functions-nested/src/distance.rs | 39 ++++++++--- datafusion/functions-nested/src/utils.rs | 69 +++++++++++++++++++ .../test_files/array/array_length.slt | 17 +++++ .../test_files/array_normalize.slt | 15 ++++ .../test_files/cosine_distance.slt | 18 +++++ 7 files changed, 198 insertions(+), 24 deletions(-) diff --git a/datafusion/functions-nested/src/array_normalize.rs b/datafusion/functions-nested/src/array_normalize.rs index 8f4342cd8389b..5790e239ece5f 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(arrays: &[ArrayRef]) -> Result(arrays: &[ArrayRef]) -> Result 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 56ca071234e1d..7eab877044748 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,14 @@ fn general_cosine_distance(arrays: &[ArrayRef]) -> Result(arrays: &[ArrayRef]) -> Result (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 d9182e871f72f..54a1405dbc729 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) { + 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 2906250ec6fa7..9a750193ca380 100644 --- a/datafusion/functions-nested/src/utils.rs +++ b/datafusion/functions-nested/src/utils.rs @@ -463,6 +463,45 @@ 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 and does not change how the result is rounded. +/// +/// Returns `None` when `values` is empty, all zero, or contains a NaN or an +/// infinity. The unscaled computation already gives the expected result for +/// those inputs. +pub(crate) fn norm_scale(values: impl IntoIterator) -> Option { + let mut max = 0.0_f64; + for value in values { + if !value.is_finite() { + return None; + } + max = max.max(value.abs()); + } + if max == 0.0 { + 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 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 less than the smallest subnormal value (about `5e-324`), so a sum of at +/// least `1e-180` is not affected by any practical number of such terms. +pub(crate) fn needs_norm_scale(sum_of_squares: f64) -> bool { + !(1e-180..f64::INFINITY).contains(&sum_of_squares) +} + #[cfg(test)] mod tests { use super::*; @@ -518,4 +557,34 @@ 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)); + assert_eq!(norm_scale([f64::MAX]), Some(2.0_f64.powi(-1023))); + assert_eq!( + norm_scale([f64::MIN_POSITIVE / 4.0]), + Some(2.0_f64.powi(1022)) + ); + } + + #[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([1.0, f64::NAN]), None); + assert_eq!(norm_scale([f64::INFINITY, 1.0]), None); + } + + #[test] + fn needs_norm_scale_only_for_sums_that_may_have_overflowed_or_underflowed() { + assert!(!needs_norm_scale(1.0)); + assert!(!needs_norm_scale(f64::MAX)); + assert!(!needs_norm_scale(1e-180)); + assert!(needs_norm_scale(1e-200)); + assert!(needs_norm_scale(0.0)); + assert!(needs_norm_scale(f64::INFINITY)); + assert!(needs_norm_scale(f64::NAN)); + } } diff --git a/datafusion/sqllogictest/test_files/array/array_length.slt b/datafusion/sqllogictest/test_files/array/array_length.slt index 7741d815bc234..bbfda578119cb 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 ba4711d02cf9d..f9fbc14144738 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 9142aac8cf684..90281214c8c99 100644 --- a/datafusion/sqllogictest/test_files/cosine_distance.slt +++ b/datafusion/sqllogictest/test_files/cosine_distance.slt @@ -165,3 +165,21 @@ 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 From efc50cf7f24a9321f8af7f7bea387856177d7c2c Mon Sep 17 00:00:00 2001 From: naman Date: Mon, 14 Sep 2026 19:33:15 +0530 Subject: [PATCH 2/5] Derive the rescaling threshold from the array length A fixed 1e-180 threshold rescaled sums that no underflowing square could have changed, such as the squares of values around 1e-100, and the rescaled path is several times slower. An underflowing square is off by at most 2^-1075, so rescale only when the sum is below len * 2^-1012. --- .../functions-nested/src/array_normalize.rs | 2 +- .../functions-nested/src/cosine_distance.rs | 3 +- datafusion/functions-nested/src/distance.rs | 2 +- datafusion/functions-nested/src/utils.rs | 36 ++++++++++++------- 4 files changed, 27 insertions(+), 16 deletions(-) diff --git a/datafusion/functions-nested/src/array_normalize.rs b/datafusion/functions-nested/src/array_normalize.rs index 5790e239ece5f..988e2679b991b 100644 --- a/datafusion/functions-nested/src/array_normalize.rs +++ b/datafusion/functions-nested/src/array_normalize.rs @@ -184,7 +184,7 @@ fn general_array_normalize(arrays: &[ArrayRef]) -> Result(arrays: &[ArrayRef]) -> Result) -> Option { Some(2.0_f64.powi(-exponent)) } -/// Returns whether a sum of squares computed without scaling may be wrong -/// because a square overflowed or underflowed, in which case it should be +/// 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 less than the smallest subnormal value (about `5e-324`), so a sum of at -/// least `1e-180` is not affected by any practical number of such terms. -pub(crate) fn needs_norm_scale(sum_of_squares: f64) -> bool { - !(1e-180..f64::INFINITY).contains(&sum_of_squares) +/// 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)] @@ -579,12 +582,19 @@ mod tests { #[test] fn needs_norm_scale_only_for_sums_that_may_have_overflowed_or_underflowed() { - assert!(!needs_norm_scale(1.0)); - assert!(!needs_norm_scale(f64::MAX)); - assert!(!needs_norm_scale(1e-180)); - assert!(needs_norm_scale(1e-200)); - assert!(needs_norm_scale(0.0)); - assert!(needs_norm_scale(f64::INFINITY)); - assert!(needs_norm_scale(f64::NAN)); + 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)); } } From b3fcbe175a19ce582c47658a7795cc233311b558 Mon Sep 17 00:00:00 2001 From: naman Date: Tue, 15 Sep 2026 19:40:27 +0530 Subject: [PATCH 3/5] Skip the rescale scan when it cannot change the result norm_scale no longer returns early on a non-finite value, so its loop vectorizes; it checks the maximum once instead. NaN values are ignored, and the scaled computation still produces NaN. cosine_distance now scans a vector only when its own sum of squares is out of range, and recomputes only when there is something to scale, so a row with one all-zero vector no longer scans both vectors and repeats the dot product. --- .../functions-nested/src/cosine_distance.rs | 31 ++++++++++++++----- datafusion/functions-nested/src/utils.rs | 26 +++++++++------- 2 files changed, 38 insertions(+), 19 deletions(-) diff --git a/datafusion/functions-nested/src/cosine_distance.rs b/datafusion/functions-nested/src/cosine_distance.rs index 9be7253fd6922..1637b2cea1340 100644 --- a/datafusion/functions-nested/src/cosine_distance.rs +++ b/datafusion/functions-nested/src/cosine_distance.rs @@ -198,14 +198,29 @@ fn general_cosine_distance(arrays: &[ArrayRef]) -> Result) -> Option { + // 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 { - if !value.is_finite() { - return None; - } max = max.max(value.abs()); } - if max == 0.0 { + if max == 0.0 || !max.is_finite() { return None; } // Unbiased exponent of `max`. Subnormal values store a biased exponent of 0, @@ -565,19 +565,23 @@ mod tests { 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)); - assert_eq!(norm_scale([f64::MAX]), Some(2.0_f64.powi(-1023))); + // 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(2.0_f64.powi(1022)) + 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([1.0, f64::NAN]), 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] From c249ef4982a807c1602741d6b2baa7d4724fc9d7 Mon Sep 17 00:00:00 2001 From: naman Date: Tue, 15 Sep 2026 21:27:31 +0530 Subject: [PATCH 4/5] Say which results rescaling can round A value that rescaling makes subnormal is rounded. That can change the last bit of an array_normalize element that is itself subnormal, so the doc no longer says it cannot change the result. --- datafusion/functions-nested/src/utils.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/datafusion/functions-nested/src/utils.rs b/datafusion/functions-nested/src/utils.rs index 9e90d3c9d9226..a6a93ed5dfc5d 100644 --- a/datafusion/functions-nested/src/utils.rs +++ b/datafusion/functions-nested/src/utils.rs @@ -470,7 +470,8 @@ where /// 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, but it is then too small to change the result. +/// 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. From 89464d57938c0029dbc78c05d759d91ec80b5200 Mon Sep 17 00:00:00 2001 From: naman Date: Wed, 16 Sep 2026 23:28:07 +0530 Subject: [PATCH 5/5] Cover one-sided scaling in cosine_distance The existing cases scale both vectors or neither. These two scale only the first: its sum of squares is infinite while the second's is 1, which is in range, so the dot product multiplies a scaled value by an unscaled one. --- datafusion/sqllogictest/test_files/cosine_distance.slt | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/datafusion/sqllogictest/test_files/cosine_distance.slt b/datafusion/sqllogictest/test_files/cosine_distance.slt index 90281214c8c99..12bc5621ed744 100644 --- a/datafusion/sqllogictest/test_files/cosine_distance.slt +++ b/datafusion/sqllogictest/test_files/cosine_distance.slt @@ -183,3 +183,12 @@ select 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