jayzhan211 commented on code in PR #25289:
URL: https://github.com/apache/datafusion/pull/25289#discussion_r4026960442
##########
datafusion/functions-nested/src/cosine_distance.rs:
##########
@@ -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) {
Review Comment:
The new tests only scale both vectors or neither. The one-sided path
(`scale1 = Some`, `scale2 = None`), where the dot product mixes scaled and
unscaled values, has no coverage. I checked it by hand and it returns the
correct results; please add a case to `cosine_distance.slt` so it stays that
way. Fine to do in a follow-up ("Add one-sided scaling tests for
cosine_distance").
```sql
# only one vector's sum of squares is out of range
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
```
--
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]
