feat(array): add scalar probes with lazy child contexts

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

@@ -426,6 +426,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,8 @@ 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::ProbeUsage;
 use crate::arrays::Constant;
 use crate::arrays::DictArray;
 use crate::arrays::FilterArray;
@@ -273,31 +275,37 @@ impl ArrayRef {
     }
 
     /// Execute the array to extract a scalar at the given index.
+    // TODO(joe): deprecate this.
     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(ProbeUsage::Once).execute_scalar(index, ctx)
+    }
+
+    /// Create an accessor over this array with the requested policy for retaining state between
+    /// scalar lookups. The probe borrows this handle; use [`ArrayProbe::new`] for an owned probe.
+    ///
+    /// ```
+    /// use vortex_array::{IntoArray, ProbeUsage, 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.probe(ProbeUsage::Repeated);
+    /// 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 probe(&self, usage: ProbeUsage) -> ArrayProbe<'_> {
+        ArrayProbe::borrowed(self, usage)
     }
 
     /// Returns whether the item at `index` is valid.
+    // TODO(joe): deprecate this.
     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(ProbeUsage::Once).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,13 +222,12 @@ pub(crate) trait DynArrayData: 'static + private::Sealed + Send + Sync + Debug {
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<ExecutionResult>;
 
-    /// Execute the scalar at the given index.
-    ///
-    /// This method panics if the index is out of bounds for the array.
-    fn execute_scalar(
+    /// Read a non-null scalar at `index`, using `storage` for retained state if given.
+    fn probe_scalar(
         &self,
         this: &ArrayRef,
         index: usize,
+        storage: Option<&mut ProbeStorage>,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<Scalar>;
 }
@@ -490,14 +491,23 @@ impl<V: VTable> DynArrayData for ArrayData<V> {
         V::execute(typed, ctx)
     }
 
-    fn execute_scalar(
+    fn probe_scalar(
         &self,
         this: &ArrayRef,
         index: usize,
+        storage: Option<&mut ProbeStorage>,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<Scalar> {
+        type Ops<V> = <V as VTable>::OperationsVTable;
+        // 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)
+        match storage {
+            Some(storage) => Ops::<V>::probe_scalar(view, index, storage.get_or_init()?, ctx),
+            // One-off access: the state lives on the stack for this call only.
+            None => {
+                Ops::<V>::probe_scalar(view, index, &mut ProbeState::new(ProbeUsage::Once), ctx)
+            }
+        }
     }
 }