From 264e0d44cd6e002630eabcd45d547a830d52fcc7 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 15 Sep 2026 18:51:55 +0100 Subject: [PATCH 1/2] feat(array): add probe_scalar to OperationsVTable with a ProbeState shell Adds the vtable hook the scalar probe API builds on, without any of the probe machinery. `OperationsVTable` gains `type ProbeState`, the state an encoding may keep across repeated reads, and `probe_scalar(state, index, ctx)`, defaulting to `scalar_at` so no encoding changes behaviour. `ProbeState<'_, V>` for now only carries the `ArrayView`; the erased dispatch routes one-off reads through it. Every encoding declares `type ProbeState = ()`. `scalar_at` stays until encodings migrate to `probe_scalar`. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_016CqrLKgPqFYGZK5sjk1qe7 --- encodings/alp/src/alp/ops.rs | 2 ++ encodings/alp/src/alp_rd/ops.rs | 2 ++ encodings/bytebool/src/array.rs | 2 ++ encodings/datetime-parts/src/ops.rs | 2 ++ .../src/decimal_byte_parts/mod.rs | 2 ++ .../src/bitpacking/vtable/operations.rs | 2 ++ .../fastlanes/src/delta/vtable/operations.rs | 2 ++ .../fastlanes/src/for/vtable/operations.rs | 2 ++ .../fastlanes/src/rle/vtable/operations.rs | 2 ++ encodings/fastlanes/src/transposed_bool.rs | 2 ++ encodings/fsst/src/ops.rs | 2 ++ encodings/onpair/src/ops.rs | 2 ++ encodings/parquet-variant/src/operations.rs | 2 ++ encodings/pco/src/array.rs | 2 ++ encodings/runend/src/ops.rs | 2 ++ encodings/sequence/src/array.rs | 2 ++ encodings/sparse/src/ops.rs | 2 ++ encodings/zigzag/src/array.rs | 2 ++ encodings/zstd/src/array.rs | 2 ++ encodings/zstd/src/zstd_buffers.rs | 2 ++ vortex-array/src/array/mod.rs | 8 +++++- vortex-array/src/array/probe/array.rs | 28 +++++++++++++++++++ vortex-array/src/array/probe/mod.rs | 8 ++++++ vortex-array/src/array/vtable/operations.rs | 26 +++++++++++++++++ .../src/arrays/bool/vtable/operations.rs | 2 ++ .../src/arrays/chunked/vtable/operations.rs | 2 ++ .../src/arrays/constant/vtable/operations.rs | 2 ++ .../src/arrays/decimal/vtable/operations.rs | 2 ++ .../src/arrays/dict/vtable/operations.rs | 2 ++ .../src/arrays/extension/vtable/operations.rs | 2 ++ vortex-array/src/arrays/filter/vtable.rs | 2 ++ .../fixed_size_list/vtable/operations.rs | 2 ++ vortex-array/src/arrays/interleave/mod.rs | 2 ++ .../src/arrays/list/vtable/operations.rs | 2 ++ .../src/arrays/listview/vtable/operations.rs | 2 ++ .../src/arrays/map/vtable/operations.rs | 2 ++ .../src/arrays/masked/vtable/operations.rs | 2 ++ vortex-array/src/arrays/null/mod.rs | 2 ++ .../src/arrays/patched/vtable/operations.rs | 2 ++ .../src/arrays/piecewise_sequence/vtable.rs | 2 ++ .../src/arrays/primitive/vtable/operations.rs | 2 ++ .../src/arrays/scalar_fn/vtable/operations.rs | 2 ++ vortex-array/src/arrays/shared/vtable.rs | 2 ++ vortex-array/src/arrays/slice/vtable.rs | 2 ++ .../src/arrays/struct_/vtable/operations.rs | 2 ++ .../src/arrays/union/vtable/operations.rs | 2 ++ .../src/arrays/varbin/vtable/operations.rs | 2 ++ .../arrays/varbinview/vtable/operations.rs | 2 ++ .../src/arrays/variant/vtable/operations.rs | 2 ++ vortex-python/src/arrays/py/vtable.rs | 2 ++ 50 files changed, 161 insertions(+), 1 deletion(-) create mode 100644 vortex-array/src/array/probe/array.rs create mode 100644 vortex-array/src/array/probe/mod.rs 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 a36a1f6edab..f3f44e64f6e 100644 --- a/encodings/sequence/src/array.rs +++ b/encodings/sequence/src/array.rs @@ -425,6 +425,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/mod.rs b/vortex-array/src/array/mod.rs index 8b9b2812bfa..9b0b7ab8f95 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::*; @@ -497,7 +499,11 @@ impl DynArrayData for ArrayData { ctx: &mut ExecutionCtx, ) -> VortexResult { let view = unsafe { ArrayView::new_unchecked(this, &self.data) }; - >::scalar_at(view, index, ctx) + >::probe_scalar( + &mut ProbeState::once(view), + index, + 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..acb45fdaf3b --- /dev/null +++ b/vortex-array/src/array/probe/array.rs @@ -0,0 +1,28 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use crate::array::ArrayView; +use crate::array::VTable; + +/// Everything an encoding's `probe_scalar` runs with. +/// +/// Passed to [`OperationsVTable::probe_scalar`](crate::vtable::OperationsVTable::probe_scalar). +/// Holds the typed view of the array being read; it owns nothing, so building one per read is +/// free. +pub struct ProbeState<'a, V: VTable> { + array: ArrayView<'a, V>, +} + +impl<'a, V: VTable> ProbeState<'a, V> { + /// State for a single read of `array`. + #[inline] + pub fn once(array: ArrayView<'a, V>) -> Self { + Self { array } + } + + /// The typed view of the array being read. + #[inline] + pub fn array(&self) -> ArrayView<'a, V> { + self.array + } +} diff --git a/vortex-array/src/array/probe/mod.rs b/vortex-array/src/array/probe/mod.rs new file mode 100644 index 00000000000..c13332f115f --- /dev/null +++ b/vortex-array/src/array/probe/mod.rs @@ -0,0 +1,8 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! State passed to an encoding's +//! [`probe_scalar`](crate::vtable::OperationsVTable::probe_scalar). + +mod array; +pub use array::*; diff --git a/vortex-array/src/array/vtable/operations.rs b/vortex-array/src/array/vtable/operations.rs index 7f49e683640..96540169e14 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,29 @@ use crate::vtable::NotSupported; /// [`ArrayRef`](crate::ArrayRef) /// methods perform common checks before dispatching here. pub trait OperationsVTable { + /// Encoding-specific state retained across repeated scalar reads. + /// + /// Built once per retained probe and never for a one-off read. Use `()` when no state is + /// needed. + type ProbeState: Default + 'static; + + /// Read the non-null scalar at `index` of the array in `state`. + /// + /// Bounds and validity have been checked; the row is non-null. `state` carries the typed + /// view of the array being read. The scalar must retain the source's logical dtype, + /// including nullability. + /// + /// The default preserves the existing scalar path. + fn probe_scalar( + state: &mut ProbeState<'_, V>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + // FIXME: Remove this default once all encodings have migrated to probe_scalar. + Self::scalar_at(state.array(), index, ctx) + } + + // FIXME: Deprecate scalar_at once encodings have migrated to probe_scalar. /// Fetch the scalar at the given index. /// /// ## Preconditions @@ -35,6 +59,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 260abe0046c..f5a61ba79d7 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-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, From 8a4beb1a68732be5e230066846adab90dc4fad4d Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 15 Sep 2026 18:54:58 +0100 Subject: [PATCH 2/2] feat(array): scalar probes with one-off and repeated probe types Row access over arrays on top of the `probe_scalar` hook: - `ArrayProbe<'a>` is a bare `&ArrayRef` with no destructor. One-off reads, including every child read while recursing through nested arrays, own nothing, so there is no unwind cleanup, no out-of-line drop glue and no pinned temporary on the per-row path. This is what the CodSpeed regression on the earlier design came down to. - `RepeatedArrayProbe` owns its handle, encoding state, validity probe and child probes, created on first use and reused across rows. - Both implement `Probe`. `ArrayRef::probe()` and `ArrayRef::repeated_probe()` are the entry points; `execute_scalar` and `is_valid` are shims over `probe()`. - `ProbeState<'_, V>` carries the `ArrayView` and, for a repeated read, a borrow of the retained `RepeatedState`; it has no destructor. `state.array()`, `state.retained()`, `state.slot(i) -> impl Probe`, and the direct `child_scalar` / `child_is_valid`. Encodings never branch on policy for child reads. - `Validity::probe()` gives the retained validity accessor. - `Struct` and `Primitive` implement `probe_scalar`; `scalar_at` delegates through `ProbeState::once`. Not generic over the retention policy: a generic vtable method is instantiated in every crate that constructs arrays of the encoding, so the hot body compiles downstream without this crate's inlining; measured at 20-30 points on struct reads and rejected. The policy is a runtime `Option` check, folded into the one-off arm by `inline(always)` on the probe bodies. Local divan medians vs the develop tip, interleaved runs of 1000 samples: is_valid_per_element[1024] -10%, [256] -12%, execute_scalar_struct_wide -39%, execute_scalar_struct_simple -29%. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_016CqrLKgPqFYGZK5sjk1qe7 --- vortex-array/src/array/erased.rs | 64 ++- vortex-array/src/array/mod.rs | 33 +- vortex-array/src/array/probe/array.rs | 494 +++++++++++++++++- vortex-array/src/array/probe/mod.rs | 11 +- vortex-array/src/array/probe/validity.rs | 45 ++ vortex-array/src/array/vtable/operations.rs | 17 +- .../src/arrays/primitive/vtable/operations.rs | 14 +- .../src/arrays/struct_/vtable/operations.rs | 27 +- vortex-array/src/validity.rs | 17 + 9 files changed, 679 insertions(+), 43 deletions(-) create mode 100644 vortex-array/src/array/probe/validity.rs diff --git a/vortex-array/src/array/erased.rs b/vortex-array/src/array/erased.rs index 76c0b33cd18..2b8434303da 100644 --- a/vortex-array/src/array/erased.rs +++ b/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 { - 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 { - 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 { Ok(!self.is_valid(index, ctx)?) } diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs index 9b0b7ab8f95..25351e63af6 100644 --- a/vortex-array/src/array/mod.rs +++ b/vortex-array/src/array/mod.rs @@ -222,15 +222,26 @@ pub(crate) trait DynArrayData: 'static + private::Sealed + Send + Sync + Debug { ctx: &mut ExecutionCtx, ) -> VortexResult; - /// 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; + + /// 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; } /// Trait for converting a type into a Vortex [`ArrayRef`]. @@ -492,12 +503,13 @@ impl DynArrayData for ArrayData { V::execute(typed, ctx) } - fn execute_scalar( + fn probe_scalar_once( &self, this: &ArrayRef, index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { + // SAFETY: this adapter belongs to the ArrayData stored in `this`. let view = unsafe { ArrayView::new_unchecked(this, &self.data) }; >::probe_scalar( &mut ProbeState::once(view), @@ -505,6 +517,19 @@ impl DynArrayData for ArrayData { ctx, ) } + + fn probe_scalar_retained( + &self, + this: &ArrayRef, + index: usize, + storage: &mut ProbeStorage, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + // SAFETY: this adapter belongs to the ArrayData stored in `this`. + let view = unsafe { ArrayView::new_unchecked(this, &self.data) }; + let mut state = ProbeState::repeated(view, storage.get_or_init()?); + >::probe_scalar(&mut state, index, ctx) + } } /// Wrapper around `&mut dyn Hasher` that implements `Hasher` (and is `Sized`). diff --git a/vortex-array/src/array/probe/array.rs b/vortex-array/src/array/probe/array.rs index acb45fdaf3b..12f2a097676 100644 --- a/vortex-array/src/array/probe/array.rs +++ b/vortex-array/src/array/probe/array.rs @@ -1,23 +1,230 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::Any; + +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::ExecutionCtx; use crate::array::ArrayView; use crate::array::VTable; +use crate::array::probe::ProbeValidity; +use crate::scalar::Scalar; +use crate::vtable::OperationsVTable; + +/// Row access shared by [`ArrayProbe`] and [`RepeatedArrayProbe`]. +/// +/// Both probe kinds expose the same reads; the trait lets code that recurses through child +/// slots treat them uniformly. +pub trait Probe { + /// The array this probe reads from. + fn array(&self) -> &ArrayRef; + + /// Read the scalar at `index`, including its nullness. + fn execute_scalar(&mut self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult; + + /// Whether the row at `index` is valid. + fn execute_is_valid(&mut self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult; + + /// Whether the row at `index` is null. + fn execute_is_invalid(&mut self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(!self.execute_is_valid(index, ctx)?) + } +} + +/// A one-off row accessor over a borrowed array. +/// +/// It is a bare reference: it retains nothing between reads and has no destructor, so building +/// one per read, including for every child slot of a nested array, costs nothing. Use +/// [`RepeatedArrayProbe`] when the same array is read many times. +/// +/// This is split is a performance concern -- If we merge these there was a performance regression +/// for the one-off array probe. +#[derive(Clone, Copy)] +pub struct ArrayProbe<'a> { + array: &'a ArrayRef, +} + +impl<'a> ArrayProbe<'a> { + /// Borrow an array for one-off reads. + #[inline] + pub fn new(array: &'a ArrayRef) -> Self { + Self { array } + } +} + +impl Probe for ArrayProbe<'_> { + #[inline] + fn array(&self) -> &ArrayRef { + self.array + } + + // Always inlined: this is the per-element read path, and callers that read once per row + // must get a single body with one dynamic call rather than an extra frame. Measured at + // about 10% on the per-element benchmarks against plain `#[inline]`. + #[allow(clippy::inline_always)] + #[inline(always)] + fn execute_scalar(&mut self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { + let array = self.array; + if !self.execute_is_valid(index, ctx)? { + return Ok(Scalar::null(array.dtype().clone())); + } + check_dtype( + array, + array.dyn_array().probe_scalar_once(array, index, ctx), + ) + } + + // Always inlined, see `execute_scalar`. + #[allow(clippy::inline_always)] + #[inline(always)] + fn execute_is_valid(&mut self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { + let array = self.array; + check_bounds(array, index)?; + if !array.dtype().is_nullable() { + return Ok(true); + } + // Matching the validity directly keeps this path free of any retained temporary. + array.validity()?.execute_is_valid(index, ctx) + } +} + +/// A row accessor that owns its array and keeps preparation between reads. +/// +/// The encoding's state, the validity probe and the probes over child slots are created on +/// first use and reused by every following read. Dropping the probe drops all of them. Array +/// handles share their buffers, so the probe can outlive the handle it was built from. Probes +/// are local to a thread. +pub struct RepeatedArrayProbe { + array: ArrayRef, + storage: ProbeStorage, + /// Created on the first read of a nullable array. Boxed because the validity probe holds + /// another `RepeatedArrayProbe`. + validity: Option>, +} + +impl RepeatedArrayProbe { + /// Own an array for repeated reads. + pub fn new(array: ArrayRef) -> Self { + Self { + array, + storage: ProbeStorage::default(), + validity: None, + } + } +} + +impl Probe for RepeatedArrayProbe { + fn array(&self) -> &ArrayRef { + &self.array + } + + 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 = &self.array; + let result = array + .dyn_array() + .probe_scalar_retained(array, index, &mut self.storage, ctx); + check_dtype(array, result) + } -/// Everything an encoding's `probe_scalar` runs with. + fn execute_is_valid(&mut self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { + check_bounds(&self.array, index)?; + if !self.array.dtype().is_nullable() { + return Ok(true); + } + if self.validity.is_none() { + self.validity = Some(Box::new(self.array.validity()?.probe())); + } + self.validity + .as_mut() + .ok_or_else(|| vortex_err!("validity probe was just initialized"))? + .execute_is_valid(index, ctx) + } +} + +/// Pass an encoding's result through, checking its dtype in debug builds. +/// +/// Unwrapping and re-wrapping the result here would cost an extra copy of the scalar on every +/// read. +#[inline] +fn check_dtype(array: &ArrayRef, result: VortexResult) -> VortexResult { + result.inspect(|scalar| { + debug_assert_eq!(scalar.dtype(), array.dtype(), "Scalar dtype mismatch"); + }) +} + +#[inline] +fn check_bounds(array: &ArrayRef, index: usize) -> VortexResult<()> { + if index >= array.len() { + return Err(out_of_bounds(index, array.len())); + } + Ok(()) +} + +/// Kept out of line so the error path, which captures a backtrace, does not count against the +/// hot probe bodies when the inliner sizes them. +#[cold] +#[inline(never)] +fn out_of_bounds(index: usize, len: usize) -> VortexError { + vortex_err!(OutOfBounds: index, 0, len) +} + +/// The encoding state type of `V`'s operations vtable. +pub type EncodingProbeState = + <::OperationsVTable as OperationsVTable>::ProbeState; + +/// Everything an encoding's `probe_scalar` runs with: the typed view of the array being read +/// and, for a repeated read, a borrow of the state its [`RepeatedArrayProbe`] keeps. /// /// Passed to [`OperationsVTable::probe_scalar`](crate::vtable::OperationsVTable::probe_scalar). -/// Holds the typed view of the array being read; it owns nothing, so building one per read is -/// free. +/// A one-off read gets [`ProbeState::once`], which holds only the view. A repeated read borrows +/// the [`RepeatedState`]: the encoding's own state and the lazily created child probes. Neither +/// owns anything, so building one per read is free. +/// +/// Encodings never inspect the policy: [`ProbeState::slot`] hands out a probe over a child under +/// whichever policy is in force, [`ProbeState::child_scalar`] and [`ProbeState::child_is_valid`] +/// read one directly, and [`ProbeState::retained`] hands out the encoding state only when it is +/// kept. pub struct ProbeState<'a, V: VTable> { array: ArrayView<'a, V>, + retained: Option<&'a mut RepeatedState>>, +} + +/// What a [`RepeatedArrayProbe`] keeps for its encoding between reads. +pub struct RepeatedState { + state: S, + /// Probes over the source's child slots, allocated on first use. Slots that are never + /// requested stay empty. + slots: Vec>, } impl<'a, V: VTable> ProbeState<'a, V> { - /// State for a single read of `array`. + /// State for a single read of `array`. Encodings use this to run their `probe_scalar` path + /// from `scalar_at`. #[inline] pub fn once(array: ArrayView<'a, V>) -> Self { - Self { array } + Self { + array, + retained: None, + } + } + + /// State for a read through a [`RepeatedArrayProbe`], borrowing what it keeps. + #[inline] + pub(crate) fn repeated( + array: ArrayView<'a, V>, + retained: &'a mut RepeatedState>, + ) -> Self { + Self { + array, + retained: Some(retained), + } } /// The typed view of the array being read. @@ -25,4 +232,281 @@ impl<'a, V: VTable> ProbeState<'a, V> { pub fn array(&self) -> ArrayView<'a, V> { self.array } + + /// The encoding's retained state, or `None` for a one-off read. + /// + /// Encodings whose repeated algorithm differs from their one-off one branch on this. + #[inline] + pub fn retained(&mut self) -> Option<&mut EncodingProbeState> { + self.retained + .as_deref_mut() + .map(|repeated| &mut repeated.state) + } + + /// A probe over the array's child in `slot`, under this state's policy. + /// + /// For a one-off read this is a borrowed [`ArrayProbe`] built for the call; for a repeated + /// read it is the [`RepeatedArrayProbe`] kept in the slot table, created on first use. + /// Errors for an out-of-bounds or absent slot. + #[inline] + pub fn slot(&mut self, slot: usize) -> VortexResult { + let parent = self.array.array(); + Ok(match &mut self.retained { + None => ChildProbe::Once(ArrayProbe::new(child_of(parent, slot)?)), + Some(repeated) => ChildProbe::Repeated(repeated.child_probe(parent, slot)?), + }) + } + + /// Read the scalar at `index` of the array's child in `slot`, including its nullness. + /// + /// Equivalent to [`Self::slot`] followed by [`Probe::execute_scalar`], without building + /// the intermediate probe: use this for the per-row read of a child, and `slot` when a + /// probe needs to be held or passed on. + #[inline] + pub fn child_scalar( + &mut self, + slot: usize, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let parent = self.array.array(); + match &mut self.retained { + None => ArrayProbe::new(child_of(parent, slot)?).execute_scalar(index, ctx), + Some(repeated) => repeated + .child_probe(parent, slot)? + .execute_scalar(index, ctx), + } + } + + /// Whether the row at `index` of the array's child in `slot` is valid. + /// + /// Same relationship to [`Self::slot`] as [`Self::child_scalar`]. + #[inline] + pub fn child_is_valid( + &mut self, + slot: usize, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let parent = self.array.array(); + match &mut self.retained { + None => ArrayProbe::new(child_of(parent, slot)?).execute_is_valid(index, ctx), + Some(repeated) => repeated + .child_probe(parent, slot)? + .execute_is_valid(index, ctx), + } + } +} + +/// The probe [`ProbeState::slot`] hands out: a bare one-off probe or the retained one. +enum ChildProbe<'a> { + Once(ArrayProbe<'a>), + Repeated(&'a mut RepeatedArrayProbe), +} + +impl Probe for ChildProbe<'_> { + #[inline] + fn array(&self) -> &ArrayRef { + match self { + Self::Once(probe) => probe.array(), + Self::Repeated(probe) => probe.array(), + } + } + + // Always inlined for the same reason as `ArrayProbe::execute_scalar`: the one-off arm is + // the per-element path through nested arrays. + #[allow(clippy::inline_always)] + #[inline(always)] + fn execute_scalar(&mut self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { + match self { + Self::Once(probe) => probe.execute_scalar(index, ctx), + Self::Repeated(probe) => probe.execute_scalar(index, ctx), + } + } + + #[allow(clippy::inline_always)] + #[inline(always)] + fn execute_is_valid(&mut self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { + match self { + Self::Once(probe) => probe.execute_is_valid(index, ctx), + Self::Repeated(probe) => probe.execute_is_valid(index, ctx), + } + } +} + +impl Default for RepeatedState { + fn default() -> Self { + Self { + state: S::default(), + slots: Vec::new(), + } + } +} + +impl RepeatedState { + /// Get or create the retained probe over `parent`'s child in `slot`. + fn child_probe( + &mut self, + parent: &ArrayRef, + slot: usize, + ) -> VortexResult<&mut RepeatedArrayProbe> { + let child = child_of(parent, slot)?; + if self.slots.is_empty() { + self.slots.resize_with(parent.slots().len(), || None); + } + Ok(self.slots[slot].get_or_insert_with(|| RepeatedArrayProbe::new(child.clone()))) + } +} + +/// The child of `parent` in `slot`, as an error if the slot is out of bounds or absent. +#[inline] +fn child_of(parent: &ArrayRef, slot: usize) -> VortexResult<&ArrayRef> { + parent + .slots() + .get(slot) + .ok_or_else(|| vortex_err!("Probe slot {slot} is out of bounds"))? + .as_ref() + .ok_or_else(|| vortex_err!("Probe slot {slot} is absent")) +} + +/// Type-erased, lazily initialized [`RepeatedState`] owned by a [`RepeatedArrayProbe`]. +#[derive(Default)] +pub(crate) struct ProbeStorage(Option>); + +impl ProbeStorage { + pub(crate) fn get_or_init( + &mut self, + ) -> VortexResult<&mut RepeatedState> { + self.0 + .get_or_insert_with(|| Box::new(RepeatedState::::default())) + .downcast_mut::>() + .ok_or_else(|| vortex_err!("Probe state type mismatch")) + } +} + +#[cfg(test)] +mod tests { + use std::mem::needs_drop; + use std::mem::size_of; + + 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::Struct; + use crate::arrays::StructArray; + + fn nullable_ints() -> ArrayRef { + PrimitiveArray::from_option_iter([Some(10i32), None, Some(30)]).into_array() + } + + fn check_reads(probe: &mut dyn Probe, ctx: &mut ExecutionCtx) -> VortexResult<()> { + assert!(probe.execute_scalar(3, ctx).is_err()); + assert!(probe.execute_is_valid(3, ctx).is_err()); + assert!(probe.execute_scalar(1, ctx)?.is_null()); + assert!(!probe.execute_is_valid(1, ctx)?); + assert_eq!(probe.execute_scalar(2, ctx)?, Scalar::from(Some(30i32))); + assert_eq!(probe.execute_scalar(0, ctx)?, Scalar::from(Some(10i32))); + Ok(()) + } + + #[test] + fn once_probe_checks_bounds_and_nulls() -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let array = nullable_ints(); + check_reads(&mut array.probe(), &mut ctx) + } + + #[test] + fn once_probe_is_a_bare_reference() { + assert_eq!(size_of::>(), size_of::<&ArrayRef>()); + assert!(!needs_drop::>()); + } + + #[test] + fn repeated_probe_initializes_lazily_and_outlives_handle() -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let array = nullable_ints(); + let mut probe = RepeatedArrayProbe::new(array.clone()); + drop(array); + assert!(probe.storage.0.is_none()); + assert!(probe.validity.is_none()); + + check_reads(&mut probe, &mut ctx)?; + + assert!(probe.storage.0.is_some()); + assert!(probe.validity.is_some()); + Ok(()) + } + + #[test] + fn once_state_reads_children_without_retaining() -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let array = struct_of_two_fields()?; + let typed = array + .as_opt::() + .ok_or_else(|| vortex_err!("expected a struct"))?; + let mut state = ProbeState::once(typed); + + assert!(state.slot(5).is_err()); + // Slot 0 is the struct's absent validity. + assert!(state.slot(0).is_err()); + { + let mut field = state.slot(2)?; + assert_eq!(field.array().len(), 2); + assert_eq!(field.execute_scalar(1, &mut ctx)?, Scalar::from(4i64)); + assert!(field.execute_is_valid(0, &mut ctx)?); + } + assert_eq!(state.child_scalar(2, 0, &mut ctx)?, Scalar::from(3i64)); + assert!(state.child_is_valid(1, 0, &mut ctx)?); + assert!(state.retained().is_none()); + assert!(!needs_drop::>()); + Ok(()) + } + + #[test] + fn repeated_state_creates_children_on_demand() -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let array = struct_of_two_fields()?; + let typed = array + .as_opt::() + .ok_or_else(|| vortex_err!("expected a struct"))?; + let mut repeated = RepeatedState::<()>::default(); + { + let mut state = ProbeState::repeated(typed, &mut repeated); + assert!(state.slot(5).is_err()); + assert!(state.slot(0).is_err()); + assert!(state.retained().is_some()); + let mut field = state.slot(2)?; + assert_eq!(field.execute_scalar(1, &mut ctx)?, Scalar::from(4i64)); + assert_eq!(field.execute_scalar(0, &mut ctx)?, Scalar::from(3i64)); + } + + assert_eq!(repeated.slots.len(), array.slots().len()); + assert!(repeated.slots[1].is_none()); + let child = repeated.slots[2] + .as_ref() + .ok_or_else(|| vortex_err!("missing child probe"))?; + assert_eq!(child.array().len(), 2); + Ok(()) + } + + fn struct_of_two_fields() -> VortexResult { + Ok(StructArray::from_fields(&[ + ("a", PrimitiveArray::from_iter([1i32, 2]).into_array()), + ("b", PrimitiveArray::from_iter([3i64, 4]).into_array()), + ])? + .into_array()) + } + + #[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 index c13332f115f..0060153ae5f 100644 --- a/vortex-array/src/array/probe/mod.rs +++ b/vortex-array/src/array/probe/mod.rs @@ -1,8 +1,15 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! State passed to an encoding's -//! [`probe_scalar`](crate::vtable::OperationsVTable::probe_scalar). +//! Row access over arrays. +//! +//! [`ArrayProbe`] reads a borrowed array once and retains nothing. [`RepeatedArrayProbe`] owns +//! its array and keeps encoding state, its validity probe and child probes between reads. Both +//! implement [`Probe`]. Encodings implement a single +//! [`probe_scalar`](crate::vtable::OperationsVTable::probe_scalar) that serves both through +//! [`ProbeState`]. mod array; pub use array::*; +mod validity; +pub use validity::*; diff --git a/vortex-array/src/array/probe/validity.rs b/vortex-array/src/array/probe/validity.rs new file mode 100644 index 00000000000..1ce85178b8d --- /dev/null +++ b/vortex-array/src/array/probe/validity.rs @@ -0,0 +1,45 @@ +// 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::Probe; +use crate::array::probe::RepeatedArrayProbe; + +/// A retained validity accessor built by [`Validity::probe`](crate::validity::Validity::probe). +/// +/// Uniform validity retains nothing; array-backed validity keeps a [`RepeatedArrayProbe`] over +/// the boolean array. One-off reads use +/// [`Validity::execute_is_valid`](crate::validity::Validity::execute_is_valid) directly. +pub enum ProbeValidity { + /// Validity is uniform, so no lookup is needed. + Constant(bool), + /// Validity backed by a boolean array. + Array(RepeatedArrayProbe), +} + +impl ProbeValidity { + /// Returns whether the row at `index` is valid. + #[inline] + 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 96540169e14..2d4e8ae0ca1 100644 --- a/vortex-array/src/array/vtable/operations.rs +++ b/vortex-array/src/array/vtable/operations.rs @@ -18,19 +18,24 @@ use crate::vtable::NotSupported; /// [`ArrayRef`](crate::ArrayRef) /// methods perform common checks before dispatching here. pub trait OperationsVTable { - /// Encoding-specific state retained across repeated scalar reads. + /// Encoding-specific state retained by repeated scalar access. /// - /// Built once per retained probe and never for a one-off read. Use `()` when no state is - /// needed. + /// Built once per repeated probe and never for a one-off read. 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 the non-null scalar at `index` of the array in `state`. /// /// Bounds and validity have been checked; the row is non-null. `state` carries the typed - /// view of the array being read. The scalar must retain the source's logical dtype, - /// including nullability. + /// view of the array and, for a read through a + /// [`RepeatedArrayProbe`](crate::RepeatedArrayProbe), the state that probe keeps. Read + /// children through [`ProbeState::child_scalar`], or hold a child probe from + /// [`ProbeState::slot`]; both follow the read's policy without the encoding having to know + /// it. Take encoding state from [`ProbeState::retained`]. The scalar must retain the source's + /// logical dtype, including nullability. /// - /// The default preserves the existing scalar path. + /// The default preserves the existing scalar path without adding caching. fn probe_scalar( state: &mut ProbeState<'_, V>, index: usize, diff --git a/vortex-array/src/arrays/primitive/vtable/operations.rs b/vortex-array/src/arrays/primitive/vtable/operations.rs index 9501513fb1c..0f77ed22d9d 100644 --- a/vortex-array/src/arrays/primitive/vtable/operations.rs +++ b/vortex-array/src/arrays/primitive/vtable/operations.rs @@ -6,6 +6,7 @@ use vortex_error::VortexResult; use crate::ExecutionCtx; use crate::array::ArrayView; use crate::array::OperationsVTable; +use crate::array::ProbeState; use crate::arrays::Primitive; use crate::match_each_native_ptype; use crate::scalar::Scalar; @@ -13,13 +14,22 @@ use crate::scalar::Scalar; impl OperationsVTable for Primitive { type ProbeState = (); - fn scalar_at( - array: ArrayView<'_, Primitive>, + fn probe_scalar( + state: &mut ProbeState<'_, Primitive>, index: usize, _ctx: &mut ExecutionCtx, ) -> VortexResult { + let array = state.array(); Ok(match_each_native_ptype!(array.ptype(), |T| { Scalar::primitive(array.as_slice::()[index], array.dtype().nullability()) })) } + + fn scalar_at( + array: ArrayView<'_, Primitive>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Self::probe_scalar(&mut ProbeState::once(array), index, ctx) + } } diff --git a/vortex-array/src/arrays/struct_/vtable/operations.rs b/vortex-array/src/arrays/struct_/vtable/operations.rs index 54d086bf7c9..f315b781ccd 100644 --- a/vortex-array/src/arrays/struct_/vtable/operations.rs +++ b/vortex-array/src/arrays/struct_/vtable/operations.rs @@ -6,25 +6,30 @@ use vortex_error::VortexResult; use crate::ExecutionCtx; use crate::array::ArrayView; use crate::array::OperationsVTable; +use crate::array::ProbeState; use crate::arrays::Struct; use crate::arrays::struct_::StructArrayExt; +use crate::arrays::struct_::StructSlots; use crate::scalar::Scalar; use crate::scalar::ScalarValue; impl OperationsVTable for Struct { type ProbeState = (); - fn scalar_at( - array: ArrayView<'_, Struct>, + fn probe_scalar( + state: &mut ProbeState<'_, Struct>, index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - let field_values = array - .iter_unmasked_fields() - .map(|field| field.execute_scalar(index, ctx).map(Scalar::into_value)) - .collect::>>()?; + let array = state.array(); + let nfields = array.iter_unmasked_fields().len(); + let mut field_values = Vec::with_capacity(nfields); + for field in 0..nfields { + let slot = StructSlots::FIELDS_OFFSET + field; + field_values.push(state.child_scalar(slot, index, ctx)?.into_value()); + } // SAFETY: The vtable guarantees index is in-bounds and non-null before this is called. - // Each field's scalar_at returns a value with the field's own dtype. + // Each field read returns a value with the field's own dtype. Ok(unsafe { Scalar::new_unchecked( array.dtype().clone(), @@ -32,4 +37,12 @@ impl OperationsVTable for Struct { ) }) } + + fn scalar_at( + array: ArrayView<'_, Struct>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Self::probe_scalar(&mut ProbeState::once(array), index, ctx) + } } diff --git a/vortex-array/src/validity.rs b/vortex-array/src/validity.rs index a2c10e30ca7..03d970b89a2 100644 --- a/vortex-array/src/validity.rs +++ b/vortex-array/src/validity.rs @@ -24,6 +24,8 @@ use crate::ArrayRef; use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; +use crate::ProbeValidity; +use crate::RepeatedArrayProbe; use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; @@ -162,7 +164,21 @@ impl Validity { } } + /// Create a retained accessor for repeated lookups. + /// + /// Mirrors [`ArrayRef::repeated_probe`]. Uniform validity retains nothing; array-backed + /// validity keeps a probe over the underlying boolean array. For a single lookup use + /// [`Self::execute_is_valid`]. + pub fn probe(self) -> ProbeValidity { + match self { + Self::NonNullable | Self::AllValid => ProbeValidity::Constant(true), + Self::AllInvalid => ProbeValidity::Constant(false), + Self::Array(array) => ProbeValidity::Array(RepeatedArrayProbe::new(array)), + } + } + /// 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)?)