From a0496405aace25d96bdf103da048a35fe1ff14cb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=B6rn=20Horstmann?= Date: Wed, 9 Feb 2022 17:33:36 +0100 Subject: [PATCH] Fix bitmask creation also for simd comparisons with scalar (#1290) --- arrow/src/compute/kernels/comparison.rs | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/arrow/src/compute/kernels/comparison.rs b/arrow/src/compute/kernels/comparison.rs index d0a20daed330..930c275c6418 100644 --- a/arrow/src/compute/kernels/comparison.rs +++ b/arrow/src/compute/kernels/comparison.rs @@ -1799,7 +1799,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; } @@ -2478,6 +2478,18 @@ mod tests { let a = Int64Array::from($A_VEC); let c = $KERNEL(&a, $B).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 e = vec![]; + for _i in 0..10 { + a.extend($A_VEC); + e.extend($EXPECTED); + } + let a = Int64Array::from(a); + let c = $KERNEL(&a, $B).unwrap(); + assert_eq!(BooleanArray::from(e), c); + }; }