Use PCO logical-page indexes for retained scalar probes

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<usize> {
+        self.last.as_ref().map(|(page, _)| *page)
+    }
+    pub fn cached_position(&self, row: usize) -> Option<Option<(usize, usize)>> {
+        let (page, cached) = self.last.as_ref()?;
+        if !cached.rows.contains(&row) {
+            return None;
+        }
+        Some(
+            cached
+                .mask
+                .value_index::<true>(row - cached.rows.start)
+                .map(|value| (*page, value)),
+        )
+    }
 }
 pub(crate) struct PageSpan {
     pub page: usize,

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<Pco> for Pco {
 }
 
 impl OperationsVTable<Pco> for Pco {
-    type ProbeState = crate::probe::PcoProbeState;
+    type ProbeState = PcoProbeState;
 
     fn probe_scalar(
         state: &mut ProbeState<'_, Pco>,

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<T: Number> {
     decoder: PageDecoder<T>,
     chunks: ChunkCursor,
     pages: BTreeMap<usize, PrimitiveArray>,
+    last: Option<(usize, PrimitiveArray)>,
 }
 
 impl<T: Number + NativePType> DecodedPages<T> {
@@ -91,18 +93,26 @@ impl<T: Number + NativePType> DecodedPages<T> {
             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<'_>,
         array: ArrayView<'_, Pco>,
         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<T: Number + NativePType> DecodedPages<T> {
                     .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)
     }
 }

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<Box<dyn BoundProbe>>,
+}
+
+trait BoundProbe: Send + Sync {
+    fn scalar(
+        &mut self,
+        array: ArrayView<'_, Pco>,
+        row: usize,
+        ranked: bool,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Scalar>;
+}
+
+struct TypedProbe<T: Number> {
+    owner: Arc<IndexedIndex>,
+    access: AccessState,
+    prepared: BTreeMap<usize, AccessState>,
+    decoded: Option<DecodedPages<T>>,
+}
+
+impl<T> BoundProbe for TypedProbe<T>
+where
+    T: Number + NativePType + Into<PValue>,
+{
+    fn scalar(
+        &mut self,
+        array: ArrayView<'_, Pco>,
+        row: usize,
+        ranked: bool,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Scalar> {
+        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::<T>()[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::<true>(view, array, row, ctx)?
+        } else {
+            self.access.resolve::<false>(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::<T>()[value],
+            array.dtype().nullability(),
+        ))
+    }
+}
+
+fn bind<T: Number + NativePType + Into<PValue>>(owner: Arc<IndexedIndex>) -> Box<dyn BoundProbe> {
+    Box::new(TypedProbe::<T> {
+        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<Scalar> {
+    // 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<T> => {
+                bind::<T>(Arc::clone(owner))
+            }),
+        );
+    }
+    state.bound.as_mut().vortex_expect("bound above").scalar(
+        array,
+        array.slice_start() + index,
+        ranked,
+        ctx,
+    )
+}

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;

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<Mask>,
     rank: Vec<usize>,
     pages: Vec<Page>,
@@ -56,6 +59,9 @@ pub(crate) fn scalar_at(
     state: Option<&mut PcoProbeState>,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<Scalar> {
+    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