[arrow-rs] branch master updated: Fix nullif when existing array has no nulls (#3034)

[tu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new 12f0ef4ac Fix nullif when existing array has no nulls (#3034)
12f0ef4ac is described below

commit 12f0ef4ac424e035e480a140088914e631607cae
Author: Raphael Taylor-Davies <17...@users.noreply.github.com>
AuthorDate: Mon Nov 7 11:56:10 2022 +1300

    Fix nullif when existing array has no nulls (#3034)
---
 arrow/src/compute/kernels/boolean.rs | 32 ++++++++++++++++++++++++++------
 1 file changed, 26 insertions(+), 6 deletions(-)

diff --git a/arrow/src/compute/kernels/boolean.rs b/arrow/src/compute/kernels/boolean.rs
index ea3b49e8c..dee5d0d1b 100644
--- a/arrow/src/compute/kernels/boolean.rs
+++ b/arrow/src/compute/kernels/boolean.rs
@@ -472,7 +472,7 @@ pub fn is_not_null(input: &dyn Array) -> Result<BooleanArray> {
 }
 
 /// Copies original array, setting validity bit to false if a secondary comparison
-/// boolean array is set to true or null
+/// boolean array is set to true
 ///
 /// Typically used to implement NULLIF.
 pub fn nullif(left: &dyn Array, right: &BooleanArray) -> Result<ArrayRef> {
@@ -522,11 +522,19 @@ pub fn nullif(left: &dyn Array, right: &BooleanArray) -> Result<ArrayRef> {
                 t
             })
         }
-        None => bitwise_unary_op_helper(&right, right_offset, len, |b| {
-            let t = !b;
-            valid_count += t.count_ones() as usize;
-            t
-        }),
+        None => {
+            let buffer = bitwise_unary_op_helper(&right, right_offset, len, |b| {
+                let t = !b;
+                valid_count += t.count_ones() as usize;
+                t
+            });
+            // We need to compensate for the additional bits read from the end
+            let remainder_len = len % 64;
+            if remainder_len != 0 {
+                valid_count -= 64 - remainder_len
+            }
+            buffer
+        }
     };
 
     // Need to construct null buffer with offset of left
@@ -1411,4 +1419,16 @@ mod tests {
 
         assert_eq!(&expected, res);
     }
+
+    #[test]
+    fn test_nullif_no_nulls() {
+        let a = Int32Array::from(vec![Some(15), Some(7), Some(8), Some(1), Some(9)]);
+        let comp =
+            BooleanArray::from(vec![Some(false), None, Some(true), Some(false), None]);
+        let res = nullif(&a, &comp).unwrap();
+        let res = as_primitive_array::<Int32Type>(res.as_ref());
+
+        let expected = Int32Array::from(vec![Some(15), Some(7), None, Some(1), Some(9)]);
+        assert_eq!(res, &expected);
+    }
 }