perf(array): reuse probe state in primitive and compressed arrays

encodings/fastlanes/src/rle/mod.rs

@@ -10,6 +10,8 @@ pub use array::RLESlots;
 mod compute;
 mod kernel;
 
+mod probe;
+
 mod vtable;
 pub use vtable::RLE;
 pub use vtable::RLEArray;

encodings/fastlanes/src/rle/probe.rs

@@ -0,0 +1,210 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! RLE probes route through the indices, offsets, and values slots.
+//!
+//! A slot that already holds a materialized primitive array is read directly from its typed
+//! buffer. Any other child, including one with a lazy validity expression, is read through a
+//! child probe that keeps its own state, so nesting works to any depth.
+
+use num_traits::ToPrimitive;
+use vortex_array::ArrayRef;
+use vortex_array::ArrayView;
+use vortex_array::ExecutionCtx;
+use vortex_array::ProbeChildren;
+use vortex_array::ProbeCtx;
+use vortex_array::arrays::Bool;
+use vortex_array::arrays::Primitive;
+use vortex_array::arrays::primitive::PrimitiveArrayExt;
+use vortex_array::match_each_native_ptype;
+use vortex_array::match_each_unsigned_integer_ptype;
+use vortex_array::scalar::Scalar;
+use vortex_array::validity::Validity;
+use vortex_error::VortexResult;
+use vortex_error::vortex_ensure;
+use vortex_error::vortex_err;
+use vortex_mask::Mask;
+
+use crate::FL_CHUNK_SIZE;
+use crate::RLE;
+use crate::rle::RLEArrayExt;
+use crate::rle::RLEArraySlotsExt;
+use crate::rle::RLESlots;
+
+/// State for repeated RLE probes: one reader per child slot plus the slice's base value offset.
+#[derive(Default)]
+pub struct RleProbeState<'a> {
+    children: Option<Children<'a>>,
+}
+
+struct Children<'a> {
+    indices: Child<'a>,
+    offsets: Child<'a>,
+    values: Child<'a>,
+    base: usize,
+}
+
+enum Child<'a> {
+    /// A materialized primitive array with its validity resolved once.
+    Primitive {
+        view: ArrayView<'a, Primitive>,
+        validity: Option<Mask>,
+    },
+    /// A slot read through the context's retained child probe.
+    Probe(usize),
+}
+
+impl<'a> Child<'a> {
+    fn new(
+        slots: &'a [Option<ArrayRef>],
+        slot: usize,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Self> {
+        let child = slots[slot]
+            .as_ref()
+            .ok_or_else(|| vortex_err!("RLE slot {slot} is missing"))?;
+        let Some(view) = child.as_opt::<Primitive>() else {
+            return Ok(Self::Probe(slot));
+        };
+        let validity = match PrimitiveArrayExt::validity(&view) {
+            Validity::NonNullable | Validity::AllValid => None,
+            // A lazy validity expression may fail on rows this probe never requests, so the
+            // primitive probe evaluates it per row instead.
+            Validity::Array(array) if !array.is::<Bool>() => {
+                return Ok(Self::Probe(slot));
+            }
+            validity => Some(validity.execute_mask(view.len(), ctx)?),
+        };
+        Ok(Self::Primitive { view, validity })
+    }
+
+    /// Read an unsigned routing value, or `None` when the row is null.
+    fn index_at(
+        &mut self,
+        index: usize,
+        children: &mut ProbeChildren<'a>,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<usize>> {
+        match self {
+            Self::Primitive { view, validity } => {
+                vortex_ensure!(index < view.len(), OutOfBounds: index, 0, view.len());
+                if validity.as_ref().is_some_and(|mask| !mask.value(index)) {
+                    return Ok(None);
+                }
+                match_each_unsigned_integer_ptype!(view.ptype(), |T| {
+                    view.as_slice::<T>()[index]
+                        .to_usize()
+                        .map(Some)
+                        .ok_or_else(|| vortex_err!("RLE index does not fit usize"))
+                })
+            }
+            Self::Probe(slot) => Ok(children
+                .child(*slot)?
+                .scalar_at(index, ctx)?
+                .as_primitive()
+                .as_::<usize>()),
+        }
+    }
+
+    /// Read a value, tagged with the RLE array's dtype.
+    fn value_at(
+        &mut self,
+        index: usize,
+        array: ArrayView<'_, RLE>,
+        children: &mut ProbeChildren<'a>,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Scalar> {
+        match self {
+            Self::Primitive { view, .. } => {
+                vortex_ensure!(index < view.len(), OutOfBounds: index, 0, view.len());
+                Ok(match_each_native_ptype!(view.ptype(), |T| {
+                    Scalar::primitive(view.as_slice::<T>()[index], array.dtype().nullability())
+                }))
+            }
+            Self::Probe(slot) => Scalar::try_new(
+                array.dtype().clone(),
+                children.child(*slot)?.scalar_at(index, ctx)?.into_value(),
+            ),
+        }
+    }
+}
+
+pub(crate) fn scalar_at<'a>(
+    array: ArrayView<'a, RLE>,
+    index: usize,
+    probe: Option<&mut ProbeCtx<'a, RleProbeState<'a>>>,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<Scalar> {
+    let logical_index = array.offset() + index;
+    let chunk = logical_index / FL_CHUNK_SIZE;
+    let Some(probe) = probe else {
+        // The index scalar supplies RLE's nullness and routing value in one lookup.
+        let code = array.indices().execute_scalar(logical_index, ctx)?;
+        let Some(code) = code.as_primitive().as_::<usize>() else {
+            return Ok(Scalar::null(array.dtype().clone()));
+        };
+        let offset = if chunk == 0 {
+            0
+        } else {
+            let offsets = array.values_idx_offsets();
+            read_offset(offsets.execute_scalar(chunk, ctx)?)?
+                .checked_sub(read_offset(offsets.execute_scalar(0, ctx)?)?)
+                .ok_or_else(|| vortex_err!("RLE offsets precede the slice base"))?
+        };
+        let value_index = value_index(offset, code)?;
+        return Scalar::try_new(
+            array.dtype().clone(),
+            array
+                .values()
+                .execute_scalar(value_index, ctx)?
+                .into_value(),
+        );
+    };
+
+    let (state, probes) = probe.parts();
+    let children = match &mut state.children {
+        Some(children) => children,
+        slot @ None => {
+            // Borrow from the source slots, whose lifetime is independent of the temporary view.
+            let slots = array.slots();
+            let mut offsets = Child::new(slots, RLESlots::VALUES_IDX_OFFSETS, ctx)?;
+            let base = offsets
+                .index_at(0, probes, ctx)?
+                .ok_or_else(|| vortex_err!("RLE offset must be a non-null usize"))?;
+            slot.insert(Children {
+                indices: Child::new(slots, RLESlots::INDICES, ctx)?,
+                offsets,
+                values: Child::new(slots, RLESlots::VALUES, ctx)?,
+                base,
+            })
+        }
+    };
+    let Some(code) = children.indices.index_at(logical_index, probes, ctx)? else {
+        return Ok(Scalar::null(array.dtype().clone()));
+    };
+    let offset = children
+        .offsets
+        .index_at(chunk, probes, ctx)?
+        .ok_or_else(|| vortex_err!("RLE offset must be a non-null usize"))?
+        .checked_sub(children.base)
+        .ok_or_else(|| vortex_err!("RLE offsets precede the slice base"))?;
+    children
+        .values
+        .value_at(value_index(offset, code)?, array, probes, ctx)
+}
+
+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"))
+}
+
+fn value_index(offset: usize, code: usize) -> VortexResult<usize> {
+    offset
+        .checked_add(code)
+        .ok_or_else(|| vortex_err!("RLE value index overflow"))
+}
+
+#[cfg(test)]
+mod tests;

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

