Simplified traits.

jorgecarleitao · Dec 21, 2021 · b1186a8 · b1186a8
1 parent 4bbbe62
commit b1186a8
Show file tree

Hide file tree

Showing 38 changed files with 588 additions and 648 deletions.
diff --git a/src/array/binary/mod.rs b/src/array/binary/mod.rs
@@ -1,8 +1,9 @@
 use crate::{bitmap::Bitmap, buffer::Buffer, datatypes::DataType};
 
 use super::{
-    display_fmt, display_helper, specification::check_offsets,
-    specification::check_offsets_minimal, specification::Offset, Array, GenericBinaryArray,
+    display_fmt, display_helper,
+    specification::{check_offsets, check_offsets_minimal},
+    Array, GenericBinaryArray, Offset,
 };
 
 mod ffi;

diff --git a/src/array/dictionary/ffi.rs b/src/array/dictionary/ffi.rs
@@ -30,7 +30,8 @@ impl<K: DictionaryKey, A: ffi::ArrowArrayRef> FromFfi<A> for DictionaryArray<K>
         let validity = unsafe { array.validity() }?;
         let values = unsafe { array.buffer::<K>(1) }?;
 
-        let keys = PrimitiveArray::<K>::from_data(K::DATA_TYPE, values, validity);
+        let data_type = K::PRIMITIVE.into();
+        let keys = PrimitiveArray::<K>::from_data(data_type, values, validity);
         let values = array.dictionary()?.unwrap();
         let values = ffi::try_from(values)?.into();
 

