bjchambers commented on a change in pull request #521:
URL: https://github.com/apache/arrow-rs/pull/521#discussion_r670118830
##########
File path: arrow/src/compute/kernels/boolean.rs
##########
@@ -458,101 +458,113 @@ pub fn is_not_null(input: &Array) ->
Result<BooleanArray> {
Ok(BooleanArray::from(data))
}
-/// Copies original array, setting null bit to true if a secondary comparison
boolean array is set to true.
-/// Typically used to implement NULLIF.
-// NOTE: For now this only supports Primitive Arrays. Although the code could
be made generic, the issue
-// is that currently the bitmap operations result in a final bitmap which is
aligned to bit 0, and thus
-// the left array's data needs to be sliced to a new offset, and for
non-primitive arrays shifting the
-// data might be too complicated. In the future, to avoid shifting left
array's data, we could instead
-// shift the final bitbuffer to the right, prepending with 0's instead.
-pub fn nullif<T>(
- left: &PrimitiveArray<T>,
- right: &BooleanArray,
-) -> Result<PrimitiveArray<T>>
-where
- T: ArrowNumericType,
-{
+/// Copies original array, setting null bit to true if a secondary comparison
+/// boolean array is set to true. Typically used to implement NULLIF.
+pub fn nullif(left: &ArrayRef, right: &BooleanArray) -> Result<ArrayRef> {
if left.len() != right.len() {
return Err(ArrowError::ComputeError(
"Cannot perform comparison operation on arrays of different length"
.to_string(),
));
}
let left_data = left.data();
- let right_data = right.data();
-
- // If left has no bitmap, create a new one with all values set for nullity
op later
- // left=0 (null) right=null output bitmap=null
- // left=0 right=1 output bitmap=null
- // left=1 (set) right=null output bitmap=set (passthrough)
- // left=1 right=1 & comp=true output bitmap=null
- // left=1 right=1 & comp=false output bitmap=set
- //
- // Thus: result = left null bitmap & (!right_values | !right_bitmap)
- // OR left null bitmap & !(right_values & right_bitmap)
+
+ // If right has no trues and no nulls, then it will pass the left array
+ // unmodified.
+ if right.null_count() == 0
+ && right
+ .values()
+ .count_set_bits_offset(right.offset(), right.len())
+ == 0
+ {
+ return Ok(left.clone());
Review comment:
Added.
##########
File path: arrow/src/compute/kernels/boolean.rs
##########
@@ -1148,12 +1222,215 @@ mod tests {
let comp = comp.slice(2, 3); // Some(false), None, Some(true)
let comp = comp.as_any().downcast_ref::<BooleanArray>().unwrap();
let res = nullif(&a, &comp).unwrap();
+ let res = res.as_any().downcast_ref::<Int32Array>().unwrap();
let expected = Int32Array::from(vec![
Some(15), // False => keep it
Some(8), // None => keep it
None, // true => None
]);
- assert_eq!(&expected, &res)
+ assert_eq!(&expected, res)
+ }
+
+ #[test]
+ fn test_nullif_int_large_right_offset() {
+ let a: ArrayRef = Arc::new(Int32Array::from(vec![
+ None, // 0
+ Some(15), // 1
+ Some(8),
+ Some(1),
+ Some(9),
+ ]));
+ let a = a.slice(1, 3); // Some(15), Some(8), Some(1)
+
+ let comp = BooleanArray::from(vec![
+ Some(false), // 0
+ Some(false),
+ Some(false),
+ Some(false),
+ Some(false),
+ Some(false),
+ Some(false),
+ Some(false),
+ Some(false), // 8
+ Some(false),
+ Some(false),
+ Some(false),
+ Some(false),
+ Some(false),
+ Some(false),
+ Some(false),
+ Some(false), // 16
+ Some(false), // 17
+ Some(false), // 18
+ None,
+ Some(true),
+ Some(false),
+ None,
+ ]);
+ let comp = comp.slice(18, 3); // Some(false), None, Some(true)
+ let comp = comp.as_any().downcast_ref::<BooleanArray>().unwrap();
+ let res = nullif(&a, &comp).unwrap();
+ let res = res.as_any().downcast_ref::<Int32Array>().unwrap();
+
+ let expected = Int32Array::from(vec![
+ Some(15), // False => keep it
+ Some(8), // None => keep it
+ None, // true => None
+ ]);
+ assert_eq!(&expected, res)
+ }
+
+ #[test]
+ fn test_nullif_boolean_offset() {
+ let a: ArrayRef = Arc::new(BooleanArray::from(vec![
+ None, // 0
+ Some(true), // 1
+ Some(false),
+ Some(true),
+ Some(true),
+ ]));
+ let a = a.slice(1, 3); // Some(true), Some(false), Some(true)
+
+ let comp = BooleanArray::from(vec![
+ Some(false), // 0
+ Some(false), // 1
+ Some(false), // 2
+ None,
+ Some(true),
+ Some(false),
+ None,
+ ]);
+ let comp = comp.slice(2, 3); // Some(false), None, Some(true)
+ let comp = comp.as_any().downcast_ref::<BooleanArray>().unwrap();
+ let res = nullif(&a, &comp).unwrap();
+ let res = res.as_any().downcast_ref::<BooleanArray>().unwrap();
+
+ let expected = BooleanArray::from(vec![
+ Some(true), // False => keep it
+ Some(false), // None => keep it
+ None, // true => None
+ ]);
+ assert_eq!(&expected, res)
+ }
+
+ #[test]
+ fn test_nullif_passthrough_all_nonnull() {
+ // DO NOT SUBMIT: Write this test
Review comment:
Added.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
