You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by al...@apache.org on 2022/11/20 11:53:13 UTC

[arrow-datafusion] branch master updated: replace the comparison op for decimal array op using the arrow-rs kernel (#4290)

This is an automated email from the ASF dual-hosted git repository.

alamb pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/master by this push:
     new 60d9bbe87 replace the comparison op for decimal array op using the arrow-rs kernel (#4290)
60d9bbe87 is described below

commit 60d9bbe87a4e38e7b01229cf0690e7689ff321f4
Author: Kun Liu <li...@apache.org>
AuthorDate: Sun Nov 20 19:53:07 2022 +0800

    replace the comparison op for decimal array op using the arrow-rs kernel (#4290)
    
    * use the arrow-rs kernel: decimal comparison op
    
    * move the kernel test to the binary
    
    * refine the macro code
---
 datafusion/physical-expr/src/expressions/binary.rs |  87 +++++++++++++++---
 .../src/expressions/binary/adapter.rs              |  18 +---
 .../src/expressions/binary/kernels_arrow.rs        | 102 +--------------------
 3 files changed, 78 insertions(+), 129 deletions(-)

diff --git a/datafusion/physical-expr/src/expressions/binary.rs b/datafusion/physical-expr/src/expressions/binary.rs
index e067e34b5..2d5f0854a 100644
--- a/datafusion/physical-expr/src/expressions/binary.rs
+++ b/datafusion/physical-expr/src/expressions/binary.rs
@@ -376,8 +376,6 @@ macro_rules! binary_string_array_op {
 macro_rules! binary_primitive_array_op {
     ($LEFT:expr, $RIGHT:expr, $OP:ident) => {{
         match $LEFT.data_type() {
-            // TODO support decimal type
-            // which is not the primitive type
             DataType::Decimal128(_,_) => compute_decimal_op!($LEFT, $RIGHT, $OP, Decimal128Array),
             DataType::Int8 => compute_op!($LEFT, $RIGHT, $OP, Int8Array),
             DataType::Int16 => compute_op!($LEFT, $RIGHT, $OP, Int16Array),
@@ -2374,25 +2372,90 @@ mod tests {
         )
         .unwrap();
 
-        // compare decimal array with other array type
+        let value: i128 = 123;
+        let decimal_array = Arc::new(create_decimal_array(
+            &[Some(value), None, Some(value - 1), Some(value + 1)],
+            10,
+            0,
+        )) as ArrayRef;
+
+        // comparison array op for decimal array
         let schema = Arc::new(Schema::new(vec![
-            Field::new("a", DataType::Int64, true),
+            Field::new("a", DataType::Decimal128(10, 0), true),
             Field::new("b", DataType::Decimal128(10, 0), true),
         ]));
-
-        let value: i64 = 123;
-
-        let decimal_array = Arc::new(create_decimal_array(
+        let right_decimal_array = Arc::new(create_decimal_array(
             &[
-                Some(value as i128),
-                None,
-                Some((value - 1) as i128),
-                Some((value + 1) as i128),
+                Some(value - 1),
+                Some(value),
+                Some(value + 1),
+                Some(value + 1),
             ],
             10,
             0,
         )) as ArrayRef;
 
+        apply_logic_op(
+            &schema,
+            &decimal_array,
+            &right_decimal_array,
+            Operator::Eq,
+            BooleanArray::from(vec![Some(false), None, Some(false), Some(true)]),
+        )
+        .unwrap();
+
+        apply_logic_op(
+            &schema,
+            &decimal_array,
+            &right_decimal_array,
+            Operator::NotEq,
+            BooleanArray::from(vec![Some(true), None, Some(true), Some(false)]),
+        )
+        .unwrap();
+
+        apply_logic_op(
+            &schema,
+            &decimal_array,
+            &right_decimal_array,
+            Operator::Lt,
+            BooleanArray::from(vec![Some(false), None, Some(true), Some(false)]),
+        )
+        .unwrap();
+
+        apply_logic_op(
+            &schema,
+            &decimal_array,
+            &right_decimal_array,
+            Operator::LtEq,
+            BooleanArray::from(vec![Some(false), None, Some(true), Some(true)]),
+        )
+        .unwrap();
+
+        apply_logic_op(
+            &schema,
+            &decimal_array,
+            &right_decimal_array,
+            Operator::Gt,
+            BooleanArray::from(vec![Some(true), None, Some(false), Some(false)]),
+        )
+        .unwrap();
+
+        apply_logic_op(
+            &schema,
+            &decimal_array,
+            &right_decimal_array,
+            Operator::GtEq,
+            BooleanArray::from(vec![Some(true), None, Some(false), Some(true)]),
+        )
+        .unwrap();
+
+        // compare decimal array with other array type
+        let value: i64 = 123;
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("a", DataType::Int64, true),
+            Field::new("b", DataType::Decimal128(10, 0), true),
+        ]));
+
         let int64_array = Arc::new(Int64Array::from(vec![
             Some(value),
             Some(value - 1),
diff --git a/datafusion/physical-expr/src/expressions/binary/adapter.rs b/datafusion/physical-expr/src/expressions/binary/adapter.rs
index a4726ba2c..ec0eda392 100644
--- a/datafusion/physical-expr/src/expressions/binary/adapter.rs
+++ b/datafusion/physical-expr/src/expressions/binary/adapter.rs
@@ -20,10 +20,7 @@
 
 use std::sync::Arc;
 
-use super::kernels_arrow::*;
 use arrow::array::*;
-use arrow::datatypes::DataType;
-use datafusion_common::cast::as_decimal128_array;
 use datafusion_common::Result;
 
 /// create a `dyn_op` wrapper function for the specified operation
@@ -35,20 +32,9 @@ macro_rules! make_dyn_comp_op {
             /// wrapper over arrow compute kernel that maps Error types and
             /// patches missing support in arrow
             pub(crate) fn [<$OP _dyn>] (left: &dyn Array, right: &dyn Array) -> Result<ArrayRef> {
-                match (left.data_type(), right.data_type()) {
-                    // Call `op_decimal` (e.g. `eq_decimal) until
-                    // arrow has native support
-                    // https://github.com/apache/arrow-rs/issues/1200
-                    (DataType::Decimal128(_, _), DataType::Decimal128(_, _)) => {
-                        [<$OP _decimal>](as_decimal128_array(left).unwrap(), as_decimal128_array(right).unwrap())
-                    },
-                    // By default call the arrow kernel
-                    _ => {
-                    arrow::compute::kernels::comparison::[<$OP _dyn>](left, right)
+                arrow::compute::kernels::comparison::[<$OP _dyn>](left, right)
                             .map_err(|e| e.into())
-                    }
-                }
-                .map(|a| Arc::new(a) as ArrayRef)
+                            .map(|a| Arc::new(a) as ArrayRef)
             }
         }
     };
diff --git a/datafusion/physical-expr/src/expressions/binary/kernels_arrow.rs b/datafusion/physical-expr/src/expressions/binary/kernels_arrow.rs
index d73302d7e..0edb7f162 100644
--- a/datafusion/physical-expr/src/expressions/binary/kernels_arrow.rs
+++ b/datafusion/physical-expr/src/expressions/binary/kernels_arrow.rs
@@ -137,29 +137,6 @@ where
         .collect())
 }
 
-/// Creates an BooleanArray the same size as `left`,
-/// by applying `op` to all non-null elements of left and right
-pub(crate) fn compare_decimal<F>(
-    left: &Decimal128Array,
-    right: &Decimal128Array,
-    op: F,
-) -> Result<BooleanArray>
-where
-    F: Fn(i128, i128) -> bool,
-{
-    Ok(left
-        .iter()
-        .zip(right.iter())
-        .map(|(left, right)| {
-            if let (Some(left), Some(right)) = (left, right) {
-                Some(op(left, right))
-            } else {
-                None
-            }
-        })
-        .collect())
-}
-
 pub(crate) fn eq_decimal_scalar(
     left: &Decimal128Array,
     right: i128,
@@ -167,13 +144,6 @@ pub(crate) fn eq_decimal_scalar(
     compare_decimal_scalar(left, right, |left, right| left == right)
 }
 
-pub(crate) fn eq_decimal(
-    left: &Decimal128Array,
-    right: &Decimal128Array,
-) -> Result<BooleanArray> {
-    compare_decimal(left, right, |left, right| left == right)
-}
-
 pub(crate) fn neq_decimal_scalar(
     left: &Decimal128Array,
     right: i128,
@@ -181,13 +151,6 @@ pub(crate) fn neq_decimal_scalar(
     compare_decimal_scalar(left, right, |left, right| left != right)
 }
 
-pub(crate) fn neq_decimal(
-    left: &Decimal128Array,
-    right: &Decimal128Array,
-) -> Result<BooleanArray> {
-    compare_decimal(left, right, |left, right| left != right)
-}
-
 pub(crate) fn lt_decimal_scalar(
     left: &Decimal128Array,
     right: i128,
@@ -195,13 +158,6 @@ pub(crate) fn lt_decimal_scalar(
     compare_decimal_scalar(left, right, |left, right| left < right)
 }
 
-pub(crate) fn lt_decimal(
-    left: &Decimal128Array,
-    right: &Decimal128Array,
-) -> Result<BooleanArray> {
-    compare_decimal(left, right, |left, right| left < right)
-}
-
 pub(crate) fn lt_eq_decimal_scalar(
     left: &Decimal128Array,
     right: i128,
@@ -209,13 +165,6 @@ pub(crate) fn lt_eq_decimal_scalar(
     compare_decimal_scalar(left, right, |left, right| left <= right)
 }
 
-pub(crate) fn lt_eq_decimal(
-    left: &Decimal128Array,
-    right: &Decimal128Array,
-) -> Result<BooleanArray> {
-    compare_decimal(left, right, |left, right| left <= right)
-}
-
 pub(crate) fn gt_decimal_scalar(
     left: &Decimal128Array,
     right: i128,
@@ -223,13 +172,6 @@ pub(crate) fn gt_decimal_scalar(
     compare_decimal_scalar(left, right, |left, right| left > right)
 }
 
-pub(crate) fn gt_decimal(
-    left: &Decimal128Array,
-    right: &Decimal128Array,
-) -> Result<BooleanArray> {
-    compare_decimal(left, right, |left, right| left > right)
-}
-
 pub(crate) fn gt_eq_decimal_scalar(
     left: &Decimal128Array,
     right: i128,
@@ -237,13 +179,6 @@ pub(crate) fn gt_eq_decimal_scalar(
     compare_decimal_scalar(left, right, |left, right| left >= right)
 }
 
-pub(crate) fn gt_eq_decimal(
-    left: &Decimal128Array,
-    right: &Decimal128Array,
-) -> Result<BooleanArray> {
-    compare_decimal(left, right, |left, right| left >= right)
-}
-
 pub(crate) fn is_distinct_from_decimal(
     left: &Decimal128Array,
     right: &Decimal128Array,
@@ -516,42 +451,7 @@ mod tests {
             25,
             3,
         );
-        // eq: left == right
-        let result = eq_decimal(&left_decimal_array, &right_decimal_array)?;
-        assert_eq!(
-            BooleanArray::from(vec![Some(false), None, Some(false), Some(true)]),
-            result
-        );
-        // neq: left != right
-        let result = neq_decimal(&left_decimal_array, &right_decimal_array)?;
-        assert_eq!(
-            BooleanArray::from(vec![Some(true), None, Some(true), Some(false)]),
-            result
-        );
-        // lt: left < right
-        let result = lt_decimal(&left_decimal_array, &right_decimal_array)?;
-        assert_eq!(
-            BooleanArray::from(vec![Some(false), None, Some(true), Some(false)]),
-            result
-        );
-        // lt_eq: left <= right
-        let result = lt_eq_decimal(&left_decimal_array, &right_decimal_array)?;
-        assert_eq!(
-            BooleanArray::from(vec![Some(false), None, Some(true), Some(true)]),
-            result
-        );
-        // gt: left > right
-        let result = gt_decimal(&left_decimal_array, &right_decimal_array)?;
-        assert_eq!(
-            BooleanArray::from(vec![Some(true), None, Some(false), Some(false)]),
-            result
-        );
-        // gt_eq: left >= right
-        let result = gt_eq_decimal(&left_decimal_array, &right_decimal_array)?;
-        assert_eq!(
-            BooleanArray::from(vec![Some(true), None, Some(false), Some(true)]),
-            result
-        );
+
         // is_distinct: left distinct right
         let result = is_distinct_from_decimal(&left_decimal_array, &right_decimal_array)?;
         assert_eq!(