[arrow-rs] branch master updated: feat: additional data type groups (#4057)

[tu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new eec499db9 feat: additional data type groups (#4057)
eec499db9 is described below

commit eec499db94a5794e2e4fd979177f2d63b112e41c
Author: Igor Izvekov <iz...@gmail.com>
AuthorDate: Wed Apr 12 06:46:36 2023 +0300

    feat: additional data type groups (#4057)
---
 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 3ec5597b2..3f684285c 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);