Convert trivial filters into slices during reduction, MaskValues::last uses BitBuffer::last_set_index

vortex-array/src/arrays/filter/execute/mod.rs

@@ -181,6 +181,7 @@ mod tests {
     use super::*;
     use crate::VortexSessionExecute;
     use crate::array_session;
+    use crate::arrays::Primitive;
     use crate::arrays::PrimitiveArray;
 
     #[test]
@@ -194,10 +195,10 @@ mod tests {
             values.cached_indices().is_none() && values.cached_slices().is_none()
         }));
 
-        let filtered = array
-            .into_array()
-            .filter(mask)?
-            .execute::<PrimitiveArray>(&mut array_session().create_execution_ctx())?;
+        let filtered = array.into_array().filter(mask)?;
+        assert!(filtered.is::<Primitive>());
+        let filtered =
+            filtered.execute::<PrimitiveArray>(&mut array_session().create_execution_ctx())?;
         let filtered_values = filtered.to_buffer::<i32>();
 
         assert_eq!(filtered_values.as_slice(), &(37..91).collect::<Vec<_>>());

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

@@ -15,6 +15,7 @@ use crate::arrays::filter::FilterArraySlotsExt;
 use crate::arrays::filter::FilterReduce;
 use crate::arrays::filter::FilterReduceAdaptor;
 use crate::arrays::filter::execute::buffer::prepare_mask_for_reuse;
+use crate::arrays::filter::execute::contiguous_filter_range;
 use crate::arrays::scalar_fn::ExactScalarFn;
 use crate::arrays::scalar_fn::ScalarFnArrayView;
 use crate::arrays::struct_::StructDataParts;
@@ -67,7 +68,9 @@ impl ArrayReduceRule<Filter> for TrivialFilterRule {
         match array.filter_mask() {
             Mask::AllTrue(_) => Ok(Some(array.child().clone())),
             Mask::AllFalse(_) => Ok(Some(Canonical::empty(array.dtype()).into_array())),
-            Mask::Values(_) => Ok(None),
+            Mask::Values(_) => contiguous_filter_range(array.filter_mask())
+                .map(|range| array.child().slice(range))
+                .transpose(),
         }
     }
 }

vortex-array/src/test_harness/trace/tests.rs

@@ -40,9 +40,10 @@ use crate::arrays::Filter;
 use crate::arrays::FilterArray;
 use crate::arrays::Primitive;
 use crate::arrays::PrimitiveArray;
+use crate::arrays::Slice;
 use crate::arrays::StructArray;
 use crate::arrays::VarBinViewArray;
-use crate::arrays::filter::FilterArraySlotsExt;
+use crate::arrays::slice::SliceArraySlotsExt;
 use crate::assert_arrays_eq;
 use crate::buffer::BufferHandle;
 use crate::dtype::DType;
@@ -385,7 +386,7 @@ optimize root=vortex.filter(i32, len=4) session=false
 }
 
 #[test]
-fn trace_optimize_parent_reduce_fixpoint_attempts() -> VortexResult<()> {
+fn trace_optimize_contiguous_filter() -> VortexResult<()> {
     let values = PrimitiveArray::from_iter([0i32, 1, 2, 3, 4, 5]).into_array();
     let inner = FilterArray::try_new(
         values,
@@ -402,17 +403,17 @@ fn trace_optimize_parent_reduce_fixpoint_attempts() -> VortexResult<()> {
         || outer.optimize(),
     )?;
 
-    let optimized_filter = traced.output.as_::<Filter>();
-    assert!(optimized_filter.child().is::<Primitive>());
+    let optimized_slice = traced.output.as_::<Slice>();
+    assert!(optimized_slice.child().is::<Filter>());
     assert_arrays_eq!(
         traced.output,
         PrimitiveArray::from_iter([2i32, 3]),
         &mut execution_ctx()
     );
     insta::assert_snapshot!(traced.trace.to_string(), @r"
     optimize root=vortex.filter(i32, len=2) session=false
-      reduce_parent static:FilterReduceAdaptor(Filter) slot=0 parent=vortex.filter(i32, len=2) child=vortex.filter(i32, len=4) -> vortex.filter(i32, len=2)
-      done output=vortex.filter(i32, len=2)
+      reduce TrivialFilterRule: vortex.filter(i32, len=2) -> vortex.slice(i32, len=2)
+      done output=vortex.slice(i32, len=2)
     ");
 
     let mut ctx = ExecutionCtx::new(VortexSession::empty().with::<ArraySession>());
@@ -457,9 +458,6 @@ fn trace_optimize_parent_reduce_fixpoint_attempts() -> VortexResult<()> {
 
 /// A filter whose mask is not one contiguous run cannot be answered as a slice, so it has to
 /// execute through its child.
-///
-/// The test above happens to build a contiguous combined mask, which short-circuits before the
-/// child is reached; without this case no trace would cover an executed filter at all.
 #[test]
 fn trace_execute_filter_with_scattered_mask() -> VortexResult<()> {
     let values = PrimitiveArray::from_iter([0i32, 1, 2, 3, 4, 5]).into_array();

vortex-mask/src/lib.rs

@@ -483,22 +483,7 @@ impl Mask {
                     return slices.last().map(|(_, end)| end - 1);
                 }
 
-                if values.true_count == 0 {
-                    return None;
-                }
-
-                Some(
-                    values
-                        .buffer
-                        .select(values.true_count - 1)
-                        .unwrap_or_else(|| {
-                            vortex_panic!(
-                                "Rank {} out of bounds for mask with true count {}",
-                                values.true_count - 1,
-                                values.true_count
-                            )
-                        }),
-                )
+                values.buffer.last_set_index()
             }
         }
     }

vortex-mask/src/tests.rs

@@ -104,6 +104,22 @@ fn test_mask_first() {
     assert_eq!(values_slices.first(), Some(2));
 }
 
+#[test]
+fn test_mask_last() {
+    assert_eq!(Mask::new_true(5).last(), Some(4));
+    assert_eq!(Mask::new_false(5).last(), None);
+
+    let buffer = BitBuffer::from_iter([true, false, true, false, true, false, true]);
+    let values = Mask::from_buffer(buffer.slice(1..6));
+    assert_eq!(values.last(), Some(3));
+
+    let values_indices = Mask::from_indices(5, vec![1, 3]);
+    assert_eq!(values_indices.last(), Some(3));
+
+    let values_slices = Mask::from_slices(5, vec![(1, 2), (3, 4)]);
+    assert_eq!(values_slices.last(), Some(3));
+}
+
 #[test]
 fn test_mask_false_count() {
     assert_eq!(Mask::new_true(5).false_count(), 0);