jhorstmann commented on a change in pull request #8688:
URL: https://github.com/apache/arrow/pull/8688#discussion_r530215752
##########
File path: rust/arrow/src/compute/kernels/boolean.rs
##########
@@ -223,6 +224,101 @@ pub fn is_not_null(input: &Array) -> Result<BooleanArray>
{
Ok(BooleanArray::from(Arc::new(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,
+{
+ 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)
+ //
+ // Do the right expression !(right_values & right_bitmap) first since
there are two steps
+ // TRICK: convert BooleanArray buffer as a bitmap for faster operation
+ let right_combo_buffer = match right.data().null_bitmap() {
+ Some(right_bitmap) => {
+ // NOTE: right values and bitmaps are combined and stay at bit
offset right.offset()
+ (&right.values() & &right_bitmap.bits).ok().map(|b| b.not())
+ }
+ None => Some(!&right.values()),
+ };
+
+ // AND of original left null bitmap with right expression
+ // Here we take care of the possible offsets of the left and right arrays
all at once.
+ let modified_null_buffer = match left_data.null_bitmap() {
+ Some(left_null_bitmap) => match right_combo_buffer {
+ Some(rcb) => Some(buffer_bin_and(
+ &left_null_bitmap.bits,
+ left_data.offset(),
+ &rcb,
+ right_data.offset(),
+ left_data.len(),
+ )),
+ None => Some(
+ left_null_bitmap
+ .bits
+ .bit_slice(left_data.offset(), left.len()),
+ ),
+ },
+ None => right_combo_buffer,
Review comment:
I think the right side also needs to be sliced because in the ArrayData
later it will be accessed with an offset of 0. It could be more efficient to do
the slicing before computing the `!(values & bits)`, but it probably would not
be a big difference.
Otherwise this looks good.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
