Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Commit

Permalink
DRY (#332)
Browse files Browse the repository at this point in the history
  • Loading branch information
@jorgecarleitao
jorgecarleitao authored Aug 24, 2021
1 parent d3464d1 commit a9932a6
Show file tree
Hide file tree
Showing 15 changed files with 138 additions and 446 deletions.
46 changes: 5 additions & 41 deletions src/array/equal/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -278,49 +278,13 @@ pub fn equal(lhs: &dyn Array, rhs: &dyn Array) -> bool {
let rhs = rhs.as_any().downcast_ref::<StructArray>().unwrap();
struct_::equal(lhs, rhs)
}
DataType::Dictionary(key_type, _) => match key_type.as_ref() {
DataType::Int8 => {
DataType::Dictionary(key_type, _) => {
with_match_dictionary_key_type!(key_type.as_ref(), |$T| {
let lhs = lhs.as_any().downcast_ref().unwrap();
let rhs = rhs.as_any().downcast_ref().unwrap();
dictionary::equal::<i8>(lhs, rhs)
}
DataType::Int16 => {
let lhs = lhs.as_any().downcast_ref().unwrap();
let rhs = rhs.as_any().downcast_ref().unwrap();
dictionary::equal::<i16>(lhs, rhs)
}
DataType::Int32 => {
let lhs = lhs.as_any().downcast_ref().unwrap();
let rhs = rhs.as_any().downcast_ref().unwrap();
dictionary::equal::<i32>(lhs, rhs)
}
DataType::Int64 => {
let lhs = lhs.as_any().downcast_ref().unwrap();
let rhs = rhs.as_any().downcast_ref().unwrap();
dictionary::equal::<i64>(lhs, rhs)
}
DataType::UInt8 => {
let lhs = lhs.as_any().downcast_ref().unwrap();
let rhs = rhs.as_any().downcast_ref().unwrap();
dictionary::equal::<u8>(lhs, rhs)
}
DataType::UInt16 => {
let lhs = lhs.as_any().downcast_ref().unwrap();
let rhs = rhs.as_any().downcast_ref().unwrap();
dictionary::equal::<u16>(lhs, rhs)
}
DataType::UInt32 => {
let lhs = lhs.as_any().downcast_ref().unwrap();
let rhs = rhs.as_any().downcast_ref().unwrap();
dictionary::equal::<u32>(lhs, rhs)
}
DataType::UInt64 => {
let lhs = lhs.as_any().downcast_ref().unwrap();
let rhs = rhs.as_any().downcast_ref().unwrap();
dictionary::equal::<u64>(lhs, rhs)
}
_ => unreachable!(),
},
dictionary::equal::<$T>(lhs, rhs)
})
}
DataType::FixedSizeBinary(_) => {
let lhs = lhs.as_any().downcast_ref().unwrap();
let rhs = rhs.as_any().downcast_ref().unwrap();
Expand Down
32 changes: 10 additions & 22 deletions src/array/ffi.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,17 +33,6 @@ macro_rules! ffi_dyn {
}};
}

macro_rules! ffi_dict_dyn {
($array:expr, $ty:ty) => {{
let array = $array.as_any().downcast_ref::<$ty>().unwrap();
(
array.buffers(),
array.children(),
Some(array.values().clone()),
)
}};
}

