Compute gcd with u64 instead of i64 because of overflows

datafusion/functions/src/math/gcd.rs

@@ -88,16 +88,16 @@ fn gcd(args: &[ArrayRef]) -> Result<ArrayRef> {
 
 /// Computes greatest common divisor using Binary GCD algorithm.
 pub fn compute_gcd(x: i64, y: i64) -> i64 {
-    let mut a = x.wrapping_abs();
-    let mut b = y.wrapping_abs();
-
-    if a == 0 {
-        return b;
+    if x == 0 {
+        return y;
     }
-    if b == 0 {
-        return a;
+    if y == 0 {
+        return x;
     }
 
+    let mut a = x.unsigned_abs();
+    let mut b = y.unsigned_abs();
+
     let shift = (a | b).trailing_zeros();
     a >>= shift;
     b >>= shift;
@@ -112,7 +112,8 @@ pub fn compute_gcd(x: i64, y: i64) -> i64 {
         b -= a;
 
         if b == 0 {
-            return a << shift;
+            // because the input values are i64, casting this back to i64 is safe
+            return (a << shift) as i64;
         }
     }
 }
@@ -142,4 +143,22 @@ mod test {
         assert_eq!(ints.value(2), 5);
         assert_eq!(ints.value(3), 8);
     }
+
+    // from issue https://github.com/apache/datafusion/issues/11031 we know that the previous implementation could
+    // not handle cases were one or both of the inputs were an i64::MAX or i64::MIN coupled with other values (neg or pos)
+    #[test]
+    fn test_gcd_i64_maxes() {
+        let args: Vec<ArrayRef> = vec![
+            Arc::new(Int64Array::from(vec![i64::MAX, i64::MIN, i64::MAX])), // x
+            Arc::new(Int64Array::from(vec![i64::MIN, -1, -1])),             // y
+        ];
+
+        let result = gcd(&args).expect("failed to initialize function gcd");
+        let ints = as_int64_array(&result).expect("failed to initialize function gcd");
+
+        assert_eq!(ints.len(), 3);
+        assert_eq!(ints.value(0), 1);
+        assert_eq!(ints.value(1), 1);
+        assert_eq!(ints.value(1), 1);
+    }
 }