Decode selected PCO pages for primitive-index takes

encodings/pco/Cargo.toml

@@ -16,6 +16,7 @@ version = { workspace = true }
 workspace = true
 
 [dependencies]
+once_cell = { workspace = true }
 pco = { workspace = true }
 prost = { workspace = true }
 vortex-array = { workspace = true }

encodings/pco/benches/selection.rs

@@ -27,6 +27,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;
@@ -53,7 +54,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,144 @@
+// 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::LogicalPageIndex;
+
+const RANK_STRIDE: usize = 512;
+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: &LogicalPageIndex,
+        array: ArrayView<'_, Pco>,
+        row: usize,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<(usize, &PreparedPage)> {
+        if !self
+            .last
+            .as_ref()
+            .is_some_and(|(_, page)| page.rows.contains(&row))
+        {
+            let page_index = view.find_row(row);
+            let rows = view.rows(page_index);
+            let rows = usize::try_from(rows.start)
+                .vortex_expect("page row start must fit the validity array")
+                ..usize::try_from(rows.end)
+                    .vortex_expect("page row end must fit the validity array");
+            let mask = array
+                .unsliced_validity()
+                .slice(rows.clone())?
+                .execute_mask(rows.len(), ctx)?;
+            self.last = Some((
+                page_index,
+                PreparedPage {
+                    rows,
+                    mask: RankMask::new(mask),
+                },
+            ));
+        }
+        let (page_index, page) = self.last.as_ref().vortex_expect("page prepared above");
+        Ok((*page_index, page))
+    }
+    pub fn resolve<const RANKED: bool>(
+        &mut self,
+        view: &LogicalPageIndex,
+        array: ArrayView<'_, Pco>,
+        row: usize,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<(usize, usize)>> {
+        let (page_index, page) = self.page(view, array, row, ctx)?;
+        Ok(page
+            .mask
+            .value_index::<RANKED>(row - page.rows.start)
+            .map(|value| (page_index, value)))
+    }
+}

encodings/pco/src/array.rs

@@ -9,6 +9,7 @@ use std::hash::Hash;
 use std::hash::Hasher;
 use std::sync::Arc;
 
+use once_cell::sync::OnceCell;
 use pco::ChunkConfig;
 use pco::PagingSpec;
 use pco::data_types::Number;
@@ -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<Vec<ByteBuffer>>,
     pub(crate) pages: Arc<Vec<ByteBuffer>>,
     pub(crate) metadata: Arc<PcoMetadata>,
+    page_index: Arc<OnceCell<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(OnceCell::new()),
             ptype,
             unsliced_n_rows: len,
             slice_start: 0,
@@ -781,6 +786,11 @@ impl PcoData {
             .into_byte_buffer())
     }
 
+    pub(crate) fn page_index(&self) -> VortexResult<&PageIndex> {
+        self.page_index
+            .get_or_try_init(|| PageIndex::try_new(Arc::clone(&self.metadata)))
+    }
+
     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,100 @@
+// 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::ArrayView;
+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::VortexExpect;
+use vortex_error::VortexResult;
+use vortex_error::vortex_err;
+
+use crate::Pco;
+use crate::array::vortex_err_from_pco;
+use crate::page_index::LogicalPageIndex;
+use crate::page_index::PhysicalPage;
+
+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>::with_capacity_in(n_values, ctx.allocator().clone());
+        unsafe {
+            values.set_len(n_values);
+        }
+        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(&mut values).map_err(vortex_err_from_pco)?;
+        Ok(PrimitiveArray::new(values.freeze(), Validity::NonNullable))
+    }
+}
+
+#[derive(Default)]
+pub(crate) struct PageLookup {
+    last: Option<PhysicalPage>,
+}
+
+impl PageLookup {
+    pub fn resolve(
+        &mut self,
+        index: &LogicalPageIndex,
+        array: ArrayView<'_, Pco>,
+        page: usize,
+    ) -> VortexResult<(usize, usize)> {
+        if !self
+            .last
+            .as_ref()
+            .is_some_and(|last| last.pages.contains(&page))
+        {
+            let info = index.chunk_for_page(page)?;
+            let result = (info.chunk, info.n_values);
+            self.last = Some(info);
+            return Ok(result);
+        }
+        let last = self.last.as_ref().vortex_expect("chunk lookup initialized");
+        let info = array.metadata.chunks[last.chunk]
+            .pages
+            .get(page - last.pages.start)
+            .ok_or_else(|| vortex_err!("Pco page metadata is missing"))?;
+        Ok((last.chunk, info.n_values as usize))
+    }
+}

encodings/pco/src/lib.rs

@@ -9,24 +9,33 @@
 //!
 //! Pco supports integer and floating-point primitive dtypes handled by the upstream `pco` crate.
 //! It is normally selected through the BtrBlocks compressor when the `pco` feature is enabled.
-//! To deserialize arrays manually, register the encoding in the array session:
+//! To register the encoding and its selective take kernel in a session:
 //!
 //! ```rust
-//! use vortex_array::session::ArraySessionExt;
+//! use vortex_pco::initialize;
 //! use vortex_session::VortexSession;
-//! use vortex_pco::Pco;
 //!
 //! let session = VortexSession::empty();
-//! session.arrays().register(Pco);
+//! initialize(&session);
 //! ```
 
 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::*;
 
@@ -77,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;