Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add dyn boolean kernels #1131

Merged
merged 6 commits into from
Jan 5, 2022
Merged
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
138 changes: 134 additions & 4 deletions arrow/src/compute/kernels/comparison.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1425,30 +1425,105 @@ pub fn neq_dyn_utf8_scalar(left: Arc<dyn Array>, right: &str) -> Result<BooleanA
dyn_compare_utf8_scalar!(&left, right, key_type, neq_utf8_scalar)
}
_ => 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 => {
let left = as_string_array(&left);
neq_utf8_scalar(left, right)
}
_ => 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
}

/// Perform `left == right` operation on an array and a numeric scalar
/// value. Supports BooleanArrays, and DictionaryArrays that have string values
/// value. Supports BooleanArrays.
pub fn eq_dyn_bool_scalar(left: Arc<dyn Array>, right: bool) -> Result<BooleanArray> {
let result = match left.data_type() {
DataType::Boolean => {
let left = as_boolean_array(&left);
eq_bool_scalar(left, right)
}
_ => Err(ArrowError::ComputeError(
"Kernel only supports BooleanArray".to_string(),
"eq_dyn_bool_scalar only supports BooleanArray".to_string(),
)),
};
result
}

/// Perform `left < right` operation on an array and a numeric scalar
/// value. Supports BooleanArrays.
pub fn lt_dyn_bool_scalar(left: Arc<dyn Array>, right: bool) -> Result<BooleanArray> {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we should update this to

Suggested change
pub fn lt_dyn_bool_scalar(left: Arc<dyn Array>, right: bool) -> Result<BooleanArray> {
pub fn lt_dyn_bool_scalar(left: &dyn Array, right: bool) -> Result<BooleanArray> {

Rather than taking an owned Arc to be consistent with the work in #1127

I am happy to make the change as a follow on PR as well.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

since this PR is for bool i think best to update the signatures in a follow on PR. to confirm though, this would be for all dyn scalar kernels right? scalar, utf8, and bool?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

since this PR is for bool i think best to update the signatures in a follow on PR. to confirm though, this would be for all dyn scalar kernels right? scalar, utf8, and bool?

Yes

However, I made the change for scalar and utf8 already in #1127 -- so how about I merge this PR and then I can make the updates to that PR as well

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sure sounds good

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated API in e97f588

let result = match left.data_type() {
DataType::Boolean => {
let left = as_boolean_array(&left);
lt_bool_scalar(left, right)
}
_ => Err(ArrowError::ComputeError(
"lt_dyn_bool_scalar only supports BooleanArray".to_string(),
)),
};
result
}

/// Perform `left > right` operation on an array and a numeric scalar
/// value. Supports BooleanArrays.
pub fn gt_dyn_bool_scalar(left: Arc<dyn Array>, right: bool) -> Result<BooleanArray> {
let result = match left.data_type() {
DataType::Boolean => {
let left = as_boolean_array(&left);
gt_bool_scalar(left, right)
}
_ => Err(ArrowError::ComputeError(
"gt_dyn_bool_scalar only supports BooleanArray".to_string(),
)),
};
result
}

/// Perform `left <= right` operation on an array and a numeric scalar
/// value. Supports BooleanArrays.
pub fn lt_eq_dyn_bool_scalar(left: Arc<dyn Array>, right: bool) -> Result<BooleanArray> {
let result = match left.data_type() {
DataType::Boolean => {
let left = as_boolean_array(&left);
lt_eq_bool_scalar(left, right)
}
_ => Err(ArrowError::ComputeError(
"lt_eq_dyn_bool_scalar only supports BooleanArray".to_string(),
)),
};
result
}

/// Perform `left >= right` operation on an array and a numeric scalar
/// value. Supports BooleanArrays.
pub fn gt_eq_dyn_bool_scalar(left: Arc<dyn Array>, right: bool) -> Result<BooleanArray> {
let result = match left.data_type() {
DataType::Boolean => {
let left = as_boolean_array(&left);
gt_eq_bool_scalar(left, right)
}
_ => Err(ArrowError::ComputeError(
"gt_eq_dyn_bool_scalar only supports BooleanArray".to_string(),
)),
};
result
}

/// Perform `left != right` operation on an array and a numeric scalar
/// value. Supports BooleanArrays.
pub fn neq_dyn_bool_scalar(left: Arc<dyn Array>, right: bool) -> Result<BooleanArray> {
let result = match left.data_type() {
DataType::Boolean => {
let left = as_boolean_array(&left);
neq_bool_scalar(left, right)
}
_ => Err(ArrowError::ComputeError(
"neq_dyn_bool_scalar only supports BooleanArray".to_string(),
)),
};
result
Expand Down Expand Up @@ -3599,4 +3674,59 @@ mod tests {
BooleanArray::from(vec![Some(false), Some(true), Some(false)])
);
}

#[test]
fn test_lt_dyn_bool_scalar() {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you add the null element in your test?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sure - but can you expand on why? is there a specific concern on this test with nulls?

Copy link
Contributor

@liukun4515 liukun4515 Jan 5, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

NULL > FALSE -> NULL.
If left or right is NULL, the result must be NULL.
We should make sure this case is right.

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), None])
);
}

#[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]);
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_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]);
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)])
);
}
}