From 4394b23aeb26080ed7c406a104f523f7b60ee7f7 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Wed, 16 Sep 2026 10:48:31 +0000 Subject: [PATCH] perf(array): forward probe_scalar through pass-through encodings Slice, Dict, Chunked, Shared, Masked, Extension, FoR, ZigZag, ALP, FSST and DateTimeParts cost almost nothing per row themselves, but each still served reads from the default `probe_scalar`, which forwards to `scalar_at` and reads children with a one-off `execute_scalar`. Any retained preparation below them was rebuilt on every row, so an encoding that keeps decoded state only paid off at the root of a tree. Each now implements `probe_scalar` once and reads its children through `state.slot(..)`, with `scalar_at` delegating via `ProbeState::once` so one-off and repeated reads share a body. Two need more than a slot read: Shared resolves to either its source or its computed cache and rebuilds its probe when that changes, and FSST rebuilds its codes array from a buffer and the offsets slot, so it retains one probe over that array. ALP's patched path still reads through `Patches`; only its unpatched read is probed. Adds a conformance test asserting a retained probe agrees with one-off reads over a backwards, forwards and sparse visit order, so every encoding in the consistency suite covers its probe path. Signed-off-by: "Joe Isaacs" Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Ui7E5mtZLgh4j41ugrLXbS Signed-off-by: Joe Isaacs --- encodings/alp/src/alp/ops.rs | 25 ++++++-- encodings/datetime-parts/src/ops.rs | 64 ++++++++++++------- encodings/fastlanes/src/for_/tests.rs | 17 +++++ .../fastlanes/src/for_/vtable/operations.rs | 27 ++++++-- encodings/fsst/src/ops.rs | 36 +++++++++-- encodings/zigzag/src/array.rs | 20 +++++- .../src/arrays/chunked/vtable/operations.rs | 19 +++++- .../src/arrays/dict/vtable/operations.rs | 37 +++++++---- .../src/arrays/extension/vtable/operations.rs | 21 ++++-- vortex-array/src/arrays/masked/tests.rs | 24 +++++++ .../src/arrays/masked/vtable/operations.rs | 22 ++++++- vortex-array/src/arrays/shared/vtable.rs | 32 +++++++++- vortex-array/src/arrays/slice/vtable.rs | 16 ++++- .../src/compute/conformance/consistency.rs | 58 +++++++++++++++++ 14 files changed, 355 insertions(+), 63 deletions(-) diff --git a/encodings/alp/src/alp/ops.rs b/encodings/alp/src/alp/ops.rs index 58b5ee99360..b24fd5b8456 100644 --- a/encodings/alp/src/alp/ops.rs +++ b/encodings/alp/src/alp/ops.rs @@ -3,32 +3,41 @@ use vortex_array::ArrayView; use vortex_array::ExecutionCtx; +use vortex_array::ProbeState; use vortex_array::scalar::Scalar; use vortex_array::vtable::OperationsVTable; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_err; use crate::ALP; use crate::ALPArrayExt; -use crate::ALPArraySlotsExt; use crate::ALPFloat; +use crate::ALPSlots; use crate::match_each_alp_float_ptype; impl OperationsVTable for ALP { type ProbeState = (); - fn scalar_at( - array: ArrayView<'_, ALP>, + fn probe_scalar( + state: &mut ProbeState<'_, ALP>, index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { + if !state.is_valid(index, ctx)? { + return Ok(Scalar::null(state.array().dtype().clone())); + } + let array = state.array(); if let Some(patches) = array.patches() && let Some(patch) = patches.get_patched(index)? { return patch.cast(array.dtype()); } - let encoded_val = array.encoded().execute_scalar(index, ctx)?; + let encoded_val = state + .slot(ALPSlots::ENCODED)? + .ok_or_else(|| vortex_err!("ALP encoded slot is missing"))? + .execute_scalar(index, ctx)?; Ok(match_each_alp_float_ptype!(array.dtype().as_ptype(), |T| { let encoded_val: ::ALPInt = @@ -39,4 +48,12 @@ impl OperationsVTable for ALP { ) })) } + + fn scalar_at( + array: ArrayView<'_, ALP>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Self::probe_scalar(&mut ProbeState::once(array), index, ctx) + } } diff --git a/encodings/datetime-parts/src/ops.rs b/encodings/datetime-parts/src/ops.rs index 0ee4f9e315c..4fc0c0e3638 100644 --- a/encodings/datetime-parts/src/ops.rs +++ b/encodings/datetime-parts/src/ops.rs @@ -3,27 +3,30 @@ use vortex_array::ArrayView; use vortex_array::ExecutionCtx; +use vortex_array::ProbeState; use vortex_array::dtype::DType; use vortex_array::extension::datetime::Timestamp; use vortex_array::scalar::Scalar; use vortex_array::vtable::OperationsVTable; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_err; use vortex_error::vortex_panic; use crate::DateTimeParts; -use crate::array::DateTimePartsArraySlotsExt; +use crate::DateTimePartsSlots; use crate::timestamp; use crate::timestamp::TimestampParts; impl OperationsVTable for DateTimeParts { type ProbeState = (); - fn scalar_at( - array: ArrayView<'_, DateTimeParts>, + fn probe_scalar( + state: &mut ProbeState<'_, DateTimeParts>, index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { + let array = state.array(); let DType::Extension(ext) = array.dtype().clone() else { vortex_panic!( "DateTimePartsArray must have extension dtype, found {}", @@ -35,28 +38,19 @@ impl OperationsVTable for DateTimeParts { vortex_panic!(Compute: "must decode TemporalMetadata from extension metadata"); }; - if !array.as_ref().is_valid(index, ctx)? { + if !state.is_valid(index, ctx)? { return Ok(Scalar::null(DType::Extension(ext))); } - let days: i32 = array - .days() - .execute_scalar(index, ctx)? - .as_primitive() - .as_::() - .vortex_expect("days fits in i32"); - let seconds: i32 = array - .seconds() - .execute_scalar(index, ctx)? - .as_primitive() - .as_::() - .vortex_expect("seconds fits in i32"); - let subseconds: i32 = array - .subseconds() - .execute_scalar(index, ctx)? - .as_primitive() - .as_::() - .vortex_expect("subseconds fits in i32"); + let days = part_at(state, DateTimePartsSlots::DAYS, "days", index, ctx)?; + let seconds = part_at(state, DateTimePartsSlots::SECONDS, "seconds", index, ctx)?; + let subseconds = part_at( + state, + DateTimePartsSlots::SUBSECONDS, + "subseconds", + index, + ctx, + )?; let ts = timestamp::combine( TimestampParts { @@ -72,4 +66,30 @@ impl OperationsVTable for DateTimeParts { Scalar::primitive(ts, ext.storage_dtype().nullability()), )) } + + fn scalar_at( + array: ArrayView<'_, DateTimeParts>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Self::probe_scalar(&mut ProbeState::once(array), index, ctx) + } +} + +/// Reads one timestamp part out of `slot`, through the probe so a repeated read keeps the +/// child's preparation. +fn part_at( + state: &mut ProbeState<'_, DateTimeParts>, + slot: usize, + name: &'static str, + index: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult { + Ok(state + .slot(slot)? + .ok_or_else(|| vortex_err!("DateTimeParts {name} slot is missing"))? + .execute_scalar(index, ctx)? + .as_primitive() + .as_::() + .vortex_expect("timestamp part fits in i32")) } diff --git a/encodings/fastlanes/src/for_/tests.rs b/encodings/fastlanes/src/for_/tests.rs index b930446575d..8fce5d69b0a 100644 --- a/encodings/fastlanes/src/for_/tests.rs +++ b/encodings/fastlanes/src/for_/tests.rs @@ -146,6 +146,23 @@ fn drifting_u32(len: u32) -> PrimitiveArray { PrimitiveArray::from_iter((0..len).map(|i| (i / 1024) * 1_000_000 + i % 100)) } +#[test] +fn repeated_probe_across_sliced_chunks() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = drifting_u32(3000); + let encoded = FoR::encode_chunked(values.clone(), &mut ctx)?; + let values = values.into_array(); + let sliced = encoded.into_array().slice(1000..2100)?; + let mut probe = sliced.repeated_probe(); + for index in [1048, 24, 0, 1099, 23, 1047, 24] { + assert_eq!( + probe.execute_scalar(index, &mut ctx)?, + values.execute_scalar(index + 1000, &mut ctx)? + ); + } + Ok(()) +} + #[rstest] #[case::empty(PrimitiveArray::from_iter(Vec::::new()))] #[case::one(PrimitiveArray::from_iter([7u32]))] diff --git a/encodings/fastlanes/src/for_/vtable/operations.rs b/encodings/fastlanes/src/for_/vtable/operations.rs index a7e1f197567..db5fa8e38d1 100644 --- a/encodings/fastlanes/src/for_/vtable/operations.rs +++ b/encodings/fastlanes/src/for_/vtable/operations.rs @@ -3,28 +3,37 @@ use vortex_array::ArrayView; use vortex_array::ExecutionCtx; +use vortex_array::ProbeState; use vortex_array::match_each_integer_ptype; use vortex_array::scalar::Scalar; use vortex_array::vtable::OperationsVTable; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_err; use super::FoR; use crate::FL_CHUNK_SIZE; use crate::for_::array::FoRArrayExt; -use crate::for_::array::FoRArraySlotsExt; +use crate::for_::array::FoRSlots; impl OperationsVTable for FoR { type ProbeState = (); - fn scalar_at( - array: ArrayView<'_, FoR>, + fn probe_scalar( + state: &mut ProbeState<'_, FoR>, index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - let encoded_pvalue = array.encoded().execute_scalar(index, ctx)?; + let array = state.array(); + let encoded_pvalue = state + .slot(FoRSlots::ENCODED)? + .ok_or_else(|| vortex_err!("FoR encoded slot is missing"))? + .execute_scalar(index, ctx)?; let encoded_pvalue = encoded_pvalue.as_primitive(); let chunk = (usize::from(array.offset()) + index) / FL_CHUNK_SIZE; - let reference = array.references().execute_scalar(chunk, ctx)?; + let reference = state + .slot(FoRSlots::REFERENCES)? + .ok_or_else(|| vortex_err!("FoR references slot is missing"))? + .execute_scalar(chunk, ctx)?; let reference = reference.as_primitive(); Ok(match_each_integer_ptype!(array.ptype(), |P| { @@ -41,6 +50,14 @@ impl OperationsVTable for FoR { .unwrap_or_else(|| Scalar::null(array.dtype().clone())) })) } + + fn scalar_at( + array: ArrayView<'_, FoR>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Self::probe_scalar(&mut ProbeState::once(array), index, ctx) + } } #[cfg(test)] diff --git a/encodings/fsst/src/ops.rs b/encodings/fsst/src/ops.rs index 03560b2c8b8..31e270a097c 100644 --- a/encodings/fsst/src/ops.rs +++ b/encodings/fsst/src/ops.rs @@ -3,6 +3,9 @@ use vortex_array::ArrayView; use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::ProbeState; +use vortex_array::RepeatedArrayProbe; use vortex_array::arrays::varbin::varbin_scalar; use vortex_array::scalar::Scalar; use vortex_array::vtable::OperationsVTable; @@ -13,18 +16,43 @@ use vortex_error::VortexResult; use crate::FSST; use crate::FSSTArrayExt; +/// The codes array is rebuilt from the codes buffer and the offsets slot on every read, so a +/// repeated probe keeps one probe over it and the offsets child keeps its preparation. +#[derive(Default)] +pub struct FsstProbeState { + codes: Option, +} + impl OperationsVTable for FSST { - type ProbeState = (); + type ProbeState = FsstProbeState; - fn scalar_at( - array: ArrayView<'_, FSST>, + fn probe_scalar( + state: &mut ProbeState<'_, FSST>, index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - let compressed = array.codes().execute_scalar(index, ctx)?; + if !state.is_valid(index, ctx)? { + return Ok(Scalar::null(state.array().dtype().clone())); + } + let array = state.array(); + let compressed = match state.retained() { + None => array.codes().into_array().execute_scalar(index, ctx)?, + Some(retained) => retained + .codes + .get_or_insert_with(|| RepeatedArrayProbe::new(array.codes().into_array())) + .execute_scalar(index, ctx)?, + }; let binary_datum = compressed.as_binary().value().vortex_expect("non-null"); let decoded_buffer = ByteBuffer::from(array.decompressor().decompress(binary_datum)); Ok(varbin_scalar(decoded_buffer, array.dtype())) } + + fn scalar_at( + array: ArrayView<'_, FSST>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Self::probe_scalar(&mut ProbeState::once(array), index, ctx) + } } diff --git a/encodings/zigzag/src/array.rs b/encodings/zigzag/src/array.rs index 09076fcbbdf..187b2309803 100644 --- a/encodings/zigzag/src/array.rs +++ b/encodings/zigzag/src/array.rs @@ -16,6 +16,7 @@ use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; use vortex_array::IntoArray; +use vortex_array::ProbeState; use vortex_array::TypedArrayRef; use vortex_array::array_slots; use vortex_array::buffer::BufferHandle; @@ -33,6 +34,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -233,12 +235,16 @@ impl Default for ZigZagData { impl OperationsVTable for ZigZag { type ProbeState = (); - fn scalar_at( - array: ArrayView<'_, ZigZag>, + fn probe_scalar( + state: &mut ProbeState<'_, ZigZag>, index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - let scalar = array.encoded().execute_scalar(index, ctx)?; + let array = state.array(); + let scalar = state + .slot(ZigZagSlots::ENCODED)? + .ok_or_else(|| vortex_err!("ZigZag encoded slot is missing"))? + .execute_scalar(index, ctx)?; if scalar.is_null() { return scalar.primitive_reinterpret_cast(ZigZagArrayExt::ptype(&array)); } @@ -255,6 +261,14 @@ impl OperationsVTable for ZigZag { ) })) } + + fn scalar_at( + array: ArrayView<'_, ZigZag>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Self::probe_scalar(&mut ProbeState::once(array), index, ctx) + } } impl ValidityChild for ZigZag { diff --git a/vortex-array/src/arrays/chunked/vtable/operations.rs b/vortex-array/src/arrays/chunked/vtable/operations.rs index 395e8aff70a..3edca4df82f 100644 --- a/vortex-array/src/arrays/chunked/vtable/operations.rs +++ b/vortex-array/src/arrays/chunked/vtable/operations.rs @@ -2,24 +2,39 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; +use vortex_error::vortex_err; use crate::ExecutionCtx; use crate::array::ArrayView; use crate::array::OperationsVTable; +use crate::array::ProbeState; use crate::arrays::Chunked; use crate::arrays::chunked::ChunkedArrayExt; +use crate::arrays::chunked::ChunkedSlots; use crate::scalar::Scalar; impl OperationsVTable for Chunked { type ProbeState = (); + fn probe_scalar( + state: &mut ProbeState<'_, Chunked>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let (chunk_index, chunk_offset) = state.array().find_chunk_idx(index)?; + let slot = ChunkedSlots::CHUNKS_OFFSET + chunk_index; + state + .slot(slot)? + .ok_or_else(|| vortex_err!("Chunked chunk slot {slot} is missing"))? + .execute_scalar(chunk_offset, ctx) + } + fn scalar_at( array: ArrayView<'_, Chunked>, index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - let (chunk_index, chunk_offset) = array.find_chunk_idx(index)?; - array.chunk(chunk_index).execute_scalar(chunk_offset, ctx) + Self::probe_scalar(&mut ProbeState::once(array), index, ctx) } } diff --git a/vortex-array/src/arrays/dict/vtable/operations.rs b/vortex-array/src/arrays/dict/vtable/operations.rs index d497db0f2f8..139a23ef05e 100644 --- a/vortex-array/src/arrays/dict/vtable/operations.rs +++ b/vortex-array/src/arrays/dict/vtable/operations.rs @@ -3,35 +3,46 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_err; use super::Dict; use crate::ExecutionCtx; use crate::array::ArrayView; use crate::array::OperationsVTable; -use crate::arrays::dict::DictArraySlotsExt; +use crate::array::ProbeState; +use crate::arrays::dict::DictSlots; use crate::scalar::Scalar; impl OperationsVTable for Dict { type ProbeState = (); - fn scalar_at( - array: ArrayView<'_, Dict>, + fn probe_scalar( + state: &mut ProbeState<'_, Dict>, index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - let Some(dict_index) = array - .codes() - .execute_scalar(index, ctx)? - .as_primitive() - .as_::() - else { - return Ok(Scalar::null(array.dtype().clone())); + let dtype = state.array().dtype().clone(); + let code = state + .slot(DictSlots::CODES)? + .ok_or_else(|| vortex_err!("Dict codes slot is missing"))? + .execute_scalar(index, ctx)?; + let Some(dict_index) = code.as_primitive().as_::() else { + return Ok(Scalar::null(dtype)); }; - Ok(array - .values() + Ok(state + .slot(DictSlots::VALUES)? + .ok_or_else(|| vortex_err!("Dict values slot is missing"))? .execute_scalar(dict_index, ctx)? - .cast(array.dtype()) + .cast(&dtype) .vortex_expect("Array dtype will only differ by nullability")) } + + fn scalar_at( + array: ArrayView<'_, Dict>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Self::probe_scalar(&mut ProbeState::once(array), index, ctx) + } } diff --git a/vortex-array/src/arrays/extension/vtable/operations.rs b/vortex-array/src/arrays/extension/vtable/operations.rs index 519ef6088f7..14d21988dd7 100644 --- a/vortex-array/src/arrays/extension/vtable/operations.rs +++ b/vortex-array/src/arrays/extension/vtable/operations.rs @@ -2,25 +2,38 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; +use vortex_error::vortex_err; use crate::ExecutionCtx; use crate::array::ArrayView; use crate::array::OperationsVTable; +use crate::array::ProbeState; use crate::arrays::Extension; use crate::arrays::extension::ExtensionArrayExt; +use crate::arrays::extension::ExtensionSlots; use crate::scalar::Scalar; impl OperationsVTable for Extension { type ProbeState = (); + fn probe_scalar( + state: &mut ProbeState<'_, Extension>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let ext_dtype = state.array().ext_dtype().clone(); + let storage = state + .slot(ExtensionSlots::STORAGE)? + .ok_or_else(|| vortex_err!("Extension storage slot is missing"))? + .execute_scalar(index, ctx)?; + Ok(Scalar::extension_ref(ext_dtype, storage)) + } + fn scalar_at( array: ArrayView<'_, Extension>, index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - Ok(Scalar::extension_ref( - array.ext_dtype().clone(), - array.storage_array().execute_scalar(index, ctx)?, - )) + Self::probe_scalar(&mut ProbeState::once(array), index, ctx) } } diff --git a/vortex-array/src/arrays/masked/tests.rs b/vortex-array/src/arrays/masked/tests.rs index b326380bbf4..93a61e9f06f 100644 --- a/vortex-array/src/arrays/masked/tests.rs +++ b/vortex-array/src/arrays/masked/tests.rs @@ -18,8 +18,32 @@ use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::dtype::DType; use crate::dtype::Nullability; +use crate::scalar::Scalar; use crate::validity::Validity; +#[test] +fn test_repeated_probe_preserves_masked_nulls() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = MaskedArray::try_new( + PrimitiveArray::from_iter([1i32, 2, 3]).into_array(), + Validity::from_iter([true, false, true]), + )? + .into_array(); + let expected = [Some(1i32), None, Some(3)]; + let mut probe = array.repeated_probe(); + for index in [1, 2, 0, 1] { + assert_eq!( + probe.execute_scalar(index, &mut ctx)?, + Scalar::from(expected[index]) + ); + assert_eq!( + array.execute_scalar(index, &mut ctx)?, + Scalar::from(expected[index]) + ); + } + Ok(()) +} + #[rstest] #[case(Validity::AllValid, Nullability::Nullable)] #[case(Validity::from_iter([true, false, true]), Nullability::Nullable)] diff --git a/vortex-array/src/arrays/masked/vtable/operations.rs b/vortex-array/src/arrays/masked/vtable/operations.rs index a418e23f217..6215bdcc7f2 100644 --- a/vortex-array/src/arrays/masked/vtable/operations.rs +++ b/vortex-array/src/arrays/masked/vtable/operations.rs @@ -2,23 +2,39 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; +use vortex_error::vortex_err; use crate::ExecutionCtx; use crate::array::ArrayView; use crate::array::OperationsVTable; +use crate::array::ProbeState; use crate::arrays::Masked; -use crate::arrays::masked::MaskedArraySlotsExt; +use crate::arrays::masked::MaskedSlots; use crate::scalar::Scalar; impl OperationsVTable for Masked { type ProbeState = (); + fn probe_scalar( + state: &mut ProbeState<'_, Masked>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + if !state.is_valid(index, ctx)? { + return Ok(Scalar::null(state.array().dtype().clone())); + } + Ok(state + .slot(MaskedSlots::CHILD)? + .ok_or_else(|| vortex_err!("Masked child slot is missing"))? + .execute_scalar(index, ctx)? + .into_nullable()) + } + fn scalar_at( array: ArrayView<'_, Masked>, index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - // Invalid indices are handled by the entrypoint function. - Ok(array.child().execute_scalar(index, ctx)?.into_nullable()) + Self::probe_scalar(&mut ProbeState::once(array), index, ctx) } } diff --git a/vortex-array/src/arrays/shared/vtable.rs b/vortex-array/src/arrays/shared/vtable.rs index b1a9a4d7b15..9df66667062 100644 --- a/vortex-array/src/arrays/shared/vtable.rs +++ b/vortex-array/src/arrays/shared/vtable.rs @@ -21,6 +21,8 @@ use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayView; use crate::array::OperationsVTable; +use crate::array::ProbeState; +use crate::array::RepeatedArrayProbe; use crate::array::VTable; use crate::array::ValidityVTable; use crate::array::with_empty_buffers; @@ -124,15 +126,41 @@ impl VTable for Shared { .map(ExecutionResult::done) } } + +/// A probe over whichever array [`SharedArrayExt::current_array_ref`] resolves to. The source +/// can be replaced by the computed result between reads, so the retained probe is rebuilt when +/// the array it was built over is no longer the current one. +#[derive(Default)] +pub struct SharedProbeState { + current: Option, +} + impl OperationsVTable for Shared { - type ProbeState = (); + type ProbeState = SharedProbeState; + + fn probe_scalar( + state: &mut ProbeState<'_, Shared>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let array = state.array(); + let current = array.current_array_ref(); + let Some(retained) = state.retained() else { + return current.probe().execute_scalar(index, ctx); + }; + let probe = match &mut retained.current { + Some(probe) if ArrayRef::ptr_eq(probe.array(), current) => probe, + slot => slot.insert(RepeatedArrayProbe::new(current.clone())), + }; + probe.execute_scalar(index, ctx) + } fn scalar_at( array: ArrayView<'_, Shared>, index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - array.current_array_ref().execute_scalar(index, ctx) + Self::probe_scalar(&mut ProbeState::once(array), index, ctx) } } diff --git a/vortex-array/src/arrays/slice/vtable.rs b/vortex-array/src/arrays/slice/vtable.rs index c09d81ee4f1..de87931ab5f 100644 --- a/vortex-array/src/arrays/slice/vtable.rs +++ b/vortex-array/src/arrays/slice/vtable.rs @@ -11,6 +11,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -25,6 +26,7 @@ use crate::array::Array; use crate::array::ArrayId; use crate::array::ArrayView; use crate::array::OperationsVTable; +use crate::array::ProbeState; use crate::array::VTable; use crate::array::ValidityVTable; use crate::array::with_empty_buffers; @@ -171,12 +173,24 @@ impl VTable for Slice { impl OperationsVTable for Slice { type ProbeState = (); + fn probe_scalar( + state: &mut ProbeState<'_, Slice>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let child_index = state.array().range.start + index; + state + .slot(SliceSlots::CHILD)? + .ok_or_else(|| vortex_err!("Slice child slot is missing"))? + .execute_scalar(child_index, ctx) + } + fn scalar_at( array: ArrayView<'_, Slice>, index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - array.child().execute_scalar(array.range.start + index, ctx) + Self::probe_scalar(&mut ProbeState::once(array), index, ctx) } } diff --git a/vortex-array/src/compute/conformance/consistency.rs b/vortex-array/src/compute/conformance/consistency.rs index ed6b42f3514..1c18be97f55 100644 --- a/vortex-array/src/compute/conformance/consistency.rs +++ b/vortex-array/src/compute/conformance/consistency.rs @@ -1390,6 +1390,61 @@ fn test_cast_slice_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { } } +/// Tests that a repeated probe agrees with one-off scalar reads. +/// +/// # Invariant +/// `array.repeated_probe().execute_scalar(i)` equals `array.execute_scalar(i)` for every `i`, +/// whatever order the rows are visited in. +/// +/// # Test Details +/// - Reads every row one-off to build the expected values +/// - Reads the same rows through a single retained probe, backwards, then forwards, then +/// sparsely, so a probe that caches per-row state has to stay correct across jumps and +/// revisits +/// - Checks nullness through the probe as well as the value +/// +/// # Why This Matters +/// Encodings serve one-off and repeated reads from the same `probe_scalar` body, and a parent +/// encoding passes its children's probes down. A retained cache that outlives the row it was +/// built for, or a slot read through the wrong index, shows up here as a wrong value. +fn test_repeated_probe_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { + // Long arrays are covered by their prefix; the three passes below already cost 3n reads. + let len = array.len().min(1024); + if len == 0 { + return; + } + + let expected: Vec<_> = (0..len) + .map(|i| { + array + .execute_scalar(i, ctx) + .vortex_expect("scalar_at should succeed in conformance test") + }) + .collect(); + + let mut probe = array.repeated_probe(); + let order = (0..len).rev().chain(0..len).chain((0..len).step_by(7)); + for i in order { + let actual = probe + .execute_scalar(i, ctx) + .vortex_expect("repeated probe read should succeed in conformance test"); + assert_eq!( + actual, expected[i], + "Repeated probe and one-off read disagree at index {i}. \ + Probe value: {actual:?}, one-off value: {:?}", + expected[i] + ); + let valid = probe + .execute_is_valid(i, ctx) + .vortex_expect("repeated probe validity read should succeed in conformance test"); + assert_eq!( + valid, + !expected[i].is_null(), + "Repeated probe validity disagrees with the value read at index {i}" + ); + } +} + /// Run all consistency tests on an array. /// /// This function executes a comprehensive suite of consistency tests that verify @@ -1454,6 +1509,9 @@ pub fn test_array_consistency(array: &ArrayRef, ctx: &mut ExecutionCtx) { test_mask_identity(array, ctx); test_take_preserves_properties(array, ctx); + // Row access + test_repeated_probe_consistency(array, ctx); + // Ordering and correctness test_filter_preserves_order(array, ctx); test_take_repeated_indices(array, ctx);