diff --git a/encodings/alp/src/alp/ops.rs b/encodings/alp/src/alp/ops.rs index a8850744056..58b5ee99360 100644 --- a/encodings/alp/src/alp/ops.rs +++ b/encodings/alp/src/alp/ops.rs @@ -15,6 +15,8 @@ use crate::ALPFloat; use crate::match_each_alp_float_ptype; impl OperationsVTable for ALP { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, ALP>, index: usize, diff --git a/encodings/alp/src/alp_rd/ops.rs b/encodings/alp/src/alp_rd/ops.rs index edb2fb21186..1433966a01e 100644 --- a/encodings/alp/src/alp_rd/ops.rs +++ b/encodings/alp/src/alp_rd/ops.rs @@ -14,6 +14,8 @@ use crate::ALPRDArrayExt; use crate::ALPRDArraySlotsExt; impl OperationsVTable for ALPRD { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, ALPRD>, index: usize, diff --git a/encodings/bytebool/src/array.rs b/encodings/bytebool/src/array.rs index faa1fea81fa..02bb4fcb2b8 100644 --- a/encodings/bytebool/src/array.rs +++ b/encodings/bytebool/src/array.rs @@ -312,6 +312,8 @@ impl ValidityVTable for ByteBool { } impl OperationsVTable for ByteBool { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, ByteBool>, index: usize, diff --git a/encodings/datetime-parts/src/ops.rs b/encodings/datetime-parts/src/ops.rs index d99e55c7542..0ee4f9e315c 100644 --- a/encodings/datetime-parts/src/ops.rs +++ b/encodings/datetime-parts/src/ops.rs @@ -17,6 +17,8 @@ use crate::timestamp; use crate::timestamp::TimestampParts; impl OperationsVTable for DateTimeParts { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, DateTimeParts>, index: usize, diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs index d5b0024f5b7..e1d081ddaca 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -289,6 +289,8 @@ fn to_canonical_decimal( } impl OperationsVTable for DecimalByteParts { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, DecimalByteParts>, index: usize, diff --git a/encodings/fastlanes/src/bitpacking/vtable/operations.rs b/encodings/fastlanes/src/bitpacking/vtable/operations.rs index e14b27323c1..2816407ac03 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/operations.rs +++ b/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 for BitPacked { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, BitPacked>, index: usize, diff --git a/encodings/fastlanes/src/delta/vtable/operations.rs b/encodings/fastlanes/src/delta/vtable/operations.rs index 7ed57a0886d..b37760621ea 100644 --- a/encodings/fastlanes/src/delta/vtable/operations.rs +++ b/encodings/fastlanes/src/delta/vtable/operations.rs @@ -11,6 +11,8 @@ use vortex_error::VortexResult; use super::Delta; impl OperationsVTable for Delta { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, Delta>, index: usize, diff --git a/encodings/fastlanes/src/for/vtable/operations.rs b/encodings/fastlanes/src/for/vtable/operations.rs index 36dac998cbe..361020824d5 100644 --- a/encodings/fastlanes/src/for/vtable/operations.rs +++ b/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 { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, FoR>, index: usize, diff --git a/encodings/fastlanes/src/rle/vtable/operations.rs b/encodings/fastlanes/src/rle/vtable/operations.rs index ca4d2d39545..ba6a624f1f5 100644 --- a/encodings/fastlanes/src/rle/vtable/operations.rs +++ b/encodings/fastlanes/src/rle/vtable/operations.rs @@ -14,6 +14,8 @@ use crate::rle::RLEArrayExt; use crate::rle::RLEArraySlotsExt; impl OperationsVTable for RLE { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, RLE>, index: usize, diff --git a/encodings/fastlanes/src/transposed_bool.rs b/encodings/fastlanes/src/transposed_bool.rs index efb3443cdfc..92da3b3b79c 100644 --- a/encodings/fastlanes/src/transposed_bool.rs +++ b/encodings/fastlanes/src/transposed_bool.rs @@ -256,6 +256,8 @@ impl VTable for TransposedBool { } impl OperationsVTable for TransposedBool { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, TransposedBool>, index: usize, diff --git a/encodings/fsst/src/ops.rs b/encodings/fsst/src/ops.rs index b630508ed9e..03560b2c8b8 100644 --- a/encodings/fsst/src/ops.rs +++ b/encodings/fsst/src/ops.rs @@ -14,6 +14,8 @@ use crate::FSST; use crate::FSSTArrayExt; impl OperationsVTable for FSST { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, FSST>, index: usize, diff --git a/encodings/onpair/src/ops.rs b/encodings/onpair/src/ops.rs index 728e5a0e6f2..082bc55e199 100644 --- a/encodings/onpair/src/ops.rs +++ b/encodings/onpair/src/ops.rs @@ -18,6 +18,8 @@ use crate::decode::code_boundary_at; use crate::decode::collect_widened; impl OperationsVTable for OnPair { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, OnPair>, index: usize, diff --git a/encodings/parquet-variant/src/operations.rs b/encodings/parquet-variant/src/operations.rs index 8517df27f22..eab80bf63e1 100644 --- a/encodings/parquet-variant/src/operations.rs +++ b/encodings/parquet-variant/src/operations.rs @@ -31,6 +31,8 @@ use crate::ParquetVariantArraySlotsExt; use crate::vtable::ParquetVariant; impl OperationsVTable 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 diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array.rs index 44a6a8e4045..97a17fe6a3c 100644 --- a/encodings/pco/src/array.rs +++ b/encodings/pco/src/array.rs @@ -778,6 +778,8 @@ impl ValidityVTable for Pco { } impl OperationsVTable for Pco { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, Pco>, index: usize, diff --git a/encodings/runend/src/ops.rs b/encodings/runend/src/ops.rs index e2c2e3fc99b..46949d84e1b 100644 --- a/encodings/runend/src/ops.rs +++ b/encodings/runend/src/ops.rs @@ -18,6 +18,8 @@ use crate::array::RunEndArrayExt; use crate::array::RunEndArraySlotsExt; impl OperationsVTable for RunEnd { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, RunEnd>, index: usize, diff --git a/encodings/sequence/src/array.rs b/encodings/sequence/src/array.rs index 10a222544b6..e61eb54e4eb 100644 --- a/encodings/sequence/src/array.rs +++ b/encodings/sequence/src/array.rs @@ -426,6 +426,8 @@ impl VTable for Sequence { } impl OperationsVTable for Sequence { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, Sequence>, index: usize, diff --git a/encodings/sparse/src/ops.rs b/encodings/sparse/src/ops.rs index 568d8d377d1..acabbb103d5 100644 --- a/encodings/sparse/src/ops.rs +++ b/encodings/sparse/src/ops.rs @@ -11,6 +11,8 @@ use crate::Sparse; use crate::SparseExt as _; impl OperationsVTable for Sparse { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, Sparse>, index: usize, diff --git a/encodings/zigzag/src/array.rs b/encodings/zigzag/src/array.rs index ef3165132f4..09076fcbbdf 100644 --- a/encodings/zigzag/src/array.rs +++ b/encodings/zigzag/src/array.rs @@ -231,6 +231,8 @@ impl Default for ZigZagData { } impl OperationsVTable for ZigZag { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, ZigZag>, index: usize, diff --git a/encodings/zstd/src/array.rs b/encodings/zstd/src/array.rs index fb3551e539b..dafe8e820a9 100644 --- a/encodings/zstd/src/array.rs +++ b/encodings/zstd/src/array.rs @@ -1588,6 +1588,8 @@ impl ValidityVTable for Zstd { } impl OperationsVTable for Zstd { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, Zstd>, index: usize, diff --git a/encodings/zstd/src/zstd_buffers.rs b/encodings/zstd/src/zstd_buffers.rs index f21deee6f64..51d338777ca 100644 --- a/encodings/zstd/src/zstd_buffers.rs +++ b/encodings/zstd/src/zstd_buffers.rs @@ -520,6 +520,8 @@ impl VTable for ZstdBuffers { } impl OperationsVTable for ZstdBuffers { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, ZstdBuffers>, index: usize, diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index 76c0b33cd18..5307568a73c 100644 --- a/vortex-array/src/array/erased.rs +++ b/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 { - 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 { - 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 { Ok(!self.is_valid(index, ctx)?) } diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs index 8b9b2812bfa..5d5f3211c8e 100644 --- a/vortex-array/src/array/mod.rs +++ b/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; - /// 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; } @@ -490,14 +491,23 @@ impl DynArrayData for ArrayData { V::execute(typed, ctx) } - fn execute_scalar( + fn probe_scalar( &self, this: &ArrayRef, index: usize, + storage: Option<&mut ProbeStorage>, ctx: &mut ExecutionCtx, ) -> VortexResult { + type Ops = ::OperationsVTable; + // SAFETY: this adapter belongs to the ArrayData stored in `this`. let view = unsafe { ArrayView::new_unchecked(this, &self.data) }; - >::scalar_at(view, index, ctx) + match storage { + Some(storage) => Ops::::probe_scalar(view, index, storage.get_or_init()?, ctx), + // One-off access: the state lives on the stack for this call only. + None => { + Ops::::probe_scalar(view, index, &mut ProbeState::new(ProbeUsage::Once), ctx) + } + } } } diff --git a/vortex-array/src/array/probe/array.rs b/vortex-array/src/array/probe/array.rs new file mode 100644 index 00000000000..7b070cbcff9 --- /dev/null +++ b/vortex-array/src/array/probe/array.rs @@ -0,0 +1,310 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::any::Any; +use std::borrow::Cow; + +use smallvec::SmallVec; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::array::probe::ProbeUsage; +use crate::array::probe::ProbeValidity; +use crate::scalar::Scalar; + +/// A scalar accessor over a borrowed or owned array that optionally retains preparation +/// between lookups. +/// +/// Array handles share their buffers. [`ProbeUsage::Repeated`] initializes encoding state and +/// child probes lazily and reuses them for following lookups; [`ProbeUsage::Once`] retains +/// nothing. A probe built with [`ArrayProbe::new`] owns its handle and can outlive the original; +/// [`ArrayRef::probe`] borrows it. Probes are local to a thread. +pub struct ArrayProbe<'a> { + array: Cow<'a, ArrayRef>, + retained: Option, +} + +/// Everything a repeated-access probe keeps between lookups. +#[derive(Default)] +pub(crate) struct Retained { + state: ProbeStorage, + // Boxed because `ProbeValidity` itself contains an `ArrayProbe`. + validity: Option>, +} + +impl ArrayProbe<'static> { + /// Own an array and choose whether to retain preparation. + pub fn new(array: ArrayRef, usage: ProbeUsage) -> Self { + Self { + array: Cow::Owned(array), + retained: Retained::for_usage(usage), + } + } +} + +impl<'a> ArrayProbe<'a> { + /// Borrow an array and choose whether to retain preparation. + pub(crate) fn borrowed(array: &'a ArrayRef, usage: ProbeUsage) -> Self { + Self { + array: Cow::Borrowed(array), + retained: Retained::for_usage(usage), + } + } + + /// The array this probe reads from. + pub fn array(&self) -> &ArrayRef { + &self.array + } + + /// Read a scalar, including its nullness, preparing and reusing state as appropriate. + pub fn execute_scalar(&mut self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { + if !self.execute_is_valid(index, ctx)? { + return Ok(Scalar::null(self.array.dtype().clone())); + } + let array: &ArrayRef = &self.array; + let storage = self.retained.as_mut().map(|retained| &mut retained.state); + let scalar = array.dyn_array().probe_scalar(array, index, storage, ctx)?; + debug_assert_eq!(scalar.dtype(), array.dtype(), "Scalar dtype mismatch"); + Ok(scalar) + } + + /// Check bounds and read validity using this probe's retention policy. + pub fn execute_is_valid(&mut self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { + let array: &ArrayRef = &self.array; + vortex_ensure!(index < array.len(), OutOfBounds: index, 0, array.len()); + if !array.dtype().is_nullable() { + return Ok(true); + } + let Some(retained) = &mut self.retained else { + return array + .validity()? + .probe(ProbeUsage::Once) + .execute_is_valid(index, ctx); + }; + if retained.validity.is_none() { + let validity = array.validity()?.probe(ProbeUsage::Repeated); + retained.validity = Some(Box::new(validity)); + } + retained + .validity + .as_mut() + .ok_or_else(|| vortex_err!("validity probe was just initialized"))? + .execute_is_valid(index, ctx) + } +} + +impl Retained { + fn for_usage(usage: ProbeUsage) -> Option { + match usage { + ProbeUsage::Once => None, + ProbeUsage::Repeated => Some(Self::default()), + } + } +} + +/// Encoding state and lazy child probes for one source array. +/// +/// Passed to [`OperationsVTable::probe_scalar`](crate::vtable::OperationsVTable::probe_scalar). +/// For repeated access the framework boxes this once per probe and reuses it; for one-off +/// access it lives on the stack for the duration of the call, so `S::default()` must be cheap. +pub struct ProbeState { + /// The encoding's own retained state. + pub state: S, + /// Lazily created probes over the source's child slots. + pub children: ProbeChildren, +} + +impl ProbeState { + /// Create empty state with the given retention policy. + /// + /// The framework builds this for every probe; encodings only need it to run their + /// `probe_scalar` path from `scalar_at` with [`ProbeUsage::Once`]. + pub fn new(usage: ProbeUsage) -> Self { + Self { + state: S::default(), + children: ProbeChildren { + usage, + slots: SmallVec::new(), + }, + } + } +} + +impl ProbeState { + /// Whether this state is discarded after one lookup or reused across many. + /// + /// Encodings whose preparation is expensive should skip it for [`ProbeUsage::Once`]. + pub fn usage(&self) -> ProbeUsage { + self.children.usage + } +} + +/// Lazy child probes over the slots of one source array. +/// +/// Child probes inherit the parent's retention policy. The slot table is created on the first +/// valid request; slots that are never requested stay empty. Dropping the parent state drops +/// every child probe. +pub struct ProbeChildren { + usage: ProbeUsage, + slots: SmallVec<[Option>; 4]>, +} + +impl ProbeChildren { + /// Get or create a probe over `parent`'s given child slot. + /// + /// `parent` must be the array this state was created for. Returns `None` for an absent + /// slot and an error for an out-of-bounds slot; neither case creates the slot table. + pub fn slot( + &mut self, + parent: &ArrayRef, + slot: usize, + ) -> VortexResult>> { + let slots = parent.slots(); + let child = slots + .get(slot) + .ok_or_else(|| vortex_err!("Probe slot {slot} is out of bounds"))?; + let Some(child) = child else { + return Ok(None); + }; + if self.slots.is_empty() { + self.slots.resize_with(slots.len(), || None); + } + let usage = self.usage; + Ok(Some(self.slots[slot].get_or_insert_with(|| { + ArrayProbe::new(child.clone(), usage) + }))) + } +} + +/// Type-erased, lazily initialized [`ProbeState`] for a retained probe. +#[derive(Default)] +pub(crate) struct ProbeStorage(Option>); + +impl ProbeStorage { + pub(crate) fn get_or_init(&mut self) -> VortexResult<&mut ProbeState> { + self.0 + .get_or_insert_with(|| Box::new(ProbeState::::new(ProbeUsage::Repeated))) + .downcast_mut::>() + .ok_or_else(|| vortex_err!("Probe state type mismatch")) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + + use super::*; + use crate::VortexSessionExecute; + use crate::array::IntoArray; + use crate::arrays::PrimitiveArray; + use crate::arrays::StructArray; + + fn nullable_ints() -> ArrayRef { + PrimitiveArray::from_option_iter([Some(10i32), None, Some(30)]).into_array() + } + + #[rstest] + fn checks_bounds_and_nulls( + #[values(ProbeUsage::Once, ProbeUsage::Repeated)] usage: ProbeUsage, + ) -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let array = nullable_ints(); + let mut probe = array.probe(usage); + assert!(probe.execute_scalar(3, &mut ctx).is_err()); + assert!(probe.execute_is_valid(3, &mut ctx).is_err()); + assert!(probe.execute_scalar(1, &mut ctx)?.is_null()); + assert!(!probe.execute_is_valid(1, &mut ctx)?); + assert_eq!( + probe.execute_scalar(2, &mut ctx)?, + Scalar::from(Some(30i32)) + ); + assert_eq!( + probe.execute_scalar(0, &mut ctx)?, + Scalar::from(Some(10i32)) + ); + Ok(()) + } + + #[test] + fn once_retains_nothing() -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let array = nullable_ints(); + let mut probe = array.probe(ProbeUsage::Once); + probe.execute_scalar(0, &mut ctx)?; + assert!(probe.retained.is_none()); + Ok(()) + } + + #[test] + fn repeated_initializes_lazily_and_outlives_handle() -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let array = nullable_ints(); + let mut probe = ArrayProbe::new(array.clone(), ProbeUsage::Repeated); + drop(array); + + let retained = probe + .retained + .as_ref() + .ok_or_else(|| vortex_err!("missing retained state"))?; + assert!(retained.state.0.is_none()); + assert!(retained.validity.is_none()); + + assert!(probe.execute_scalar(3, &mut ctx).is_err()); + assert!(probe.execute_scalar(1, &mut ctx)?.is_null()); + assert_eq!( + probe.execute_scalar(0, &mut ctx)?, + Scalar::from(Some(10i32)) + ); + + let retained = probe + .retained + .as_ref() + .ok_or_else(|| vortex_err!("missing retained state"))?; + assert!(retained.state.0.is_some()); + assert!(retained.validity.is_some()); + Ok(()) + } + + #[rstest] + fn children_are_created_on_demand( + #[values(ProbeUsage::Once, ProbeUsage::Repeated)] usage: ProbeUsage, + ) -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let array = StructArray::from_fields(&[ + ("a", PrimitiveArray::from_iter([1i32, 2]).into_array()), + ("b", PrimitiveArray::from_iter([3i64, 4]).into_array()), + ])? + .into_array(); + let mut state = ProbeState::<()>::new(usage); + + assert!(state.children.slot(&array, 5).is_err()); + // Slot 0 is the struct's absent validity. + assert!(state.children.slot(&array, 0)?.is_none()); + assert!(state.children.slots.is_empty()); + + let child = state + .children + .slot(&array, 2)? + .ok_or_else(|| vortex_err!("missing slot"))?; + assert_eq!(child.execute_scalar(1, &mut ctx)?, Scalar::from(4i64)); + assert_eq!(child.retained.is_some(), usage == ProbeUsage::Repeated); + + assert_eq!(state.children.slots.len(), array.slots().len()); + assert!(state.children.slots[1].is_none()); + assert!(state.children.slots[2].is_some()); + Ok(()) + } + + #[test] + fn storage_rejects_mismatched_state_type() -> VortexResult<()> { + let mut storage = ProbeStorage::default(); + storage.get_or_init::<()>()?; + assert!(storage.get_or_init::().is_err()); + Ok(()) + } +} diff --git a/vortex-array/src/array/probe/mod.rs b/vortex-array/src/array/probe/mod.rs new file mode 100644 index 00000000000..095af0079e3 --- /dev/null +++ b/vortex-array/src/array/probe/mod.rs @@ -0,0 +1,16 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +mod array; +pub use array::*; +mod validity; +pub use validity::*; + +/// Whether scalar access should retain preparation for subsequent lookups. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ProbeUsage { + /// Use temporary resources only, without initializing retained probe state. + Once, + /// Allow the encoding to retain preparation and decoded data between lookups. + Repeated, +} diff --git a/vortex-array/src/array/probe/validity.rs b/vortex-array/src/array/probe/validity.rs new file mode 100644 index 00000000000..9f937f21a74 --- /dev/null +++ b/vortex-array/src/array/probe/validity.rs @@ -0,0 +1,42 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::ExecutionCtx; +use crate::array::probe::ArrayProbe; + +/// A validity accessor built by [`Validity::probe`](crate::validity::Validity::probe). +/// +/// Uniform validity retains nothing; array-backed validity delegates to an [`ArrayProbe`] +/// over the boolean array with the requested retention policy. +pub enum ProbeValidity { + /// Validity is uniform, so no lookup is needed. + Constant(bool), + /// Validity backed by a boolean array. + Array(ArrayProbe<'static>), +} + +impl ProbeValidity { + /// Returns whether the row at `index` is valid. + pub fn execute_is_valid(&mut self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { + match self { + Self::Constant(valid) => Ok(*valid), + Self::Array(probe) => probe + .execute_scalar(index, ctx)? + .as_bool() + .value() + .ok_or_else(|| vortex_err!("validity value at index {index} is null")), + } + } + + /// Returns whether the row at `index` is null. + pub fn execute_is_invalid( + &mut self, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(!self.execute_is_valid(index, ctx)?) + } +} diff --git a/vortex-array/src/array/vtable/operations.rs b/vortex-array/src/array/vtable/operations.rs index 7f49e683640..6771798f6ec 100644 --- a/vortex-array/src/array/vtable/operations.rs +++ b/vortex-array/src/array/vtable/operations.rs @@ -7,6 +7,7 @@ use vortex_error::vortex_bail; use crate::ExecutionCtx; use crate::array::ArrayView; use crate::array::VTable; +use crate::array::probe::ProbeState; use crate::scalar::Scalar; use crate::vtable::NotSupported; @@ -17,6 +18,32 @@ use crate::vtable::NotSupported; /// [`ArrayRef`](crate::ArrayRef) /// methods perform common checks before dispatching here. pub trait OperationsVTable { + /// Encoding-specific state retained by repeated scalar access. + /// + /// Default construction must be cheap and allocation-free: one-off reads build it on the + /// stack for every call. Preparation belongs in [`Self::probe_scalar`]. State owns its + /// preparation and may hold shared buffer or array handles. Use `()` when no state is needed. + type ProbeState: Default + 'static; + + /// Read a non-null scalar, retaining preparation in `state` for subsequent reads. + /// + /// Bounds and validity have been checked; the row is non-null. Whether `state` is reused + /// across calls or discarded after this one is decided by the caller; child probes obtained + /// through [`ProbeState::children`] follow the same policy. The scalar must retain the + /// source's logical dtype, including nullability. + /// + /// The default preserves the existing scalar path without adding caching. + fn probe_scalar( + array: ArrayView<'_, V>, + index: usize, + _state: &mut ProbeState, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + // FIXME: Remove this default once all encodings have migrated to probe_scalar. + Self::scalar_at(array, index, ctx) + } + + // FIXME: Deprecate scalar_at once encodings have migrated to probe_scalar. /// Fetch the scalar at the given index. /// /// ## Preconditions @@ -35,6 +62,8 @@ pub trait OperationsVTable { } impl OperationsVTable for NotSupported { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, V>, _index: usize, diff --git a/vortex-array/src/arrays/bool/vtable/operations.rs b/vortex-array/src/arrays/bool/vtable/operations.rs index c29ab20331b..e1ec02ecbc8 100644 --- a/vortex-array/src/arrays/bool/vtable/operations.rs +++ b/vortex-array/src/arrays/bool/vtable/operations.rs @@ -11,6 +11,8 @@ use crate::arrays::bool::BoolArrayExt; use crate::scalar::Scalar; impl OperationsVTable for Bool { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, Bool>, index: usize, diff --git a/vortex-array/src/arrays/chunked/vtable/operations.rs b/vortex-array/src/arrays/chunked/vtable/operations.rs index 8f9e0867a88..395e8aff70a 100644 --- a/vortex-array/src/arrays/chunked/vtable/operations.rs +++ b/vortex-array/src/arrays/chunked/vtable/operations.rs @@ -11,6 +11,8 @@ use crate::arrays::chunked::ChunkedArrayExt; use crate::scalar::Scalar; impl OperationsVTable for Chunked { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, Chunked>, index: usize, diff --git a/vortex-array/src/arrays/constant/vtable/operations.rs b/vortex-array/src/arrays/constant/vtable/operations.rs index e3568a9c39f..9c49094f324 100644 --- a/vortex-array/src/arrays/constant/vtable/operations.rs +++ b/vortex-array/src/arrays/constant/vtable/operations.rs @@ -10,6 +10,8 @@ use crate::arrays::Constant; use crate::scalar::Scalar; impl OperationsVTable for Constant { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, Constant>, _index: usize, diff --git a/vortex-array/src/arrays/decimal/vtable/operations.rs b/vortex-array/src/arrays/decimal/vtable/operations.rs index 257f26127ae..aecf6638bd2 100644 --- a/vortex-array/src/arrays/decimal/vtable/operations.rs +++ b/vortex-array/src/arrays/decimal/vtable/operations.rs @@ -12,6 +12,8 @@ use crate::scalar::DecimalValue; use crate::scalar::Scalar; impl OperationsVTable for Decimal { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, Decimal>, index: usize, diff --git a/vortex-array/src/arrays/dict/vtable/operations.rs b/vortex-array/src/arrays/dict/vtable/operations.rs index 1982a1e0870..d497db0f2f8 100644 --- a/vortex-array/src/arrays/dict/vtable/operations.rs +++ b/vortex-array/src/arrays/dict/vtable/operations.rs @@ -12,6 +12,8 @@ use crate::arrays::dict::DictArraySlotsExt; use crate::scalar::Scalar; impl OperationsVTable for Dict { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, Dict>, index: usize, diff --git a/vortex-array/src/arrays/extension/vtable/operations.rs b/vortex-array/src/arrays/extension/vtable/operations.rs index 66de94b596a..519ef6088f7 100644 --- a/vortex-array/src/arrays/extension/vtable/operations.rs +++ b/vortex-array/src/arrays/extension/vtable/operations.rs @@ -11,6 +11,8 @@ use crate::arrays::extension::ExtensionArrayExt; use crate::scalar::Scalar; impl OperationsVTable for Extension { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, Extension>, index: usize, diff --git a/vortex-array/src/arrays/filter/vtable.rs b/vortex-array/src/arrays/filter/vtable.rs index 56fc112aa16..7f0e0bfd014 100644 --- a/vortex-array/src/arrays/filter/vtable.rs +++ b/vortex-array/src/arrays/filter/vtable.rs @@ -200,6 +200,8 @@ impl VTable for Filter { } } impl OperationsVTable for Filter { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, Filter>, index: usize, diff --git a/vortex-array/src/arrays/fixed_size_list/vtable/operations.rs b/vortex-array/src/arrays/fixed_size_list/vtable/operations.rs index 9f4cf02fbf8..eca0e45ad95 100644 --- a/vortex-array/src/arrays/fixed_size_list/vtable/operations.rs +++ b/vortex-array/src/arrays/fixed_size_list/vtable/operations.rs @@ -11,6 +11,8 @@ use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::scalar::Scalar; impl OperationsVTable for FixedSizeList { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, FixedSizeList>, index: usize, diff --git a/vortex-array/src/arrays/interleave/mod.rs b/vortex-array/src/arrays/interleave/mod.rs index d29981249d4..fb649cf6e19 100644 --- a/vortex-array/src/arrays/interleave/mod.rs +++ b/vortex-array/src/arrays/interleave/mod.rs @@ -389,6 +389,8 @@ impl VTable for Interleave { } impl OperationsVTable for Interleave { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, Interleave>, index: usize, diff --git a/vortex-array/src/arrays/list/vtable/operations.rs b/vortex-array/src/arrays/list/vtable/operations.rs index 02c686cd1f1..668cc18a32f 100644 --- a/vortex-array/src/arrays/list/vtable/operations.rs +++ b/vortex-array/src/arrays/list/vtable/operations.rs @@ -13,6 +13,8 @@ use crate::arrays::list::ListArrayExt; use crate::scalar::Scalar; impl OperationsVTable for List { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, List>, index: usize, diff --git a/vortex-array/src/arrays/listview/vtable/operations.rs b/vortex-array/src/arrays/listview/vtable/operations.rs index f0cb9539cc3..8608985463d 100644 --- a/vortex-array/src/arrays/listview/vtable/operations.rs +++ b/vortex-array/src/arrays/listview/vtable/operations.rs @@ -13,6 +13,8 @@ use crate::arrays::listview::ListViewArrayExt; use crate::scalar::Scalar; impl OperationsVTable for ListView { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, ListView>, index: usize, diff --git a/vortex-array/src/arrays/map/vtable/operations.rs b/vortex-array/src/arrays/map/vtable/operations.rs index d6e8fe87f12..66b8e47a75a 100644 --- a/vortex-array/src/arrays/map/vtable/operations.rs +++ b/vortex-array/src/arrays/map/vtable/operations.rs @@ -13,6 +13,8 @@ use crate::arrays::struct_::StructArrayExt; use crate::scalar::Scalar; impl OperationsVTable for Map { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, Map>, index: usize, diff --git a/vortex-array/src/arrays/masked/vtable/operations.rs b/vortex-array/src/arrays/masked/vtable/operations.rs index c82d0bf03ed..a418e23f217 100644 --- a/vortex-array/src/arrays/masked/vtable/operations.rs +++ b/vortex-array/src/arrays/masked/vtable/operations.rs @@ -11,6 +11,8 @@ use crate::arrays::masked::MaskedArraySlotsExt; use crate::scalar::Scalar; impl OperationsVTable for Masked { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, Masked>, index: usize, diff --git a/vortex-array/src/arrays/null/mod.rs b/vortex-array/src/arrays/null/mod.rs index 8479b7137cc..b4d4536c9cf 100644 --- a/vortex-array/src/arrays/null/mod.rs +++ b/vortex-array/src/arrays/null/mod.rs @@ -175,6 +175,8 @@ impl Array { } impl OperationsVTable for Null { + type ProbeState = (); + fn scalar_at( _array: ArrayView<'_, Null>, _index: usize, diff --git a/vortex-array/src/arrays/patched/vtable/operations.rs b/vortex-array/src/arrays/patched/vtable/operations.rs index 51dd1fc9e3c..bd7440965ab 100644 --- a/vortex-array/src/arrays/patched/vtable/operations.rs +++ b/vortex-array/src/arrays/patched/vtable/operations.rs @@ -14,6 +14,8 @@ use crate::optimizer::ArrayOptimizer; use crate::scalar::Scalar; impl OperationsVTable for Patched { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, Patched>, index: usize, diff --git a/vortex-array/src/arrays/piecewise_sequence/vtable.rs b/vortex-array/src/arrays/piecewise_sequence/vtable.rs index d68a00815f3..d39bd6eccac 100644 --- a/vortex-array/src/arrays/piecewise_sequence/vtable.rs +++ b/vortex-array/src/arrays/piecewise_sequence/vtable.rs @@ -145,6 +145,8 @@ impl VTable for PiecewiseSequence { } impl OperationsVTable for PiecewiseSequence { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, PiecewiseSequence>, index: usize, diff --git a/vortex-array/src/arrays/primitive/vtable/operations.rs b/vortex-array/src/arrays/primitive/vtable/operations.rs index ddeaa386485..9501513fb1c 100644 --- a/vortex-array/src/arrays/primitive/vtable/operations.rs +++ b/vortex-array/src/arrays/primitive/vtable/operations.rs @@ -11,6 +11,8 @@ use crate::match_each_native_ptype; use crate::scalar::Scalar; impl OperationsVTable for Primitive { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, Primitive>, index: usize, diff --git a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs index 40d75906356..21af5728572 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs @@ -15,6 +15,8 @@ use crate::scalar::Scalar; use crate::scalar_fn::VecExecutionArgs; impl OperationsVTable for ScalarFn { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, ScalarFn>, index: usize, diff --git a/vortex-array/src/arrays/shared/vtable.rs b/vortex-array/src/arrays/shared/vtable.rs index 758d091b557..b1a9a4d7b15 100644 --- a/vortex-array/src/arrays/shared/vtable.rs +++ b/vortex-array/src/arrays/shared/vtable.rs @@ -125,6 +125,8 @@ impl VTable for Shared { } } impl OperationsVTable for Shared { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, Shared>, index: usize, diff --git a/vortex-array/src/arrays/slice/vtable.rs b/vortex-array/src/arrays/slice/vtable.rs index c88c28a9a4d..c09d81ee4f1 100644 --- a/vortex-array/src/arrays/slice/vtable.rs +++ b/vortex-array/src/arrays/slice/vtable.rs @@ -169,6 +169,8 @@ impl VTable for Slice { } } impl OperationsVTable for Slice { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, Slice>, index: usize, diff --git a/vortex-array/src/arrays/struct_/vtable/operations.rs b/vortex-array/src/arrays/struct_/vtable/operations.rs index b491e231520..54d086bf7c9 100644 --- a/vortex-array/src/arrays/struct_/vtable/operations.rs +++ b/vortex-array/src/arrays/struct_/vtable/operations.rs @@ -12,6 +12,8 @@ use crate::scalar::Scalar; use crate::scalar::ScalarValue; impl OperationsVTable for Struct { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, Struct>, index: usize, diff --git a/vortex-array/src/arrays/union/vtable/operations.rs b/vortex-array/src/arrays/union/vtable/operations.rs index 39ba95ed4ee..78b9759ab1c 100644 --- a/vortex-array/src/arrays/union/vtable/operations.rs +++ b/vortex-array/src/arrays/union/vtable/operations.rs @@ -14,6 +14,8 @@ use crate::arrays::union::UnionArraySlotsExt; use crate::scalar::Scalar; impl OperationsVTable for Union { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, Union>, index: usize, diff --git a/vortex-array/src/arrays/varbin/vtable/operations.rs b/vortex-array/src/arrays/varbin/vtable/operations.rs index e11043e605a..cc5090d8912 100644 --- a/vortex-array/src/arrays/varbin/vtable/operations.rs +++ b/vortex-array/src/arrays/varbin/vtable/operations.rs @@ -12,6 +12,8 @@ use crate::arrays::varbin::varbin_scalar; use crate::scalar::Scalar; impl OperationsVTable for VarBin { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, VarBin>, index: usize, diff --git a/vortex-array/src/arrays/varbinview/vtable/operations.rs b/vortex-array/src/arrays/varbinview/vtable/operations.rs index 1a1f20a0dbe..53866d29fb5 100644 --- a/vortex-array/src/arrays/varbinview/vtable/operations.rs +++ b/vortex-array/src/arrays/varbinview/vtable/operations.rs @@ -11,6 +11,8 @@ use crate::arrays::varbin::varbin_scalar; use crate::scalar::Scalar; impl OperationsVTable for VarBinView { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, VarBinView>, index: usize, diff --git a/vortex-array/src/arrays/variant/vtable/operations.rs b/vortex-array/src/arrays/variant/vtable/operations.rs index c7f7d36fd97..ed61223d637 100644 --- a/vortex-array/src/arrays/variant/vtable/operations.rs +++ b/vortex-array/src/arrays/variant/vtable/operations.rs @@ -12,6 +12,8 @@ use crate::arrays::variant::VariantArraySlotsExt; use crate::scalar::Scalar; impl OperationsVTable for Variant { + type ProbeState = (); + fn scalar_at( array: ArrayView<'_, Variant>, index: usize, diff --git a/vortex-array/src/validity.rs b/vortex-array/src/validity.rs index a2c10e30ca7..c9b44045b6c 100644 --- a/vortex-array/src/validity.rs +++ b/vortex-array/src/validity.rs @@ -20,10 +20,13 @@ use vortex_error::vortex_err; use vortex_mask::Mask; use vortex_mask::MaskValues; +use crate::ArrayProbe; use crate::ArrayRef; use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; +use crate::ProbeUsage; +use crate::ProbeValidity; use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; @@ -162,7 +165,20 @@ impl Validity { } } + /// Create an accessor with the requested policy for retaining state between lookups. + /// + /// Mirrors [`ArrayRef::probe`]. Uniform validity retains nothing; array-backed validity + /// probes the underlying boolean array. + pub fn probe(self, usage: ProbeUsage) -> ProbeValidity { + match self { + Self::NonNullable | Self::AllValid => ProbeValidity::Constant(true), + Self::AllInvalid => ProbeValidity::Constant(false), + Self::Array(array) => ProbeValidity::Array(ArrayProbe::new(array, usage)), + } + } + /// Returns whether the `index` item is valid, using `ctx` to execute the validity array. + // Todo(joe): deprecate this #[inline] pub fn execute_is_valid(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { Ok(match self { @@ -177,6 +193,7 @@ impl Validity { } /// Returns whether the `index` item is null, using `ctx` to execute the validity array. + // Todo(joe): deprecate this #[inline] pub fn execute_is_null(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { Ok(!self.execute_is_valid(index, ctx)?) diff --git a/vortex-python/src/arrays/py/vtable.rs b/vortex-python/src/arrays/py/vtable.rs index 5ec749d3a5f..5697d5a7574 100644 --- a/vortex-python/src/arrays/py/vtable.rs +++ b/vortex-python/src/arrays/py/vtable.rs @@ -122,6 +122,8 @@ impl VTable for PythonVTable { } impl OperationsVTable for PythonVTable { + type ProbeState = (); + fn scalar_at( _array: ArrayView<'_, PythonVTable>, _index: usize,