perf(array): reuse probe state in RLE, RunEnd and PCO

encodings/fastlanes/src/rle/mod.rs

@@ -9,6 +9,7 @@ pub use array::RLESlots;
 
 mod compute;
 mod kernel;
+mod probe;
 
 mod vtable;
 pub use vtable::RLE;

encodings/fastlanes/src/rle/probe.rs

@@ -0,0 +1,86 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! RLE scalar access through the indices, offsets, and values slots.
+
+use vortex_array::ExecutionCtx;
+use vortex_array::ProbeState;
+use vortex_array::scalar::Scalar;
+use vortex_error::VortexResult;
+use vortex_error::vortex_err;
+
+use crate::FL_CHUNK_SIZE;
+use crate::RLE;
+use crate::rle::RLEArrayExt;
+use crate::rle::RLESlots;
+
+/// The slice's base value offset, retained across repeated lookups.
+#[derive(Default)]
+pub struct RleProbeState {
+    base: Option<usize>,
+}
+
+pub(crate) fn scalar_at(
+    state: &mut ProbeState<'_, RLE>,
+    index: usize,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<Scalar> {
+    let array = state.array();
+    let logical_index = array.offset() + index;
+    // Hold the retained state while reading children: the base offset is cached across rows.
+    let (mut retained, mut children) = state.split();
+    let code = children
+        .slot(RLESlots::INDICES)?
+        .ok_or_else(|| vortex_err!("RLE indices slot is missing"))?
+        .execute_scalar(logical_index, ctx)?;
+    let Some(code) = code.as_primitive().as_::<usize>() else {
+        return Ok(Scalar::null(array.dtype().clone()));
+    };
+
+    let chunk = logical_index / FL_CHUNK_SIZE;
+    let offset = if chunk == 0 {
+        0
+    } else {
+        let base = match retained.as_deref().and_then(|state| state.base) {
+            Some(base) => base,
+            None => {
+                let value = read_offset(
+                    children
+                        .slot(RLESlots::VALUES_IDX_OFFSETS)?
+                        .ok_or_else(|| vortex_err!("RLE offsets slot is missing"))?
+                        .execute_scalar(0, ctx)?,
+                )?;
+                if let Some(state) = retained.as_mut() {
+                    state.base = Some(value);
+                }
+                value
+            }
+        };
+        read_offset(
+            children
+                .slot(RLESlots::VALUES_IDX_OFFSETS)?
+                .ok_or_else(|| vortex_err!("RLE offsets slot is missing"))?
+                .execute_scalar(chunk, ctx)?,
+        )?
+        .checked_sub(base)
+        .ok_or_else(|| vortex_err!("RLE offsets precede the slice base"))?
+    };
+    let value_index = offset
+        .checked_add(code)
+        .ok_or_else(|| vortex_err!("RLE value index overflow"))?;
+    let scalar = children
+        .slot(RLESlots::VALUES)?
+        .ok_or_else(|| vortex_err!("RLE values slot is missing"))?
+        .execute_scalar(value_index, ctx)?;
+    Scalar::try_new(array.dtype().clone(), scalar.into_value())
+}
+
+fn read_offset(scalar: Scalar) -> VortexResult<usize> {
+    scalar
+        .as_primitive()
+        .as_::<usize>()
+        .ok_or_else(|| vortex_err!("RLE offset must be a non-null usize"))
+}
+
+#[cfg(test)]
+mod tests;

encodings/fastlanes/src/rle/probe/tests.rs

@@ -0,0 +1,94 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use rstest::rstest;
+use vortex_array::ArrayProbe;
+use vortex_array::ArrayRef;
+use vortex_array::IntoArray;
+use vortex_array::RepeatedArrayProbe;
+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::builtins::ArrayBuiltins;
+use vortex_array::scalar_fn::fns::operators::Operator;
+use vortex_array::validity::Validity;
+use vortex_error::VortexResult;
+
+use crate::RLE;
+use crate::RLEData;
+
+/// A one-off or retained probe over `array`, the retained one living in `retained`.
+fn probe_for<'a>(
+    array: &'a ArrayRef,
+    retained: &'a mut Option<RepeatedArrayProbe>,
+    repeated: bool,
+) -> ArrayProbe<'a> {
+    if repeated {
+        retained.insert(array.repeated_probe()).as_probe()
+    } else {
+        array.probe()
+    }
+}
+
+#[rstest]
+fn random_access_across_chunks_and_nulls(
+    #[values(false, true)] repeated: bool,
+    #[values(false, true)] sliced: bool,
+) -> VortexResult<()> {
+    let mut ctx = crate::test::SESSION.create_execution_ctx();
+    let input =
+        PrimitiveArray::from_option_iter((0..8192u32).map(|i| (i % 11 != 0).then_some(i / 16)));
+    let encoded = RLEData::encode(input.as_view(), &mut ctx)?.into_array();
+    let range = if sliced { 1777..7333 } else { 0..8192 };
+    let source = if sliced {
+        encoded
+            .slice(range.clone())?
+            .execute::<ArrayRef>(&mut ctx)?
+    } else {
+        encoded
+    };
+    assert!(source.is::<RLE>());
+    let input = input.slice(range)?;
+    let indices = [0u32, 1, 1023, 1024, 2048, 2047, 4097, 11, 33, 17, 0];
+    let mut actual = builder_with_capacity_in(source.dtype(), indices.len(), ctx.allocator());
+    let mut retained = None;
+    let mut probe = probe_for(&source, &mut retained, repeated);
+    for index in indices {
+        actual.append_scalar(&probe.execute_scalar(index as usize, &mut ctx)?)?;
+    }
+    assert_arrays_eq!(
+        actual.finish(),
+        input.take(PrimitiveArray::from_iter(indices).into_array())?,
+        &mut ctx
+    );
+    assert!(probe.execute_scalar(source.len(), &mut ctx).is_err());
+    Ok(())
+}
+
+#[rstest]
+fn lazy_validity_does_not_evaluate_unrequested_rows(
+    #[values(false, true)] repeated: bool,
+) -> VortexResult<()> {
+    let mut ctx = crate::test::SESSION.create_execution_ctx();
+    let numerators = PrimitiveArray::from_iter(vec![1u32; 1024]).into_array();
+    let denominators =
+        PrimitiveArray::from_iter((0..1024).map(|i| u32::from(i != 1023))).into_array();
+    let validity = numerators
+        .binary(denominators, Operator::Div)?
+        .binary(numerators, Operator::Eq)?;
+    let array = RLE::try_new(
+        PrimitiveArray::from_iter([42u32]).into_array(),
+        PrimitiveArray::new(vec![0u16; 1024], Validity::Array(validity)).into_array(),
+        PrimitiveArray::from_iter([0u64]).into_array(),
+        0,
+        1024,
+    )?
+    .into_array();
+    let expected = array.execute_scalar(0, &mut ctx)?;
+    let mut retained = None;
+    let mut probe = probe_for(&array, &mut retained, repeated);
+    assert_eq!(probe.execute_scalar(0, &mut ctx)?, expected);
+    assert!(probe.execute_scalar(1023, &mut ctx).is_err());
+    Ok(())
+}

