Decode selected PCO pages for primitive-index takes

encodings/pco/benches/selection.rs

@@ -26,6 +26,7 @@ use vortex_array::session::ArraySessionExt;
 use vortex_array::validity::Validity;
 use vortex_error::VortexExpect;
 use vortex_pco::Pco;
+use vortex_pco::initialize as initialize_pco;
 use vortex_session::VortexSession;
 
 const ROWS: usize = 4_194_304;
@@ -35,7 +36,7 @@ static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
     let session = VortexSession::empty()
         .with::<KernelSession>()
         .with::<MemorySession>();
-    session.arrays().register(Pco);
+    initialize_pco(&session);
     session
 });
 

encodings/pco/src/access.rs

@@ -0,0 +1,167 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::ops::Range;
+use std::sync::OnceLock;
+
+use vortex_array::ArrayView;
+use vortex_array::ExecutionCtx;
+use vortex_error::VortexExpect;
+use vortex_error::VortexResult;
+use vortex_mask::Mask;
+
+use crate::Pco;
+use crate::PcoArrayExt;
+use crate::page_index::IndexedView;
+
+const RANK_STRIDE: usize = 512;
+pub(crate) struct RankMask {
+    mask: Mask,
+    ranks: OnceLock<Vec<usize>>,
+}
+
+impl RankMask {
+    fn new(mask: Mask) -> Self {
+        Self {
+            mask,
+            ranks: OnceLock::new(),
+        }
+    }
+
+    fn boundary_rank(&self, row: usize) -> usize {
+        match &self.mask {
+            Mask::AllTrue(_) => row,
+            Mask::AllFalse(_) => 0,
+            Mask::Values(values) => {
+                let bits = values.bit_buffer();
+                if row <= bits.len() - row {
+                    bits.count_range(0, row)
+                } else {
+                    self.mask.true_count() - bits.count_range(row, bits.len())
+                }
+            }
+        }
+    }
+
+    fn rank(&self, row: usize) -> usize {
+        match &self.mask {
+            Mask::AllTrue(_) => row,
+            Mask::AllFalse(_) => 0,
+            Mask::Values(values) => {
+                if row == self.mask.len() {
+                    return self.mask.true_count();
+                }
+                let ranks = self.ranks.get_or_init(|| {
+                    let bits = values.bit_buffer();
+                    let mut ranks = Vec::with_capacity(bits.len().div_ceil(RANK_STRIDE));
+                    let mut count = 0;
+                    for start in (0..bits.len()).step_by(RANK_STRIDE) {
+                        ranks.push(count);
+                        count += bits.count_range(start, (start + RANK_STRIDE).min(bits.len()));
+                    }
+                    ranks
+                });
+                let block = row / RANK_STRIDE;
+                ranks[block] + values.bit_buffer().count_range(block * RANK_STRIDE, row)
+            }
+        }
+    }
+
+    fn value_index<const INDEXED: bool>(&self, row: usize) -> Option<usize> {
+        match &self.mask {
+            Mask::AllTrue(_) => Some(row),
+            Mask::AllFalse(_) => None,
+            Mask::Values(values) => {
+                let bits = values.bit_buffer();
+                bits.value(row).then(|| {
+                    if INDEXED {
+                        self.rank(row)
+                    } else {
+                        self.boundary_rank(row)
+                    }
+                })
+            }
+        }
+    }
+}
+
+/// Only new-prefix arrays use this operation-owned last-page validity preparation.
+#[derive(Default)]
+pub(crate) struct AccessState {
+    last: Option<(usize, PreparedPage)>,
+}
+struct PreparedPage {
+    rows: Range<usize>,
+    mask: RankMask,
+}
+impl AccessState {
+    fn page(
+        &mut self,
+        view: IndexedView<'_>,
+        array: ArrayView<'_, Pco>,
+        row: usize,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<(usize, &PreparedPage)> {
+        if !self
+            .last
+            .as_ref()
+            .is_some_and(|(_, cached)| cached.rows.contains(&row))
+        {
+            let page = view.find_row(row);
+            let rows = view.rows(page)?;
+            let mask = array
+                .unsliced_validity()
+                .slice(rows.clone())?
+                .execute_mask(rows.len(), ctx)?;
+            self.last = Some((
+                page,
+                PreparedPage {
+                    rows,
+                    mask: RankMask::new(mask),
+                },
+            ));
+        }
+        let (page, cached) = self.last.as_ref().vortex_expect("page prepared above");
+        Ok((*page, cached))
+    }
+    pub fn resolve<const RANKED: bool>(
+        &mut self,
+        view: IndexedView<'_>,
+        array: ArrayView<'_, Pco>,
+        row: usize,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<(usize, usize)>> {
+        let (page, cached) = self.page(view, array, row, ctx)?;
+        Ok(cached
+            .mask
+            .value_index::<RANKED>(row - cached.rows.start)
+            .map(|value| (page, value)))
+    }
+    pub fn span(
+        &mut self,
+        view: IndexedView<'_>,
+        array: ArrayView<'_, Pco>,
+        row: usize,
+        stop: usize,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<PageSpan> {
+        let (page, cached) = self.page(view, array, row, ctx)?;
+        let end = stop.min(cached.rows.end);
+        let start_in_page = row - cached.rows.start;
+        let end_in_page = end - cached.rows.start;
+        Ok(PageSpan {
+            page,
+            stop: end,
+            values: cached.mask.boundary_rank(start_in_page)
+                ..cached.mask.boundary_rank(end_in_page),
+            mask: cached.mask.mask.slice(start_in_page..end_in_page),
+        })
+    }
+}
+pub(crate) struct PageSpan {
+    pub page: usize,
+    pub stop: usize,
+    /// Dense-value positions within this physical page, never global value ordinals.
+    pub values: Range<usize>,
+    pub mask: Mask,
+}

encodings/pco/src/array.rs

@@ -8,6 +8,7 @@ use std::fmt::Formatter;
 use std::hash::Hash;
 use std::hash::Hasher;
 use std::sync::Arc;
+use std::sync::OnceLock;
 
 use pco::ChunkConfig;
 use pco::PagingSpec;
@@ -50,6 +51,7 @@ use vortex_array::vtable::validity_to_child;
 use vortex_buffer::BufferMut;
 use vortex_buffer::ByteBuffer;
 use vortex_error::VortexError;
+use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure;
@@ -63,6 +65,7 @@ use crate::PcoChunkInfo;
 use crate::PcoLogicalRowPrefix;
 use crate::PcoMetadata;
 use crate::PcoPageInfo;
+use crate::page_index::PageIndex;
 
 // Overall approach here:
 // Chunk the array into Pco chunks (currently using the default recommended size
@@ -406,6 +409,7 @@ pub struct PcoData {
     pub(crate) chunk_metas: Arc<[ByteBuffer]>,
     pub(crate) pages: Arc<[ByteBuffer]>,
     pub(crate) metadata: Arc<PcoMetadata>,
+    page_index: Arc<OnceLock<PageIndex>>,
     ptype: PType,
     unsliced_n_rows: usize,
     slice_start: usize,
@@ -505,6 +509,7 @@ impl PcoData {
             chunk_metas: chunk_metas.into(),
             pages: pages.into(),
             metadata: Arc::new(metadata),
+            page_index: Arc::new(OnceLock::new()),
             ptype,
             unsliced_n_rows: len,
             slice_start: 0,
@@ -781,6 +786,19 @@ impl PcoData {
             .into_byte_buffer())
     }
 
+    pub(crate) fn page_index(&self) -> VortexResult<&PageIndex> {
+        if self.page_index.get().is_none() {
+            let index = PageIndex::try_new(Arc::clone(&self.metadata))?;
+            // Concurrent first readers may prepare equivalent immutable descriptors. Only one
+            // index is retained, and no decoded values or validity are cached in the array.
+            drop(self.page_index.set(index));
+        }
+        Ok(self
+            .page_index
+            .get()
+            .vortex_expect("page index was initialized"))
+    }
+
     pub(crate) fn _slice(&self, start: usize, stop: usize) -> Self {
         PcoData {
             slice_start: self.slice_start + start,

encodings/pco/src/decoder.rs

@@ -0,0 +1,73 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! File and chunk decoding setup amortized over all selected pages in an operation.
+
+use std::collections::BTreeMap;
+use std::collections::btree_map::Entry;
+
+use pco::data_types::Number;
+use pco::wrapped::ChunkDecompressor;
+use pco::wrapped::FileDecompressor;
+use vortex_array::ExecutionCtx;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_array::dtype::NativePType;
+use vortex_array::validity::Validity;
+use vortex_buffer::BufferMut;
+use vortex_error::VortexResult;
+
+use crate::array::vortex_err_from_pco;
+
+pub(crate) struct PageDecoder<T: Number> {
+    file: FileDecompressor,
+    chunks: BTreeMap<usize, ChunkDecompressor<T>>,
+}
+
+impl<T: Number + NativePType> PageDecoder<T> {
+    pub fn new(header: &[u8]) -> VortexResult<Self> {
+        let (file, _) = FileDecompressor::new(header).map_err(vortex_err_from_pco)?;
+        Ok(Self {
+            file,
+            chunks: BTreeMap::new(),
+        })
+    }
+
+    pub fn read(
+        &mut self,
+        chunk: usize,
+        chunk_meta: &[u8],
+        n_values: usize,
+        buffer: &[u8],
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<PrimitiveArray> {
+        let mut values = BufferMut::<T>::zeroed_in(n_values, ctx.allocator().clone());
+        self.read_into(chunk, chunk_meta, n_values, buffer, &mut values)?;
+        Ok(PrimitiveArray::new(values.freeze(), Validity::NonNullable))
+    }
+    /// Decode into initialized storage and propagate codec errors. No additional semantic
+    /// completion check is needed to keep a partial/corrupt result memory-safe.
+    pub fn read_into(
+        &mut self,
+        chunk: usize,
+        chunk_meta: &[u8],
+        n_values: usize,
+        buffer: &[u8],
+        output: &mut [T],
+    ) -> VortexResult<()> {
+        let chunk = match self.chunks.entry(chunk) {
+            Entry::Occupied(entry) => entry.into_mut(),
+            Entry::Vacant(entry) => {
+                let (decoder, _) = self
+                    .file
+                    .chunk_decompressor::<T, _>(chunk_meta)
+                    .map_err(vortex_err_from_pco)?;
+                entry.insert(decoder)
+            }
+        };
+        let mut decoder = chunk
+            .page_decompressor(buffer, n_values)
+            .map_err(vortex_err_from_pco)?;
+        decoder.read(output).map_err(vortex_err_from_pco)?;
+        Ok(())
+    }
+}

encodings/pco/src/lib.rs

@@ -21,11 +21,22 @@
 //! ```
 
 use prost::Message;
+use vortex_array::ArrayVTable;
+use vortex_array::arrays::Dict;
+use vortex_array::arrays::dict::TakeExecuteAdaptor;
+use vortex_array::optimizer::kernels::ArrayKernelsExt;
+use vortex_array::session::ArraySessionExt;
+use vortex_session::VortexSession;
+
+mod access;
 mod array;
 mod compute;
+mod decoder;
+mod page_index;
 mod probe;
 mod rules;
 mod slice;
+mod take;
 
 pub use array::*;
 
@@ -75,6 +86,17 @@ pub struct PcoLogicalRowPrefix {
     pub logical_row_ends: Vec<u64>,
 }
 
+/// Register PCO encoding and its selective take execution kernel in a session.
+///
+/// Registering only `Pco` in the array registry is sufficient for legacy reads, but does not
+/// enable the batched take kernel.
+pub fn initialize(session: &VortexSession) {
+    session.arrays().register(Pco);
+    session
+        .kernels()
+        .register_execute_parent_kernel(Dict.id(), Pco, TakeExecuteAdaptor(Pco));
+}
+
 #[cfg(test)]
 mod tests;