perf: speed up runtime pruning and run-end execution

encodings/pco/src/probe/tests.rs

@@ -10,11 +10,51 @@ use vortex_array::VortexSessionExecute;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::assert_arrays_eq;
 use vortex_array::builders::builder_with_capacity_in;
+use vortex_array::search_sorted::SearchResult;
+use vortex_array::search_sorted::SearchSorted;
+use vortex_array::search_sorted::SearchSortedPrimitiveArray;
+use vortex_array::search_sorted::SearchSortedSide;
 use vortex_array::validity::Validity;
 use vortex_error::VortexResult;
 
 use crate::Pco;
 
+#[rstest]
+fn search_sorted_across_pages(
+    #[values(false, true)] nullable: bool,
+    #[values(false, true)] sliced: bool,
+) -> VortexResult<()> {
+    let mut ctx = vortex_array::array_session().create_execution_ctx();
+    let expected: Vec<Option<u32>> = (0..4096u32)
+        .map(|i| (!nullable || i >= 32).then_some(i / 3))
+        .collect();
+    let input = if nullable {
+        PrimitiveArray::from_option_iter(expected.iter().copied())
+    } else {
+        PrimitiveArray::from_iter((0..4096u32).map(|i| i / 3))
+    };
+    let encoded = Pco::from_primitive(input.as_view(), 3, 128, &mut ctx)?.into_array();
+    let range = if sliced { 17..3333 } else { 0..4096 };
+    let source = encoded.slice(range.clone())?;
+    let expected = &expected[range];
+    let searcher = SearchSortedPrimitiveArray::<u32>::new_repeated(&source, &mut ctx);
+    for needle in [None, Some(0), Some(511), Some(1365), Some(100), Some(2048)] {
+        for side in [SearchSortedSide::Left, SearchSortedSide::Right] {
+            let index = expected.partition_point(|value| match side {
+                SearchSortedSide::Left => value < &needle,
+                SearchSortedSide::Right => value <= &needle,
+            });
+            let result = if expected.contains(&needle) {
+                SearchResult::Found(index)
+            } else {
+                SearchResult::NotFound(index)
+            };
+            assert_eq!(searcher.search_sorted(&needle, side)?, result);
+        }
+    }
+    Ok(())
+}
+
 /// A one-off or retained probe over `array`, the retained one living in `retained`.
 fn probe_for<'a>(
     array: &'a ArrayRef,

encodings/runend/benches/run_end_decode.rs

@@ -11,9 +11,16 @@ use mimalloc::MiMalloc;
 use vortex_array::VortexSessionExecute;
 use vortex_array::arrays::BoolArray;
 use vortex_array::arrays::PrimitiveArray;
+use vortex_array::dtype::NativePType;
+use vortex_array::dtype::Nullability;
 use vortex_array::validity::Validity;
 use vortex_buffer::BitBuffer;
+use vortex_buffer::Buffer;
 use vortex_buffer::BufferMut;
+use vortex_error::VortexExpect;
+use vortex_mask::Mask;
+use vortex_runend::compress::runend_decode_primitive;
+use vortex_runend::compress::runend_decode_typed_primitive;
 use vortex_runend::decompress_bool::runend_decode_bools;
 use vortex_session::VortexSession;
 
@@ -389,3 +396,101 @@ fn decode_bool_nullable(bencher: Bencher, args: NullableBoolBenchArgs) {
             runend_decode_bools(ends.clone(), values.clone(), 0, total_length, ctx)
         });
 }
