vortex_array/array/probe/
repeated.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::Any;
5
6use vortex_error::VortexResult;
7use vortex_error::vortex_err;
8
9use crate::ArrayRef;
10use crate::ExecutionCtx;
11use crate::IntoArray;
12use crate::array::probe::ArrayProbe;
13use crate::array::probe::array::check_bounds;
14use crate::array::probe::array::check_dtype;
15use crate::array::probe::array::child_of;
16use crate::arrays::BoolArray;
17use crate::arrays::ScalarFn;
18use crate::scalar::Scalar;
19use crate::validity::Validity;
20
21/// A row accessor that owns its array and keeps preparation between reads.
22///
23/// The encoding's state, its validity probe and the probes over child slots are created on
24/// first use and reused by every following read. Dropping the probe drops all of them. Array
25/// handles share their buffers, so the probe can outlive the handle it was built from. Probes
26/// are local to a thread.
27pub struct RepeatedArrayProbe {
28    array: ArrayRef,
29    /// The encoding's [`RepeatedState`], type-erased, created on the first read.
30    state: Option<Box<dyn Any>>,
31}
32
33impl RepeatedArrayProbe {
34    /// Own an array for repeated reads.
35    pub fn new(array: ArrayRef) -> Self {
36        Self { array, state: None }
37    }
38
39    /// The array this probe reads from.
40    pub fn array(&self) -> &ArrayRef {
41        &self.array
42    }
43
44    /// Borrow this probe as an [`ArrayProbe`].
45    #[inline]
46    pub fn as_probe(&mut self) -> ArrayProbe<'_> {
47        ArrayProbe::Repeated(self)
48    }
49
50    /// Read the scalar at `index`, including its nullness, reusing retained preparation.
51    pub fn execute_scalar(&mut self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult<Scalar> {
52        check_bounds(&self.array, index)?;
53        let result =
54            self.array
55                .dyn_array()
56                .probe_scalar_retained(&self.array, index, &mut self.state, ctx);
57        check_dtype(&self.array, result)
58    }
59
60    /// Whether the row at `index` is valid, through the retained validity.
61    pub fn execute_is_valid(&mut self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
62        check_bounds(&self.array, index)?;
63        if !self.array.dtype().is_nullable() {
64            return Ok(true);
65        }
66        self.array
67            .dyn_array()
68            .probe_is_valid_retained(&self.array, index, &mut self.state, ctx)
69    }
70
71    /// Whether the row at `index` is null.
72    pub fn execute_is_invalid(
73        &mut self,
74        index: usize,
75        ctx: &mut ExecutionCtx,
76    ) -> VortexResult<bool> {
77        Ok(!self.execute_is_valid(index, ctx)?)
78    }
79}
80
81/// The validity of a [`RepeatedState`], resolved on first use and kept between reads.
82enum ResolvedValidity {
83    Uniform(bool),
84    Array(RepeatedArrayProbe),
85}
86
87impl ResolvedValidity {
88    /// Whether row `index` is valid.
89    #[inline]
90    fn is_valid(&mut self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
91        match self {
92            Self::Uniform(valid) => Ok(*valid),
93            Self::Array(probe) => is_valid_scalar(probe.execute_scalar(index, ctx)?, index),
94        }
95    }
96
97    /// The resolved validity of `array`.
98    fn resolve(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
99        if !array.dtype().is_nullable() {
100            return Ok(Self::Uniform(true));
101        }
102        Ok(match array.validity()? {
103            Validity::NonNullable | Validity::AllValid => Self::Uniform(true),
104            Validity::AllInvalid => Self::Uniform(false),
105            Validity::Array(validity) => {
106                // ScalarFn validity is lazy; materialize it once instead of once per row.
107                // TODO(myrrc): remove this once the probing API can specify validity.
108                let validity = if array.is::<ScalarFn>() {
109                    validity.execute::<BoolArray>(ctx)?.into_array()
110                } else {
111                    validity
112                };
113                Self::Array(RepeatedArrayProbe::new(validity))
114            }
115        })
116    }
117}
118
119/// The validity bit read out of a validity array's scalar.
120#[inline]
121pub(super) fn is_valid_scalar(scalar: Scalar, index: usize) -> VortexResult<bool> {
122    scalar
123        .as_bool()
124        .value()
125        .ok_or_else(|| vortex_err!("validity value at index {index} is null"))
126}
127
128/// Get or create the [`RepeatedState`] for encoding state `S` in a probe's erased slot.
129pub(crate) fn repeated_state<S: Default + 'static>(
130    slot: &mut Option<Box<dyn Any>>,
131) -> VortexResult<&mut RepeatedState<S>> {
132    slot.get_or_insert_with(|| Box::new(RepeatedState::<S>::default()))
133        .downcast_mut::<RepeatedState<S>>()
134        .ok_or_else(|| vortex_err!("Probe state type mismatch"))
135}
136
137/// What a [`RepeatedArrayProbe`] keeps for its encoding between reads.
138pub(crate) struct RepeatedState<S> {
139    state: S,
140    /// Probes over the source's child slots, allocated on first use. Slots that are never
141    /// requested stay empty.
142    slots: Vec<Option<RepeatedArrayProbe>>,
143    /// The source's validity, resolved on first use.
144    validity: Option<ResolvedValidity>,
145}
146
147impl<S: Default> Default for RepeatedState<S> {
148    fn default() -> Self {
149        Self {
150            state: S::default(),
151            slots: Vec::new(),
152            validity: None,
153        }
154    }
155}
156
157impl<S> RepeatedState<S> {
158    #[inline]
159    pub(super) fn state_mut(&mut self) -> &mut S {
160        &mut self.state
161    }
162
163    /// The encoding state and the child slot table as disjoint borrows.
164    #[inline]
165    pub(super) fn split_mut(&mut self) -> (&mut S, &mut Vec<Option<RepeatedArrayProbe>>) {
166        (&mut self.state, &mut self.slots)
167    }
168
169    /// Whether row `index` of `array` is valid, through the kept validity.
170    #[inline]
171    pub(crate) fn is_valid(
172        &mut self,
173        array: &ArrayRef,
174        index: usize,
175        ctx: &mut ExecutionCtx,
176    ) -> VortexResult<bool> {
177        let validity = match &mut self.validity {
178            Some(validity) => validity,
179            slot @ None => slot.insert(ResolvedValidity::resolve(array, ctx)?),
180        };
181        validity.is_valid(index, ctx)
182    }
183}
184
185/// Get or create the retained probe over `parent`'s child in `slot`.
186pub(super) fn child_probe<'s>(
187    slots: &'s mut Vec<Option<RepeatedArrayProbe>>,
188    parent: &ArrayRef,
189    slot: usize,
190) -> VortexResult<Option<&'s mut RepeatedArrayProbe>> {
191    let Some(child) = child_of(parent, slot)? else {
192        return Ok(None);
193    };
194    if slots.is_empty() {
195        slots.resize_with(parent.slots().len(), || None);
196    }
197    Ok(Some(slots[slot].get_or_insert_with(|| {
198        RepeatedArrayProbe::new(child.clone())
199    })))
200}
201
202#[cfg(test)]
203mod tests {
204    use vortex_error::VortexResult;
205    use vortex_error::vortex_err;
206
207    use super::*;
208    use crate::VortexSessionExecute;
209    use crate::array::IntoArray;
210    use crate::array::probe::ProbeState;
211    use crate::arrays::PrimitiveArray;
212    use crate::arrays::Struct;
213    use crate::arrays::StructArray;
214
215    fn nullable_ints() -> ArrayRef {
216        PrimitiveArray::from_option_iter([Some(10i32), None, Some(30)]).into_array()
217    }
218
219    fn check_reads(mut probe: ArrayProbe<'_>, ctx: &mut ExecutionCtx) -> VortexResult<()> {
220        assert!(probe.execute_scalar(3, ctx).is_err());
221        assert!(probe.execute_is_valid(3, ctx).is_err());
222        assert!(probe.execute_scalar(1, ctx)?.is_null());
223        assert!(!probe.execute_is_valid(1, ctx)?);
224        assert!(probe.execute_is_invalid(1, ctx)?);
225        assert_eq!(probe.execute_scalar(2, ctx)?, Scalar::from(Some(30i32)));
226        assert_eq!(probe.execute_scalar(0, ctx)?, Scalar::from(Some(10i32)));
227        Ok(())
228    }
229
230    #[test]
231    fn repeated_probe_initializes_lazily_and_outlives_handle() -> VortexResult<()> {
232        let mut ctx = crate::array_session().create_execution_ctx();
233        let array = nullable_ints();
234        let mut probe = RepeatedArrayProbe::new(array.clone());
235        drop(array);
236        assert!(probe.state.is_none());
237
238        check_reads(probe.as_probe(), &mut ctx)?;
239
240        let state = repeated_state::<()>(&mut probe.state)?;
241        assert!(matches!(state.validity, Some(ResolvedValidity::Array(_))));
242        Ok(())
243    }
244
245    #[test]
246    fn repeated_probe_resolves_uniform_validity_once() -> VortexResult<()> {
247        let mut ctx = crate::array_session().create_execution_ctx();
248        let array = PrimitiveArray::from_option_iter([Some(1i32), Some(2)]).into_array();
249        let mut probe = array.repeated_probe();
250        assert!(probe.execute_is_valid(1, &mut ctx)?);
251        let state = repeated_state::<()>(&mut probe.state)?;
252        assert!(matches!(
253            state.validity,
254            Some(ResolvedValidity::Uniform(true))
255        ));
256        assert_eq!(probe.execute_scalar(1, &mut ctx)?, Scalar::from(Some(2i32)));
257        Ok(())
258    }
259
260    #[test]
261    fn validity_and_scalar_reads_share_the_validity_slot() -> VortexResult<()> {
262        let mut ctx = crate::array_session().create_execution_ctx();
263        let mut probe = nullable_ints().repeated_probe();
264        assert!(!probe.execute_is_valid(1, &mut ctx)?);
265        let Some(ResolvedValidity::Array(validity)) =
266            &repeated_state::<()>(&mut probe.state)?.validity
267        else {
268            return Err(vortex_err!("validity slot should hold a probe"));
269        };
270        assert!(validity.state.is_some());
271        assert!(probe.execute_scalar(1, &mut ctx)?.is_null());
272        assert_eq!(
273            probe.execute_scalar(2, &mut ctx)?,
274            Scalar::from(Some(30i32))
275        );
276        Ok(())
277    }
278
279    #[test]
280    fn repeated_probe_reads_null_rows_through_encoding() -> VortexResult<()> {
281        let mut ctx = crate::array_session().create_execution_ctx();
282        let array = PrimitiveArray::from_option_iter([Some(1i32), None]).into_array();
283        let mut probe = array.repeated_probe();
284        assert!(probe.execute_scalar(1, &mut ctx)?.is_null());
285        assert_eq!(probe.execute_scalar(0, &mut ctx)?, Scalar::from(Some(1i32)));
286        assert!(!probe.execute_is_valid(1, &mut ctx)?);
287        Ok(())
288    }
289
290    #[test]
291    fn repeated_state_creates_children_on_demand() -> VortexResult<()> {
292        let mut ctx = crate::array_session().create_execution_ctx();
293        let array = struct_of_two_fields()?;
294        let typed = array
295            .as_opt::<Struct>()
296            .ok_or_else(|| vortex_err!("expected a struct"))?;
297        let mut repeated = RepeatedState::<()>::default();
298        {
299            let mut state = ProbeState::repeated(typed, &mut repeated);
300            assert!(state.slot(5).is_err());
301            assert!(state.slot(0)?.is_none());
302            assert!(state.retained().is_some());
303            let mut field = state.slot(2)?.ok_or_else(|| vortex_err!("missing field"))?;
304            assert!(matches!(field, ArrayProbe::Repeated(_)));
305            assert_eq!(field.execute_scalar(1, &mut ctx)?, Scalar::from(4i64));
306            assert_eq!(field.execute_scalar(0, &mut ctx)?, Scalar::from(3i64));
307        }
308
309        assert_eq!(repeated.slots.len(), array.slots().len());
310        assert!(repeated.slots[1].is_none());
311        let child = repeated.slots[2]
312            .as_ref()
313            .ok_or_else(|| vortex_err!("missing child probe"))?;
314        assert_eq!(child.array().len(), 2);
315        Ok(())
316    }
317
318    fn struct_of_two_fields() -> VortexResult<ArrayRef> {
319        Ok(StructArray::from_fields(&[
320            ("a", PrimitiveArray::from_iter([1i32, 2]).into_array()),
321            ("b", PrimitiveArray::from_iter([3i64, 4]).into_array()),
322        ])?
323        .into_array())
324    }
325
326    #[test]
327    fn state_slot_rejects_mismatched_type() -> VortexResult<()> {
328        let mut slot: Option<Box<dyn Any>> = None;
329        repeated_state::<()>(&mut slot)?;
330        assert!(repeated_state::<u8>(&mut slot).is_err());
331        Ok(())
332    }
333}
vortex_array/array/probe/repeated.rs

vortex_array/array/probe/
repeated.rs