diff --git a/src/array/dictionary/mod.rs b/src/array/dictionary/mod.rs
@@ -62,15 +62,17 @@ impl<K: DictionaryKey> DictionaryArray<K> {
     pub fn new_empty(data_type: DataType) -> Self {
         let values = Self::get_child(&data_type);
         let values = new_empty_array(values.clone()).into();
-        Self::from_data(PrimitiveArray::<K>::new_empty(K::DATA_TYPE), values)
+        let data_type = K::PRIMITIVE.into();
+        Self::from_data(PrimitiveArray::<K>::new_empty(data_type), values)
     }
 
     /// Returns an [`DictionaryArray`] whose all elements are null
     #[inline]
     pub fn new_null(data_type: DataType, length: usize) -> Self {
         let values = Self::get_child(&data_type);
+        let data_type = K::PRIMITIVE.into();
         Self::from_data(
-            PrimitiveArray::<K>::new_null(K::DATA_TYPE, length),
+            PrimitiveArray::<K>::new_null(data_type, length),
             new_empty_array(values.clone()).into(),
         )
     }

diff --git a/src/array/growable/dictionary.rs b/src/array/growable/dictionary.rs
@@ -83,7 +83,8 @@ impl<'a, T: DictionaryKey> GrowableDictionary<'a, T> {
         let validity = std::mem::take(&mut self.key_validity);
         let values = std::mem::take(&mut self.key_values);
 
-        let keys = PrimitiveArray::<T>::from_data(T::DATA_TYPE, values.into(), validity.into());
+        let data_type = T::PRIMITIVE.into();
+        let keys = PrimitiveArray::<T>::from_data(data_type, values.into(), validity.into());
 
         DictionaryArray::<T>::from_data(keys, self.values.clone())
     }
@@ -125,8 +126,9 @@ impl<'a, T: DictionaryKey> Growable<'a> for GrowableDictionary<'a, T> {
 impl<'a, T: DictionaryKey> From<GrowableDictionary<'a, T>> for DictionaryArray<T> {
     #[inline]
     fn from(val: GrowableDictionary<'a, T>) -> Self {
+        let data_type = T::PRIMITIVE.into();
         let keys = PrimitiveArray::<T>::from_data(
-            T::DATA_TYPE,
+            data_type,
             val.key_values.into(),
             val.key_validity.into(),
         );

diff --git a/src/array/list/ffi.rs b/src/array/list/ffi.rs
@@ -2,7 +2,7 @@ use std::sync::Arc;
 
 use crate::{array::FromFfi, bitmap::align, error::Result, ffi};
 
-use super::super::{ffi::ToFfi, specification::Offset, Array};
+use super::super::{ffi::ToFfi, Array, Offset};
 use super::ListArray;
 
 unsafe impl<O: Offset> ToFfi for ListArray<O> {

diff --git a/src/array/list/mod.rs b/src/array/list/mod.rs
@@ -6,11 +6,7 @@ use crate::{
     datatypes::{DataType, Field},
 };
 
-use super::{
-    debug_fmt, new_empty_array,
-    specification::{check_offsets, Offset},
-    Array,
-};
+use super::{debug_fmt, new_empty_array, specification::check_offsets, Array, Offset};
 
 mod ffi;
 mod iterator;

diff --git a/src/array/mod.rs b/src/array/mod.rs
@@ -361,6 +361,7 @@ pub mod ord;
 pub use display::get_display;
 pub use equal::equal;
 
+pub use crate::types::Offset;
 pub use binary::{BinaryArray, BinaryValueIter, MutableBinaryArray};
 pub use boolean::{BooleanArray, MutableBooleanArray};
 pub use dictionary::{DictionaryArray, DictionaryKey, MutableDictionaryArray};
@@ -370,7 +371,6 @@ pub use list::{ListArray, MutableListArray};
 pub use map::MapArray;
 pub use null::NullArray;
 pub use primitive::*;
-pub use specification::Offset;
 pub use struct_::StructArray;
 pub use union::UnionArray;
 pub use utf8::{MutableUtf8Array, Utf8Array, Utf8ValuesIter};

diff --git a/src/array/primitive/from_natural.rs b/src/array/primitive/from_natural.rs
@@ -1,43 +1,42 @@
 use std::iter::FromIterator;
 
-use crate::{
-    trusted_len::TrustedLen,
-    types::{NativeType, NaturalDataType},
-};
+use crate::{trusted_len::TrustedLen, types::NativeType};
 
 use super::{MutablePrimitiveArray, PrimitiveArray};
 
-impl<T: NativeType + NaturalDataType, P: AsRef<[Option<T>]>> From<P> for PrimitiveArray<T> {
+impl<T: NativeType, P: AsRef<[Option<T>]>> From<P> for PrimitiveArray<T> {
     fn from(slice: P) -> Self {
         MutablePrimitiveArray::<T>::from(slice).into()
     }
 }
 
-impl<T: NativeType + NaturalDataType, Ptr: std::borrow::Borrow<Option<T>>> FromIterator<Ptr>
-    for PrimitiveArray<T>
-{
+impl<T: NativeType, Ptr: std::borrow::Borrow<Option<T>>> FromIterator<Ptr> for PrimitiveArray<T> {
     fn from_iter<I: IntoIterator<Item = Ptr>>(iter: I) -> Self {
         MutablePrimitiveArray::<T>::from_iter(iter).into()
     }
 }
 
-impl<T: NativeType + NaturalDataType> PrimitiveArray<T> {
+impl<T: NativeType> PrimitiveArray<T> {
     /// Creates a (non-null) [`PrimitiveArray`] from an iterator of values.
     /// # Implementation
     /// This does not assume that the iterator has a known length.
     pub fn from_values<I: IntoIterator<Item = T>>(iter: I) -> Self {
-        Self::from_data(T::DATA_TYPE, Vec::<T>::from_iter(iter).into(), None)
+        Self::from_data(T::PRIMITIVE.into(), Vec::<T>::from_iter(iter).into(), None)
     }
 
     /// Creates a (non-null) [`PrimitiveArray`] from a slice of values.
     /// # Implementation
     /// This is essentially a memcopy and is the fastest way to create a [`PrimitiveArray`].
     pub fn from_slice<P: AsRef<[T]>>(slice: P) -> Self {
-        Self::from_data(T::DATA_TYPE, Vec::<T>::from(slice.as_ref()).into(), None)
+        Self::from_data(
+            T::PRIMITIVE.into(),
+            Vec::<T>::from(slice.as_ref()).into(),
+            None,
+        )
     }
 }
 
-impl<T: NativeType + NaturalDataType> PrimitiveArray<T> {
+impl<T: NativeType> PrimitiveArray<T> {
     /// Creates a (non-null) [`PrimitiveArray`] from a [`TrustedLen`] of values.
     /// # Implementation
     /// This does not assume that the iterator has a known length.

diff --git a/src/array/primitive/mod.rs b/src/array/primitive/mod.rs
@@ -59,7 +59,7 @@ impl<T: NativeType> PrimitiveArray<T> {
     /// * `data_type` is not supported by the physical type
     /// * The validity is not `None` and its length is different from the `values`'s length
     pub fn from_data(data_type: DataType, values: Buffer<T>, validity: Option<Bitmap>) -> Self {
-        if !T::is_valid(&data_type) {
+        if !data_type.to_physical_type().eq_primitive(T::PRIMITIVE) {
             Err(ArrowError::InvalidArgumentError(format!(
                 "Type {} does not support logical type {:?}",
                 std::any::type_name::<T>(),
@@ -170,7 +170,7 @@ impl<T: NativeType> PrimitiveArray<T> {
     /// Panics iff the data_type is not supported for the physical type.
     #[inline]
     pub fn to(self, data_type: DataType) -> Self {
-        if !T::is_valid(&data_type) {
+        if !data_type.to_physical_type().eq_primitive(T::PRIMITIVE) {
             Err(ArrowError::InvalidArgumentError(format!(
                 "Type {} does not support logical type {:?}",
                 std::any::type_name::<T>(),

diff --git a/src/array/primitive/mutable.rs b/src/array/primitive/mutable.rs
@@ -6,7 +6,7 @@ use crate::{
     datatypes::DataType,
     error::{ArrowError, Result},
     trusted_len::TrustedLen,
-    types::{NativeType, NaturalDataType},
+    types::NativeType,
 };
 
 use super::PrimitiveArray;
@@ -32,21 +32,21 @@ impl<T: NativeType> From<MutablePrimitiveArray<T>> for PrimitiveArray<T> {
     }
 }
 
-impl<T: NativeType + NaturalDataType, P: AsRef<[Option<T>]>> From<P> for MutablePrimitiveArray<T> {
+impl<T: NativeType, P: AsRef<[Option<T>]>> From<P> for MutablePrimitiveArray<T> {
     fn from(slice: P) -> Self {
         Self::from_trusted_len_iter(slice.as_ref().iter().map(|x| x.as_ref()))
     }
 }
 
-impl<T: NativeType + NaturalDataType> MutablePrimitiveArray<T> {
+impl<T: NativeType> MutablePrimitiveArray<T> {
     /// Creates a new empty [`MutablePrimitiveArray`].
     pub fn new() -> Self {
         Self::with_capacity(0)
     }
 
     /// Creates a new [`MutablePrimitiveArray`] with a capacity.
     pub fn with_capacity(capacity: usize) -> Self {
-        Self::with_capacity_from(capacity, T::DATA_TYPE)
+        Self::with_capacity_from(capacity, T::PRIMITIVE.into())
     }
 
     /// Create a [`MutablePrimitiveArray`] out of low-end APIs.
@@ -55,7 +55,7 @@ impl<T: NativeType + NaturalDataType> MutablePrimitiveArray<T> {
     /// * `data_type` is not supported by the physical type
     /// * The validity is not `None` and its length is different from the `values`'s length
     pub fn from_data(data_type: DataType, values: Vec<T>, validity: Option<MutableBitmap>) -> Self {
-        if !T::is_valid(&data_type) {
+        if !data_type.to_physical_type().eq_primitive(T::PRIMITIVE) {
             Err(ArrowError::InvalidArgumentError(format!(
                 "Type {} does not support logical type {:?}",
                 std::any::type_name::<T>(),
@@ -79,15 +79,15 @@ impl<T: NativeType + NaturalDataType> MutablePrimitiveArray<T> {
     }
 }
 
-impl<T: NativeType + NaturalDataType> Default for MutablePrimitiveArray<T> {
+impl<T: NativeType> Default for MutablePrimitiveArray<T> {
     fn default() -> Self {
         Self::new()
     }
 }
 
 impl<T: NativeType> From<DataType> for MutablePrimitiveArray<T> {
     fn from(data_type: DataType) -> Self {
-        assert!(T::is_valid(&data_type));
+        assert!(data_type.to_physical_type().eq_primitive(T::PRIMITIVE));
         Self {
             data_type,
             values: Vec::<T>::new(),
@@ -99,7 +99,7 @@ impl<T: NativeType> From<DataType> for MutablePrimitiveArray<T> {
 impl<T: NativeType> MutablePrimitiveArray<T> {
     /// Creates a new [`MutablePrimitiveArray`] from a capacity and [`DataType`].
     pub fn with_capacity_from(capacity: usize, data_type: DataType) -> Self {
-        assert!(T::is_valid(&data_type));
+        assert!(data_type.to_physical_type().eq_primitive(T::PRIMITIVE));
         Self {
             data_type,
             values: Vec::<T>::with_capacity(capacity),
@@ -388,7 +388,7 @@ impl<T: NativeType> MutableArray for MutablePrimitiveArray<T> {
     }
 }
 
-impl<T: NativeType + NaturalDataType> MutablePrimitiveArray<T> {
+impl<T: NativeType> MutablePrimitiveArray<T> {
     /// Creates a [`MutablePrimitiveArray`] from a slice of values.
     pub fn from_slice<P: AsRef<[T]>>(slice: P) -> Self {
         Self::from_trusted_len_values_iter(slice.as_ref().iter().copied())
@@ -407,7 +407,7 @@ impl<T: NativeType + NaturalDataType> MutablePrimitiveArray<T> {
         let (validity, values) = trusted_len_unzip(iterator);
 
         Self {
-            data_type: T::DATA_TYPE,
+            data_type: T::PRIMITIVE.into(),
             values,
             validity,
         }
@@ -440,7 +440,7 @@ impl<T: NativeType + NaturalDataType> MutablePrimitiveArray<T> {
         let (validity, values) = try_trusted_len_unzip(iterator)?;
 
         Ok(Self {
-            data_type: T::DATA_TYPE,
+            data_type: T::PRIMITIVE.into(),
             values,
             validity,
         })
@@ -459,7 +459,7 @@ impl<T: NativeType + NaturalDataType> MutablePrimitiveArray<T> {
     /// Creates a new [`MutablePrimitiveArray`] out an iterator over values
     pub fn from_trusted_len_values_iter<I: TrustedLen<Item = T>>(iter: I) -> Self {
         Self {
-            data_type: T::DATA_TYPE,
+            data_type: T::PRIMITIVE.into(),
             values: iter.collect(),
             validity: None,
         }
@@ -471,14 +471,14 @@ impl<T: NativeType + NaturalDataType> MutablePrimitiveArray<T> {
     /// I.e. that `size_hint().1` correctly reports its length.
     pub unsafe fn from_trusted_len_values_iter_unchecked<I: Iterator<Item = T>>(iter: I) -> Self {
         Self {
-            data_type: T::DATA_TYPE,
+            data_type: T::PRIMITIVE.into(),
             values: iter.collect(),
             validity: None,
         }
     }
 }
 
-impl<T: NativeType + NaturalDataType, Ptr: std::borrow::Borrow<Option<T>>> FromIterator<Ptr>
+impl<T: NativeType, Ptr: std::borrow::Borrow<Option<T>>> FromIterator<Ptr>
     for MutablePrimitiveArray<T>
 {
     fn from_iter<I: IntoIterator<Item = Ptr>>(iter: I) -> Self {
@@ -506,7 +506,7 @@ impl<T: NativeType + NaturalDataType, Ptr: std::borrow::Borrow<Option<T>>> FromI
         };
 
         Self {
-            data_type: T::DATA_TYPE,
+            data_type: T::PRIMITIVE.into(),
             values,
             validity,
         }

diff --git a/src/array/specification.rs b/src/array/specification.rs
@@ -1,61 +1,4 @@
-use std::convert::TryFrom;
-
-use num_traits::Num;
-
-use crate::types::Index;
-
-mod private {
-    pub trait Sealed {}
-
-    impl Sealed for i32 {}
-    impl Sealed for i64 {}
-}
-
-/// Sealed trait describing types that can be used as offsets in Arrow (`i32` and `i64`).
-pub trait Offset: private::Sealed + Index + Num + Ord + num_traits::CheckedAdd {
-    /// Whether it is `i32` or `i64`
-    fn is_large() -> bool;
-
-    /// converts itself to `isize`
-    fn to_isize(&self) -> isize;
-
-    /// converts from `isize`
-    fn from_isize(value: isize) -> Option<Self>;
-}
-
-impl Offset for i32 {
-    #[inline]
-    fn is_large() -> bool {
-        false
-    }
-
-    #[inline]
-    fn from_isize(value: isize) -> Option<Self> {
-        Self::try_from(value).ok()
-    }
-
-    #[inline]
-    fn to_isize(&self) -> isize {
-        *self as isize
-    }
-}
-
-impl Offset for i64 {
-    #[inline]
-    fn is_large() -> bool {
-        true
-    }
-
-    #[inline]
-    fn from_isize(value: isize) -> Option<Self> {
-        Self::try_from(value).ok()
-    }
-
-    #[inline]
-    fn to_isize(&self) -> isize {
-        *self as isize
-    }
-}
+use crate::types::Offset;
 
 pub fn check_offsets_minimal<O: Offset>(offsets: &[O], values_len: usize) -> usize {
     assert!(