feat: cache FilterMask iterators (#1351)

This PR uses a cute trick to pre-cache filter indices when a FilterMask
is cloned. Unfortunately... we have to clone one every time we pass it
into with_dyn!!

Author: gatesn

encodings/alp/src/alp/compute.rs

@@ -86,9 +86,9 @@ impl SliceFn for ALPArray {
 }
 
 impl FilterFn for ALPArray {
-    fn filter(&self, mask: &FilterMask) -> VortexResult<ArrayData> {
+    fn filter(&self, mask: FilterMask) -> VortexResult<ArrayData> {
         Ok(Self::try_new(
-            filter(&self.encoded(), mask)?,
+            filter(&self.encoded(), mask.clone())?,
             self.exponents(),
             self.patches().map(|p| filter(&p, mask)).transpose()?,
         )?

encodings/alp/src/alp_rd/compute/filter.rs

@@ -5,15 +5,15 @@ use vortex_error::VortexResult;
 use crate::ALPRDArray;
 
 impl FilterFn for ALPRDArray {
-    fn filter(&self, mask: &FilterMask) -> VortexResult<ArrayData> {
+    fn filter(&self, mask: FilterMask) -> VortexResult<ArrayData> {
         let left_parts_exceptions = self
             .left_parts_exceptions()
-            .map(|array| filter(&array, mask))
+            .map(|array| filter(&array, mask.clone()))
             .transpose()?;
 
         Ok(ALPRDArray::try_new(
             self.dtype().clone(),
-            filter(&self.left_parts(), mask)?,
+            filter(&self.left_parts(), mask.clone())?,
             self.left_parts_dict(),
             filter(&self.right_parts(), mask)?,
             self.right_bit_width(),
@@ -43,13 +43,10 @@ mod test {
         assert!(encoded.left_parts_exceptions().is_some());
 
         // The first two values need no patching
-        let filtered = filter(
-            encoded.as_ref(),
-            &FilterMask::from_iter([true, false, true]),
-        )
-        .unwrap()
-        .into_primitive()
-        .unwrap();
+        let filtered = filter(encoded.as_ref(), FilterMask::from_iter([true, false, true]))
+            .unwrap()
+            .into_primitive()
+            .unwrap();
         assert_eq!(filtered.maybe_null_slice::<T>(), &[a, outlier]);
     }
 }

encodings/dict/src/compute.rs

@@ -86,7 +86,7 @@ impl TakeFn for DictArray {
 }
 
 impl FilterFn for DictArray {
-    fn filter(&self, mask: &FilterMask) -> VortexResult<ArrayData> {
+    fn filter(&self, mask: FilterMask) -> VortexResult<ArrayData> {
         let codes = filter(&self.codes(), mask)?;
         Self::try_new(codes, self.values()).map(|a| a.into_array())
     }

encodings/fastlanes/src/for/compute.rs

@@ -48,7 +48,7 @@ impl TakeFn for FoRArray {
 }
 
 impl FilterFn for FoRArray {
-    fn filter(&self, mask: &FilterMask) -> VortexResult<ArrayData> {
+    fn filter(&self, mask: FilterMask) -> VortexResult<ArrayData> {
         Self::try_new(
             filter(&self.encoded(), mask)?,
             self.owned_reference_scalar(),

encodings/fsst/src/compute.rs

@@ -151,12 +151,12 @@ impl ScalarAtFn for FSSTArray {
 
 impl FilterFn for FSSTArray {
     // Filtering an FSSTArray filters the codes array, leaving the symbols array untouched
-    fn filter(&self, mask: &FilterMask) -> VortexResult<ArrayData> {
+    fn filter(&self, mask: FilterMask) -> VortexResult<ArrayData> {
         Ok(Self::try_new(
             self.dtype().clone(),
             self.symbols(),
             self.symbol_lengths(),
-            filter(&self.codes(), mask)?,
+            filter(&self.codes(), mask.clone())?,
             filter(&self.uncompressed_lengths(), mask)?,
         )?
         .into_array())

encodings/fsst/tests/fsst_tests.rs

@@ -87,7 +87,7 @@ fn test_fsst_array_ops() {
     // test filter
     let mask = FilterMask::from_iter([false, true, false]);
 
-    let fsst_filtered = filter(&fsst_array, &mask).unwrap();
+    let fsst_filtered = filter(&fsst_array, mask).unwrap();
     assert_eq!(fsst_filtered.encoding().id(), FSST::ENCODING.id());
     assert_eq!(fsst_filtered.len(), 1);
     assert_nth_scalar!(

encodings/runend/src/compute.rs

@@ -77,6 +77,7 @@ impl ScalarAtFn for RunEndArray {
 }
 
 impl TakeFn for RunEndArray {
+    #[allow(deprecated)]
     fn take(&self, indices: &ArrayData, options: TakeOptions) -> VortexResult<ArrayData> {
         let primitive_indices = indices.clone().into_primitive()?;
         let u64_indices = match_each_integer_ptype!(primitive_indices.ptype(), |$P| {
@@ -111,7 +112,6 @@ impl TakeFn for RunEndArray {
             Validity::Array(original_validity) => {
                 let dense_validity =
                     FilterMask::try_from(take(&original_validity, indices, options)?)?;
-                let filtered_values = filter(&dense_values, &dense_validity)?;
                 let length = dense_validity.len();
                 let dense_nonnull_indices = PrimitiveArray::from(
                     dense_validity
@@ -120,6 +120,7 @@ impl TakeFn for RunEndArray {
                         .collect::<Vec<_>>(),
                 )
                 .into_array();
+                let filtered_values = filter(&dense_values, dense_validity)?;
 
                 SparseArray::try_new(
                     dense_nonnull_indices,
@@ -150,13 +151,13 @@ impl SliceFn for RunEndArray {
 }
 
 impl FilterFn for RunEndArray {
-    fn filter(&self, mask: &FilterMask) -> VortexResult<ArrayData> {
+    fn filter(&self, mask: FilterMask) -> VortexResult<ArrayData> {
         let primitive_run_ends = self.ends().into_primitive()?;
         let (run_ends, mask) = match_each_unsigned_integer_ptype!(primitive_run_ends.ptype(), |$P| {
             filter_run_ends(primitive_run_ends.maybe_null_slice::<$P>(), mask)?
         });
-        let values = filter(&self.values(), &mask)?;
         let validity = self.validity().filter(&mask)?;
+        let values = filter(&self.values(), mask)?;
 
         RunEndArray::try_new(run_ends.into_array(), values, validity).map(|a| a.into_array())
     }
@@ -165,7 +166,7 @@ impl FilterFn for RunEndArray {
 // Code adapted from apache arrow-rs https://github.com/apache/arrow-rs/blob/b1f5c250ebb6c1252b4e7c51d15b8e77f4c361fa/arrow-select/src/filter.rs#L425
 fn filter_run_ends<R: NativePType + AddAssign + From<bool> + AsPrimitive<u64>>(
     run_ends: &[R],
-    mask: &FilterMask,
+    mask: FilterMask,
 ) -> VortexResult<(PrimitiveArray, FilterMask)> {
     let mut new_run_ends = vec![R::zero(); run_ends.len()];
 
@@ -444,7 +445,7 @@ mod test {
         let arr = ree_array();
         let filtered = filter(
             arr.as_ref(),
-            &FilterMask::from_iter([
+            FilterMask::from_iter([
                 true, true, false, false, false, false, false, false, false, false, true, true,
             ]),
         )

fuzz/fuzz_targets/array_ops.rs

@@ -58,7 +58,7 @@ fuzz_target!(|fuzz_action: FuzzArrayAction| -> Corpus {
                 assert_search_sorted(sorted, s, side, expected.search(), i)
             }
             Action::Filter(mask) => {
-                current_array = filter(&current_array, &mask).unwrap();
+                current_array = filter(&current_array, mask).unwrap();
                 assert_array_eq(&expected.array(), &current_array, i);
             }
         }

pyvortex/src/array.rs

@@ -268,7 +268,7 @@ impl PyArray {
     fn filter(&self, filter: &Bound<PyArray>) -> PyResult<PyArray> {
         let filter = filter.borrow();
         let inner =
-            vortex::compute::filter(&self.inner, &FilterMask::try_from(filter.inner.clone())?)?;
+            vortex::compute::filter(&self.inner, FilterMask::try_from(filter.inner.clone())?)?;
         Ok(PyArray { inner })
     }
 

vortex-array/src/array/bool/compute/filter.rs

@@ -1,67 +1,93 @@
-use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder};
-use vortex_error::VortexResult;
+use arrow_buffer::{bit_util, BooleanBuffer, BooleanBufferBuilder};
+use vortex_error::{VortexExpect, VortexResult};
 
 use crate::array::BoolArray;
-use crate::compute::{FilterFn, FilterMask};
+use crate::compute::{FilterFn, FilterIter, FilterMask};
 use crate::{ArrayData, IntoArrayData};
 
 impl FilterFn for BoolArray {
-    fn filter(&self, mask: &FilterMask) -> VortexResult<ArrayData> {
-        filter_select_bool(self, mask).map(|a| a.into_array())
+    fn filter(&self, mask: FilterMask) -> VortexResult<ArrayData> {
+        let validity = self.validity().filter(&mask)?;
+
+        let buffer = match mask.iter()? {
+            FilterIter::Indices(indices) => filter_indices_slice(&self.boolean_buffer(), indices),
+            FilterIter::IndicesIter(iter) => {
+                filter_indices(&self.boolean_buffer(), mask.true_count(), iter)
+            }
+            FilterIter::Slices(slices) => filter_slices(
+                &self.boolean_buffer(),
+                mask.true_count(),
+                slices.iter().copied(),
+            ),
+            FilterIter::SlicesIter(iter) => {
+                filter_slices(&self.boolean_buffer(), mask.true_count(), iter)
+            }
+        };
+
+        Ok(Self::try_new(buffer, validity)?.into_array())
     }
 }
 
-fn filter_select_bool(arr: &BoolArray, mask: &FilterMask) -> VortexResult<BoolArray> {
-    let validity = arr.validity().filter(mask)?;
-
-    let selection_count = mask.true_count();
-    let out = if selection_count * 2 > arr.len() {
-        filter_select_bool_by_slice(&arr.boolean_buffer(), mask, selection_count)
-    } else {
-        filter_select_bool_by_index(&arr.boolean_buffer(), mask, selection_count)
-    };
-    BoolArray::try_new(out?, validity)
+/// Select indices from a boolean buffer.
+/// NOTE: it was benchmarked to be faster using collect_bool to index into a slice than to
+///  pass the indices as an iterator of usize. So we keep this alternate implementation.
+fn filter_indices_slice(buffer: &BooleanBuffer, indices: &[usize]) -> BooleanBuffer {
+    let src = buffer.values().as_ptr();
+    let offset = buffer.offset();
+    BooleanBuffer::collect_bool(indices.len(), |idx| unsafe {
+        bit_util::get_bit_raw(src, *indices.get_unchecked(idx) + offset)
+    })
 }
 
-fn filter_select_bool_by_slice(
-    values: &BooleanBuffer,
-    mask: &FilterMask,
-    selection_count: usize,
-) -> VortexResult<BooleanBuffer> {
-    let mut out_buf = BooleanBufferBuilder::new(selection_count);
-    mask.iter_slices()?.for_each(|(start, end)| {
-        out_buf.append_buffer(&values.slice(start, end - start));
-    });
-    Ok(out_buf.finish())
+pub fn filter_indices(
+    buffer: &BooleanBuffer,
+    indices_len: usize,
+    mut indices: impl Iterator<Item = usize>,
+) -> BooleanBuffer {
+    let src = buffer.values().as_ptr();
+    let offset = buffer.offset();
+
+    BooleanBuffer::collect_bool(indices_len, |_idx| {
+        let idx = indices
+            .next()
+            .vortex_expect("iterator is guaranteed to be within the length of the array.");
+        unsafe { bit_util::get_bit_raw(src, idx + offset) }
+    })
 }
 
-fn filter_select_bool_by_index(
-    values: &BooleanBuffer,
-    mask: &FilterMask,
-    selection_count: usize,
-) -> VortexResult<BooleanBuffer> {
-    let mut out_buf = BooleanBufferBuilder::new(selection_count);
-    mask.iter_indices()?
-        .for_each(|idx| out_buf.append(values.value(idx)));
-    Ok(out_buf.finish())
+pub fn filter_slices(
+    buffer: &BooleanBuffer,
+    indices_len: usize,
+    slices: impl Iterator<Item = (usize, usize)>,
+) -> BooleanBuffer {
+    let src = buffer.values();
+    let offset = buffer.offset();
+
+    let mut builder = BooleanBufferBuilder::new(indices_len);
+    for (start, end) in slices {
+        builder.append_packed_range(start + offset..end + offset, src)
+    }
+    builder.into()
 }
 
 #[cfg(test)]
 mod test {
     use itertools::Itertools;
 
-    use crate::array::bool::compute::filter::{
-        filter_select_bool, filter_select_bool_by_index, filter_select_bool_by_slice,
-    };
+    use crate::array::bool::compute::filter::{filter_indices, filter_slices};
     use crate::array::BoolArray;
-    use crate::compute::FilterMask;
+    use crate::compute::{filter, FilterMask};
+    use crate::{IntoArrayData, IntoArrayVariant};
 
     #[test]
     fn filter_bool_test() {
         let arr = BoolArray::from_iter([true, true, false]);
         let mask = FilterMask::from_iter([true, false, true]);
 
-        let filtered = filter_select_bool(&arr, &mask).unwrap();
+        let filtered = filter(&arr.into_array(), mask)
+            .unwrap()
+            .into_bool()
+            .unwrap();
         assert_eq!(2, filtered.len());
 
         assert_eq!(
@@ -73,9 +99,8 @@ mod test {
     #[test]
     fn filter_bool_by_slice_test() {
         let arr = BoolArray::from_iter([true, true, false]);
-        let mask = FilterMask::from_iter([true, false, true]);
 
-        let filtered = filter_select_bool_by_slice(&arr.boolean_buffer(), &mask, 2).unwrap();
+        let filtered = filter_slices(&arr.boolean_buffer(), 2, [(0, 1), (2, 3)].into_iter());
         assert_eq!(2, filtered.len());
 
         assert_eq!(vec![true, false], filtered.iter().collect_vec())
@@ -84,9 +109,8 @@ mod test {
     #[test]
     fn filter_bool_by_index_test() {
         let arr = BoolArray::from_iter([true, true, false]);
-        let mask = FilterMask::from_iter([true, false, true]);
 
-        let filtered = filter_select_bool_by_index(&arr.boolean_buffer(), &mask, 2).unwrap();
+        let filtered = filter_indices(&arr.boolean_buffer(), 2, [0, 2].into_iter());
         assert_eq!(2, filtered.len());
 
         assert_eq!(vec![true, false], filtered.iter().collect_vec())