type BuffersChildren = (
Vec<Option<std::ptr::NonNull<u8>>>,
Vec<Arc<dyn Array>>,
Expand Down Expand Up @@ -86,16 +75,15 @@ pub fn buffers_children_dictionary(array: &dyn Array) -> BuffersChildren {
DataType::FixedSizeList(_, _) => ffi_dyn!(array, FixedSizeListArray),
DataType::Struct(_) => ffi_dyn!(array, StructArray),
DataType::Union(_, _, _) => ffi_dyn!(array, UnionArray),
DataType::Dictionary(key_type, _) => match key_type.as_ref() {
DataType::Int8 => ffi_dict_dyn!(array, DictionaryArray::<i8>),
DataType::Int16 => ffi_dict_dyn!(array, DictionaryArray::<i16>),
DataType::Int32 => ffi_dict_dyn!(array, DictionaryArray::<i32>),
DataType::Int64 => ffi_dict_dyn!(array, DictionaryArray::<i64>),
DataType::UInt8 => ffi_dict_dyn!(array, DictionaryArray::<u8>),
DataType::UInt16 => ffi_dict_dyn!(array, DictionaryArray::<u16>),
DataType::UInt32 => ffi_dict_dyn!(array, DictionaryArray::<u32>),
DataType::UInt64 => ffi_dict_dyn!(array, DictionaryArray::<u64>),
_ => unreachable!(),
},
DataType::Dictionary(key_type, _) => {
with_match_dictionary_key_type!(key_type.as_ref(), |$T| {
let array = array.as_any().downcast_ref::<DictionaryArray<$T>>().unwrap();
(
array.buffers(),
array.children(),
Some(array.values().clone()),
)
})
}
}
}
16 changes: 5 additions & 11 deletions src/array/growable/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -226,16 +226,10 @@ pub fn make_growable<'a>(
}
DataType::FixedSizeList(_, _) => todo!(),
DataType::Union(_, _, _) => todo!(),
DataType::Dictionary(key, _) => match key.as_ref() {
DataType::UInt8 => dyn_dict_growable!(u8, arrays, use_validity, capacity),
DataType::UInt16 => dyn_dict_growable!(u16, arrays, use_validity, capacity),
DataType::UInt32 => dyn_dict_growable!(u32, arrays, use_validity, capacity),
DataType::UInt64 => dyn_dict_growable!(u64, arrays, use_validity, capacity),
DataType::Int8 => dyn_dict_growable!(i8, arrays, use_validity, capacity),
DataType::Int16 => dyn_dict_growable!(i16, arrays, use_validity, capacity),
DataType::Int32 => dyn_dict_growable!(i32, arrays, use_validity, capacity),
DataType::Int64 => dyn_dict_growable!(i64, arrays, use_validity, capacity),
_ => unreachable!(),
},
DataType::Dictionary(key_type, _) => {
with_match_dictionary_key_type!(key_type.as_ref(), |$T| {
dyn_dict_growable!($T, arrays, use_validity, capacity)
})
}
}
}
81 changes: 37 additions & 44 deletions src/array/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -146,6 +146,23 @@ macro_rules! fmt_dyn {
}};
}

macro_rules! with_match_dictionary_key_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"),
}
})}