+
+#[derive(Clone, Copy)]
+struct PrimitiveBenchArgs {
+    run_length: usize,
+    valid_percent: usize,
+    validity_offset: usize,
+}
+
+impl fmt::Display for PrimitiveBenchArgs {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "run_{}_valid_{}_offset_{}",
+            self.run_length, self.valid_percent, self.validity_offset
+        )
+    }
+}
+
+const PRIMITIVE_ARGS: [PrimitiveBenchArgs; 30] = {
+    let mut args = [PrimitiveBenchArgs {
+        run_length: 1,
+        valid_percent: 99,
+        validity_offset: 0,
+    }; 30];
+    let lengths = [1, 2, 16, 128, 4096];
+    let validities = [10, 50, 99];
+    let mut i = 0;
+    while i < args.len() {
+        args[i] = PrimitiveBenchArgs {
+            run_length: lengths[i / 6],
+            valid_percent: validities[(i / 2) % 3],
+            validity_offset: (i % 2) * 3,
+        };
+        i += 1;
+    }
+    args
+};
+
+#[divan::bench(types = [u16, i64], args = PRIMITIVE_ARGS)]
+fn decode_primitive_nullable<T: NativePType>(bencher: Bencher, args: PrimitiveBenchArgs) {
+    let length = 65_536usize;
+    let run_count = length.div_ceil(args.run_length);
+    let values: Vec<T> = (0..run_count)
+        .map(|i| T::from_usize(i % 1024).vortex_expect("benchmark value fits"))
+        .collect();
+    let validity = BitBuffer::from_iter(
+        (0..run_count + args.validity_offset).map(|i| (i * 37) % 100 < args.valid_percent),
+    )
+    .slice(args.validity_offset..run_count + args.validity_offset);
+    let validity = Mask::from(validity);
+    bencher.bench(|| {
+        runend_decode_typed_primitive(
+            (1..=run_count).map(|i| (i * args.run_length).min(length)),
+            &values,
+            validity.clone(),
+            Nullability::Nullable,
+            length,
+        )
+    });
+}
+
+#[divan::bench(types = [u16, i64], args = [1usize, 2, 16, 128, 4096])]
+fn decode_primitive_non_nullable<T: NativePType>(bencher: Bencher, run_length: usize) {
+    let length = 65_536usize;
+    let run_count = length.div_ceil(run_length);
+    let values: Vec<T> = (0..run_count)
+        .map(|i| T::from_usize(i % 1024).vortex_expect("benchmark value fits"))
+        .collect();
+    bencher.bench(|| {
+        runend_decode_typed_primitive(
+            (1..=run_count).map(|i| (i * run_length).min(length)),
+            &values,
+            Mask::AllTrue(run_count),
+            Nullability::NonNullable,
+            length,
+        )
+    });
+}
+
+#[divan::bench(types = [u16, i64], args = [1usize, 2, 16, 128, 4096])]
+fn decode_primitive_array<T: NativePType>(bencher: Bencher, run_length: usize) {
+    let length = 65_536usize;
+    let run_count = length.div_ceil(run_length);
+    let ends =
+        PrimitiveArray::from_iter((1..=run_count).map(|i| (i * run_length).min(length) as u32));
+    let values = PrimitiveArray::new(
+        Buffer::from_iter(
+            (0..run_count).map(|i| T::from_usize(i % 1024).vortex_expect("benchmark value fits")),
+        ),
+        Validity::NonNullable,
+    );
+    bencher
+        .with_inputs(|| SESSION.create_execution_ctx())
+        .bench_refs(|ctx| {
+            runend_decode_primitive(ends.clone(), values.clone(), 0, length, ctx)
+                .vortex_expect("benchmark decode succeeds")
+        });
+}

encodings/runend/benches/run_end_filter.rs

@@ -14,14 +14,21 @@
 #![expect(clippy::expect_used)]
 
 use std::fmt;
+use std::sync::LazyLock;
 
 use divan::Bencher;
 use mimalloc::MiMalloc;
 use rand::SeedableRng;
 use rand::rngs::StdRng;
 use rand::seq::SliceRandom;
+use vortex_array::IntoArray;
+use vortex_array::VortexSessionExecute;
+use vortex_array::arrays::PrimitiveArray;
 use vortex_buffer::BitBuffer;
+use vortex_mask::Mask;
 use vortex_runend::_benchmarking::filter_run_end_primitive;
+use vortex_runend::RunEnd;
+use vortex_session::VortexSession;
 
 #[global_allocator]
 static GLOBAL: MiMalloc = MiMalloc;
@@ -30,6 +37,12 @@ fn main() {
     divan::main();
 }
 
