Simplified traits.

jorgecarleitao · Dec 21, 2021 · 9d60670 · 9d60670
1 parent 8cfcb01
commit 9d60670
Show file tree

Hide file tree

Showing 38 changed files with 584 additions and 645 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
@@ -84,7 +84,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
@@ -370,7 +370,7 @@ pub use list::{ListArray, MutableListArray};
 pub use map::MapArray;
 pub use null::NullArray;
 pub use primitive::*;
-pub use specification::Offset;
+pub use crate::types::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
@@ -3,32 +3,30 @@ use std::iter::FromIterator;
 use crate::{
     buffer::{Buffer, MutableBuffer},
     trusted_len::TrustedLen,
-    types::{NativeType, NaturalDataType},
+    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,
+            T::PRIMITIVE.into(),
             MutableBuffer::<T>::from_iter(iter).into(),
             None,
         )
@@ -38,11 +36,11 @@ impl<T: NativeType + NaturalDataType> PrimitiveArray<T> {
     /// # 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, Buffer::<T>::from(slice), None)
+        Self::from_data(T::PRIMITIVE.into(), Buffer::<T>::from(slice), 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
@@ -7,7 +7,7 @@ use crate::{
     datatypes::DataType,
     error::{ArrowError, Result},
     trusted_len::TrustedLen,
-    types::{NativeType, NaturalDataType},
+    types::NativeType,
 };
 
 use super::PrimitiveArray;
@@ -33,21 +33,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.
@@ -60,7 +60,7 @@ impl<T: NativeType + NaturalDataType> MutablePrimitiveArray<T> {
         values: MutableBuffer<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>(),
@@ -84,15 +84,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: MutableBuffer::<T>::new(),
@@ -104,7 +104,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: MutableBuffer::<T>::with_capacity(capacity),
@@ -392,7 +392,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())
@@ -411,7 +411,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,
         }
@@ -444,7 +444,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,
         })
@@ -463,7 +463,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: MutableBuffer::<T>::from_trusted_len_iter(iter),
             validity: None,
         }
@@ -475,14 +475,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: MutableBuffer::<T>::from_trusted_len_iter_unchecked(iter),
             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 {
@@ -510,7 +510,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!(