impl Display for dyn Array {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self.data_type() {
Expand Down Expand Up @@ -185,17 +202,11 @@ impl Display for dyn Array {
DataType::FixedSizeList(_, _) => fmt_dyn!(self, FixedSizeListArray, f),
DataType::Struct(_) => fmt_dyn!(self, StructArray, f),
DataType::Union(_, _, _) => fmt_dyn!(self, UnionArray, f),
DataType::Dictionary(key_type, _) => match key_type.as_ref() {
DataType::Int8 => fmt_dyn!(self, DictionaryArray::<i8>, f),
DataType::Int16 => fmt_dyn!(self, DictionaryArray::<i16>, f),
DataType::Int32 => fmt_dyn!(self, DictionaryArray::<i32>, f),
DataType::Int64 => fmt_dyn!(self, DictionaryArray::<i64>, f),
DataType::UInt8 => fmt_dyn!(self, DictionaryArray::<u8>, f),
DataType::UInt16 => fmt_dyn!(self, DictionaryArray::<u16>, f),
DataType::UInt32 => fmt_dyn!(self, DictionaryArray::<u32>, f),
DataType::UInt64 => fmt_dyn!(self, DictionaryArray::<u64>, f),
_ => unreachable!(),
},
DataType::Dictionary(key_type, _) => {
with_match_dictionary_key_type!(key_type.as_ref(), |$T| {
fmt_dyn!(self, DictionaryArray::<$T>, f)
})
}
}
}
}
Expand Down Expand Up @@ -239,17 +250,11 @@ pub fn new_empty_array(data_type: DataType) -> Box<dyn Array> {
DataType::FixedSizeList(_, _) => Box::new(FixedSizeListArray::new_empty(data_type)),
DataType::Struct(fields) => Box::new(StructArray::new_empty(&fields)),
DataType::Union(_, _, _) => Box::new(UnionArray::new_empty(data_type)),
DataType::Dictionary(key_type, value_type) => match key_type.as_ref() {
DataType::Int8 => Box::new(DictionaryArray::<i8>::new_empty(*value_type)),
DataType::Int16 => Box::new(DictionaryArray::<i16>::new_empty(*value_type)),
DataType::Int32 => Box::new(DictionaryArray::<i32>::new_empty(*value_type)),
DataType::Int64 => Box::new(DictionaryArray::<i64>::new_empty(*value_type)),
DataType::UInt8 => Box::new(DictionaryArray::<u8>::new_empty(*value_type)),
DataType::UInt16 => Box::new(DictionaryArray::<u16>::new_empty(*value_type)),
DataType::UInt32 => Box::new(DictionaryArray::<u32>::new_empty(*value_type)),
DataType::UInt64 => Box::new(DictionaryArray::<u64>::new_empty(*value_type)),
_ => unreachable!(),
},
DataType::Dictionary(key_type, value_type) => {
with_match_dictionary_key_type!(key_type.as_ref(), |$T| {
Box::new(DictionaryArray::<$T>::new_empty(*value_type))
})
}
}
}

Expand Down Expand Up @@ -294,17 +299,11 @@ pub fn new_null_array(data_type: DataType, length: usize) -> Box<dyn Array> {
DataType::FixedSizeList(_, _) => Box::new(FixedSizeListArray::new_null(data_type, length)),
DataType::Struct(fields) => Box::new(StructArray::new_null(&fields, length)),
DataType::Union(_, _, _) => Box::new(UnionArray::new_null(data_type, length)),
DataType::Dictionary(key_type, value_type) => match key_type.as_ref() {
DataType::Int8 => Box::new(DictionaryArray::<i8>::new_null(*value_type, length)),
DataType::Int16 => Box::new(DictionaryArray::<i16>::new_null(*value_type, length)),
DataType::Int32 => Box::new(DictionaryArray::<i32>::new_null(*value_type, length)),
DataType::Int64 => Box::new(DictionaryArray::<i64>::new_null(*value_type, length)),
DataType::UInt8 => Box::new(DictionaryArray::<u8>::new_null(*value_type, length)),
DataType::UInt16 => Box::new(DictionaryArray::<u16>::new_null(*value_type, length)),
DataType::UInt32 => Box::new(DictionaryArray::<u32>::new_null(*value_type, length)),
DataType::UInt64 => Box::new(DictionaryArray::<u64>::new_null(*value_type, length)),
_ => unreachable!(),
},
DataType::Dictionary(key_type, value_type) => {
with_match_dictionary_key_type!(key_type.as_ref(), |$T| {
Box::new(DictionaryArray::<$T>::new_null(*value_type, length))
})
}
}
}

Expand Down Expand Up @@ -355,17 +354,11 @@ pub fn clone(array: &dyn Array) -> Box<dyn Array> {
DataType::FixedSizeList(_, _) => clone_dyn!(array, FixedSizeListArray),
DataType::Struct(_) => clone_dyn!(array, StructArray),
DataType::Union(_, _, _) => clone_dyn!(array, UnionArray),
DataType::Dictionary(key_type, _) => match key_type.as_ref() {
DataType::Int8 => clone_dyn!(array, DictionaryArray::<i8>),
DataType::Int16 => clone_dyn!(array, DictionaryArray::<i16>),
DataType::Int32 => clone_dyn!(array, DictionaryArray::<i32>),
DataType::Int64 => clone_dyn!(array, DictionaryArray::<i64>),
DataType::UInt8 => clone_dyn!(array, DictionaryArray::<u8>),
DataType::UInt16 => clone_dyn!(array, DictionaryArray::<u16>),
DataType::UInt32 => clone_dyn!(array, DictionaryArray::<u32>),
DataType::UInt64 => clone_dyn!(array, DictionaryArray::<u64>),
_ => unreachable!(),
},
DataType::Dictionary(key_type, _) => {
with_match_dictionary_key_type!(key_type.as_ref(), |$T| {
clone_dyn!(array, DictionaryArray::<$T>)
})
}
}
}

