From 7da4edd20f83d7c23f88a7d7b417aa350d3312d9 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Mon, 26 Oct 2026 17:08:05 -0400 Subject: [PATCH] Use an operation-owned indexed probe cache for new-prefix arrays Signed-off-by: Daniel King --- encodings/pco/src/access.rs | 15 +++ encodings/pco/src/array.rs | 3 +- encodings/pco/src/decoder.rs | 30 ++++-- encodings/pco/src/indexed_probe.rs | 152 +++++++++++++++++++++++++++++ encodings/pco/src/lib.rs | 1 + encodings/pco/src/probe.rs | 6 ++ 6 files changed, 200 insertions(+), 7 deletions(-) create mode 100644 encodings/pco/src/indexed_probe.rs diff --git a/encodings/pco/src/access.rs b/encodings/pco/src/access.rs index 78fde096525..11379ae2012 100644 --- a/encodings/pco/src/access.rs +++ b/encodings/pco/src/access.rs @@ -157,6 +157,21 @@ impl AccessState { mask: cached.mask.mask.slice(start_in_page..end_in_page), }) } + pub fn page_id(&self) -> Option { + self.last.as_ref().map(|(page, _)| *page) + } + pub fn cached_position(&self, row: usize) -> Option> { + let (page, cached) = self.last.as_ref()?; + if !cached.rows.contains(&row) { + return None; + } + Some( + cached + .mask + .value_index::(row - cached.rows.start) + .map(|value| (*page, value)), + ) + } } pub(crate) struct PageSpan { pub page: usize, diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array.rs index 355893afa6e..2aaf2b68bc6 100644 --- a/encodings/pco/src/array.rs +++ b/encodings/pco/src/array.rs @@ -66,6 +66,7 @@ use crate::PcoLogicalRowPrefix; use crate::PcoMetadata; use crate::PcoPageInfo; use crate::page_index::PageIndex; +use crate::probe::PcoProbeState; use crate::take::slice as decode_slice; // Overall approach here: @@ -847,7 +848,7 @@ impl ValidityVTable for Pco { } impl OperationsVTable for Pco { - type ProbeState = crate::probe::PcoProbeState; + type ProbeState = PcoProbeState; fn probe_scalar( state: &mut ProbeState<'_, Pco>, diff --git a/encodings/pco/src/decoder.rs b/encodings/pco/src/decoder.rs index 88f32488e26..32e1afeca8d 100644 --- a/encodings/pco/src/decoder.rs +++ b/encodings/pco/src/decoder.rs @@ -15,6 +15,7 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::NativePType; use vortex_array::validity::Validity; use vortex_buffer::BufferMut; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -83,6 +84,7 @@ pub(crate) struct DecodedPages { decoder: PageDecoder, chunks: ChunkCursor, pages: BTreeMap, + last: Option<(usize, PrimitiveArray)>, } impl DecodedPages { @@ -91,8 +93,15 @@ impl DecodedPages { decoder: PageDecoder::new(array.metadata.header.as_slice())?, chunks: ChunkCursor::default(), pages: BTreeMap::new(), + last: None, }) } + pub fn last_values(&self, page: usize) -> Option<&PrimitiveArray> { + self.last + .as_ref() + .filter(|(last, _)| *last == page) + .map(|(_, values)| values) + } pub fn values( &mut self, view: IndexedView<'_>, @@ -100,9 +109,10 @@ impl DecodedPages { page: usize, ctx: &mut ExecutionCtx, ) -> VortexResult<&PrimitiveArray> { - let entry = match self.pages.entry(page) { - Entry::Occupied(entry) => entry.into_mut(), - Entry::Vacant(entry) => { + if !self.last.as_ref().is_some_and(|(last, _)| *last == page) { + let decoded = if let Some(decoded) = self.pages.remove(&page) { + decoded + } else { let info = view.physical(page, &mut self.chunks)?; let buffer = array .pages @@ -112,15 +122,23 @@ impl DecodedPages { .chunk_metas .get(info.chunk) .ok_or_else(|| vortex_err!("Pco chunk metadata is missing"))?; - entry.insert(self.decoder.read( + self.decoder.read( info.chunk, chunk.as_slice(), info.n_values, buffer.as_slice(), ctx, - )?) + )? + }; + if let Some((page, previous)) = self.last.replace((page, decoded)) { + self.pages.insert(page, previous); } - }; + } + let entry = &self + .last + .as_ref() + .vortex_expect("decoded page prepared above") + .1; Ok(entry) } } diff --git a/encodings/pco/src/indexed_probe.rs b/encodings/pco/src/indexed_probe.rs new file mode 100644 index 00000000000..649cbc92a07 --- /dev/null +++ b/encodings/pco/src/indexed_probe.rs @@ -0,0 +1,152 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Operation-owned probes bind indexed or legacy addressing on first use. + +use std::collections::BTreeMap; +use std::mem::replace; +use std::sync::Arc; + +use pco::data_types::Number; +use pco::data_types::NumberType; +use pco::match_number_enum; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::half; +use vortex_array::scalar::PValue; +use vortex_array::scalar::Scalar; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; + +use crate::Pco; +use crate::access::AccessState; +use crate::array::number_type_from_ptype; +use crate::decoder::DecodedPages; +use crate::page_index::IndexedIndex; +use crate::page_index::PageIndex; + +/// Construction is empty; first use binds an owned immutable page-index variant and +/// typed decoder. Touched decoded pages can be retained up to the entire source array. +#[derive(Default)] +pub(crate) struct IndexedProbeState { + bound: Option>, +} + +trait BoundProbe: Send + Sync { + fn scalar( + &mut self, + array: ArrayView<'_, Pco>, + row: usize, + ranked: bool, + ctx: &mut ExecutionCtx, + ) -> VortexResult; +} + +struct TypedProbe { + owner: Arc, + access: AccessState, + prepared: BTreeMap, + decoded: Option>, +} + +impl BoundProbe for TypedProbe +where + T: Number + NativePType + Into, +{ + fn scalar( + &mut self, + array: ArrayView<'_, Pco>, + row: usize, + ranked: bool, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + if ranked && let Some(position) = self.access.cached_position(row) { + match position { + None => return Ok(Scalar::null(array.dtype().clone())), + Some((page, value)) => { + if let Some(values) = self + .decoded + .as_ref() + .and_then(|decoded| decoded.last_values(page)) + { + return Ok(Scalar::primitive( + values.as_slice::()[value], + array.dtype().nullability(), + )); + } + } + } + } + let view = self.owner.view(); + let page = view.find_row(row); + if self.access.page_id() != Some(page) { + let next = self.prepared.remove(&page).unwrap_or_default(); + let previous = replace(&mut self.access, next); + if let Some(previous_page) = previous.page_id() { + self.prepared.insert(previous_page, previous); + } + } + let position = if ranked { + self.access.resolve::(view, array, row, ctx)? + } else { + self.access.resolve::(view, array, row, ctx)? + }; + let Some((page, value)) = position else { + return Ok(Scalar::null(array.dtype().clone())); + }; + if self.decoded.is_none() { + self.decoded = Some(DecodedPages::new(array)?); + } + let decoded = self + .decoded + .as_mut() + .vortex_expect("decoder prepared above"); + let values = decoded.values(view, array, page, ctx)?; + Ok(Scalar::primitive( + values.as_slice::()[value], + array.dtype().nullability(), + )) + } +} + +fn bind>(owner: Arc) -> Box { + Box::new(TypedProbe:: { + owner, + access: AccessState::default(), + prepared: BTreeMap::new(), + decoded: None, + }) +} + +pub(crate) fn scalar_at( + array: ArrayView<'_, Pco>, + index: usize, + state: Option<&mut IndexedProbeState>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + // The public probe dispatcher already checks visible bounds. A zero-page array + // stores no defined values, so no validity/index/decoder preparation is needed. + if array.pages.is_empty() { + return Ok(Scalar::null(array.dtype().clone())); + } + let mut once = IndexedProbeState::default(); + let ranked = state.is_some(); + let state = state.unwrap_or(&mut once); + if state.bound.is_none() { + let PageIndex::IndexedFast(owner) = array.page_index()? else { + unreachable!("indexed probes are dispatched only for a present prefix") + }; + state.bound = Some( + match_number_enum!(number_type_from_ptype(array.dtype().as_ptype()), NumberType => { + bind::(Arc::clone(owner)) + }), + ); + } + state.bound.as_mut().vortex_expect("bound above").scalar( + array, + array.slice_start() + index, + ranked, + ctx, + ) +} diff --git a/encodings/pco/src/lib.rs b/encodings/pco/src/lib.rs index 51dfcf55223..02147301e15 100644 --- a/encodings/pco/src/lib.rs +++ b/encodings/pco/src/lib.rs @@ -32,6 +32,7 @@ mod access; mod array; mod compute; mod decoder; +mod indexed_probe; mod page_index; mod probe; mod rules; diff --git a/encodings/pco/src/probe.rs b/encodings/pco/src/probe.rs index 00aa801254a..d0a6480a41a 100644 --- a/encodings/pco/src/probe.rs +++ b/encodings/pco/src/probe.rs @@ -27,6 +27,8 @@ use crate::Pco; use crate::PcoArrayExt; use crate::array::number_type_from_ptype; use crate::array::vortex_err_from_pco; +use crate::indexed_probe::IndexedProbeState; +use crate::indexed_probe::scalar_at as indexed_scalar_at; const RANK_STRIDE: usize = 512; @@ -37,6 +39,7 @@ const RANK_STRIDE: usize = 512; /// and never beyond the decompressed array. #[derive(Default)] pub struct PcoProbeState { + indexed: IndexedProbeState, validity: Option, rank: Vec, pages: Vec, @@ -56,6 +59,9 @@ pub(crate) fn scalar_at( state: Option<&mut PcoProbeState>, ctx: &mut ExecutionCtx, ) -> VortexResult { + if array.metadata.logical_row_prefix.is_some() { + return indexed_scalar_at(array, index, state.map(|state| &mut state.indexed), ctx); + } let Some(state) = state else { let validity = array.unsliced_validity(); return array