Filter on lists (#5775)

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

vortex-array/src/arrays/dict/vtable/rules.rs

@@ -13,18 +13,45 @@ use crate::arrays::ConstantArray;
 use crate::arrays::ConstantVTable;
 use crate::arrays::DictArray;
 use crate::arrays::DictVTable;
+use crate::arrays::FilterArray;
+use crate::arrays::FilterVTable;
 use crate::arrays::ScalarFnArray;
 use crate::builtins::ArrayBuiltins;
 use crate::expr::Pack;
+use crate::matchers::Exact;
 use crate::optimizer::ArrayOptimizer;
 use crate::optimizer::rules::ArrayParentReduceRule;
 use crate::optimizer::rules::ParentRuleSet;
 
 pub(super) const PARENT_RULES: ParentRuleSet<DictVTable> = ParentRuleSet::new(&[
+    ParentRuleSet::lift(&DictionaryFilterPushDownRule),
     ParentRuleSet::lift(&DictionaryScalarFnValuesPushDownRule),
     ParentRuleSet::lift(&DictionaryScalarFnCodesPullUpRule),
 ]);
 
+#[derive(Debug)]
+struct DictionaryFilterPushDownRule;
+
+impl ArrayParentReduceRule<DictVTable> for DictionaryFilterPushDownRule {
+    type Parent = Exact<FilterVTable>;
+
+    fn parent(&self) -> Self::Parent {
+        Exact::from(&FilterVTable)
+    }
+
+    fn reduce_parent(
+        &self,
+        array: &DictArray,
+        parent: &FilterArray,
+        _child_idx: usize,
+    ) -> VortexResult<Option<ArrayRef>> {
+        let new_codes = array.codes().filter(parent.filter_mask().clone())?;
+        let new_dict =
+            unsafe { DictArray::new_unchecked(new_codes, array.values().clone()) }.into_array();
+        Ok(Some(new_dict))
+    }
+}
+
 /// Push down a scalar function to run only over the values of a dictionary array.
 #[derive(Debug)]
 struct DictionaryScalarFnValuesPushDownRule;

vortex-array/src/arrays/filter/rules.rs

@@ -3,6 +3,7 @@
 
 use vortex_error::VortexResult;
 
+use crate::Array;
 use crate::ArrayRef;
 use crate::IntoArray;
 use crate::arrays::FilterArray;
@@ -32,7 +33,7 @@ impl ArrayParentReduceRule<FilterVTable> for FilterFilterRule {
         _child_idx: usize,
     ) -> VortexResult<Option<ArrayRef>> {
         let combined_mask = child.mask.intersect_by_rank(&parent.mask);
-        let new_array = FilterArray::new(child.child.clone(), combined_mask);
+        let new_array = child.child.filter(combined_mask)?;
         Ok(Some(new_array.into_array()))
     }
 }

vortex-array/src/arrays/list/compute/filter.rs

@@ -1,6 +1,8 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::sync::Arc;
+
 use vortex_buffer::BitBufferMut;
 use vortex_buffer::BufferMut;
 use vortex_dtype::IntegerPType;
@@ -9,19 +11,18 @@ use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_mask::Mask;
 use vortex_mask::MaskIter;
+use vortex_mask::MaskValues;
 
 use crate::Array;
 use crate::ArrayRef;
 use crate::IntoArray;
 use crate::ToCanonical;
 use crate::arrays::ListArray;
 use crate::arrays::ListVTable;
-use crate::arrays::PrimitiveArray;
 use crate::compute::FilterKernel;
 use crate::compute::FilterKernelAdapter;
 use crate::compute::filter;
 use crate::register_kernel;
-use crate::validity::Validity;
 use crate::vtable::ValidityHelper;
 
 /// Density threshold for choosing between indices and slices representation when expanding masks.
@@ -44,21 +45,22 @@ impl FilterKernel for ListVTable {
         );
 
         let (new_elements, new_offsets) = match_each_integer_ptype!(offsets.ptype(), |O| {
-            compute_filtered_elements_and_offsets::<O>(
+            let (new_elements, new_offsets) = compute_filtered_elements_and_offsets::<O>(
                 elements.as_ref(),
                 offsets.as_slice::<O>(),
                 selection_mask,
-            )?
+            )?;
+            (new_elements, new_offsets.into_array())
         });
 
         // SAFETY: Filter operation maintains all ListArray invariants:
         // - Offsets are monotonically increasing (built correctly above).
         // - Elements are properly filtered to match the offsets.
         // - Validity matches the original array's nullability.
-        Ok(unsafe {
-            ListArray::new_unchecked(new_elements, new_offsets.into_array(), new_validity)
-        }
-        .into_array())
+        Ok(
+            unsafe { ListArray::new_unchecked(new_elements, new_offsets, new_validity) }
+                .into_array(),
+        )
     }
 }
 
