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

Refactored comparison Operator #604

Merged
merged 1 commit into from
Nov 14, 2021
Merged

Refactored comparison Operator #604

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
2 changes: 1 addition & 1 deletion src/array/utf8/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -178,7 +178,7 @@ impl<O: Offset> Utf8Array<O> {
impl<O: Offset> Utf8Array<O> {
/// Returns the length of this array
#[inline]
fn len(&self) -> usize {
pub fn len(&self) -> usize {
self.offsets.len() - 1
}

Expand Down
124 changes: 28 additions & 96 deletions src/compute/comparison/binary.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,39 +1,20 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! Comparison functions for [`BinaryArray`]
use crate::{
array::{BinaryArray, BooleanArray, Offset},
bitmap::Bitmap,
datatypes::DataType,
};

use crate::datatypes::DataType;
use crate::error::{ArrowError, Result};
use crate::scalar::{BinaryScalar, Scalar};
use crate::{array::*, bitmap::Bitmap};

use super::{super::utils::combine_validities, Operator};
use super::super::utils::combine_validities;

/// Evaluate `op(lhs, rhs)` for [`BinaryArray`]s using a specified
/// comparison function.
fn compare_op<O, F>(lhs: &BinaryArray<O>, rhs: &BinaryArray<O>, op: F) -> Result<BooleanArray>
fn compare_op<O, F>(lhs: &BinaryArray<O>, rhs: &BinaryArray<O>, op: F) -> BooleanArray
where
O: Offset,
F: Fn(&[u8], &[u8]) -> bool,
{
if lhs.len() != rhs.len() {
return Err(ArrowError::InvalidArgumentError(
"Cannot perform comparison operation on arrays of different length".to_string(),
));
}
assert_eq!(lhs.len(), rhs.len());

let validity = combine_validities(lhs.validity(), rhs.validity());

Expand All @@ -43,7 +24,7 @@ where
.map(|(lhs, rhs)| op(lhs, rhs));
let values = Bitmap::from_trusted_len_iter(values);

Ok(BooleanArray::from_data(DataType::Boolean, values, validity))
BooleanArray::from_data(DataType::Boolean, values, validity)
}

/// Evaluate `op(lhs, rhs)` for [`BinaryArray`] and scalar using
Expand All @@ -62,123 +43,74 @@ where
}

/// Perform `lhs == rhs` operation on [`BinaryArray`].
fn eq<O: Offset>(lhs: &BinaryArray<O>, rhs: &BinaryArray<O>) -> Result<BooleanArray> {
/// # Panic
/// iff the arrays do not have the same length.
pub fn eq<O: Offset>(lhs: &BinaryArray<O>, rhs: &BinaryArray<O>) -> BooleanArray {
compare_op(lhs, rhs, |a, b| a == b)
}

/// Perform `lhs == rhs` operation on [`BinaryArray`] and a scalar.
fn eq_scalar<O: Offset>(lhs: &BinaryArray<O>, rhs: &[u8]) -> BooleanArray {
pub fn eq_scalar<O: Offset>(lhs: &BinaryArray<O>, rhs: &[u8]) -> BooleanArray {
compare_op_scalar(lhs, rhs, |a, b| a == b)
}

/// Perform `lhs != rhs` operation on [`BinaryArray`].
fn neq<O: Offset>(lhs: &BinaryArray<O>, rhs: &BinaryArray<O>) -> Result<BooleanArray> {
/// # Panic
/// iff the arrays do not have the same length.
pub fn neq<O: Offset>(lhs: &BinaryArray<O>, rhs: &BinaryArray<O>) -> BooleanArray {
compare_op(lhs, rhs, |a, b| a != b)
}

/// Perform `lhs != rhs` operation on [`BinaryArray`] and a scalar.
fn neq_scalar<O: Offset>(lhs: &BinaryArray<O>, rhs: &[u8]) -> BooleanArray {
pub fn neq_scalar<O: Offset>(lhs: &BinaryArray<O>, rhs: &[u8]) -> BooleanArray {
compare_op_scalar(lhs, rhs, |a, b| a != b)
}

/// Perform `lhs < rhs` operation on [`BinaryArray`].
fn lt<O: Offset>(lhs: &BinaryArray<O>, rhs: &BinaryArray<O>) -> Result<BooleanArray> {
pub fn lt<O: Offset>(lhs: &BinaryArray<O>, rhs: &BinaryArray<O>) -> BooleanArray {
compare_op(lhs, rhs, |a, b| a < b)
}

/// Perform `lhs < rhs` operation on [`BinaryArray`] and a scalar.
fn lt_scalar<O: Offset>(lhs: &BinaryArray<O>, rhs: &[u8]) -> BooleanArray {
pub fn lt_scalar<O: Offset>(lhs: &BinaryArray<O>, rhs: &[u8]) -> BooleanArray {
compare_op_scalar(lhs, rhs, |a, b| a < b)
}

/// Perform `lhs <= rhs` operation on [`BinaryArray`].
fn lt_eq<O: Offset>(lhs: &BinaryArray<O>, rhs: &BinaryArray<O>) -> Result<BooleanArray> {
pub fn lt_eq<O: Offset>(lhs: &BinaryArray<O>, rhs: &BinaryArray<O>) -> BooleanArray {
compare_op(lhs, rhs, |a, b| a <= b)
}

/// Perform `lhs <= rhs` operation on [`BinaryArray`] and a scalar.
fn lt_eq_scalar<O: Offset>(lhs: &BinaryArray<O>, rhs: &[u8]) -> BooleanArray {
pub fn lt_eq_scalar<O: Offset>(lhs: &BinaryArray<O>, rhs: &[u8]) -> BooleanArray {
compare_op_scalar(lhs, rhs, |a, b| a <= b)
}

/// Perform `lhs > rhs` operation on [`BinaryArray`].
fn gt<O: Offset>(lhs: &BinaryArray<O>, rhs: &BinaryArray<O>) -> Result<BooleanArray> {
pub fn gt<O: Offset>(lhs: &BinaryArray<O>, rhs: &BinaryArray<O>) -> BooleanArray {
compare_op(lhs, rhs, |a, b| a > b)
}

/// Perform `lhs > rhs` operation on [`BinaryArray`] and a scalar.
fn gt_scalar<O: Offset>(lhs: &BinaryArray<O>, rhs: &[u8]) -> BooleanArray {
pub fn gt_scalar<O: Offset>(lhs: &BinaryArray<O>, rhs: &[u8]) -> BooleanArray {
compare_op_scalar(lhs, rhs, |a, b| a > b)
}

/// Perform `lhs >= rhs` operation on [`BinaryArray`].
fn gt_eq<O: Offset>(lhs: &BinaryArray<O>, rhs: &BinaryArray<O>) -> Result<BooleanArray> {
pub fn gt_eq<O: Offset>(lhs: &BinaryArray<O>, rhs: &BinaryArray<O>) -> BooleanArray {
compare_op(lhs, rhs, |a, b| a >= b)
}

/// Perform `lhs >= rhs` operation on [`BinaryArray`] and a scalar.
fn gt_eq_scalar<O: Offset>(lhs: &BinaryArray<O>, rhs: &[u8]) -> BooleanArray {
pub 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>,
op: Operator,
) -> Result<BooleanArray> {
match op {
Operator::Eq => eq(lhs, rhs),
Operator::Neq => neq(lhs, rhs),
Operator::Gt => gt(lhs, rhs),
Operator::GtEq => gt_eq(lhs, rhs),
Operator::Lt => lt(lhs, rhs),
Operator::LtEq => lt_eq(lhs, rhs),
}
}

/// 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>,
op: Operator,
) -> BooleanArray {
if !rhs.is_valid() {
return BooleanArray::new_null(DataType::Boolean, lhs.len());
}
compare_scalar_non_null(lhs, rhs.value(), op)
}

pub fn compare_scalar_non_null<O: Offset>(
lhs: &BinaryArray<O>,
rhs: &[u8],
op: Operator,
) -> BooleanArray {
match op {
Operator::Eq => eq_scalar(lhs, rhs),
Operator::Neq => neq_scalar(lhs, rhs),
Operator::Gt => gt_scalar(lhs, rhs),
Operator::GtEq => gt_eq_scalar(lhs, rhs),
Operator::Lt => lt_scalar(lhs, rhs),
Operator::LtEq => lt_eq_scalar(lhs, rhs),
}
}

#[cfg(test)]
mod tests {
use super::*;

fn test_generic<O: Offset, F: Fn(&BinaryArray<O>, &BinaryArray<O>) -> Result<BooleanArray>>(
fn test_generic<O: Offset, F: Fn(&BinaryArray<O>, &BinaryArray<O>) -> BooleanArray>(
lhs: Vec<&[u8]>,
rhs: Vec<&[u8]>,
op: F,
Expand All @@ -187,7 +119,7 @@ mod tests {
let lhs = BinaryArray::<O>::from_slice(lhs);
let rhs = BinaryArray::<O>::from_slice(rhs);
let expected = BooleanArray::from_slice(expected);
assert_eq!(op(&lhs, &rhs).unwrap(), expected);
assert_eq!(op(&lhs, &rhs), expected);
}

fn test_generic_scalar<O: Offset, F: Fn(&BinaryArray<O>, &[u8]) -> BooleanArray>(
Expand Down
Loading