diff --git a/src/compute/comparison/binary.rs b/src/compute/comparison/binary.rs index eccc4189ccd..f72ca819362 100644 --- a/src/compute/comparison/binary.rs +++ b/src/compute/comparison/binary.rs @@ -121,6 +121,13 @@ fn gt_eq_scalar(lhs: &BinaryArray, rhs: &[u8]) -> BooleanArray { compare_op_scalar(lhs, rhs, |a, b| a >= b) } +/// Compare two [`BinaryArray`]s using the given [`Operator`]. +/// +/// # Errors +/// When the two arrays have different lengths. +/// +/// Check the [crate::compute::comparison](module documentation) for usage +/// examples. pub fn compare( lhs: &BinaryArray, rhs: &BinaryArray, @@ -136,6 +143,11 @@ pub fn compare( } } +/// Compare a [`BinaryArray`] and a scalar value using the given +/// [`Operator`]. +/// +/// Check the [crate::compute::comparison](module documentation) for usage +/// examples. pub fn compare_scalar( lhs: &BinaryArray, rhs: &BinaryScalar, diff --git a/src/compute/comparison/boolean.rs b/src/compute/comparison/boolean.rs index 647e3036931..845e077fdd2 100644 --- a/src/compute/comparison/boolean.rs +++ b/src/compute/comparison/boolean.rs @@ -140,6 +140,13 @@ pub fn gt_eq_scalar(lhs: &BooleanArray, rhs: bool) -> BooleanArray { compare_op_scalar(lhs, rhs, |a, b| a | !b) } +/// Compare two [`BooleanArray`]s using the given [`Operator`]. +/// +/// # Errors +/// When the two arrays have different lengths. +/// +/// Check the [crate::compute::comparison](module documentation) for usage +/// examples. pub fn compare(lhs: &BooleanArray, rhs: &BooleanArray, op: Operator) -> Result { match op { Operator::Eq => eq(lhs, rhs), @@ -151,6 +158,11 @@ pub fn compare(lhs: &BooleanArray, rhs: &BooleanArray, op: Operator) -> Result BooleanArray { if !rhs.is_valid() { return BooleanArray::new_null(DataType::Boolean, lhs.len()); diff --git a/src/compute/comparison/mod.rs b/src/compute/comparison/mod.rs index 5c3514145c4..54d12991274 100644 --- a/src/compute/comparison/mod.rs +++ b/src/compute/comparison/mod.rs @@ -15,7 +15,67 @@ // specific language governing permissions and limitations // under the License. -//! Defines basic comparison kernels. +//! Basic comparison kernels. +//! +//! The module contains functions that compare either an array and a scalar +//! or two arrays of the same [`DataType`]. The scalar-oriented functions are +//! suffixed with `_scalar`. In general, these comparison functions receive as +//! inputs the two items for comparison and an [`Operator`] which specifies the +//! type of comparison that will be conducted, such as `<=` ([`Operator::LtEq`]). +//! +//! Much like the parent module [`crate::compute`](compute), the comparison functions +//! have two variants - a statically typed one ([`primitive_compare`]) +//! which expects concrete types such as [`Int8Array`] and a dynamically typed +//! variant ([`compare`]) that compares values of type `&dyn Array` and errors +//! if the underlying concrete types mismsatch. The statically-typed functions +//! are appropriately prefixed with the concrete types they expect. +//! +//! Also note that the scalar-based comparison functions for the concrete types, +//! like [`primitive_compare_scalar`], are infallible and always return a +//! [`BooleanArray`] while the rest of the functions always wrap the returned +//! array in a [`Result`] due to their internal checks of the data types and +//! lengths. +//! +//! # Examples +//! +//! Compare two [`PrimitiveArray`]s: +//! ``` +//! use arrow2::compute::comparison::{primitive_compare, Operator}; +//! # use arrow2::array::{BooleanArray, PrimitiveArray}; +//! # use arrow2::error::{ArrowError, Result}; +//! +//! let array1 = PrimitiveArray::::from([Some(1), None, Some(2)]); +//! let array2 = PrimitiveArray::::from([Some(1), None, Some(1)]); +//! let result = primitive_compare(&array1, &array2, Operator::Gt)?; +//! assert_eq!(result, BooleanArray::from([Some(false), None, Some(true)])); +//! # Ok::<(), ArrowError>(()) +//! ``` +//! Compare two dynamically-typed arrays (trait objects): +//! ``` +//! use arrow2::compute::comparison::{compare, Operator}; +//! # use arrow2::array::{Array, BooleanArray, PrimitiveArray}; +//! # use arrow2::error::{ArrowError, Result}; +//! +//! let array1: &dyn Array = &PrimitiveArray::::from(&[Some(10.0), None, Some(20.0)]); +//! let array2: &dyn Array = &PrimitiveArray::::from(&[Some(10.0), None, Some(10.0)]); +//! let result = compare(array1, array2, Operator::LtEq)?; +//! assert_eq!(result, BooleanArray::from([Some(true), None, Some(false)])); +//! # Ok::<(), ArrowError>(()) +//! ``` +//! Compare an array of strings to a "scalar", i.e a word (note that we use +//! [`Operator::Neq`]): +//! ``` +//! use arrow2::compute::comparison::{utf8_compare_scalar, Operator}; +//! # use arrow2::array::{Array, BooleanArray, Utf8Array}; +//! # use arrow2::scalar::Utf8Scalar; +//! # use arrow2::error::{ArrowError, Result}; +//! +//! let array = Utf8Array::::from([Some("compute"), None, Some("compare")]); +//! let word = Utf8Scalar::new(Some("compare")); +//! let result = utf8_compare_scalar(&array, &word, Operator::Neq); +//! assert_eq!(result, BooleanArray::from([Some(true), None, Some(false)])); +//! # Ok::<(), ArrowError>(()) +//! ``` use crate::array::*; use crate::datatypes::{DataType, IntervalUnit}; @@ -30,6 +90,17 @@ mod utf8; mod simd; pub use simd::{Simd8, Simd8Lanes}; +pub use binary::compare as binary_compare; +pub use binary::compare_scalar as binary_compare_scalar; +pub use boolean::compare as boolean_compare; +pub use boolean::compare_scalar as boolean_compare_scalar; +pub use primitive::compare as primitive_compare; +pub use primitive::compare_scalar as primitive_compare_scalar; +pub(crate) use primitive::compare_values_op as primitive_compare_values_op; +pub use utf8::compare as utf8_compare; +pub use utf8::compare_scalar as utf8_compare_scalar; + +/// Comparison operators, such as `>` ([`Operator::Gt`]) #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum Operator { Lt, @@ -262,13 +333,6 @@ pub fn compare_scalar( } }) } - -pub use binary::compare_scalar_non_null as binary_compare_scalar; -pub use boolean::compare_scalar_non_null as boolean_compare_scalar; -pub use primitive::compare_scalar_non_null as primitive_compare_scalar; -pub(crate) use primitive::compare_values_op as primitive_compare_values_op; -pub use utf8::compare_scalar_non_null as utf8_compare_scalar; - /// Checks if an array of type `datatype` can be compared with another array of /// the same type. /// diff --git a/src/compute/comparison/primitive.rs b/src/compute/comparison/primitive.rs index 561868edeac..e4faefa5f53 100644 --- a/src/compute/comparison/primitive.rs +++ b/src/compute/comparison/primitive.rs @@ -213,6 +213,13 @@ where compare_op_scalar(lhs, rhs, |a, b| a.gt_eq(b)) } +/// Compare two [`PrimitiveArray`]s using the given [`Operator`]. +/// +/// # Errors +/// When the two arrays have different lengths. +/// +/// Check the [crate::compute::comparison](module documentation) for usage +/// examples. pub fn compare( lhs: &PrimitiveArray, rhs: &PrimitiveArray, @@ -228,6 +235,11 @@ pub fn compare( } } +/// Compare a [`PrimitiveArray`] and a scalar value using the given +/// [`Operator`]. +/// +/// Check the [crate::compute::comparison](module documentation) for usage +/// examples. pub fn compare_scalar( lhs: &PrimitiveArray, rhs: &PrimitiveScalar, diff --git a/src/compute/comparison/utf8.rs b/src/compute/comparison/utf8.rs index 7fd449e520a..11b4f335bde 100644 --- a/src/compute/comparison/utf8.rs +++ b/src/compute/comparison/utf8.rs @@ -121,6 +121,13 @@ fn gt_eq_scalar(lhs: &Utf8Array, rhs: &str) -> BooleanArray { compare_op_scalar(lhs, rhs, |a, b| a >= b) } +/// Compare two [`Utf8Array`]s using the given [`Operator`]. +/// +/// # Errors +/// When the two arrays have different lengths. +/// +/// Check the [crate::compute::comparison](module documentation) for usage +/// examples. pub fn compare( lhs: &Utf8Array, rhs: &Utf8Array, @@ -136,6 +143,11 @@ pub fn compare( } } +/// Compare a [`Utf8Array`] and a scalar value using the given +/// [`Operator`]. +/// +/// Check the [crate::compute::comparison](module documentation) for usage +/// examples. pub fn compare_scalar( lhs: &Utf8Array, rhs: &Utf8Scalar,