@@ -74,7 +76,7 @@ fn compute_filtered_elements_and_offsets<O: IntegerPType>(
     elements: &dyn Array,
     offsets: &[O],
     selection_mask: &Mask,
-) -> VortexResult<(ArrayRef, PrimitiveArray)> {
+) -> VortexResult<(ArrayRef, BufferMut<O>)> {
     let values = selection_mask
         .values()
         .vortex_expect("`AllTrue` and `AllFalse` are handled by filter entry point");
@@ -132,11 +134,50 @@ fn compute_filtered_elements_and_offsets<O: IntegerPType>(
     // The `Mask` can determine the best representation based on the buffer's density in the future.
     let new_elements = filter(elements, &Mask::from_buffer(new_mask_builder.freeze()))?;
 
-    let new_offsets = PrimitiveArray::new(new_offsets, Validity::NonNullable);
-
     Ok((new_elements, new_offsets))
 }
 
+/// Construct an element mask from contiguous list offsets and a selection mask.
+pub fn element_mask_from_offsets<O: IntegerPType>(
+    offsets: &[O],
+    selection: &Arc<MaskValues>,
+) -> Mask {
+    let first_offset = offsets.first().map_or(0, |first_offset| first_offset.as_());
+    let last_offset = offsets.last().map_or(0, |last_offset| last_offset.as_());
+    let len = last_offset - first_offset;
+
+    let mut mask_builder = BitBufferMut::with_capacity(len);
+
+    match selection.threshold_iter(MASK_EXPANSION_DENSITY_THRESHOLD) {
+        MaskIter::Slices(slices) => {
+            // Dense iteration: process ranges of consecutive selected lists.
+            for &(start, end) in slices {
+                // Optimization: for dense ranges, we can process the elements mask more efficiently.
+                let elems_start = offsets[start].as_() - first_offset;
+                let elems_end = offsets[end].as_() - first_offset;
+
+                // Process the entire range of elements at once.
+                process_element_range(elems_start, elems_end, &mut mask_builder);
+            }
+        }
+        MaskIter::Indices(indices) => {
+            // Sparse iteration: process individual selected lists.
+            for &idx in indices {
+                let list_start = offsets[idx].as_() - first_offset;
+                let list_end = offsets[idx + 1].as_() - first_offset;
+
+                // Process the elements for this list.
+                process_element_range(list_start, list_end, &mut mask_builder);
+            }
+        }
+    }
+
+    // Pad to full length if necessary.
+    mask_builder.append_n(false, last_offset - mask_builder.len());
+
+    Mask::from_buffer(mask_builder.freeze())
+}
+
 /// Process a range of elements for filtering.
 fn process_element_range(
     elems_start: usize,

vortex-array/src/arrays/list/compute/mod.rs

@@ -9,6 +9,8 @@ mod mask;
 mod min_max;
 mod take;
 
+pub(super) use filter::element_mask_from_offsets;
+
 #[cfg(test)]
 mod tests {
     use rstest::rstest;

vortex-array/src/arrays/list/vtable/kernel/filter.rs

@@ -0,0 +1,125 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::sync::Arc;
+
+use vortex_buffer::BufferMut;
+use vortex_dtype::PTypeDowncastExt;
+use vortex_dtype::match_each_integer_ptype;
+use vortex_error::VortexResult;
+use vortex_mask::Mask;
+use vortex_vector::Vector;
+use vortex_vector::VectorMutOps;
+use vortex_vector::listview::ListViewVector;
+use vortex_vector::listview::ListViewVectorMut;
+use vortex_vector::primitive::PVector;
+use vortex_vector::primitive::PrimitiveVector;
+
+use crate::Array;
+use crate::ExecutionCtx;
+use crate::VectorExecutor;
+use crate::arrays::FilterArray;
+use crate::arrays::FilterVTable;
+use crate::arrays::ListArray;
+use crate::arrays::ListVTable;
+use crate::arrays::list::compute::element_mask_from_offsets;
+use crate::kernel::ExecuteParentKernel;
+use crate::mask::MaskExecutor;
+use crate::matchers::Exact;
+use crate::validity::Validity;
+use crate::vtable::ValidityHelper;
+
+#[derive(Debug)]
+pub(super) struct ListFilterKernel;
+
+impl ExecuteParentKernel<ListVTable> for ListFilterKernel {
+    type Parent = Exact<FilterVTable>;
+
+    fn parent(&self) -> Self::Parent {
+        Exact::from(&FilterVTable)
+    }
+
+    fn execute_parent(
+        &self,
+        array: &ListArray,
+        parent: &FilterArray,
+        _child_idx: usize,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<Vector>> {
+        let selection = match parent.filter_mask() {
+            Mask::AllTrue(_) | Mask::AllFalse(_) => return Ok(None),
+            Mask::Values(v) => v,
+        };
+
+        // TODO(ngates): for ultra-sparse masks, we don't need to optimize the entire offsets.
+        let offsets = array
+            .offsets()
+            .execute_vector(ctx.session())?
+            .into_primitive();
+
+        let new_validity = match array.validity() {
+            Validity::NonNullable | Validity::AllValid => Mask::new_true(selection.true_count()),
+            Validity::AllInvalid => {
+                let mut vec = ListViewVectorMut::with_capacity(
+                    array.elements().dtype(),
+                    selection.true_count(),
+                );
+                vec.append_nulls(selection.true_count());
+                return Ok(Some(vec.freeze().into()));
+            }
+            Validity::Array(a) => a
+                .filter(parent.filter_mask().clone())?
+                .execute_mask(ctx.session())?,
+        };
+
+        let (new_offsets, new_sizes) = match_each_integer_ptype!(offsets.ptype(), |O| {
+            let offsets = (&offsets).downcast::<O>().elements().as_slice();
+            let mut new_offsets = BufferMut::<O>::with_capacity(selection.true_count());
+            let mut new_sizes = BufferMut::<O>::with_capacity(selection.true_count());
+
+            let mut offset = 0;
+            for idx in selection.indices() {
+                let start = offsets[*idx];
+                let end = offsets[idx + 1];
+                let size = end - start;
+                unsafe { new_offsets.push_unchecked(offset) };
+                unsafe { new_sizes.push_unchecked(size) };
+                offset += size;
+            }
+
+            let new_offsets = PrimitiveVector::from(PVector::<O>::new(
+                new_offsets.freeze(),
+                Mask::new_true(selection.true_count()),
+            ));
+            let new_sizes = PrimitiveVector::from(PVector::<O>::new(
+                new_sizes.freeze(),
+                Mask::new_true(selection.true_count()),
+            ));
+
+            (new_offsets, new_sizes)
+        });
+
+        // TODO(ngates): for very dense masks, there may be no point in filtering the elements,
+        //  and instead we should construct a view against the unfiltered elements.
+        let element_mask = match_each_integer_ptype!(offsets.ptype(), |O| {
+            element_mask_from_offsets::<O>((&offsets).downcast::<O>().elements(), selection)
+        });
+
+        let new_elements = array
+            .sliced_elements()
+            .filter(element_mask)?
+            .execute_vector(ctx.session())?;
+
+        Ok(Some(
+            unsafe {
+                ListViewVector::new_unchecked(
+                    Arc::new(new_elements),
+                    new_offsets,
+                    new_sizes,
+                    new_validity,
+                )
+            }
+            .into(),
+        ))
+    }
+}

vortex-array/src/arrays/list/vtable/kernel/mod.rs

@@ -0,0 +1,11 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+mod filter;
+
+use crate::arrays::ListVTable;
+use crate::arrays::list::vtable::kernel::filter::ListFilterKernel;
+use crate::kernel::ParentKernelSet;
+
+pub(super) const PARENT_KERNELS: ParentKernelSet<ListVTable> =
+    ParentKernelSet::new(&[ParentKernelSet::lift(&ListFilterKernel)]);

vortex-array/src/arrays/list/vtable/mod.rs

@@ -9,10 +9,13 @@ use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure;
+use vortex_vector::Vector;
 
 use crate::ArrayRef;
+use crate::ExecutionCtx;
 use crate::ProstMetadata;
 use crate::arrays::ListArray;
+use crate::arrays::list::vtable::kernel::PARENT_KERNELS;
 use crate::metadata::DeserializeMetadata;
 use crate::metadata::SerializeMetadata;
 use crate::serde::ArrayChildren;
@@ -27,6 +30,7 @@ use crate::vtable::ValidityVTableFromValidityHelper;
 
 mod array;
 mod canonical;
+mod kernel;
 mod operations;
 mod validity;
 mod visitor;
@@ -138,6 +142,15 @@ impl VTable for ListVTable {
         *array = new_array;
         Ok(())
     }
+
+    fn execute_parent(
+        array: &Self::Array,
+        parent: &ArrayRef,
+        child_idx: usize,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<Vector>> {
+        PARENT_KERNELS.execute(array, parent, child_idx, ctx)
+    }
 }
 
 #[derive(Debug)]

vortex-array/src/arrays/listview/vtable/mod.rs

@@ -20,6 +20,7 @@ use crate::ProstMetadata;
 use crate::SerializeMetadata;
 use crate::VectorExecutor;
 use crate::arrays::ListViewArray;
+use crate::arrays::listview::vtable::rules::PARENT_RULES;
 use crate::executor::ExecutionCtx;
 use crate::serde::ArrayChildren;
 use crate::validity::Validity;
@@ -34,6 +35,7 @@ use crate::vtable::ValidityVTableFromValidityHelper;
 mod array;
 mod canonical;
 mod operations;
+mod rules;
 mod validity;
 mod visitor;
 
@@ -187,4 +189,12 @@ impl VTable for ListViewVTable {
         }
         .into())
     }
+
+    fn reduce_parent(
+        array: &Self::Array,
+        parent: &ArrayRef,
+        child_idx: usize,
+    ) -> VortexResult<Option<ArrayRef>> {
+        PARENT_RULES.evaluate(array, parent, child_idx)
+    }
 }