encodings/fastlanes/src/rle/vtable/operations.rs

@@ -3,42 +3,32 @@
 
 use vortex_array::ArrayView;
 use vortex_array::ExecutionCtx;
+use vortex_array::ProbeState;
 use vortex_array::scalar::Scalar;
 use vortex_array::vtable::OperationsVTable;
-use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 
 use super::RLE;
-use crate::FL_CHUNK_SIZE;
-use crate::rle::RLEArrayExt;
-use crate::rle::RLEArraySlotsExt;
+use crate::rle::probe;
+use crate::rle::probe::RleProbeState;
 
 impl OperationsVTable<RLE> for RLE {
-    type ProbeState = ();
+    type ProbeState = RleProbeState;
+
+    fn probe_scalar(
+        state: &mut ProbeState<'_, RLE>,
+        index: usize,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Scalar> {
+        probe::scalar_at(state, index, ctx)
+    }
 
     fn scalar_at(
         array: ArrayView<'_, RLE>,
         index: usize,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<Scalar> {
-        let offset_in_chunk = array.offset();
-        let chunk_relative_idx = array
-            .indices()
-            .execute_scalar(offset_in_chunk + index, ctx)?;
-
-        let chunk_relative_idx = chunk_relative_idx
-            .as_primitive()
-            .as_::<usize>()
-            .vortex_expect("Index must not be null");
-
-        let chunk_id = (offset_in_chunk + index) / FL_CHUNK_SIZE;
-        let value_idx_offset = array.values_idx_offset(chunk_id, ctx);
-
-        let scalar = array
-            .values()
-            .execute_scalar(value_idx_offset + chunk_relative_idx, ctx)?;
-
-        Scalar::try_new(array.dtype().clone(), scalar.into_value())
+        probe::scalar_at(&mut ProbeState::once(array), index, ctx)
     }
 }
 
@@ -54,9 +44,9 @@ mod tests {
     use vortex_array::validity::Validity;
     use vortex_buffer::Buffer;
     use vortex_buffer::buffer;
+    use vortex_error::VortexExpect;
     use vortex_session::VortexSession;
 
-    use super::*;
     use crate::RLE;
     use crate::RLEArray;
     use crate::RLEData;

encodings/pco/benches/scalar.rs

@@ -1,8 +1,9 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-//! Scalar reads out of a PCO array. Cases are `(access_count, nullable, scattered)`; clustered
-//! indices stay inside one page so a retained decode can be reused, scattered ones cross pages.
+//! Scalar reads out of a PCO array through one repeated probe. Cases are
+//! `(access_count, nullable, scattered)`; clustered indices stay inside one page so the retained
+//! decode is reused, scattered ones cross pages.
 
 use std::sync::LazyLock;
 
@@ -65,11 +66,17 @@ fn scalar_access(bencher: Bencher, (count, nullable, scattered): (usize, bool, b
     let array = pco(nullable);
     let indices = indices(count, scattered);
     bencher
-        .with_inputs(|| (SESSION.create_execution_ctx(), Vec::with_capacity(count)))
-        .bench_refs(|(ctx, scalars)| {
+        .with_inputs(|| {
+            (
+                SESSION.create_execution_ctx(),
+                array.repeated_probe(),
+                Vec::with_capacity(count),
+            )
+        })
+        .bench_refs(|(ctx, probe, scalars)| {
             for &index in &indices {
                 scalars.push(
-                    array
+                    probe
                         .execute_scalar(index, ctx)
                         .vortex_expect("scalar access"),
                 );

encodings/pco/src/array.rs

@@ -29,6 +29,7 @@ use vortex_array::EqMode;
 use vortex_array::ExecutionCtx;
 use vortex_array::ExecutionResult;
 use vortex_array::IntoArray;
+use vortex_array::ProbeState;
 use vortex_array::TypedArrayRef;
 use vortex_array::array_slots;
 use vortex_array::arrays::Primitive;
@@ -778,19 +779,23 @@ impl ValidityVTable<Pco> for Pco {
 }
 
 impl OperationsVTable<Pco> for Pco {
-    type ProbeState = ();
+    type ProbeState = crate::probe::PcoProbeState;
+
+    fn probe_scalar(
+        state: &mut ProbeState<'_, Pco>,
+        index: usize,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Scalar> {
+        let array = state.array();
+        crate::probe::scalar_at(array, index, state.retained(), ctx)
+    }
 
     fn scalar_at(
         array: ArrayView<'_, Pco>,
         index: usize,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<Scalar> {
-        let unsliced_validity = array.unsliced_validity();
-        array
-            ._slice(index, index + 1)
-            .decompress(&unsliced_validity, ctx)?
-            .into_array()
-            .execute_scalar(0, ctx)
+        crate::probe::scalar_at(array, index, None, ctx)
     }
 }
 

encodings/pco/src/lib.rs

@@ -20,6 +20,7 @@
 
 mod array;
 mod compute;
+mod probe;
 mod rules;
 mod slice;