Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Improve comparison docs and re-export the array-comparing function #404

Merged
merged 4 commits into from
Sep 16, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions src/compute/comparison/binary.rs
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,13 @@ fn gt_eq_scalar<O: Offset>(lhs: &BinaryArray<O>, 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<O: Offset>(
lhs: &BinaryArray<O>,
rhs: &BinaryArray<O>,
Expand All @@ -136,6 +143,11 @@ pub fn compare<O: Offset>(
}
}

/// 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<O: Offset>(
lhs: &BinaryArray<O>,
rhs: &BinaryScalar<O>,
Expand Down
12 changes: 12 additions & 0 deletions src/compute/comparison/boolean.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<BooleanArray> {
match op {
Operator::Eq => eq(lhs, rhs),
Expand All @@ -151,6 +158,11 @@ pub fn compare(lhs: &BooleanArray, rhs: &BooleanArray, op: Operator) -> Result<B
}
}

/// Compare a [`BooleanArray`] and a scalar value using the given
/// [`Operator`].
///
/// Check the [crate::compute::comparison](module documentation) for usage
/// examples.
pub fn compare_scalar(lhs: &BooleanArray, rhs: &BooleanScalar, op: Operator) -> BooleanArray {
if !rhs.is_valid() {
return BooleanArray::new_null(DataType::Boolean, lhs.len());
Expand Down
80 changes: 72 additions & 8 deletions src/compute/comparison/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::<i32>::from([Some(1), None, Some(2)]);
//! let array2 = PrimitiveArray::<i32>::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::<f64>::from(&[Some(10.0), None, Some(20.0)]);
//! let array2: &dyn Array = &PrimitiveArray::<f64>::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::<i32>::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};
Expand All @@ -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,
Expand Down Expand Up @@ -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.
///
Expand Down
12 changes: 12 additions & 0 deletions src/compute/comparison/primitive.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<T: NativeType + Simd8>(
lhs: &PrimitiveArray<T>,
rhs: &PrimitiveArray<T>,
Expand All @@ -228,6 +235,11 @@ pub fn compare<T: NativeType + Simd8>(
}
}

/// 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<T: NativeType + Simd8>(
lhs: &PrimitiveArray<T>,
rhs: &PrimitiveScalar<T>,
Expand Down
12 changes: 12 additions & 0 deletions src/compute/comparison/utf8.rs
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,13 @@ fn gt_eq_scalar<O: Offset>(lhs: &Utf8Array<O>, 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<O: Offset>(
lhs: &Utf8Array<O>,
rhs: &Utf8Array<O>,
Expand All @@ -136,6 +143,11 @@ pub fn compare<O: Offset>(
}
}

/// 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<O: Offset>(
lhs: &Utf8Array<O>,
rhs: &Utf8Scalar<O>,
Expand Down