feat(array): scalar probes with one-off and repeated probe types

encodings/alp/src/alp/ops.rs

@@ -15,6 +15,8 @@ use crate::ALPFloat;
 use crate::match_each_alp_float_ptype;
 
 impl OperationsVTable<ALP> for ALP {
+    type ProbeState = ();
+
     fn scalar_at(
         array: ArrayView<'_, ALP>,
         index: usize,

encodings/alp/src/alp_rd/ops.rs

@@ -14,6 +14,8 @@ use crate::ALPRDArrayExt;
 use crate::ALPRDArraySlotsExt;
 
 impl OperationsVTable<ALPRD> for ALPRD {
+    type ProbeState = ();
+
     fn scalar_at(
         array: ArrayView<'_, ALPRD>,
         index: usize,

encodings/bytebool/src/array.rs

@@ -312,6 +312,8 @@ impl ValidityVTable<ByteBool> for ByteBool {
 }
 
 impl OperationsVTable<ByteBool> for ByteBool {
+    type ProbeState = ();
+
     fn scalar_at(
         array: ArrayView<'_, ByteBool>,
         index: usize,

encodings/datetime-parts/src/ops.rs

@@ -17,6 +17,8 @@ use crate::timestamp;
 use crate::timestamp::TimestampParts;
 
 impl OperationsVTable<DateTimeParts> for DateTimeParts {
+    type ProbeState = ();
+
     fn scalar_at(
         array: ArrayView<'_, DateTimeParts>,
         index: usize,

encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs

@@ -289,6 +289,8 @@ fn to_canonical_decimal(
 }
 
 impl OperationsVTable<DecimalByteParts> for DecimalByteParts {
+    type ProbeState = ();
+
     fn scalar_at(
         array: ArrayView<'_, DecimalByteParts>,
         index: usize,

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

@@ -11,6 +11,8 @@ use crate::BitPacked;
 use crate::bitpack_decompress;
 use crate::bitpacking::array::BitPackedArrayExt;
 impl OperationsVTable<BitPacked> for BitPacked {
+    type ProbeState = ();
+
     fn scalar_at(
         array: ArrayView<'_, BitPacked>,
         index: usize,

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

@@ -11,6 +11,8 @@ use vortex_error::VortexResult;
 
 use super::Delta;
 impl OperationsVTable<Delta> for Delta {
+    type ProbeState = ();
+
     fn scalar_at(
         array: ArrayView<'_, Delta>,
         index: usize,

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

@@ -13,6 +13,8 @@ use super::FoR;
 use crate::r#for::array::FoRArrayExt;
 use crate::r#for::array::FoRArraySlotsExt;
 impl OperationsVTable<FoR> for FoR {
+    type ProbeState = ();
+
     fn scalar_at(
         array: ArrayView<'_, FoR>,
         index: usize,

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

@@ -14,6 +14,8 @@ use crate::rle::RLEArrayExt;
 use crate::rle::RLEArraySlotsExt;
 
 impl OperationsVTable<RLE> for RLE {
+    type ProbeState = ();
+
     fn scalar_at(
         array: ArrayView<'_, RLE>,
         index: usize,

encodings/fastlanes/src/transposed_bool.rs

@@ -256,6 +256,8 @@ impl VTable for TransposedBool {
 }
 
 impl OperationsVTable<TransposedBool> for TransposedBool {
+    type ProbeState = ();
+
     fn scalar_at(
         array: ArrayView<'_, TransposedBool>,
         index: usize,

encodings/fsst/src/ops.rs

@@ -14,6 +14,8 @@ use crate::FSST;
 use crate::FSSTArrayExt;
 
 impl OperationsVTable<FSST> for FSST {
+    type ProbeState = ();
+
     fn scalar_at(
         array: ArrayView<'_, FSST>,
         index: usize,

encodings/onpair/src/ops.rs

@@ -18,6 +18,8 @@ use crate::decode::code_boundary_at;
 use crate::decode::collect_widened;
 
 impl OperationsVTable<OnPair> for OnPair {
+    type ProbeState = ();
+
     fn scalar_at(
         array: ArrayView<'_, OnPair>,
         index: usize,

encodings/parquet-variant/src/operations.rs

@@ -31,6 +31,8 @@ use crate::ParquetVariantArraySlotsExt;
 use crate::vtable::ParquetVariant;
 
 impl OperationsVTable<ParquetVariant> for ParquetVariant {
+    type ProbeState = ();
+
     /// Resolves one row according to the Parquet Variant shredding rules.
     ///
     /// For valid data, a row with both `value` and struct `typed_value` is a partially

encodings/pco/src/array.rs

@@ -778,6 +778,8 @@ impl ValidityVTable<Pco> for Pco {
 }
 
 impl OperationsVTable<Pco> for Pco {
+    type ProbeState = ();
+
     fn scalar_at(
         array: ArrayView<'_, Pco>,
         index: usize,

encodings/runend/src/ops.rs

@@ -18,6 +18,8 @@ use crate::array::RunEndArrayExt;
 use crate::array::RunEndArraySlotsExt;
 
 impl OperationsVTable<RunEnd> for RunEnd {
+    type ProbeState = ();
+
     fn scalar_at(
         array: ArrayView<'_, RunEnd>,
         index: usize,

encodings/sequence/src/array.rs

@@ -425,6 +425,8 @@ impl VTable for Sequence {
 }
 
 impl OperationsVTable<Sequence> for Sequence {
+    type ProbeState = ();
+
     fn scalar_at(
         array: ArrayView<'_, Sequence>,
         index: usize,

encodings/sparse/src/ops.rs

@@ -11,6 +11,8 @@ use crate::Sparse;
 use crate::SparseExt as _;
 
 impl OperationsVTable<Sparse> for Sparse {
+    type ProbeState = ();
+
     fn scalar_at(
         array: ArrayView<'_, Sparse>,
         index: usize,

encodings/zigzag/src/array.rs

@@ -231,6 +231,8 @@ impl Default for ZigZagData {
 }
 
 impl OperationsVTable<ZigZag> for ZigZag {
+    type ProbeState = ();
+
     fn scalar_at(
         array: ArrayView<'_, ZigZag>,
         index: usize,

encodings/zstd/src/array.rs

@@ -1588,6 +1588,8 @@ impl ValidityVTable<Zstd> for Zstd {
 }
 
 impl OperationsVTable<Zstd> for Zstd {
+    type ProbeState = ();
+
     fn scalar_at(
         array: ArrayView<'_, Zstd>,
         index: usize,

encodings/zstd/src/zstd_buffers.rs

@@ -520,6 +520,8 @@ impl VTable for ZstdBuffers {
 }
 
 impl OperationsVTable<ZstdBuffers> for ZstdBuffers {
+    type ProbeState = ();
+
     fn scalar_at(
         array: ArrayView<'_, ZstdBuffers>,
         index: usize,

vortex-array/src/array/erased.rs

@@ -34,6 +34,9 @@ use crate::array::ArrayId;
 use crate::array::ArrayInner;
 use crate::array::ArraySlots;
 use crate::array::DynArrayData;
+use crate::array::probe::ArrayProbe;
+use crate::array::probe::Probe;
+use crate::array::probe::RepeatedArrayProbe;
 use crate::arrays::Constant;
 use crate::arrays::DictArray;
 use crate::arrays::FilterArray;
@@ -273,31 +276,58 @@ impl ArrayRef {
     }
 
     /// Execute the array to extract a scalar at the given index.
+    ///
+    /// A one-off read; the same as `self.probe().execute_scalar(index, ctx)`.
+    // TODO(joe): deprecate this in favour of `probe()`.
     pub fn execute_scalar(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult<Scalar> {
-        vortex_ensure!(index < self.len(), OutOfBounds: index, 0, self.len());
-        if self.dtype().is_nullable() && self.is_invalid(index, ctx)? {
-            return Ok(Scalar::null(self.dtype().clone()));
-        }
-        let scalar = self.0.data.execute_scalar(self, index, ctx)?;
-        debug_assert_eq!(self.dtype(), scalar.dtype(), "Scalar dtype mismatch");
-        Ok(scalar)
+        self.probe().execute_scalar(index, ctx)
+    }
+
+    /// A one-off row accessor over this array. It borrows the handle and retains nothing; for
+    /// many reads of the same array use [`Self::repeated_probe`].
+    ///
+    /// ```
+    /// use vortex_array::{IntoArray, Probe, VortexSessionExecute};
+    /// use vortex_array::arrays::PrimitiveArray;
+    ///
+    /// let array = PrimitiveArray::from_iter([10i32, 20, 30]).into_array();
+    /// let mut ctx = vortex_array::array_session().create_execution_ctx();
+    /// assert_eq!(array.probe().execute_scalar(2, &mut ctx)?, 30i32.into());
+    /// # Ok::<(), vortex_error::VortexError>(())
+    /// ```
+    #[inline]
+    pub fn probe(&self) -> ArrayProbe<'_> {
+        ArrayProbe::new(self)
+    }
+
+    /// A row accessor that owns a handle to this array and keeps encoding state, its validity
+    /// probe and child probes between reads.
+    ///
+    /// ```
+    /// use vortex_array::{IntoArray, Probe, VortexSessionExecute};
+    /// use vortex_array::arrays::PrimitiveArray;
+    ///
+    /// let array = PrimitiveArray::from_iter([10i32, 20, 30]).into_array();
+    /// let mut ctx = vortex_array::array_session().create_execution_ctx();
+    /// let mut probe = array.repeated_probe();
+    /// assert_eq!(probe.execute_scalar(2, &mut ctx)?, 30i32.into());
+    /// assert_eq!(probe.execute_scalar(0, &mut ctx)?, 10i32.into());
+    /// # Ok::<(), vortex_error::VortexError>(())
+    /// ```
+    pub fn repeated_probe(&self) -> RepeatedArrayProbe {
+        RepeatedArrayProbe::new(self.clone())
     }
 
     /// Returns whether the item at `index` is valid.
+    ///
+    /// A one-off read; the same as `self.probe().execute_is_valid(index, ctx)`.
+    // TODO(joe): deprecate this in favour of `probe/repeated_probe()`.
     pub fn is_valid(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
-        vortex_ensure!(index < self.len(), OutOfBounds: index, 0, self.len());
-        match self.validity()? {
-            Validity::NonNullable | Validity::AllValid => Ok(true),
-            Validity::AllInvalid => Ok(false),
-            Validity::Array(a) => a
-                .execute_scalar(index, ctx)?
-                .as_bool()
-                .value()
-                .ok_or_else(|| vortex_err!("validity value at index {} is null", index)),
-        }
+        self.probe().execute_is_valid(index, ctx)
     }
 
     /// Returns whether the item at `index` is invalid.
+    // TODO(joe): deprecate this.
     pub fn is_invalid(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
         Ok(!self.is_valid(index, ctx)?)
     }

vortex-array/src/array/mod.rs

@@ -30,6 +30,8 @@ pub use erased::*;
 
 mod plugin;
 pub use plugin::*;
+mod probe;
+pub use probe::*;
 
 mod foreign;
 pub(crate) use foreign::*;
@@ -220,15 +222,26 @@ pub(crate) trait DynArrayData: 'static + private::Sealed + Send + Sync + Debug {
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<ExecutionResult>;
 
-    /// Execute the scalar at the given index.
+    /// Read a non-null scalar at `index` for a one-off access; nothing is retained.
     ///
-    /// This method panics if the index is out of bounds for the array.
-    fn execute_scalar(
+    /// Kept apart from [`Self::probe_scalar_retained`] so this entry is a bare trampoline into
+    /// the encoding: sharing one function made the one-off path pay the retained branch's
+    /// register saves before its tail call.
+    fn probe_scalar_once(
         &self,
         this: &ArrayRef,
         index: usize,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<Scalar>;
+
+    /// Read a non-null scalar at `index`, keeping preparation in `storage` for later reads.
+    fn probe_scalar_retained(
+        &self,
+        this: &ArrayRef,
+        index: usize,
+        storage: &mut ProbeStorage,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Scalar>;
 }
 
 /// Trait for converting a type into a Vortex [`ArrayRef`].
@@ -490,14 +503,32 @@ impl<V: VTable> DynArrayData for ArrayData<V> {
         V::execute(typed, ctx)
     }
 
-    fn execute_scalar(
+    fn probe_scalar_once(
+        &self,
+        this: &ArrayRef,
+        index: usize,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Scalar> {
+        // SAFETY: this adapter belongs to the ArrayData<V> stored in `this`.
+        let view = unsafe { ArrayView::new_unchecked(this, &self.data) };
+        <V::OperationsVTable as OperationsVTable<V>>::probe_scalar(
+            &mut ProbeState::once(view),
+            index,
+            ctx,
+        )
+    }
+
+    fn probe_scalar_retained(
         &self,
         this: &ArrayRef,
         index: usize,
+        storage: &mut ProbeStorage,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<Scalar> {
+        // SAFETY: this adapter belongs to the ArrayData<V> stored in `this`.
         let view = unsafe { ArrayView::new_unchecked(this, &self.data) };
-        <V::OperationsVTable as OperationsVTable<V>>::scalar_at(view, index, ctx)
+        let mut state = ProbeState::repeated(view, storage.get_or_init()?);
+        <V::OperationsVTable as OperationsVTable<V>>::probe_scalar(&mut state, index, ctx)
     }
 }