+static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+    let session = vortex_array::array_session();
+    vortex_runend::initialize(&session);
+    session
+});
+
 #[derive(Clone, Copy)]
 struct FilterBenchArgs {
     /// Total logical length of the decoded array.
@@ -38,50 +51,41 @@ struct FilterBenchArgs {
     run_length: usize,
     /// Fraction of mask bits that are set to `true`.
     density: f64,
+    /// Bit offset of a sliced predicate mask.
+    mask_offset: usize,
 }
 
 impl fmt::Display for FilterBenchArgs {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(
             f,
-            "len={}_run={}_density={:.1}",
-            self.length, self.run_length, self.density
+            "len={}_run={}_density={:.1}_offset={}",
+            self.length, self.run_length, self.density, self.mask_offset
         )
     }
 }
 
-const FILTER_ARGS: &[FilterBenchArgs] = &[
-    FilterBenchArgs {
-        length: 4_096,
-        run_length: 16,
-        density: 0.1,
-    },
-    FilterBenchArgs {
-        length: 4_096,
-        run_length: 16,
-        density: 0.5,
-    },
-    FilterBenchArgs {
-        length: 4_096,
-        run_length: 16,
-        density: 0.9,
-    },
-    FilterBenchArgs {
+const FILTER_ARGS: [FilterBenchArgs; 36] = {
+    let mut args = [FilterBenchArgs {
         length: 16_384,
-        run_length: 16,
+        run_length: 1,
         density: 0.1,
-    },
-    FilterBenchArgs {
-        length: 16_384,
-        run_length: 16,
-        density: 0.5,
-    },
-    FilterBenchArgs {
-        length: 16_384,
-        run_length: 16,
-        density: 0.9,
-    },
-];
+        mask_offset: 0,
+    }; 36];
+    let lengths = [1, 4, 16, 64, 256, 1024];
+    let densities = [0.1, 0.5, 0.9];
+    let mut i = 0;
+    while i < args.len() {
+        args[i] = FilterBenchArgs {
+            length: 16_384,
+            run_length: lengths[i / 6],
+            density: densities[(i / 2) % 3],
+            mask_offset: (i % 2) * 3,
+        };
+        i += 1;
+    }
+    args
+};
 
 /// Build the run-end boundaries (cumulative run lengths) for `length` rows.
 fn build_run_ends(length: usize, run_length: usize) -> Vec<u32> {
@@ -107,11 +111,37 @@ fn build_mask(length: usize, density: f64) -> BitBuffer {
 #[divan::bench(args = FILTER_ARGS)]
 fn filter_run_end(bencher: Bencher, args: FilterBenchArgs) {
     let run_ends = build_run_ends(args.length, args.run_length);
-    let mask = build_mask(args.length, args.density);
+    let mask = build_mask(args.length + args.mask_offset, args.density)
+        .slice(args.mask_offset..args.length + args.mask_offset);
     let length = args.length as u64;
     bencher
         .with_inputs(|| (run_ends.clone(), mask.clone()))
         .bench_refs(|(run_ends, mask)| {
             filter_run_end_primitive::<u32>(run_ends, 0, length, mask).expect("filter")
         });
 }
+
+#[divan::bench(args = FILTER_ARGS)]
+fn filter_run_end_array(bencher: Bencher, args: FilterBenchArgs) {
+    let ends = build_run_ends(args.length, args.run_length);
+    let values = PrimitiveArray::from_iter(0..ends.len() as i64).into_array();
+    let array = RunEnd::new(
+        PrimitiveArray::from_iter(ends).into_array(),
+        values,
+        &mut SESSION.create_execution_ctx(),
+    )
+    .into_array();
+    let mask = Mask::from(
+        build_mask(args.length + args.mask_offset, args.density)
+            .slice(args.mask_offset..args.length + args.mask_offset),
+    );
+    bencher
+        .with_inputs(|| (array.clone(), mask.clone(), SESSION.create_execution_ctx()))
+        .bench_values(|(array, mask, mut ctx)| {
+            array
+                .filter(mask)
+                .expect("filter")
+                .execute::<PrimitiveArray>(&mut ctx)
+                .expect("execute")
+        });
+}