From e44cb5b478257751916d1674292123eaf5d80a7c Mon Sep 17 00:00:00 2001
From: Raphael Taylor-Davies <1781103+tustvold@users.noreply.github.com>
Date: Thu, 10 Nov 2022 16:56:20 +1300
Subject: [PATCH] Add compare to ArrowNativeTypeOp (#3070)

* Add total_cmp to ArrowNativeTypeOp

* Format
---
 arrow-array/src/arithmetic.rs | 94 ++++++++++++++---------------------
 1 file changed, 36 insertions(+), 58 deletions(-)

diff --git a/arrow-array/src/arithmetic.rs b/arrow-array/src/arithmetic.rs
index f21532364f65..566f3742e93d 100644
--- a/arrow-array/src/arithmetic.rs
+++ b/arrow-array/src/arithmetic.rs
@@ -19,6 +19,7 @@ use arrow_buffer::{i256, ArrowNativeType};
 use arrow_schema::ArrowError;
 use half::f16;
 use num::complex::ComplexFloat;
+use std::cmp::Ordering;
 
 /// Trait for [`ArrowNativeType`] that adds checked and unchecked arithmetic operations,
 /// and totally ordered comparison operations
@@ -74,17 +75,34 @@ pub trait ArrowNativeTypeOp: ArrowNativeType {
 
     fn is_zero(self) -> bool;
 
-    fn is_eq(self, rhs: Self) -> bool;
-
-    fn is_ne(self, rhs: Self) -> bool;
-
-    fn is_lt(self, rhs: Self) -> bool;
+    fn compare(self, rhs: Self) -> Ordering;
 
-    fn is_le(self, rhs: Self) -> bool;
-
-    fn is_gt(self, rhs: Self) -> bool;
+    fn is_eq(self, rhs: Self) -> bool;
 
-    fn is_ge(self, rhs: Self) -> bool;
+    #[inline]
+    fn is_ne(self, rhs: Self) -> bool {
+        !self.is_eq(rhs)
+    }
+
+    #[inline]
+    fn is_lt(self, rhs: Self) -> bool {
+        self.compare(rhs).is_lt()
+    }
+
+    #[inline]
+    fn is_le(self, rhs: Self) -> bool {
+        self.compare(rhs).is_le()
+    }
+
+    #[inline]
+    fn is_gt(self, rhs: Self) -> bool {
+        self.compare(rhs).is_gt()
+    }
+
+    #[inline]
+    fn is_ge(self, rhs: Self) -> bool {
+        self.compare(rhs).is_ge()
+    }
 }
 
 macro_rules! native_type_op {
@@ -209,33 +227,13 @@ macro_rules! native_type_op {
             }
 
             #[inline]
-            fn is_eq(self, rhs: Self) -> bool {
-                self == rhs
-            }
-
-            #[inline]
-            fn is_ne(self, rhs: Self) -> bool {
-                self != rhs
+            fn compare(self, rhs: Self) -> Ordering {
+                self.cmp(&rhs)
             }
 
             #[inline]
-            fn is_lt(self, rhs: Self) -> bool {
-                self < rhs
-            }
-
-            #[inline]
-            fn is_le(self, rhs: Self) -> bool {
-                self <= rhs
-            }
-
-            #[inline]
-            fn is_gt(self, rhs: Self) -> bool {
-                self > rhs
-            }
-
-            #[inline]
-            fn is_ge(self, rhs: Self) -> bool {
-                self >= rhs
+            fn is_eq(self, rhs: Self) -> bool {
+                self == rhs
             }
         }
     };
@@ -341,6 +339,11 @@ macro_rules! native_type_float_op {
                 self == $zero
             }
 
+            #[inline]
+            fn compare(self, rhs: Self) -> Ordering {
+                <$t>::total_cmp(&self, &rhs)
+            }
+
             #[inline]
             fn is_eq(self, rhs: Self) -> bool {
                 // Equivalent to `self.total_cmp(&rhs).is_eq()`
@@ -348,31 +351,6 @@ macro_rules! native_type_float_op {
                 // https://rust.godbolt.org/z/347nWGxoW
                 self.to_bits() == rhs.to_bits()
             }
-
-            #[inline]
-            fn is_ne(self, rhs: Self) -> bool {
-                !self.is_eq(rhs)
-            }
-
-            #[inline]
-            fn is_lt(self, rhs: Self) -> bool {
-                self.total_cmp(&rhs).is_lt()
-            }
-
-            #[inline]
-            fn is_le(self, rhs: Self) -> bool {
-                self.total_cmp(&rhs).is_le()
-            }
-
-            #[inline]
-            fn is_gt(self, rhs: Self) -> bool {
-                self.total_cmp(&rhs).is_gt()
-            }
-
-            #[inline]
-            fn is_ge(self, rhs: Self) -> bool {
-                self.total_cmp(&rhs).is_ge()
-            }
         }
     };
 }