Add compare to ArrowNativeTypeOp (#3070)

* Add total_cmp to ArrowNativeTypeOp * Format
apache · Nov 10, 2022 · e44cb5b · e44cb5b
1 parent 5fb3033
commit e44cb5b
Showing 1 changed file with 36 additions and 58 deletions.
diff --git 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,38 +339,18 @@ 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()`
                 // but LLVM isn't able to realise this is bitwise equality
                 // 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()
-            }
         }
     };
 }