[arrow-rs] branch master updated: Fix bitmask creation in chunked part of simd comparison (#1286)

[al...@apache.org]

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

alamb 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 936ed5e  Fix bitmask creation in chunked part of simd comparison (#1286)
936ed5e is described below

commit 936ed5e4762871e5fa889e76c95acc07b955f0c2
Author: Jörn Horstmann <gi...@jhorstmann.net>
AuthorDate: Mon Feb 7 18:04:42 2022 +0100

    Fix bitmask creation in chunked part of simd comparison (#1286)
---
 arrow/src/compute/kernels/comparison.rs | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/arrow/src/compute/kernels/comparison.rs b/arrow/src/compute/kernels/comparison.rs
index 37a4b16..d0a20da 100644
--- a/arrow/src/compute/kernels/comparison.rs
+++ b/arrow/src/compute/kernels/comparison.rs
@@ -1708,7 +1708,7 @@ where
                 let simd_result = simd_op(simd_left, simd_right);
 
                 let m = T::mask_to_u64(&simd_result);
-                bitmask |= m << (i / lanes);
+                bitmask |= m << i;
 
                 i += lanes;
             }
@@ -2442,6 +2442,20 @@ mod tests {
             let b = b.slice(0, b.len());
             let c = $DYN_KERNEL(a.as_ref(), b.as_ref()).unwrap();
             assert_eq!(BooleanArray::from($EXPECTED), c);
+
+            // test with a larger version of the same data to ensure we cover the chunked part of the comparison
+            let mut a = vec![];
+            let mut b = vec![];
+            let mut e = vec![];
+            for _i in 0..10 {
+                a.extend($A_VEC);
+                b.extend($B_VEC);
+                e.extend($EXPECTED);
+            }
+            let a = $ARRAY::from(a);
+            let b = $ARRAY::from(b);
+            let c = $KERNEL(&a, &b).unwrap();
+            assert_eq!(BooleanArray::from(e), c);
         };
     }