diff --git a/datafusion/functions-nested/src/array_normalize.rs b/datafusion/functions-nested/src/array_normalize.rs index 8f4342cd8389b..988e2679b991b 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..1637b2cea1340 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(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..36b263870d8e6 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 2906250ec6fa7..a6a93ed5dfc5d 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) -> 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 { + 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 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..12bc5621ed744 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