Simplified dictionary indexes

src/array/dictionary/mod.rs

@@ -2,9 +2,9 @@ use std::sync::Arc;
 
 use crate::{
     bitmap::Bitmap,
-    datatypes::DataType,
+    datatypes::{DataType, IntegerType},
     scalar::{new_scalar, Scalar},
-    types::{NativeType, NaturalDataType},
+    types::NativeType,
 };
 
 mod ffi;
@@ -16,19 +16,35 @@ pub use mutable::*;
 use super::{new_empty_array, primitive::PrimitiveArray, Array};
 
 /// Trait denoting [`NativeType`]s that can be used as keys of a dictionary.
-pub trait DictionaryKey:
-    NativeType + NaturalDataType + num_traits::NumCast + num_traits::FromPrimitive
-{
+pub trait DictionaryKey: NativeType + num_traits::NumCast + num_traits::FromPrimitive {
+    /// The corresponding [`IntegerType`] of this key
+    const KEY_TYPE: IntegerType;
 }
 
-impl DictionaryKey for i8 {}
-impl DictionaryKey for i16 {}
-impl DictionaryKey for i32 {}
-impl DictionaryKey for i64 {}
-impl DictionaryKey for u8 {}
-impl DictionaryKey for u16 {}
-impl DictionaryKey for u32 {}
-impl DictionaryKey for u64 {}
+impl DictionaryKey for i8 {
+    const KEY_TYPE: IntegerType = IntegerType::Int8;
+}
+impl DictionaryKey for i16 {
+    const KEY_TYPE: IntegerType = IntegerType::Int16;
+}
+impl DictionaryKey for i32 {
+    const KEY_TYPE: IntegerType = IntegerType::Int32;
+}
+impl DictionaryKey for i64 {
+    const KEY_TYPE: IntegerType = IntegerType::Int64;
+}
+impl DictionaryKey for u8 {
+    const KEY_TYPE: IntegerType = IntegerType::UInt8;
+}
+impl DictionaryKey for u16 {
+    const KEY_TYPE: IntegerType = IntegerType::UInt16;
+}
+impl DictionaryKey for u32 {
+    const KEY_TYPE: IntegerType = IntegerType::UInt32;
+}
+impl DictionaryKey for u64 {
+    const KEY_TYPE: IntegerType = IntegerType::UInt64;
+}
 
 /// An [`Array`] whose values are encoded by keys. This [`Array`] is useful when the cardinality of
 /// values is low compared to the length of the [`Array`].
@@ -59,10 +75,7 @@ impl<K: DictionaryKey> DictionaryArray<K> {
 
     /// The canonical method to create a new [`DictionaryArray`].
     pub fn from_data(keys: PrimitiveArray<K>, values: Arc<dyn Array>) -> Self {
-        let data_type = DataType::Dictionary(
-            Box::new(keys.data_type().clone()),
-            Box::new(values.data_type().clone()),
-        );
+        let data_type = DataType::Dictionary(K::KEY_TYPE, Box::new(values.data_type().clone()));
 
         Self {
             data_type,

src/array/dictionary/mutable.rs

@@ -31,10 +31,7 @@ impl<K: DictionaryKey, M: MutableArray> From<MutableDictionaryArray<K, M>> for D
 impl<K: DictionaryKey, M: MutableArray> From<M> for MutableDictionaryArray<K, M> {
     fn from(values: M) -> Self {
         Self {
-            data_type: DataType::Dictionary(
-                Box::new(K::DATA_TYPE),
-                Box::new(values.data_type().clone()),
-            ),
+            data_type: DataType::Dictionary(K::KEY_TYPE, Box::new(values.data_type().clone())),
             keys: MutablePrimitiveArray::<K>::new(),
             map: HashedMap::default(),
             values,
@@ -47,10 +44,7 @@ impl<K: DictionaryKey, M: MutableArray + Default> MutableDictionaryArray<K, M> {
     pub fn new() -> Self {
         let values = M::default();
         Self {
-            data_type: DataType::Dictionary(
-                Box::new(K::DATA_TYPE),
-                Box::new(values.data_type().clone()),
-            ),
+            data_type: DataType::Dictionary(K::KEY_TYPE, Box::new(values.data_type().clone())),
             keys: MutablePrimitiveArray::<K>::new(),
             map: HashedMap::default(),
             values,

src/array/display.rs

@@ -17,18 +17,6 @@ macro_rules! dyn_primitive {
     }};
 }
 
-macro_rules! dyn_dict {
-    ($array:expr, $ty:ty) => {{
-        let a = $array
-            .as_any()
-            .downcast_ref::<DictionaryArray<$ty>>()
-            .unwrap();
-        let keys = a.keys();
-        let display = get_display(a.values().as_ref());
-        Box::new(move |row: usize| display(keys.value(row) as usize))
-    }};
-}
-
 /// Returns a function of index returning the string representation of the _value_ of `array`.
 /// This does not take nulls into account.
 pub fn get_value_display<'a>(array: &'a dyn Array) -> Box<dyn Fn(usize) -> String + 'a> {
@@ -170,17 +158,15 @@ pub fn get_value_display<'a>(array: &'a dyn Array) -> Box<dyn Fn(usize) -> Strin
             };
             dyn_display!(array, ListArray<i64>, f)
         }
-        Dictionary(key_type, _) => match key_type.as_ref() {
-            DataType::Int8 => dyn_dict!(array, i8),
-            DataType::Int16 => dyn_dict!(array, i16),
-            DataType::Int32 => dyn_dict!(array, i32),
-            DataType::Int64 => dyn_dict!(array, i64),
-            DataType::UInt8 => dyn_dict!(array, u8),
-            DataType::UInt16 => dyn_dict!(array, u16),
-            DataType::UInt32 => dyn_dict!(array, u32),
-            DataType::UInt64 => dyn_dict!(array, u64),
-            _ => unreachable!(),
-        },
+        Dictionary(key_type, _) => match_integer_type!(key_type, |$T| {
+            let a = array
+                .as_any()
+                .downcast_ref::<DictionaryArray<$T>>()
+                .unwrap();
+            let keys = a.keys();
+            let display = get_display(a.values().as_ref());
+            Box::new(move |row: usize| display(keys.value(row) as usize))
+        }),
         Map(_, _) => todo!(),
         Struct(_) => {
             let a = array.as_any().downcast_ref::<StructArray>().unwrap();

src/array/equal/mod.rs

@@ -215,7 +215,7 @@ pub fn equal(lhs: &dyn Array, rhs: &dyn Array) -> bool {
             struct_::equal(lhs, rhs)
         }
         Dictionary(key_type) => {
-            with_match_physical_dictionary_key_type!(key_type, |$T| {
+            match_integer_type!(key_type, |$T| {
                 let lhs = lhs.as_any().downcast_ref().unwrap();
                 let rhs = rhs.as_any().downcast_ref().unwrap();
                 dictionary::equal::<$T>(lhs, rhs)

src/array/ffi.rs

@@ -74,7 +74,7 @@ pub fn offset_buffers_children_dictionary(array: &dyn Array) -> BuffersChildren
         Union => ffi_dyn!(array, UnionArray),
         Map => ffi_dyn!(array, MapArray),
         Dictionary(key_type) => {
-            with_match_physical_dictionary_key_type!(key_type, |$T| {
+            match_integer_type!(key_type, |$T| {
                 let array = array.as_any().downcast_ref::<DictionaryArray<$T>>().unwrap();
                 (
                     array.offset().unwrap(),

src/array/growable/mod.rs

@@ -61,25 +61,6 @@ macro_rules! dyn_growable {
     }};
 }
 
-macro_rules! dyn_dict_growable {
-    ($ty:ty, $arrays:expr, $use_validity:expr, $capacity:expr) => {{
-        let arrays = $arrays
-            .iter()
-            .map(|array| {
-                array
-                    .as_any()
-                    .downcast_ref::<DictionaryArray<$ty>>()
-                    .unwrap()
-            })
-            .collect::<Vec<_>>();
-        Box::new(dictionary::GrowableDictionary::<$ty>::new(
-            &arrays,
-            $use_validity,
-            $capacity,
-        ))
-    }};
-}
-
 /// Creates a new [`Growable`] from an arbitrary number of [`Array`]s.
 /// # Panics
 /// This function panics iff
@@ -132,8 +113,21 @@ pub fn make_growable<'a>(
         ),
         Union | Map => todo!(),
         Dictionary(key_type) => {
-            with_match_physical_dictionary_key_type!(key_type, |$T| {
-                dyn_dict_growable!($T, arrays, use_validity, capacity)
+            match_integer_type!(key_type, |$T| {
+                let arrays = arrays
+                    .iter()
+                    .map(|array| {
+                        array
+                            .as_any()
+                            .downcast_ref::<DictionaryArray<$T>>()
+                            .unwrap()
+                    })
+                    .collect::<Vec<_>>();
+                Box::new(dictionary::GrowableDictionary::<$T>::new(
+                    &arrays,
+                    use_validity,
+                    capacity,
+                ))
             })
         }
     }

src/array/mod.rs

@@ -174,28 +174,11 @@ macro_rules! fmt_dyn {
     }};
 }
 
-macro_rules! with_match_dictionary_key_type {(
+macro_rules! match_integer_type {(
     $key_type:expr, | $_:tt $T:ident | $($body:tt)*
 ) => ({
     macro_rules! __with_ty__ {( $_ $T:ident ) => ( $($body)* )}
-    match $key_type {
-        DataType::Int8 => __with_ty__! { i8 },
-        DataType::Int16 => __with_ty__! { i16 },
-        DataType::Int32 => __with_ty__! { i32 },
-        DataType::Int64 => __with_ty__! { i64 },
-        DataType::UInt8 => __with_ty__! { u8 },
-        DataType::UInt16 => __with_ty__! { u16 },
-        DataType::UInt32 => __with_ty__! { u32 },
-        DataType::UInt64 => __with_ty__! { u64 },
-        _ => ::core::unreachable!("A dictionary key type can only be of integer types"),
-    }
-})}
-
-macro_rules! with_match_physical_dictionary_key_type {(
-    $key_type:expr, | $_:tt $T:ident | $($body:tt)*
-) => ({
-    macro_rules! __with_ty__ {( $_ $T:ident ) => ( $($body)* )}
-    use crate::datatypes::DictionaryIndexType::*;
+    use crate::datatypes::IntegerType::*;
     match $key_type {
         Int8 => __with_ty__! { i8 },
         Int16 => __with_ty__! { i16 },
@@ -251,7 +234,7 @@ impl Display for dyn Array {
             Struct => fmt_dyn!(self, StructArray, f),
             Union => fmt_dyn!(self, UnionArray, f),
             Dictionary(key_type) => {
-                with_match_physical_dictionary_key_type!(key_type, |$T| {
+                match_integer_type!(key_type, |$T| {
                     fmt_dyn!(self, DictionaryArray::<$T>, f)
                 })
             }
@@ -281,7 +264,7 @@ pub fn new_empty_array(data_type: DataType) -> Box<dyn Array> {
         Union => Box::new(UnionArray::new_empty(data_type)),
         Map => Box::new(MapArray::new_empty(data_type)),
         Dictionary(key_type) => {
-            with_match_physical_dictionary_key_type!(key_type, |$T| {
+            match_integer_type!(key_type, |$T| {
                 Box::new(DictionaryArray::<$T>::new_empty(data_type))
             })
         }
@@ -311,7 +294,7 @@ pub fn new_null_array(data_type: DataType, length: usize) -> Box<dyn Array> {
         Union => Box::new(UnionArray::new_null(data_type, length)),
         Map => Box::new(MapArray::new_null(data_type, length)),
         Dictionary(key_type) => {
-            with_match_physical_dictionary_key_type!(key_type, |$T| {
+            match_integer_type!(key_type, |$T| {
                 Box::new(DictionaryArray::<$T>::new_null(data_type, length))
             })
         }
@@ -349,7 +332,7 @@ pub fn clone(array: &dyn Array) -> Box<dyn Array> {
         Union => clone_dyn!(array, UnionArray),
         Map => clone_dyn!(array, MapArray),
         Dictionary(key_type) => {
-            with_match_physical_dictionary_key_type!(key_type, |$T| {
+            match_integer_type!(key_type, |$T| {
                 clone_dyn!(array, DictionaryArray::<$T>)
             })
         }

src/array/ord.rs

@@ -216,15 +216,15 @@ pub fn build_compare(left: &dyn Array, right: &dyn Array) -> Result<DynComparato
         (Binary, Binary) => compare_binary::<i32>(left, right),
         (LargeBinary, LargeBinary) => compare_binary::<i64>(left, right),
         (Dictionary(key_type_lhs, _), Dictionary(key_type_rhs, _)) => {
-            match (key_type_lhs.as_ref(), key_type_rhs.as_ref()) {
-                (UInt8, UInt8) => dyn_dict!(u8, left, right),
-                (UInt16, UInt16) => dyn_dict!(u16, left, right),
-                (UInt32, UInt32) => dyn_dict!(u32, left, right),
-                (UInt64, UInt64) => dyn_dict!(u64, left, right),
-                (Int8, Int8) => dyn_dict!(i8, left, right),
-                (Int16, Int16) => dyn_dict!(i16, left, right),
-                (Int32, Int32) => dyn_dict!(i32, left, right),
-                (Int64, Int64) => dyn_dict!(i64, left, right),
+            match (key_type_lhs, key_type_rhs) {
+                (IntegerType::UInt8, IntegerType::UInt8) => dyn_dict!(u8, left, right),
+                (IntegerType::UInt16, IntegerType::UInt16) => dyn_dict!(u16, left, right),
+                (IntegerType::UInt32, IntegerType::UInt32) => dyn_dict!(u32, left, right),
+                (IntegerType::UInt64, IntegerType::UInt64) => dyn_dict!(u64, left, right),
+                (IntegerType::Int8, IntegerType::Int8) => dyn_dict!(i8, left, right),
+                (IntegerType::Int16, IntegerType::Int16) => dyn_dict!(i16, left, right),
+                (IntegerType::Int32, IntegerType::Int32) => dyn_dict!(i32, left, right),
+                (IntegerType::Int64, IntegerType::Int64) => dyn_dict!(i64, left, right),
                 (lhs, _) => {
                     return Err(ArrowError::InvalidArgumentError(format!(
                         "Dictionaries do not support keys of type {:?}",

src/compute/aggregate/memory.rs

@@ -16,16 +16,6 @@ macro_rules! dyn_binary {
     }};
 }
 
-macro_rules! dyn_dict {
-    ($array:expr, $ty:ty) => {{
-        let array = $array
-            .as_any()
-            .downcast_ref::<DictionaryArray<$ty>>()
-            .unwrap();
-        estimated_bytes_size(array.keys()) + estimated_bytes_size(array.values().as_ref())
-    }};
-}
-
 /// Returns the total (heap) allocated size of the array in bytes.
 /// # Implementation
 /// This estimation is the sum of the size of its buffers, validity, including nested arrays.
@@ -106,8 +96,12 @@ pub fn estimated_bytes_size(array: &dyn Array) -> usize {
                 .sum::<usize>();
             types + offsets + fields
         }
-        Dictionary(key_type) => with_match_physical_dictionary_key_type!(key_type, |$T| {
-            dyn_dict!(array, $T)
+        Dictionary(key_type) => match_integer_type!(key_type, |$T| {
+            let array = array
+                .as_any()
+                .downcast_ref::<DictionaryArray<$T>>()
+                .unwrap();
+            estimated_bytes_size(array.keys()) + estimated_bytes_size(array.values().as_ref())
         }),
         Map => {
             let array = array.as_any().downcast_ref::<MapArray>().unwrap();

src/compute/cast/dictionary_to.rs

@@ -114,8 +114,9 @@ pub(super) fn dictionary_cast_dyn<K: DictionaryKey>(
             let values = cast(values.as_ref(), to_values_type, options)?.into();
 
             // create the appropriate array type
-            with_match_dictionary_key_type!(to_keys_type.as_ref(), |$T| {
-                key_cast!(keys, values, array, to_keys_type, $T)
+            let data_type = (*to_keys_type).into();
+            match_integer_type!(to_keys_type, |$T| {
+                key_cast!(keys, values, array, &data_type, $T)
             })
         }
         _ => unpack_dictionary::<K>(keys, values.as_ref(), to_type, options),

src/compute/cast/mod.rs

@@ -375,31 +375,12 @@ pub fn cast(array: &dyn Array, to_type: &DataType, options: CastOptions) -> Resu
             Ok(Box::new(list_array))
         }
 
-        (Dictionary(index_type, _), _) => match **index_type {
-            DataType::Int8 => dictionary_cast_dyn::<i8>(array, to_type, options),
-            DataType::Int16 => dictionary_cast_dyn::<i16>(array, to_type, options),
-            DataType::Int32 => dictionary_cast_dyn::<i32>(array, to_type, options),
-            DataType::Int64 => dictionary_cast_dyn::<i64>(array, to_type, options),
-            DataType::UInt8 => dictionary_cast_dyn::<u8>(array, to_type, options),
-            DataType::UInt16 => dictionary_cast_dyn::<u16>(array, to_type, options),
-            DataType::UInt32 => dictionary_cast_dyn::<u32>(array, to_type, options),
-            DataType::UInt64 => dictionary_cast_dyn::<u64>(array, to_type, options),
-            _ => unreachable!(),
-        },
-        (_, Dictionary(index_type, value_type)) => match **index_type {
-            DataType::Int8 => cast_to_dictionary::<i8>(array, value_type, options),
-            DataType::Int16 => cast_to_dictionary::<i16>(array, value_type, options),
-            DataType::Int32 => cast_to_dictionary::<i32>(array, value_type, options),
-            DataType::Int64 => cast_to_dictionary::<i64>(array, value_type, options),
-            DataType::UInt8 => cast_to_dictionary::<u8>(array, value_type, options),
-            DataType::UInt16 => cast_to_dictionary::<u16>(array, value_type, options),
-            DataType::UInt32 => cast_to_dictionary::<u32>(array, value_type, options),
-            DataType::UInt64 => cast_to_dictionary::<u64>(array, value_type, options),
-            _ => Err(ArrowError::NotYetImplemented(format!(
-                "Casting from type {:?} to dictionary type {:?} not supported",
-                from_type, to_type,
-            ))),
-        },
+        (Dictionary(index_type, _), _) => match_integer_type!(index_type, |$T| {
+            dictionary_cast_dyn::<$T>(array, to_type, options)
+        }),
+        (_, Dictionary(index_type, value_type)) => match_integer_type!(index_type, |$T| {
+            cast_to_dictionary::<$T>(array, value_type, options)
+        }),
         (_, Boolean) => match from_type {
             UInt8 => primitive_to_boolean_dyn::<u8>(array, to_type.clone()),
             UInt16 => primitive_to_boolean_dyn::<u16>(array, to_type.clone()),