diff --git a/arrow-schema/src/datatype.rs b/arrow-schema/src/datatype.rs index 4f8c8a18bd17..cf426a63357f 100644 --- a/arrow-schema/src/datatype.rs +++ b/arrow-schema/src/datatype.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::cmp::Ordering; use std::fmt; use std::sync::Arc; @@ -35,7 +36,7 @@ use crate::{Field, FieldRef, Fields, UnionFields}; /// Nested types can themselves be nested within other arrays. /// For more information on these types please see /// [the physical memory layout of Apache Arrow](https://arrow.apache.org/docs/format/Columnar.html#physical-memory-layout). -#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] +#[derive(Debug, Clone, PartialEq, Eq, Hash)] #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] pub enum DataType { /// Null type @@ -310,6 +311,38 @@ impl fmt::Display for DataType { } } +impl PartialOrd for DataType { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for DataType { + fn cmp(&self, other: &Self) -> Ordering { + match (self, other) { + (DataType::Null, DataType::Null) => Ordering::Equal, + (DataType::Null, _) => Ordering::Less, + (_, DataType::Null) => Ordering::Greater, + (DataType::Boolean, DataType::Boolean) => Ordering::Equal, + (DataType::Boolean, _) => Ordering::Less, + (_, DataType::Boolean) => Ordering::Greater, + (a, b) => { + if a.is_primitive() && b.is_primitive() { + self.primitive_width() + .unwrap_or(0) + .cmp(&other.primitive_width().unwrap_or(0)) + } else if a.is_primitive() { + Ordering::Less + } else if b.is_primitive() { + Ordering::Greater + } else { + a.size().cmp(&b.size()) + } + } + } + } +} + impl DataType { /// Returns true if the type is primitive: (numeric, temporal). #[inline] @@ -903,4 +936,19 @@ mod tests { UnionMode::Dense, ); } + + #[test] + fn test_cmp() { + assert_eq!(DataType::Null.cmp(&DataType::Boolean), Ordering::Less); + assert_eq!(DataType::Boolean.cmp(&DataType::Null), Ordering::Greater); + assert_eq!(DataType::Null.cmp(&DataType::Int16), Ordering::Less); + assert_eq!(DataType::Int8.cmp(&DataType::Int16), Ordering::Less); + assert_eq!(DataType::Int32.cmp(&DataType::UInt16), Ordering::Greater); + assert_eq!(DataType::UInt64.cmp(&DataType::UInt64), Ordering::Equal); + assert_eq!( + DataType::List(Arc::new(Field::new("A", DataType::Int64, true))) + .cmp(&DataType::Decimal256(5, 2)), + Ordering::Greater + ); + } }