fix: avoid overflow and underflow in vector functions

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));

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)
+}

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.

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));
+    }
 }

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();

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]

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