diff --git a/src/compute/filter.rs b/src/compute/filter.rs index a293695b85..0299e5045f 100644 --- a/src/compute/filter.rs +++ b/src/compute/filter.rs @@ -12,6 +12,15 @@ use crate::{array::*, types::NativeType}; /// Function that can filter arbitrary arrays pub type Filter<'a> = Box Box + 'a + Send + Sync>; +#[inline] +fn get_leading_ones(chunk: u64) -> u32 { + if cfg!(target_endian = "little") { + chunk.trailing_ones() + } else { + chunk.leading_ones() + } +} + /// # Safety /// This assumes that the `mask_chunks` contains a number of set/true items equal /// to `filter_count` @@ -29,7 +38,7 @@ where .zip(mask_chunks.by_ref()) .for_each(|(chunk, mask_chunk)| { let ones = mask_chunk.count_ones(); - let leading_ones = mask_chunk.leading_ones(); + let leading_ones = get_leading_ones(mask_chunk); if ones == leading_ones { let size = leading_ones as usize; @@ -91,7 +100,7 @@ where .zip(mask_chunks.by_ref()) .for_each(|((chunk, validity_chunk), mask_chunk)| { let ones = mask_chunk.count_ones(); - let leading_ones = mask_chunk.leading_ones(); + let leading_ones = get_leading_ones(mask_chunk); if ones == leading_ones { let size = leading_ones as usize; diff --git a/tests/it/compute/filter.rs b/tests/it/compute/filter.rs index cdec0b8297..98df9079d0 100644 --- a/tests/it/compute/filter.rs +++ b/tests/it/compute/filter.rs @@ -13,6 +13,26 @@ fn array_slice() { assert_eq!(expected, c.as_ref()); } +#[test] +fn array_large_filter_chunks() { + let len = 65usize; + let a = Int32Array::from_iter((0..(len as i32)).map(Some)); + + let init = vec![true, true, true, false, false, true]; + let remaining = len - init.len(); + let iter = init + .into_iter() + .chain(std::iter::repeat(false).take(remaining)) + .map(Some); + let b = BooleanArray::from_iter(iter); + + let c = filter(&a, &b).unwrap(); + + let expected = Int32Array::from_slice([0, 1, 2, 5]); + + assert_eq!(expected, c.as_ref()); +} + #[test] fn array_low_density() { // this test exercises the all 0's branch of the filter algorithm @@ -71,7 +91,7 @@ fn string_array_simple() { } #[test] -fn primative_array_with_null() { +fn primitive_array_with_null() { let a = Int32Array::from(&[Some(5), None]); let b = BooleanArray::from_slice(vec![false, true]); let c = filter(&a, &b).unwrap();