namanjain24-sudo commented on code in PR #25289:
URL: https://github.com/apache/datafusion/pull/25289#discussion_r4029145468


##########
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:
   Thanks, and good catch: the existing cases do scale both vectors or neither. 
Added in 89464d5 rather than a follow-up, since the checks here need approving 
again either way.
   
   I ran your two queries instead of taking the values on trust, with 
`--complete` so sqllogictest wrote what DataFusion actually returns, and it 
produced exactly what you had:
   
   ```
   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
   ```
   
   Both do take the one-sided path. In each, the first vector's sum of squares 
is infinite so it is scaled, while the second's is `1`, which is in range, so 
it is not, and the dot product multiplies a scaled value by an unscaled one. 
The two cases already in the file behave differently: the huge pair has a 
non-finite dot product and the tiny pair has both sums at zero, so each scales 
both vectors.
   
   One thing worth saying plainly: this pins the one-sided results rather than 
catching a regression, because the earlier scale-both form returns the same 
values for these inputs.



-- 
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]

Reply via email to