You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by vi...@apache.org on 2022/08/31 19:19:27 UTC
[arrow-rs] branch master updated: Support comparison between DictionaryArray to BooleanArray (#2618)
This is an automated email from the ASF dual-hosted git repository.
viirya pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/master by this push:
new f9a309445 Support comparison between DictionaryArray to BooleanArray (#2618)
f9a309445 is described below
commit f9a309445550e596e32a0dbf29879946f44c1dc6
Author: Liang-Chi Hsieh <vi...@gmail.com>
AuthorDate: Wed Aug 31 12:19:21 2022 -0700
Support comparison between DictionaryArray to BooleanArray (#2618)
---
arrow/src/compute/kernels/comparison.rs | 164 ++++++++++++++++++++++++++++++++
1 file changed, 164 insertions(+)
diff --git a/arrow/src/compute/kernels/comparison.rs b/arrow/src/compute/kernels/comparison.rs
index 7b3c293e0..dd9d4fc5d 100644
--- a/arrow/src/compute/kernels/comparison.rs
+++ b/arrow/src/compute/kernels/comparison.rs
@@ -2175,12 +2175,59 @@ macro_rules! typed_dict_string_array_cmp {
}};
}
+#[cfg(feature = "dyn_cmp_dict")]
+macro_rules! typed_dict_boolean_array_cmp {
+ ($LEFT: expr, $RIGHT: expr, $LEFT_KEY_TYPE: expr, $OP: expr) => {{
+ match $LEFT_KEY_TYPE {
+ DataType::Int8 => {
+ let left = as_dictionary_array::<Int8Type>($LEFT);
+ cmp_dict_boolean_array::<_, _>(left, $RIGHT, $OP)
+ }
+ DataType::Int16 => {
+ let left = as_dictionary_array::<Int16Type>($LEFT);
+ cmp_dict_boolean_array::<_, _>(left, $RIGHT, $OP)
+ }
+ DataType::Int32 => {
+ let left = as_dictionary_array::<Int32Type>($LEFT);
+ cmp_dict_boolean_array::<_, _>(left, $RIGHT, $OP)
+ }
+ DataType::Int64 => {
+ let left = as_dictionary_array::<Int64Type>($LEFT);
+ cmp_dict_boolean_array::<_, _>(left, $RIGHT, $OP)
+ }
+ DataType::UInt8 => {
+ let left = as_dictionary_array::<UInt8Type>($LEFT);
+ cmp_dict_boolean_array::<_, _>(left, $RIGHT, $OP)
+ }
+ DataType::UInt16 => {
+ let left = as_dictionary_array::<UInt16Type>($LEFT);
+ cmp_dict_boolean_array::<_, _>(left, $RIGHT, $OP)
+ }
+ DataType::UInt32 => {
+ let left = as_dictionary_array::<UInt32Type>($LEFT);
+ cmp_dict_boolean_array::<_, _>(left, $RIGHT, $OP)
+ }
+ DataType::UInt64 => {
+ let left = as_dictionary_array::<UInt64Type>($LEFT);
+ cmp_dict_boolean_array::<_, _>(left, $RIGHT, $OP)
+ }
+ t => Err(ArrowError::NotYetImplemented(format!(
+ "Cannot compare dictionary array of key type {}",
+ t
+ ))),
+ }
+ }};
+}
+
#[cfg(feature = "dyn_cmp_dict")]
macro_rules! typed_cmp_dict_non_dict {
($LEFT: expr, $RIGHT: expr, $OP_BOOL: expr, $OP: expr, $OP_FLOAT: expr) => {{
match ($LEFT.data_type(), $RIGHT.data_type()) {
(DataType::Dictionary(left_key_type, left_value_type), right_type) => {
match (left_value_type.as_ref(), right_type) {
+ (DataType::Boolean, DataType::Boolean) => {
+ typed_dict_boolean_array_cmp!($LEFT, $RIGHT, left_key_type.as_ref(), $OP_BOOL)
+ }
(DataType::Int8, DataType::Int8) => {
typed_dict_non_dict_cmp!($LEFT, $RIGHT, left_key_type.as_ref(), Int8Type, $OP_BOOL, $OP)
}
@@ -2609,6 +2656,25 @@ where
)
}
+/// Perform given operation on `DictionaryArray` and `BooleanArray`. The value
+/// type of `DictionaryArray` is same as `BooleanArray`'s type.
+#[cfg(feature = "dyn_cmp_dict")]
+fn cmp_dict_boolean_array<K, F>(
+ left: &DictionaryArray<K>,
+ right: &dyn Array,
+ op: F,
+) -> Result<BooleanArray>
+where
+ K: ArrowNumericType,
+ F: Fn(bool, bool) -> bool,
+{
+ compare_op(
+ left.downcast_dict::<BooleanArray>().unwrap(),
+ right.as_any().downcast_ref::<BooleanArray>().unwrap(),
+ op,
+ )
+}
+
/// Perform given operation on two `DictionaryArray`s which value type is
/// primitive type. Returns an error if the two arrays have different value
/// type
@@ -6334,4 +6400,102 @@ mod tests {
);
assert_eq!(gt_eq_dyn(&array1, &array2).unwrap(), expected);
}
+
+ #[test]
+ #[cfg(feature = "dyn_cmp_dict")]
+ fn test_eq_dyn_neq_dyn_dictionary_to_boolean_array() {
+ let test1 = vec![Some(true), None, Some(false)];
+ let test2 = vec![Some(true), None, None, Some(true)];
+
+ let values = BooleanArray::from(test1);
+ let keys = Int8Array::from_iter_values([0_i8, 0, 1, 2]);
+ let dict_array = DictionaryArray::try_new(&keys, &values).unwrap();
+
+ let array: BooleanArray = test2.iter().collect();
+
+ let result = eq_dyn(&dict_array, &array);
+ assert_eq!(
+ result.unwrap(),
+ BooleanArray::from(vec![Some(true), None, None, Some(false)])
+ );
+
+ let result = eq_dyn(&array, &dict_array);
+ assert_eq!(
+ result.unwrap(),
+ BooleanArray::from(vec![Some(true), None, None, Some(false)])
+ );
+
+ let result = neq_dyn(&dict_array, &array);
+ assert_eq!(
+ result.unwrap(),
+ BooleanArray::from(vec![Some(false), None, None, Some(true)])
+ );
+
+ let result = neq_dyn(&array, &dict_array);
+ assert_eq!(
+ result.unwrap(),
+ BooleanArray::from(vec![Some(false), None, None, Some(true)])
+ );
+ }
+
+ #[test]
+ #[cfg(feature = "dyn_cmp_dict")]
+ fn test_lt_dyn_lt_eq_dyn_gt_dyn_gt_eq_dyn_dictionary_to_boolean_array() {
+ let test1 = vec![Some(true), None, Some(false)];
+ let test2 = vec![Some(true), None, None, Some(true)];
+
+ let values = BooleanArray::from(test1);
+ let keys = Int8Array::from_iter_values([0_i8, 0, 1, 2]);
+ let dict_array = DictionaryArray::try_new(&keys, &values).unwrap();
+
+ let array: BooleanArray = test2.iter().collect();
+
+ let result = lt_dyn(&dict_array, &array);
+ assert_eq!(
+ result.unwrap(),
+ BooleanArray::from(vec![Some(false), None, None, Some(true)])
+ );
+
+ let result = lt_dyn(&array, &dict_array);
+ assert_eq!(
+ result.unwrap(),
+ BooleanArray::from(vec![Some(false), None, None, Some(false)])
+ );
+
+ let result = lt_eq_dyn(&dict_array, &array);
+ assert_eq!(
+ result.unwrap(),
+ BooleanArray::from(vec![Some(true), None, None, Some(true)])
+ );
+
+ let result = lt_eq_dyn(&array, &dict_array);
+ assert_eq!(
+ result.unwrap(),
+ BooleanArray::from(vec![Some(true), None, None, Some(false)])
+ );
+
+ let result = gt_dyn(&dict_array, &array);
+ assert_eq!(
+ result.unwrap(),
+ BooleanArray::from(vec![Some(false), None, None, Some(false)])
+ );
+
+ let result = gt_dyn(&array, &dict_array);
+ assert_eq!(
+ result.unwrap(),
+ BooleanArray::from(vec![Some(false), None, None, Some(true)])
+ );
+
+ let result = gt_eq_dyn(&dict_array, &array);
+ assert_eq!(
+ result.unwrap(),
+ BooleanArray::from(vec![Some(true), None, None, Some(false)])
+ );
+
+ let result = gt_eq_dyn(&array, &dict_array);
+ assert_eq!(
+ result.unwrap(),
+ BooleanArray::from(vec![Some(true), None, None, Some(true)])
+ );
+ }
}