From 3e30b763de6c0e548678c24ef8cffb72101b5e6f Mon Sep 17 00:00:00 2001 From: Matthew Turner Date: Fri, 31 Dec 2021 23:19:06 -0500 Subject: [PATCH 1/6] Add lt_dyn_scalar and tests --- arrow/src/compute/kernels/comparison.rs | 63 +++++++++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/arrow/src/compute/kernels/comparison.rs b/arrow/src/compute/kernels/comparison.rs index 3e7a084cf334..1ee4af92f0ca 100644 --- a/arrow/src/compute/kernels/comparison.rs +++ b/arrow/src/compute/kernels/comparison.rs @@ -1121,6 +1121,42 @@ where } } +/// Perform `left < right` operation on an array and a numeric scalar +/// value. Supports PrimitiveArrays, and DictionaryArrays that have primitive values +pub fn lt_dyn_scalar(left: Arc, right: T) -> Result +where + T: TryInto + Copy + std::fmt::Debug, +{ + match left.data_type() { + DataType::Dictionary(key_type, value_type) => match value_type.as_ref() { + DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 => {dyn_compare_scalar!(&left, right, key_type, lt_scalar)} + _ => Err(ArrowError::ComputeError( + "Kernel only supports PrimitiveArray or DictionaryArray with Primitive values".to_string(), + )) + } + DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 => { + dyn_compare_scalar!(&left, right, lt_scalar) + } + _ => Err(ArrowError::ComputeError( + "Kernel only supports PrimitiveArray or DictionaryArray with Primitive values".to_string(), + )) + } +} + /// Perform `left == right` operation on an array and a numeric scalar /// value. Supports StringArrays, and DictionaryArrays that have string values pub fn eq_dyn_utf8_scalar(left: Arc, right: &str) -> Result { @@ -2973,6 +3009,33 @@ mod tests { ); } #[test] + fn test_lt_dyn_scalar() { + let array = Int32Array::from(vec![6, 7, 8, 8, 10]); + let array = Arc::new(array); + let a_eq = lt_dyn_scalar(array, 8).unwrap(); + assert_eq!( + a_eq, + BooleanArray::from( + vec![Some(true), Some(true), Some(false), Some(false), Some(false)] + ) + ); + } + #[test] + fn test_lt_dyn_scalar_with_dict() { + let key_builder = PrimitiveBuilder::::new(3); + let value_builder = PrimitiveBuilder::::new(2); + let mut builder = PrimitiveDictionaryBuilder::new(key_builder, value_builder); + builder.append(123).unwrap(); + builder.append_null().unwrap(); + builder.append(23).unwrap(); + let array = Arc::new(builder.finish()); + let a_eq = lt_dyn_scalar(array, 123).unwrap(); + assert_eq!( + a_eq, + BooleanArray::from(vec![Some(false), None, Some(true)]) + ); + } + #[test] fn test_eq_dyn_utf8_scalar() { let array = StringArray::from(vec!["abc", "def", "xyz"]); let array = Arc::new(array); From d6d0f93ede2f1fb39430cce31600c057f0cb44d9 Mon Sep 17 00:00:00 2001 From: Matthew Turner Date: Fri, 31 Dec 2021 23:40:01 -0500 Subject: [PATCH 2/6] Add lt_eq_dyn_scalar kernel --- arrow/src/compute/kernels/comparison.rs | 63 +++++++++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/arrow/src/compute/kernels/comparison.rs b/arrow/src/compute/kernels/comparison.rs index 1ee4af92f0ca..e8cd0e5726db 100644 --- a/arrow/src/compute/kernels/comparison.rs +++ b/arrow/src/compute/kernels/comparison.rs @@ -1157,6 +1157,42 @@ where } } +/// Perform `left <= right` operation on an array and a numeric scalar +/// value. Supports PrimitiveArrays, and DictionaryArrays that have primitive values +pub fn lt_eq_dyn_scalar(left: Arc, right: T) -> Result +where + T: TryInto + Copy + std::fmt::Debug, +{ + match left.data_type() { + DataType::Dictionary(key_type, value_type) => match value_type.as_ref() { + DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 => {dyn_compare_scalar!(&left, right, key_type, lt_eq_scalar)} + _ => Err(ArrowError::ComputeError( + "Kernel only supports PrimitiveArray or DictionaryArray with Primitive values".to_string(), + )) + } + DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 => { + dyn_compare_scalar!(&left, right, lt_eq_scalar) + } + _ => Err(ArrowError::ComputeError( + "Kernel only supports PrimitiveArray or DictionaryArray with Primitive values".to_string(), + )) + } +} + /// Perform `left == right` operation on an array and a numeric scalar /// value. Supports StringArrays, and DictionaryArrays that have string values pub fn eq_dyn_utf8_scalar(left: Arc, right: &str) -> Result { @@ -3036,6 +3072,33 @@ mod tests { ); } #[test] + fn test_lt_eq_dyn_scalar() { + let array = Int32Array::from(vec![6, 7, 8, 8, 10]); + let array = Arc::new(array); + let a_eq = lt_eq_dyn_scalar(array, 8).unwrap(); + assert_eq!( + a_eq, + BooleanArray::from( + vec![Some(true), Some(true), Some(true), Some(true), Some(false)] + ) + ); + } + #[test] + fn test_lt_eq_dyn_scalar_with_dict() { + let key_builder = PrimitiveBuilder::::new(3); + let value_builder = PrimitiveBuilder::::new(2); + let mut builder = PrimitiveDictionaryBuilder::new(key_builder, value_builder); + builder.append(123).unwrap(); + builder.append_null().unwrap(); + builder.append(23).unwrap(); + let array = Arc::new(builder.finish()); + let a_eq = lt_eq_dyn_scalar(array, 23).unwrap(); + assert_eq!( + a_eq, + BooleanArray::from(vec![Some(false), None, Some(true)]) + ); + } + #[test] fn test_eq_dyn_utf8_scalar() { let array = StringArray::from(vec!["abc", "def", "xyz"]); let array = Arc::new(array); From 4d2b99c22bde8ce8d28af9bf5754c0e25e976e89 Mon Sep 17 00:00:00 2001 From: Matthew Turner Date: Fri, 31 Dec 2021 23:47:55 -0500 Subject: [PATCH 3/6] Add gt_dyn_scalar kernel --- arrow/src/compute/kernels/comparison.rs | 63 +++++++++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/arrow/src/compute/kernels/comparison.rs b/arrow/src/compute/kernels/comparison.rs index e8cd0e5726db..1c7d3529e831 100644 --- a/arrow/src/compute/kernels/comparison.rs +++ b/arrow/src/compute/kernels/comparison.rs @@ -1193,6 +1193,42 @@ where } } +/// Perform `left > right` operation on an array and a numeric scalar +/// value. Supports PrimitiveArrays, and DictionaryArrays that have primitive values +pub fn gt_dyn_scalar(left: Arc, right: T) -> Result +where + T: TryInto + Copy + std::fmt::Debug, +{ + match left.data_type() { + DataType::Dictionary(key_type, value_type) => match value_type.as_ref() { + DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 => {dyn_compare_scalar!(&left, right, key_type, gt_scalar)} + _ => Err(ArrowError::ComputeError( + "Kernel only supports PrimitiveArray or DictionaryArray with Primitive values".to_string(), + )) + } + DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 => { + dyn_compare_scalar!(&left, right, gt_scalar) + } + _ => Err(ArrowError::ComputeError( + "Kernel only supports PrimitiveArray or DictionaryArray with Primitive values".to_string(), + )) + } +} + /// Perform `left == right` operation on an array and a numeric scalar /// value. Supports StringArrays, and DictionaryArrays that have string values pub fn eq_dyn_utf8_scalar(left: Arc, right: &str) -> Result { @@ -3099,6 +3135,33 @@ mod tests { ); } #[test] + fn test_gt_dyn_scalar() { + let array = Int32Array::from(vec![6, 7, 8, 8, 10]); + let array = Arc::new(array); + let a_eq = gt_dyn_scalar(array, 8).unwrap(); + assert_eq!( + a_eq, + BooleanArray::from( + vec![Some(false), Some(false), Some(false), Some(false), Some(true)] + ) + ); + } + #[test] + fn test_gt_dyn_scalar_with_dict() { + let key_builder = PrimitiveBuilder::::new(3); + let value_builder = PrimitiveBuilder::::new(2); + let mut builder = PrimitiveDictionaryBuilder::new(key_builder, value_builder); + builder.append(123).unwrap(); + builder.append_null().unwrap(); + builder.append(23).unwrap(); + let array = Arc::new(builder.finish()); + let a_eq = gt_dyn_scalar(array, 23).unwrap(); + assert_eq!( + a_eq, + BooleanArray::from(vec![Some(true), None, Some(false)]) + ); + } + #[test] fn test_eq_dyn_utf8_scalar() { let array = StringArray::from(vec!["abc", "def", "xyz"]); let array = Arc::new(array); From 68f292e77453a83db9d113a5544e52c137477f8a Mon Sep 17 00:00:00 2001 From: Matthew Turner Date: Fri, 31 Dec 2021 23:55:08 -0500 Subject: [PATCH 4/6] Add gt_eq_dyn_scalar kernel --- arrow/src/compute/kernels/comparison.rs | 63 +++++++++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/arrow/src/compute/kernels/comparison.rs b/arrow/src/compute/kernels/comparison.rs index 1c7d3529e831..d137f093c7a2 100644 --- a/arrow/src/compute/kernels/comparison.rs +++ b/arrow/src/compute/kernels/comparison.rs @@ -1229,6 +1229,42 @@ where } } +/// Perform `left >= right` operation on an array and a numeric scalar +/// value. Supports PrimitiveArrays, and DictionaryArrays that have primitive values +pub fn gt_eq_dyn_scalar(left: Arc, right: T) -> Result +where + T: TryInto + Copy + std::fmt::Debug, +{ + match left.data_type() { + DataType::Dictionary(key_type, value_type) => match value_type.as_ref() { + DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 => {dyn_compare_scalar!(&left, right, key_type, gt_eq_scalar)} + _ => Err(ArrowError::ComputeError( + "Kernel only supports PrimitiveArray or DictionaryArray with Primitive values".to_string(), + )) + } + DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 => { + dyn_compare_scalar!(&left, right, gt_eq_scalar) + } + _ => Err(ArrowError::ComputeError( + "Kernel only supports PrimitiveArray or DictionaryArray with Primitive values".to_string(), + )) + } +} + /// Perform `left == right` operation on an array and a numeric scalar /// value. Supports StringArrays, and DictionaryArrays that have string values pub fn eq_dyn_utf8_scalar(left: Arc, right: &str) -> Result { @@ -3162,6 +3198,33 @@ mod tests { ); } #[test] + fn test_gt_eq_dyn_scalar() { + let array = Int32Array::from(vec![6, 7, 8, 8, 10]); + let array = Arc::new(array); + let a_eq = gt_eq_dyn_scalar(array, 8).unwrap(); + assert_eq!( + a_eq, + BooleanArray::from( + vec![Some(false), Some(false), Some(true), Some(true), Some(true)] + ) + ); + } + #[test] + fn test_gt_eq_dyn_scalar_with_dict() { + let key_builder = PrimitiveBuilder::::new(3); + let value_builder = PrimitiveBuilder::::new(2); + let mut builder = PrimitiveDictionaryBuilder::new(key_builder, value_builder); + builder.append(22).unwrap(); + builder.append_null().unwrap(); + builder.append(23).unwrap(); + let array = Arc::new(builder.finish()); + let a_eq = gt_eq_dyn_scalar(array, 23).unwrap(); + assert_eq!( + a_eq, + BooleanArray::from(vec![Some(false), None, Some(true)]) + ); + } + #[test] fn test_eq_dyn_utf8_scalar() { let array = StringArray::from(vec!["abc", "def", "xyz"]); let array = Arc::new(array); From d0442be707bf49089a664afb82873ac07271e828 Mon Sep 17 00:00:00 2001 From: Matthew Turner Date: Sat, 1 Jan 2022 00:12:09 -0500 Subject: [PATCH 5/6] Add neq_dyn_scalar kernel --- arrow/src/compute/kernels/comparison.rs | 63 +++++++++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/arrow/src/compute/kernels/comparison.rs b/arrow/src/compute/kernels/comparison.rs index d137f093c7a2..b7902ba0abe6 100644 --- a/arrow/src/compute/kernels/comparison.rs +++ b/arrow/src/compute/kernels/comparison.rs @@ -1265,6 +1265,42 @@ where } } +/// Perform `left != right` operation on an array and a numeric scalar +/// value. Supports PrimitiveArrays, and DictionaryArrays that have primitive values +pub fn neq_dyn_scalar(left: Arc, right: T) -> Result +where + T: TryInto + Copy + std::fmt::Debug, +{ + match left.data_type() { + DataType::Dictionary(key_type, value_type) => match value_type.as_ref() { + DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 => {dyn_compare_scalar!(&left, right, key_type, neq_scalar)} + _ => Err(ArrowError::ComputeError( + "Kernel only supports PrimitiveArray or DictionaryArray with Primitive values".to_string(), + )) + } + DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 => { + dyn_compare_scalar!(&left, right, neq_scalar) + } + _ => Err(ArrowError::ComputeError( + "Kernel only supports PrimitiveArray or DictionaryArray with Primitive values".to_string(), + )) + } +} + /// Perform `left == right` operation on an array and a numeric scalar /// value. Supports StringArrays, and DictionaryArrays that have string values pub fn eq_dyn_utf8_scalar(left: Arc, right: &str) -> Result { @@ -3225,6 +3261,33 @@ mod tests { ); } #[test] + fn test_neq_dyn_scalar() { + let array = Int32Array::from(vec![6, 7, 8, 8, 10]); + let array = Arc::new(array); + let a_eq = neq_dyn_scalar(array, 8).unwrap(); + assert_eq!( + a_eq, + BooleanArray::from( + vec![Some(true), Some(true), Some(false), Some(false), Some(true)] + ) + ); + } + #[test] + fn test_neq_dyn_scalar_with_dict() { + let key_builder = PrimitiveBuilder::::new(3); + let value_builder = PrimitiveBuilder::::new(2); + let mut builder = PrimitiveDictionaryBuilder::new(key_builder, value_builder); + builder.append(22).unwrap(); + builder.append_null().unwrap(); + builder.append(23).unwrap(); + let array = Arc::new(builder.finish()); + let a_eq = neq_dyn_scalar(array, 23).unwrap(); + assert_eq!( + a_eq, + BooleanArray::from(vec![Some(true), None, Some(false)]) + ); + } + #[test] fn test_eq_dyn_utf8_scalar() { let array = StringArray::from(vec!["abc", "def", "xyz"]); let array = Arc::new(array); From b310c580455a809d69e3d6249972f6ac1c25eca5 Mon Sep 17 00:00:00 2001 From: Matthew Turner Date: Sat, 1 Jan 2022 10:55:34 -0500 Subject: [PATCH 6/6] Add kernel to err message --- 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 b7902ba0abe6..eacdc318c794 100644 --- a/arrow/src/compute/kernels/comparison.rs +++ b/arrow/src/compute/kernels/comparison.rs @@ -1282,7 +1282,7 @@ where | DataType::UInt32 | DataType::UInt64 => {dyn_compare_scalar!(&left, right, key_type, neq_scalar)} _ => Err(ArrowError::ComputeError( - "Kernel only supports PrimitiveArray or DictionaryArray with Primitive values".to_string(), + "neq_dyn_scalar only supports PrimitiveArray or DictionaryArray with Primitive values".to_string(), )) } DataType::Int8 @@ -1296,7 +1296,7 @@ where dyn_compare_scalar!(&left, right, neq_scalar) } _ => Err(ArrowError::ComputeError( - "Kernel only supports PrimitiveArray or DictionaryArray with Primitive values".to_string(), + "neq_dyn_scalar only supports PrimitiveArray or DictionaryArray with Primitive values".to_string(), )) } }