From 824d1b0f32fc9f24c350b4d3bf1f4d74b3ad0dbb Mon Sep 17 00:00:00 2001 From: Matthew Turner Date: Mon, 3 Jan 2022 22:28:36 -0500 Subject: [PATCH 1/6] Add dyn bool kernels --- arrow/src/compute/kernels/comparison.rs | 108 ++++++++++++++++++++++++ 1 file changed, 108 insertions(+) diff --git a/arrow/src/compute/kernels/comparison.rs b/arrow/src/compute/kernels/comparison.rs index 191a47ffc9c8..bacbf1975687 100644 --- a/arrow/src/compute/kernels/comparison.rs +++ b/arrow/src/compute/kernels/comparison.rs @@ -1454,6 +1454,81 @@ pub fn eq_dyn_bool_scalar(left: Arc, right: bool) -> Result, right: bool) -> Result { + let result = match left.data_type() { + DataType::Boolean => { + let left = as_boolean_array(&left); + lt_bool_scalar(left, right) + } + _ => Err(ArrowError::ComputeError( + "Kernel only supports BooleanArray".to_string(), + )), + }; + result +} + +/// Perform `left > right` operation on an array and a numeric scalar +/// value. Supports BooleanArrays, and DictionaryArrays that have string values +pub fn gt_dyn_bool_scalar(left: Arc, right: bool) -> Result { + let result = match left.data_type() { + DataType::Boolean => { + let left = as_boolean_array(&left); + gt_bool_scalar(left, right) + } + _ => Err(ArrowError::ComputeError( + "Kernel only supports BooleanArray".to_string(), + )), + }; + result +} + +/// Perform `left <= right` operation on an array and a numeric scalar +/// value. Supports BooleanArrays, and DictionaryArrays that have string values +pub fn lt_eq_dyn_bool_scalar(left: Arc, right: bool) -> Result { + let result = match left.data_type() { + DataType::Boolean => { + let left = as_boolean_array(&left); + lt_eq_bool_scalar(left, right) + } + _ => Err(ArrowError::ComputeError( + "Kernel only supports BooleanArray".to_string(), + )), + }; + result +} + +/// Perform `left >= right` operation on an array and a numeric scalar +/// value. Supports BooleanArrays, and DictionaryArrays that have string values +pub fn gt_eq_dyn_bool_scalar(left: Arc, right: bool) -> Result { + let result = match left.data_type() { + DataType::Boolean => { + let left = as_boolean_array(&left); + gt_eq_bool_scalar(left, right) + } + _ => Err(ArrowError::ComputeError( + "Kernel only supports BooleanArray".to_string(), + )), + }; + result +} + +/// Perform `left != right` operation on an array and a numeric scalar +/// value. Supports BooleanArrays, and DictionaryArrays that have string values +pub fn neq_dyn_bool_scalar(left: Arc, right: bool) -> Result { + let result = match left.data_type() { + DataType::Boolean => { + let left = as_boolean_array(&left); + neq_bool_scalar(left, right) + } + _ => Err(ArrowError::ComputeError( + "Kernel only supports BooleanArray".to_string(), + )), + }; + result +} + /// unpacks the results of comparing left.values (as a boolean) /// /// TODO add example @@ -3599,4 +3674,37 @@ mod tests { BooleanArray::from(vec![Some(false), Some(true), Some(false)]) ); } + + #[test] + fn test_lt_dyn_bool_scalar() { + let array = BooleanArray::from(vec![true, false, true]); + let array = Arc::new(array); + let a_eq = lt_dyn_bool_scalar(array, false).unwrap(); + assert_eq!( + a_eq, + BooleanArray::from(vec![Some(false), Some(false), Some(false)]) + ); + } + + #[test] + fn test_lt_eq_dyn_bool_scalar() { + let array = BooleanArray::from(vec![true, false, true]); + let array = Arc::new(array); + let a_eq = lt_eq_dyn_bool_scalar(array, false).unwrap(); + assert_eq!( + a_eq, + BooleanArray::from(vec![Some(false), Some(true), Some(false)]) + ); + } + + #[test] + fn test_neq_dyn_bool_scalar() { + let array = BooleanArray::from(vec![true, false, true]); + let array = Arc::new(array); + let a_eq = neq_dyn_bool_scalar(array, false).unwrap(); + assert_eq!( + a_eq, + BooleanArray::from(vec![Some(true), Some(false), Some(true)]) + ); + } } From c857dcb73091689ec143ad3660e068a8e74b7a96 Mon Sep 17 00:00:00 2001 From: Matthew Turner Date: Mon, 3 Jan 2022 22:35:36 -0500 Subject: [PATCH 2/6] Add tests --- arrow/src/compute/kernels/comparison.rs | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/arrow/src/compute/kernels/comparison.rs b/arrow/src/compute/kernels/comparison.rs index bacbf1975687..1d73f56b5e56 100644 --- a/arrow/src/compute/kernels/comparison.rs +++ b/arrow/src/compute/kernels/comparison.rs @@ -3686,6 +3686,17 @@ mod tests { ); } + #[test] + fn test_gt_dyn_bool_scalar() { + let array = BooleanArray::from(vec![true, false, true]); + let array = Arc::new(array); + let a_eq = gt_dyn_bool_scalar(array, false).unwrap(); + assert_eq!( + a_eq, + BooleanArray::from(vec![Some(true), Some(false), Some(true)]) + ); + } + #[test] fn test_lt_eq_dyn_bool_scalar() { let array = BooleanArray::from(vec![true, false, true]); @@ -3697,6 +3708,17 @@ mod tests { ); } + #[test] + fn test_gt_eq_dyn_bool_scalar() { + let array = BooleanArray::from(vec![true, false, true]); + let array = Arc::new(array); + let a_eq = gt_eq_dyn_bool_scalar(array, false).unwrap(); + assert_eq!( + a_eq, + BooleanArray::from(vec![Some(true), Some(true), Some(true)]) + ); + } + #[test] fn test_neq_dyn_bool_scalar() { let array = BooleanArray::from(vec![true, false, true]); From c80e2e8712bf162e38674c3ce0a01b80f7b94251 Mon Sep 17 00:00:00 2001 From: Matthew Turner Date: Tue, 4 Jan 2022 23:12:37 -0500 Subject: [PATCH 3/6] Update error messages --- arrow/src/compute/kernels/comparison.rs | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/arrow/src/compute/kernels/comparison.rs b/arrow/src/compute/kernels/comparison.rs index 1d73f56b5e56..891fb3cbe4ce 100644 --- a/arrow/src/compute/kernels/comparison.rs +++ b/arrow/src/compute/kernels/comparison.rs @@ -1425,7 +1425,7 @@ pub fn neq_dyn_utf8_scalar(left: Arc, right: &str) -> Result Err(ArrowError::ComputeError( - "Kernel only supports Utf8 or LargeUtf8 arrays or DictionaryArray with Utf8 or LargeUtf8 values".to_string(), + "neq_dyn_utf8_scalar only supports Utf8 or LargeUtf8 arrays or DictionaryArray with Utf8 or LargeUtf8 values".to_string(), )), }, DataType::Utf8 | DataType::LargeUtf8 => { @@ -1433,7 +1433,7 @@ pub fn neq_dyn_utf8_scalar(left: Arc, right: &str) -> Result Err(ArrowError::ComputeError( - "Kernel only supports Utf8 or LargeUtf8 arrays".to_string(), + "neq_dyn_utf8_scalar only supports Utf8 or LargeUtf8 arrays".to_string(), )), }; result @@ -1448,7 +1448,7 @@ pub fn eq_dyn_bool_scalar(left: Arc, right: bool) -> Result Err(ArrowError::ComputeError( - "Kernel only supports BooleanArray".to_string(), + "eq_dyn_bool_scalar only supports BooleanArray".to_string(), )), }; result @@ -1463,7 +1463,7 @@ pub fn lt_dyn_bool_scalar(left: Arc, right: bool) -> Result Err(ArrowError::ComputeError( - "Kernel only supports BooleanArray".to_string(), + "lt_dyn_bool_scalar only supports BooleanArray".to_string(), )), }; result @@ -1478,7 +1478,7 @@ pub fn gt_dyn_bool_scalar(left: Arc, right: bool) -> Result Err(ArrowError::ComputeError( - "Kernel only supports BooleanArray".to_string(), + "gt_dyn_bool_scalar only supports BooleanArray".to_string(), )), }; result @@ -1493,7 +1493,7 @@ pub fn lt_eq_dyn_bool_scalar(left: Arc, right: bool) -> Result Err(ArrowError::ComputeError( - "Kernel only supports BooleanArray".to_string(), + "lt_eq_dyn_bool_scalar only supports BooleanArray".to_string(), )), }; result @@ -1508,7 +1508,7 @@ pub fn gt_eq_dyn_bool_scalar(left: Arc, right: bool) -> Result Err(ArrowError::ComputeError( - "Kernel only supports BooleanArray".to_string(), + "gt_eq_dyn_bool_scalar only supports BooleanArray".to_string(), )), }; result @@ -1523,7 +1523,7 @@ pub fn neq_dyn_bool_scalar(left: Arc, right: bool) -> Result Err(ArrowError::ComputeError( - "Kernel only supports BooleanArray".to_string(), + "neq_dyn_bool_scalar only supports BooleanArray".to_string(), )), }; result From 9a4f295da74d9d753ccc59ddbbc47a1e96b3872a Mon Sep 17 00:00:00 2001 From: Matthew Turner Date: Tue, 4 Jan 2022 23:21:13 -0500 Subject: [PATCH 4/6] Update test --- arrow/src/compute/kernels/comparison.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/arrow/src/compute/kernels/comparison.rs b/arrow/src/compute/kernels/comparison.rs index 891fb3cbe4ce..d9caa7204ec9 100644 --- a/arrow/src/compute/kernels/comparison.rs +++ b/arrow/src/compute/kernels/comparison.rs @@ -3699,12 +3699,12 @@ mod tests { #[test] fn test_lt_eq_dyn_bool_scalar() { - let array = BooleanArray::from(vec![true, false, true]); + let array = BooleanArray::from(vec![Some(true), Some(false), Some(true)]); let array = Arc::new(array); let a_eq = lt_eq_dyn_bool_scalar(array, false).unwrap(); assert_eq!( a_eq, - BooleanArray::from(vec![Some(false), Some(true), Some(false)]) + BooleanArray::from(vec![Some(false), Some(true), Some(false), None]) ); } From de43f64801ec116f95bd895f66e0980ff94f7813 Mon Sep 17 00:00:00 2001 From: Matthew Turner Date: Tue, 4 Jan 2022 23:57:06 -0500 Subject: [PATCH 5/6] Fix test --- arrow/src/compute/kernels/comparison.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/arrow/src/compute/kernels/comparison.rs b/arrow/src/compute/kernels/comparison.rs index d9caa7204ec9..32d79e7dc296 100644 --- a/arrow/src/compute/kernels/comparison.rs +++ b/arrow/src/compute/kernels/comparison.rs @@ -3677,12 +3677,12 @@ mod tests { #[test] fn test_lt_dyn_bool_scalar() { - let array = BooleanArray::from(vec![true, false, true]); + let array = BooleanArray::from(vec![Some(true), Some(false), Some(true), None]); let array = Arc::new(array); let a_eq = lt_dyn_bool_scalar(array, false).unwrap(); assert_eq!( a_eq, - BooleanArray::from(vec![Some(false), Some(false), Some(false)]) + BooleanArray::from(vec![Some(false), Some(false), Some(false), None]) ); } @@ -3699,12 +3699,12 @@ mod tests { #[test] fn test_lt_eq_dyn_bool_scalar() { - let array = BooleanArray::from(vec![Some(true), Some(false), Some(true)]); + let array = BooleanArray::from(vec![true, false, true]); let array = Arc::new(array); let a_eq = lt_eq_dyn_bool_scalar(array, false).unwrap(); assert_eq!( a_eq, - BooleanArray::from(vec![Some(false), Some(true), Some(false), None]) + BooleanArray::from(vec![Some(false), Some(true), Some(false)]) ); } From f9bb9dbdef9c517a43ccc94f9188889f054106f9 Mon Sep 17 00:00:00 2001 From: Matthew Turner Date: Wed, 5 Jan 2022 15:42:50 -0500 Subject: [PATCH 6/6] Update doc strings --- arrow/src/compute/kernels/comparison.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/arrow/src/compute/kernels/comparison.rs b/arrow/src/compute/kernels/comparison.rs index 32d79e7dc296..5be2605ef0dc 100644 --- a/arrow/src/compute/kernels/comparison.rs +++ b/arrow/src/compute/kernels/comparison.rs @@ -1440,7 +1440,7 @@ pub fn neq_dyn_utf8_scalar(left: Arc, right: &str) -> Result, right: bool) -> Result { let result = match left.data_type() { DataType::Boolean => { @@ -1455,7 +1455,7 @@ pub fn eq_dyn_bool_scalar(left: Arc, right: bool) -> Result, right: bool) -> Result { let result = match left.data_type() { DataType::Boolean => { @@ -1470,7 +1470,7 @@ pub fn lt_dyn_bool_scalar(left: Arc, right: bool) -> Result right` operation on an array and a numeric scalar -/// value. Supports BooleanArrays, and DictionaryArrays that have string values +/// value. Supports BooleanArrays. pub fn gt_dyn_bool_scalar(left: Arc, right: bool) -> Result { let result = match left.data_type() { DataType::Boolean => { @@ -1485,7 +1485,7 @@ pub fn gt_dyn_bool_scalar(left: Arc, right: bool) -> Result, right: bool) -> Result { let result = match left.data_type() { DataType::Boolean => { @@ -1500,7 +1500,7 @@ pub fn lt_eq_dyn_bool_scalar(left: Arc, right: bool) -> Result= right` operation on an array and a numeric scalar -/// value. Supports BooleanArrays, and DictionaryArrays that have string values +/// value. Supports BooleanArrays. pub fn gt_eq_dyn_bool_scalar(left: Arc, right: bool) -> Result { let result = match left.data_type() { DataType::Boolean => { @@ -1515,7 +1515,7 @@ pub fn gt_eq_dyn_bool_scalar(left: Arc, right: bool) -> Result, right: bool) -> Result { let result = match left.data_type() { DataType::Boolean => {