Expand Down
14 changes: 3 additions & 11 deletions src/compute/cast/dictionary_to.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -107,17 +107,9 @@ pub(super) fn dictionary_cast_dyn<K: DictionaryKey>(
let values = cast_with_options(values.as_ref(), to_values_type, options)?.into();

// create the appropriate array type
match to_keys_type.as_ref() {
DataType::Int8 => key_cast!(keys, values, array, to_keys_type, i8),
DataType::Int16 => key_cast!(keys, values, array, to_keys_type, i16),
DataType::Int32 => key_cast!(keys, values, array, to_keys_type, i32),
DataType::Int64 => key_cast!(keys, values, array, to_keys_type, i64),
DataType::UInt8 => key_cast!(keys, values, array, to_keys_type, u8),
DataType::UInt16 => key_cast!(keys, values, array, to_keys_type, u16),
DataType::UInt32 => key_cast!(keys, values, array, to_keys_type, u32),
DataType::UInt64 => key_cast!(keys, values, array, to_keys_type, u64),
_ => unreachable!(),
}
with_match_dictionary_key_type!(to_keys_type.as_ref(), |$T| {
key_cast!(keys, values, array, to_keys_type, $T)
})
}
_ => unpack_dictionary::<K>(keys, values.as_ref(), to_type, options),
}
Expand Down
63 changes: 17 additions & 46 deletions src/compute/sort/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -193,8 +193,18 @@ pub fn sort_to_indices<I: Index>(
}
}
DataType::Dictionary(key_type, value_type) => match value_type.as_ref() {
DataType::Utf8 => sort_dict::<I, i32>(values, key_type.as_ref(), options, limit),
DataType::LargeUtf8 => sort_dict::<I, i64>(values, key_type.as_ref(), options, limit),
DataType::Utf8 => Ok(sort_dict::<I, i32>(
values,
key_type.as_ref(),
options,
limit,
)),
DataType::LargeUtf8 => Ok(sort_dict::<I, i64>(
values,
key_type.as_ref(),
options,
limit,
)),
t => Err(ArrowError::NotYetImplemented(format!(
"Sort not supported for dictionary type with keys {:?}",
t
Expand All @@ -212,53 +222,14 @@ fn sort_dict<I: Index, O: Offset>(
key_type: &DataType,
options: &SortOptions,
limit: Option<usize>,
) -> Result<PrimitiveArray<I>> {
match key_type {
DataType::Int8 => Ok(utf8::indices_sorted_unstable_by_dictionary::<I, i8, O>(
values.as_any().downcast_ref().unwrap(),
options,
limit,
)),
DataType::Int16 => Ok(utf8::indices_sorted_unstable_by_dictionary::<I, i16, O>(
values.as_any().downcast_ref().unwrap(),
options,
limit,
)),
DataType::Int32 => Ok(utf8::indices_sorted_unstable_by_dictionary::<I, i32, O>(
values.as_any().downcast_ref().unwrap(),
options,
limit,
)),
DataType::Int64 => Ok(utf8::indices_sorted_unstable_by_dictionary::<I, i64, O>(
values.as_any().downcast_ref().unwrap(),
options,
limit,
)),
DataType::UInt8 => Ok(utf8::indices_sorted_unstable_by_dictionary::<I, u8, O>(
values.as_any().downcast_ref().unwrap(),
options,
limit,
)),
DataType::UInt16 => Ok(utf8::indices_sorted_unstable_by_dictionary::<I, u16, O>(
values.as_any().downcast_ref().unwrap(),
options,
limit,
)),
DataType::UInt32 => Ok(utf8::indices_sorted_unstable_by_dictionary::<I, u32, O>(
values.as_any().downcast_ref().unwrap(),
options,
limit,
)),
DataType::UInt64 => Ok(utf8::indices_sorted_unstable_by_dictionary::<I, u64, O>(
) -> PrimitiveArray<I> {
with_match_dictionary_key_type!(key_type, |$T| {
utf8::indices_sorted_unstable_by_dictionary::<I, $T, O>(
values.as_any().downcast_ref().unwrap(),
options,
limit,
)),
t => Err(ArrowError::NotYetImplemented(format!(
"Sort not supported for dictionary key type {:?}",
t
))),
}
)
})
}

/// Checks if an array of type `datatype` can be sorted
Expand Down
16 changes: 5 additions & 11 deletions src/compute/take/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -100,17 +100,11 @@ pub fn take<O: Index>(values: &dyn Array, indices: &PrimitiveArray<O>) -> Result
let values = values.as_any().downcast_ref().unwrap();
Ok(Box::new(binary::take::<i64, _>(values, indices)))
}
DataType::Dictionary(key_type, _) => match key_type.as_ref() {
DataType::Int8 => downcast_dict_take!(i8, values, indices),
DataType::Int16 => downcast_dict_take!(i16, values, indices),
DataType::Int32 => downcast_dict_take!(i32, values, indices),
DataType::Int64 => downcast_dict_take!(i64, values, indices),
DataType::UInt8 => downcast_dict_take!(u8, values, indices),
DataType::UInt16 => downcast_dict_take!(u16, values, indices),
DataType::UInt32 => downcast_dict_take!(u32, values, indices),
DataType::UInt64 => downcast_dict_take!(u64, values, indices),
_ => unreachable!(),
},
DataType::Dictionary(key_type, _) => {
with_match_dictionary_key_type!(key_type.as_ref(), |$T| {
downcast_dict_take!($T, values, indices)
})
}
DataType::Struct(_) => {
let array = values.as_any().downcast_ref().unwrap();
Ok(Box::new(structure::take::<_>(array, indices)?))
Expand Down
16 changes: 5 additions & 11 deletions src/ffi/array.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -65,17 +65,11 @@ pub fn try_from<A: ArrowArrayRef>(array: A) -> Result<Box<dyn Array>> {
DataType::List(_) => Box::new(ListArray::<i32>::try_from_ffi(array)?),
DataType::LargeList(_) => Box::new(ListArray::<i64>::try_from_ffi(array)?),
DataType::Struct(_) => Box::new(StructArray::try_from_ffi(array)?),
DataType::Dictionary(keys, _) => match keys.as_ref() {
DataType::Int8 => Box::new(DictionaryArray::<i8>::try_from_ffi(array)?),
DataType::Int16 => Box::new(DictionaryArray::<i16>::try_from_ffi(array)?),
DataType::Int32 => Box::new(DictionaryArray::<i32>::try_from_ffi(array)?),
DataType::Int64 => Box::new(DictionaryArray::<i64>::try_from_ffi(array)?),
DataType::UInt8 => Box::new(DictionaryArray::<u8>::try_from_ffi(array)?),
DataType::UInt16 => Box::new(DictionaryArray::<u16>::try_from_ffi(array)?),
DataType::UInt32 => Box::new(DictionaryArray::<u32>::try_from_ffi(array)?),
DataType::UInt64 => Box::new(DictionaryArray::<u64>::try_from_ffi(array)?),
_ => unreachable!(),
},
DataType::Dictionary(key_type, _) => {
with_match_dictionary_key_type!(key_type.as_ref(), |$T| {
Box::new(DictionaryArray::<$T>::try_from_ffi(array)?)
})
}
DataType::Union(_, _, _) => Box::new(UnionArray::try_from_ffi(array)?),
data_type => {
return Err(ArrowError::NotYetImplemented(format!(
Expand Down
Loading

0 comments on commit a9932a6

Please sign in to comment.