jayzhan211 commented on PR #25289:
URL: https://github.com/apache/datafusion/pull/25289#issuecomment-5681214558
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:
```rs
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:
```rs
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:
```rs
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))
);
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
