Use arrow kernels

apache · Nov 4, 2021 · 3cb11a4 · 3cb11a4
1 parent 2c5d434
commit 3cb11a4
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 61 deletions.
diff --git a/datafusion/Cargo.toml b/datafusion/Cargo.toml
@@ -52,8 +52,10 @@ avro = ["avro-rs", "num-traits"]
 [dependencies]
 ahash = "0.7"
 hashbrown = { version = "0.11", features = ["raw"] }
-arrow = { version = "6.0.0", features = ["prettyprint"] }
-parquet = { version = "6.0.0", features = ["arrow"] }
+#arrow = { version = "6.0.0", features = ["prettyprint"] }
+#parquet = { version = "6.0.0", features = ["arrow"] }
+arrow = { git = "https://github.com/alamb/arrow-rs.git", branch="alamb/bool_kernels_pub", features = ["prettyprint"] }
+parquet = { git = "https://github.com/alamb/arrow-rs.git", branch="alamb/bool_kernels_pub", features = ["arrow"] }
 sqlparser = "0.12"
 paste = "^1.0"
 num_cpus = "1.13.0"

diff --git a/datafusion/src/physical_plan/expressions/binary.rs b/datafusion/src/physical_plan/expressions/binary.rs
@@ -24,6 +24,9 @@ use arrow::compute::kernels::arithmetic::{
 };
 use arrow::compute::kernels::boolean::{and_kleene, not, or_kleene};
 use arrow::compute::kernels::comparison::{eq, gt, gt_eq, lt, lt_eq, neq};
+use arrow::compute::kernels::comparison::{
+    eq_bool, eq_bool_scalar, neq_bool, neq_bool_scalar,
+};
 use arrow::compute::kernels::comparison::{
     eq_scalar, gt_eq_scalar, gt_scalar, lt_eq_scalar, lt_scalar, neq_scalar,
 };
@@ -351,11 +354,8 @@ macro_rules! boolean_op_scalar {
             .downcast_ref::<BooleanArray>()
             .expect("boolean_op_scalar failed to downcast array");
 
-        let result = if let ScalarValue::Boolean(scalar) = $RIGHT {
-            Ok(
-                Arc::new(paste::expr! {[<$OP _bool_scalar>]}(&ll, scalar.as_ref())?)
-                    as ArrayRef,
-            )
+        let result = if let ScalarValue::Boolean(Some(scalar)) = $RIGHT {
+            Ok(Arc::new(paste::expr! {[<$OP _bool_scalar>]}(&ll, scalar)?) as ArrayRef)
         } else {
             Err(DataFusionError::Internal(format!(
                 "boolean_op_scalar failed to cast literal value {}",
@@ -871,60 +871,6 @@ pub fn binary(
     Ok(Arc::new(BinaryExpr::new(l, op, r)))
 }
 
-// TODO file a ticket with arrow-rs to include these kernels
-
-fn eq_bool(lhs: &BooleanArray, rhs: &BooleanArray) -> Result<BooleanArray> {
-    let arr: BooleanArray = lhs
-        .iter()
-        .zip(rhs.iter())
-        .map(|v| match v {
-            // both lhs and rhs were non null
-            (Some(lhs), Some(rhs)) => Some(lhs == rhs),
-            _ => None,
-        })
-        .collect();
-
-    Ok(arr)
-}
-
-fn eq_bool_scalar(lhs: &BooleanArray, rhs: Option<&bool>) -> Result<BooleanArray> {
-    let arr: BooleanArray = lhs
-        .iter()
-        .map(|v| match (v, rhs) {
-            // both lhs and rhs were non null
-            (Some(lhs), Some(rhs)) => Some(lhs == *rhs),
-            _ => None,
-        })
-        .collect();
-    Ok(arr)
-}
-
-fn neq_bool(lhs: &BooleanArray, rhs: &BooleanArray) -> Result<BooleanArray> {
-    let arr: BooleanArray = lhs
-        .iter()
-        .zip(rhs.iter())
-        .map(|v| match v {
-            // both lhs and rhs were non null
-            (Some(lhs), Some(rhs)) => Some(lhs != rhs),
-            _ => None,
-        })
-        .collect();
-
-    Ok(arr)
-}
-
-fn neq_bool_scalar(lhs: &BooleanArray, rhs: Option<&bool>) -> Result<BooleanArray> {
-    let arr: BooleanArray = lhs
-        .iter()
-        .map(|v| match (v, rhs) {
-            // both lhs and rhs were non null
-            (Some(lhs), Some(rhs)) => Some(lhs != *rhs),
-            _ => None,
-        })
-        .collect();
-    Ok(arr)
-}
-
 #[cfg(test)]
 mod tests {
     use arrow::datatypes::{ArrowNumericType, Field, Int32Type, SchemaRef};