@@ -0,0 +1,96 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use rstest::rstest;
+use vortex_array::ArrayRef;
+use vortex_array::IntoArray;
+use vortex_array::ProbeUsage;
+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;
+
+#[rstest]
+#[case(ProbeUsage::Once, false)]
+#[case(ProbeUsage::Repeated, false)]
+#[case(ProbeUsage::Once, true)]
+#[case(ProbeUsage::Repeated, true)]
+fn random_access_across_chunks_and_nulls(
+    #[case] usage: ProbeUsage,
+    #[case] 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 probe = source.probe(usage);
+    for index in indices {
+        actual.append_scalar(&probe.scalar_at(index as usize, &mut ctx)?)?;
+    }
+    assert_arrays_eq!(
+        actual.finish(),
+        input.take(PrimitiveArray::from_iter(indices).into_array())?,
+        &mut ctx
+    );
+    assert!(probe.scalar_at(source.len(), &mut ctx).is_err());
+    Ok(())
+}
+
+#[test]
+fn primitive_slots_use_direct_readers() -> VortexResult<()> {
+    let mut ctx = crate::test::SESSION.create_execution_ctx();
+    let input = PrimitiveArray::from_iter((0..4096u32).map(|i| i / 16));
+    let encoded = RLEData::encode(input.as_view(), &mut ctx)?;
+    for slot in [
+        super::RLESlots::INDICES,
+        super::RLESlots::VALUES_IDX_OFFSETS,
+        super::RLESlots::VALUES,
+    ] {
+        let reader = super::Child::new(encoded.as_view().slots(), slot, &mut ctx)?;
+        assert!(matches!(reader, super::Child::Primitive { .. }));
+    }
+    Ok(())
+}
+
+#[test]
+fn lazy_validity_does_not_evaluate_unrequested_rows() -> 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 probe = array.probe(ProbeUsage::Repeated);
+    assert_eq!(probe.scalar_at(0, &mut ctx)?, expected);
+    assert!(probe.scalar_at(1023, &mut ctx).is_err());
+    Ok(())
+}

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

@@ -3,6 +3,7 @@
 
 use vortex_array::ArrayView;
 use vortex_array::ExecutionCtx;
+use vortex_array::ProbeCtx;
 use vortex_array::scalar::Scalar;
 use vortex_array::vtable::OperationsVTable;
 use vortex_error::VortexExpect;
@@ -12,9 +13,20 @@ 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<'a> = ();
+    type ProbeState<'a> = RleProbeState<'a>;
+
+    fn probe_scalar<'a>(
+        array: ArrayView<'a, RLE>,
+        index: usize,
+        probe: Option<&mut ProbeCtx<'a, Self::ProbeState<'a>>>,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Scalar> {
+        probe::scalar_at(array, index, probe, ctx)
+    }
 
     fn scalar_at(
         array: ArrayView<'_, RLE>,

encodings/pco/Cargo.toml

@@ -26,7 +26,13 @@ vortex-mask = { workspace = true }
 vortex-session = { workspace = true }
 
 [dev-dependencies]
+divan = { workspace = true }
 rstest = { workspace = true }
 vortex-array = { workspace = true, features = ["_test-harness"] }
 vortex-arrow = { workspace = true }
+vortex-fastlanes = { workspace = true }
+vortex-runend = { workspace = true }
 
+[[bench]]
+name = "probe"
+harness = false