From d1bbaaaa11ba56efb9c2f069b7582627cbdbd22d Mon Sep 17 00:00:00 2001
From: izveigor <izveigor@gmail.com>
Date: Tue, 11 Apr 2023 21:43:06 +0300
Subject: [PATCH] feat: additional data type groups

---
 arrow-schema/src/datatype.rs | 58 ++++++++++++++++++++++++++++++++----
 1 file changed, 53 insertions(+), 5 deletions(-)

diff --git a/arrow-schema/src/datatype.rs b/arrow-schema/src/datatype.rs
index 3ec5597b2854..3f684285c067 100644
--- a/arrow-schema/src/datatype.rs
+++ b/arrow-schema/src/datatype.rs
@@ -355,14 +355,33 @@ impl DataType {
         )
     }
 
+    /// Returns true if this type is floating: (Float*).
+    pub fn is_floating(&self) -> bool {
+        use DataType::*;
+        matches!(self, Float16 | Float32 | Float64)
+    }
+
+    /// Returns true if this type is integer: (Int*, UInt*).
+    pub fn is_integer(&self) -> bool {
+        self.is_signed_integer() || self.is_unsigned_integer()
+    }
+
+    /// Returns true if this type is signed integer: (Int*).
+    pub fn is_signed_integer(&self) -> bool {
+        use DataType::*;
+        matches!(self, Int8 | Int16 | Int32 | Int64)
+    }
+
+    /// Returns true if this type is unsigned integer: (UInt*).
+    pub fn is_unsigned_integer(&self) -> bool {
+        use DataType::*;
+        matches!(self, UInt8 | UInt16 | UInt32 | UInt64)
+    }
+
     /// Returns true if this type is valid as a dictionary key
     #[inline]
     pub fn is_dictionary_key_type(&self) -> bool {
-        use DataType::*;
-        matches!(
-            self,
-            UInt8 | UInt16 | UInt32 | UInt64 | Int8 | Int16 | Int32 | Int64
-        )
+        self.is_integer()
     }
 
     /// Returns true if this type is valid for run-ends array in RunArray
@@ -664,6 +683,35 @@ mod tests {
         )));
     }
 
+    #[test]
+    fn test_integer() {
+        // is_integer
+        assert!(DataType::is_integer(&DataType::Int32));
+        assert!(DataType::is_integer(&DataType::UInt64));
+        assert!(!DataType::is_integer(&DataType::Float16));
+
+        // is_signed_integer
+        assert!(DataType::is_signed_integer(&DataType::Int32));
+        assert!(!DataType::is_signed_integer(&DataType::UInt64));
+        assert!(!DataType::is_signed_integer(&DataType::Float16));
+
+        // is_unsigned_integer
+        assert!(!DataType::is_unsigned_integer(&DataType::Int32));
+        assert!(DataType::is_unsigned_integer(&DataType::UInt64));
+        assert!(!DataType::is_unsigned_integer(&DataType::Float16));
+
+        // is_dictionary_key_type
+        assert!(DataType::is_dictionary_key_type(&DataType::Int32));
+        assert!(DataType::is_dictionary_key_type(&DataType::UInt64));
+        assert!(!DataType::is_dictionary_key_type(&DataType::Float16));
+    }
+
+    #[test]
+    fn test_floating() {
+        assert!(DataType::is_floating(&DataType::Float16));
+        assert!(!DataType::is_floating(&DataType::Int32));
+    }
+
     #[test]
     fn size_should_not_regress() {
         assert_eq!(std::mem::size_of::<DataType>(), 24);