|
| 1 | +// SPDX-License-Identifier: Apache-2.0 |
| 2 | +// SPDX-FileCopyrightText: Copyright the Vortex contributors |
| 3 | + |
| 4 | +//! Operation-owned probes bind indexed or legacy addressing on first use. |
| 5 | +
|
| 6 | +use std::collections::BTreeMap; |
| 7 | +use std::mem::replace; |
| 8 | +use std::sync::Arc; |
| 9 | + |
| 10 | +use pco::data_types::Number; |
| 11 | +use pco::data_types::NumberType; |
| 12 | +use pco::match_number_enum; |
| 13 | +use vortex_array::ArrayView; |
| 14 | +use vortex_array::ExecutionCtx; |
| 15 | +use vortex_array::dtype::NativePType; |
| 16 | +use vortex_array::dtype::half; |
| 17 | +use vortex_array::scalar::PValue; |
| 18 | +use vortex_array::scalar::Scalar; |
| 19 | +use vortex_error::VortexExpect; |
| 20 | +use vortex_error::VortexResult; |
| 21 | + |
| 22 | +use crate::Pco; |
| 23 | +use crate::access::AccessState; |
| 24 | +use crate::array::number_type_from_ptype; |
| 25 | +use crate::decoder::DecodedPages; |
| 26 | +use crate::page_index::LogicalPageIndex; |
| 27 | +use crate::page_index::PageIndex; |
| 28 | + |
| 29 | +/// Construction is empty; first use binds an owned immutable page-index variant and |
| 30 | +/// typed decoder. Touched decoded pages can be retained up to the entire source array. |
| 31 | +#[derive(Default)] |
| 32 | +pub(crate) struct IndexedProbeState { |
| 33 | + bound: Option<Box<dyn BoundProbe>>, |
| 34 | +} |
| 35 | + |
| 36 | +trait BoundProbe: Send + Sync { |
| 37 | + fn scalar( |
| 38 | + &mut self, |
| 39 | + array: ArrayView<'_, Pco>, |
| 40 | + row: usize, |
| 41 | + ranked: bool, |
| 42 | + ctx: &mut ExecutionCtx, |
| 43 | + ) -> VortexResult<Scalar>; |
| 44 | +} |
| 45 | + |
| 46 | +struct TypedProbe<T: Number> { |
| 47 | + owner: Arc<LogicalPageIndex>, |
| 48 | + access: AccessState, |
| 49 | + prepared: BTreeMap<usize, AccessState>, |
| 50 | + decoded: Option<DecodedPages<T>>, |
| 51 | +} |
| 52 | + |
| 53 | +impl<T> BoundProbe for TypedProbe<T> |
| 54 | +where |
| 55 | + T: Number + NativePType + Into<PValue>, |
| 56 | +{ |
| 57 | + fn scalar( |
| 58 | + &mut self, |
| 59 | + array: ArrayView<'_, Pco>, |
| 60 | + row: usize, |
| 61 | + ranked: bool, |
| 62 | + ctx: &mut ExecutionCtx, |
| 63 | + ) -> VortexResult<Scalar> { |
| 64 | + if ranked && let Some(position) = self.access.cached_position(row) { |
| 65 | + match position { |
| 66 | + None => return Ok(Scalar::null(array.dtype().clone())), |
| 67 | + Some((page, value)) => { |
| 68 | + if let Some(values) = self |
| 69 | + .decoded |
| 70 | + .as_ref() |
| 71 | + .and_then(|decoded| decoded.last_values(page)) |
| 72 | + { |
| 73 | + return Ok(Scalar::primitive( |
| 74 | + values.as_slice::<T>()[value], |
| 75 | + array.dtype().nullability(), |
| 76 | + )); |
| 77 | + } |
| 78 | + } |
| 79 | + } |
| 80 | + } |
| 81 | + let view = self.owner.as_ref(); |
| 82 | + let page = view.find_row(row); |
| 83 | + if self.access.page_id() != Some(page) { |
| 84 | + let next = self.prepared.remove(&page).unwrap_or_default(); |
| 85 | + let previous = replace(&mut self.access, next); |
| 86 | + if let Some(previous_page) = previous.page_id() { |
| 87 | + self.prepared.insert(previous_page, previous); |
| 88 | + } |
| 89 | + } |
| 90 | + let position = if ranked { |
| 91 | + self.access.resolve::<true>(view, array, row, ctx)? |
| 92 | + } else { |
| 93 | + self.access.resolve::<false>(view, array, row, ctx)? |
| 94 | + }; |
| 95 | + let Some((page, value)) = position else { |
| 96 | + return Ok(Scalar::null(array.dtype().clone())); |
| 97 | + }; |
| 98 | + if self.decoded.is_none() { |
| 99 | + self.decoded = Some(DecodedPages::new(array)?); |
| 100 | + } |
| 101 | + let decoded = self |
| 102 | + .decoded |
| 103 | + .as_mut() |
| 104 | + .vortex_expect("decoder prepared above"); |
| 105 | + let values = decoded.values(view, array, page, ctx)?; |
| 106 | + Ok(Scalar::primitive( |
| 107 | + values.as_slice::<T>()[value], |
| 108 | + array.dtype().nullability(), |
| 109 | + )) |
| 110 | + } |
| 111 | +} |
| 112 | + |
| 113 | +fn bind<T: Number + NativePType + Into<PValue>>(owner: Arc<LogicalPageIndex>) -> Box<dyn BoundProbe> { |
| 114 | + Box::new(TypedProbe::<T> { |
| 115 | + owner, |
| 116 | + access: AccessState::default(), |
| 117 | + prepared: BTreeMap::new(), |
| 118 | + decoded: None, |
| 119 | + }) |
| 120 | +} |
| 121 | + |
| 122 | +pub(crate) fn scalar_at( |
| 123 | + array: ArrayView<'_, Pco>, |
| 124 | + index: usize, |
| 125 | + state: Option<&mut IndexedProbeState>, |
| 126 | + ctx: &mut ExecutionCtx, |
| 127 | +) -> VortexResult<Scalar> { |
| 128 | + // The public probe dispatcher already checks visible bounds. A zero-page array |
| 129 | + // stores no defined values, so no validity/index/decoder preparation is needed. |
| 130 | + if array.pages.is_empty() { |
| 131 | + return Ok(Scalar::null(array.dtype().clone())); |
| 132 | + } |
| 133 | + let mut once = IndexedProbeState::default(); |
| 134 | + let ranked = state.is_some(); |
| 135 | + let state = state.unwrap_or(&mut once); |
| 136 | + if state.bound.is_none() { |
| 137 | + let PageIndex::IndexedFast(owner) = array.page_index()? else { |
| 138 | + unreachable!("indexed probes are dispatched only for a present prefix") |
| 139 | + }; |
| 140 | + state.bound = Some( |
| 141 | + match_number_enum!(number_type_from_ptype(array.dtype().as_ptype()), NumberType<T> => { |
| 142 | + bind::<T>(Arc::clone(owner)) |
| 143 | + }), |
| 144 | + ); |
| 145 | + } |
| 146 | + state.bound.as_mut().vortex_expect("bound above").scalar( |
| 147 | + array, |
| 148 | + array.slice_start() + index, |
| 149 | + ranked, |
| 150 | + ctx, |
| 151 | + ) |
| 152 | +} |
0 commit comments