From 7317076fc2f87e6128f95d4ec9088be987ccdabc Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 12:30:12 -0400 Subject: [PATCH 01/36] Add PCO logical page indexing and shared selective access Signed-off-by: Daniel King --- encodings/pco/src/access.rs | 181 +++++++++++++++++++++++++++++ encodings/pco/src/array.rs | 105 ++++++++++++++--- encodings/pco/src/compute/cast.rs | 15 +-- encodings/pco/src/lib.rs | 41 ++++++- encodings/pco/src/page_index.rs | 93 +++++++++++++++ encodings/pco/src/probe.rs | 118 ++----------------- encodings/pco/src/slice.rs | 5 + encodings/pco/src/take.rs | 187 ++++++++++++++++++++++++++++++ vortex-file/src/lib.rs | 3 +- 9 files changed, 607 insertions(+), 141 deletions(-) create mode 100644 encodings/pco/src/access.rs create mode 100644 encodings/pco/src/page_index.rs create mode 100644 encodings/pco/src/take.rs diff --git a/encodings/pco/src/access.rs b/encodings/pco/src/access.rs new file mode 100644 index 00000000000..5a54388ad1a --- /dev/null +++ b/encodings/pco/src/access.rs @@ -0,0 +1,181 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Operation-owned selected-page values and validity. No decoded state lives in array metadata. + +use std::collections::BTreeMap; + +use pco::data_types::NumberType; +use pco::match_number_enum; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::half; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; +use vortex_mask::Mask; + +use crate::Pco; +use crate::PcoArrayExt; +use crate::array::number_type_from_ptype; + +const RANK_STRIDE: usize = 512; + +#[derive(Default)] +pub(crate) struct AccessState { + legacy: Option, + pages: BTreeMap, +} + +struct CachedPage { + mask: Option, + decoded: Option, +} + +struct RankMask { + mask: Mask, + ranks: Vec, +} + +impl RankMask { + fn new(mask: Mask) -> Self { + let ranks = match &mask { + Mask::Values(values) => { + 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 + } + _ => Vec::new(), + }; + Self { mask, ranks } + } + + fn value_index(&self, row: usize) -> Option { + match &self.mask { + Mask::AllTrue(_) => Some(row), + Mask::AllFalse(_) => None, + Mask::Values(values) => { + let bits = values.bit_buffer(); + bits.value(row).then(|| { + let block = row / RANK_STRIDE; + self.ranks[block] + bits.count_range(block * RANK_STRIDE, row) + }) + } + } + } +} + +impl AccessState { + /// Resolve an unsliced logical row to a physical page and compact within-page index. + pub fn resolve( + &mut self, + array: ArrayView<'_, Pco>, + row: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + vortex_ensure!( + row < array.unsliced_n_rows(), + "Pco row exceeds array bounds" + ); + let index = array.page_index()?; + if index.logical && !index.pages.is_empty() { + let page_idx = index + .pages + .partition_point(|page| page.rows.as_ref().is_some_and(|range| range.end <= row)); + let page = index + .pages + .get(page_idx) + .ok_or_else(|| vortex_err!("Pco logical page index does not cover row"))?; + let rows = page + .rows + .as_ref() + .ok_or_else(|| vortex_err!("Pco logical page range is missing"))?; + if !self.pages.contains_key(&page_idx) { + let mask = array + .unsliced_validity() + .slice(rows.clone())? + .execute_mask(rows.len(), ctx)?; + vortex_ensure_eq!( + mask.true_count(), + page.values.len(), + "Pco logical page validity count does not match stored values" + ); + self.pages.insert( + page_idx, + CachedPage { + mask: Some(RankMask::new(mask)), + decoded: None, + }, + ); + } + return Ok(self.pages[&page_idx] + .mask + .as_ref() + .and_then(|mask| mask.value_index(row - rows.start)) + .map(|value| (page_idx, value))); + } + + // Legacy or unknown optional metadata: execute/rank the full validity once per operation, + // never once per requested index. Empty physical arrays also take this checked path. + if self.legacy.is_none() { + self.legacy = Some(RankMask::new( + array + .unsliced_validity() + .execute_mask(array.unsliced_n_rows(), ctx)?, + )); + } + let Some(value) = self.legacy.as_ref().and_then(|mask| mask.value_index(row)) else { + return Ok(None); + }; + let page_idx = index.pages.partition_point(|page| page.values.end <= value); + let page = index.pages.get(page_idx).ok_or_else(|| { + vortex_err!("Pco validity contains more valid rows than stored values") + })?; + self.pages.entry(page_idx).or_insert(CachedPage { + mask: None, + decoded: None, + }); + Ok(Some((page_idx, value - page.values.start))) + } + + pub fn values( + &mut self, + array: ArrayView<'_, Pco>, + page_idx: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult<&PrimitiveArray> { + let cached = self + .pages + .get_mut(&page_idx) + .ok_or_else(|| vortex_err!("Pco page must be resolved before decoding"))?; + if cached.decoded.is_none() { + let index = array.page_index()?; + let page = index + .pages + .get(page_idx) + .ok_or_else(|| vortex_err!("Pco page is missing"))?; + let buffer = array + .pages + .get(page_idx) + .ok_or_else(|| vortex_err!("Pco compressed page is missing"))?; + let decoded = match_number_enum!( + number_type_from_ptype(array.dtype().as_ptype()), + NumberType => { + crate::probe::decode_page::(array, page.chunk, page.values.len(), buffer.as_slice(), ctx)? + } + ); + cached.decoded = Some(decoded); + } + cached + .decoded + .as_ref() + .ok_or_else(|| vortex_err!("Pco page decoding did not produce values")) + } +} diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array.rs index c1a4347cd94..07186672181 100644 --- a/encodings/pco/src/array.rs +++ b/encodings/pco/src/array.rs @@ -7,6 +7,8 @@ use std::fmt::Display; 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; @@ -58,8 +60,10 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::PcoChunkInfo; +use crate::PcoLogicalPages; 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 @@ -89,11 +93,22 @@ impl ArrayHash for PcoData { self.unsliced_n_rows.hash(state); self.slice_start.hash(state); self.slice_stop.hash(state); + self.metadata.header.hash(state); + for chunk in &self.metadata.chunks { + chunk.pages.len().hash(state); + for page in &chunk.pages { + page.n_values.hash(state); + } + } + if let Some(index) = &self.metadata.logical_pages { + index.version.hash(state); + index.ends.hash(state); + } // Hash chunk_metas and pages using pointer-based hashing - for chunk_meta in &self.chunk_metas { + for chunk_meta in self.chunk_metas.iter() { chunk_meta.array_hash(state, accuracy); } - for page in &self.pages { + for page in self.pages.iter() { page.array_hash(state, accuracy); } } @@ -106,15 +121,16 @@ impl ArrayEq for PcoData { || self.slice_stop != other.slice_stop || self.chunk_metas.len() != other.chunk_metas.len() || self.pages.len() != other.pages.len() + || self.metadata != other.metadata { return false; } - for (a, b) in self.chunk_metas.iter().zip(&other.chunk_metas) { + for (a, b) in self.chunk_metas.iter().zip(other.chunk_metas.iter()) { if !a.array_eq(b, accuracy) { return false; } } - for (a, b) in self.pages.iter().zip(&other.pages) { + for (a, b) in self.pages.iter().zip(other.pages.iter()) { if !a.array_eq(b, accuracy) { return false; } @@ -180,11 +196,13 @@ impl VTable for Pco { data.chunk_metas = buffers[..chunk_metas_len] .iter() .map(|buffer| buffer.clone().try_to_host_sync()) - .collect::>>()?; + .collect::>>()? + .into(); data.pages = buffers[chunk_metas_len..] .iter() .map(|buffer| buffer.clone().try_to_host_sync()) - .collect::>>()?; + .collect::>>()? + .into(); Ok(ArrayParts::new( self.clone(), array.dtype().clone(), @@ -198,7 +216,11 @@ impl VTable for Pco { array: ArrayView<'_, Self>, _session: &VortexSession, ) -> VortexResult>> { - Ok(Some(array.metadata.clone().encode_to_vec())) + vortex_ensure!( + array.slice_start == 0 && array.slice_stop == array.unsliced_n_rows, + "Sliced Pco data must be executed to a primitive array (and optionally recompressed) before serialization" + ); + Ok(Some(array.metadata.encode_to_vec())) } fn deserialize( @@ -259,6 +281,11 @@ impl VTable for Pco { } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + if array.slice_start() != 0 || array.slice_stop() != array.unsliced_n_rows() { + return Ok(ExecutionResult::done( + crate::take::slice(array.as_view(), ctx)?.into_array(), + )); + } let unsliced_validity = array.unsliced_validity(); Ok(ExecutionResult::done( array @@ -384,9 +411,10 @@ impl> PcoArrayExt for T {} #[derive(Clone, Debug)] /// Encoding-specific data for a [`PcoArray`]. pub struct PcoData { - pub(crate) chunk_metas: Vec, - pub(crate) pages: Vec, - pub(crate) metadata: PcoMetadata, + pub(crate) chunk_metas: Arc<[ByteBuffer]>, + pub(crate) pages: Arc<[ByteBuffer]>, + pub(crate) metadata: Arc, + page_index: Arc>, ptype: PType, unsliced_n_rows: usize, slice_start: usize, @@ -466,6 +494,7 @@ impl PcoData { } else if validity.definitely_all_null() { vortex_ensure_eq!(n_values, 0, "Pco contains values for an all-null array"); } + self.page_index()?; Ok(()) } @@ -483,9 +512,10 @@ impl PcoData { len: usize, ) -> Self { Self { - chunk_metas, - pages, - metadata, + 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, @@ -569,9 +599,28 @@ impl PcoData { chunk_infos.push(PcoChunkInfo { pages: page_infos }) } + let mask = parray.validity()?.execute_mask(parray.len(), ctx)?; + let mut ends = Vec::with_capacity(page_buffers.len()); + let mut page_counts = chunk_infos.iter().flat_map(|chunk| &chunk.pages); + let mut remaining = page_counts.next().map(|page| page.n_values as usize); + // Compression already visits all values; build logical boundaries in the same linear + // scale, without allocating one index per valid row. + for (row, valid) in mask.iter().enumerate() { + if valid && let Some(n) = remaining.as_mut() { + *n -= 1; + if *n == 0 { + ends.push(u64::try_from(row + 1)?); + remaining = page_counts.next().map(|page| page.n_values as usize); + } + } + } + if let Some(last) = ends.last_mut() { + *last = u64::try_from(parray.len())?; + } let metadata = PcoMetadata { header, chunks: chunk_infos, + logical_pages: Some(PcoLogicalPages { version: 1, ends }), }; // SAFETY: The compressor produced matching chunk metadata and pages from `parray`, with // the same ptype, logical length, and valid-value count. @@ -637,9 +686,21 @@ impl PcoData { ctx: &mut ExecutionCtx, ) -> VortexResult { // To start, we figure out what range of values we need to decompress. - let slice_value_indices = unsliced_validity - .execute_mask(self.unsliced_n_rows, ctx)? - .valid_counts_for_indices(&[self.slice_start, self.slice_stop]); + let mask = unsliced_validity.execute_mask(self.unsliced_n_rows, ctx)?; + let index = self.page_index()?; + if index.logical { + for page in &index.pages { + if let Some(rows) = &page.rows { + vortex_ensure_eq!( + mask.slice(rows.clone()).true_count(), + page.values.len(), + "Pco logical page validity count does not match stored values" + ); + } + } + } + let slice_value_indices = + mask.valid_counts_for_indices(&[self.slice_start, self.slice_stop]); let slice_value_start = slice_value_indices[0]; let slice_value_stop = slice_value_indices[1]; let slice_n_values = slice_value_stop - slice_value_start; @@ -655,7 +716,7 @@ impl PcoData { let mut page_idx = 0; let mut page_value_start = 0usize; let mut n_skipped_values = 0; - for (chunk_info, chunk_meta) in self.metadata.chunks.iter().zip(&self.chunk_metas) { + for (chunk_info, chunk_meta) in self.metadata.chunks.iter().zip(self.chunk_metas.iter()) { // lazily initialize chunk decompressor let mut chunk_decompressor: Option> = None; for page_info in &chunk_info.pages { @@ -722,6 +783,16 @@ 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(&self.metadata, self.unsliced_n_rows)?; + // Concurrent first readers may prepare equivalent immutable descriptors. Only one + // index is retained, and no decoded values or validity are cached in the array. + let _ = self.page_index.set(index); + } + Ok(self.page_index.get().expect("page index was initialized")) + } + pub(crate) fn _slice(&self, start: usize, stop: usize) -> Self { PcoData { slice_start: self.slice_start + start, diff --git a/encodings/pco/src/compute/cast.rs b/encodings/pco/src/compute/cast.rs index af1d003d7f7..17d3835d21e 100644 --- a/encodings/pco/src/compute/cast.rs +++ b/encodings/pco/src/compute/cast.rs @@ -10,7 +10,6 @@ use vortex_error::VortexResult; use crate::Pco; use crate::PcoArrayExt; -use crate::PcoData; impl CastReduce for Pco { fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { @@ -35,18 +34,8 @@ impl CastReduce for Pco { return Ok(None); }; - // SAFETY: The buffers and metadata come from a validated Pco array, the primitive type is - // unchanged, and `Pco::try_new` validates the adjusted nullability and slice below. - let data = unsafe { - PcoData::new_unchecked( - array.chunk_metas.clone(), - array.pages.clone(), - dtype.as_ptype(), - array.metadata.clone(), - array.unsliced_n_rows(), - ) - } - ._slice(array.slice_start(), array.slice_stop()); + // Only nullability changes; retain the same immutable page index and compressed owners. + let data = array.data().clone(); Ok(Some( Pco::try_new(dtype.clone(), data, new_validity)?.into_array(), diff --git a/encodings/pco/src/lib.rs b/encodings/pco/src/lib.rs index df6581ec1d1..dc119a1c053 100644 --- a/encodings/pco/src/lib.rs +++ b/encodings/pco/src/lib.rs @@ -19,15 +19,18 @@ //! session.arrays().register(vortex_pco::Pco); //! ``` +mod access; mod array; mod compute; +mod page_index; mod probe; mod rules; mod slice; +mod take; pub use array::*; -#[derive(Clone, prost::Message)] +#[derive(Clone, PartialEq, prost::Message)] /// Metadata for one Pco page. pub struct PcoPageInfo { // Since pco limits to 2^24 values per chunk, u32 is sufficient for the @@ -39,7 +42,7 @@ pub struct PcoPageInfo { // We're calling this Info instead of Metadata because ChunkMeta refers to a specific // component of a Pco file. -#[derive(Clone, prost::Message)] +#[derive(Clone, PartialEq, prost::Message)] /// Metadata for one Pco chunk. pub struct PcoChunkInfo { /// Pages contained in this chunk. @@ -47,7 +50,7 @@ pub struct PcoChunkInfo { pub pages: Vec, } -#[derive(Clone, prost::Message)] +#[derive(Clone, PartialEq, prost::Message)] /// Serialized metadata for a [`PcoArray`]. pub struct PcoMetadata { // would be nice to reuse one header per vortex file, but it's really only 1 byte, so @@ -58,6 +61,38 @@ pub struct PcoMetadata { /// Metadata for each compressed chunk. #[prost(message, repeated, tag = "2")] pub chunks: Vec, + /// Optional logical-row boundaries for selective decoding. Older readers ignore this field. + #[prost(message, optional, tag = "3")] + pub logical_pages: Option, +} + +/// Versioned optional logical-row page index. Unknown versions retain legacy decoding. +#[derive(Clone, PartialEq, prost::Message)] +pub struct PcoLogicalPages { + /// Index interpretation version; currently one. + #[prost(uint32, tag = "1")] + pub version: u32, + /// Exclusive unsliced logical-row ends in physical page order. + /// + /// The final page includes trailing null rows. An empty index has no valid values. + #[prost(uint64, repeated, tag = "2")] + pub ends: Vec, +} + +/// 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: &vortex_session::VortexSession) { + 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; + session.arrays().register(Pco); + session + .kernels() + .register_execute_parent_kernel(Dict.id(), Pco, TakeExecuteAdaptor(Pco)); } #[cfg(test)] diff --git a/encodings/pco/src/page_index.rs b/encodings/pco/src/page_index.rs new file mode 100644 index 00000000000..8fe0faab72f --- /dev/null +++ b/encodings/pco/src/page_index.rs @@ -0,0 +1,93 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Immutable physical-page descriptors shared by sliced owners and access operations. + +use std::ops::Range; + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; + +use crate::PcoMetadata; + +#[derive(Debug)] +pub(crate) struct Page { + pub chunk: usize, + pub values: Range, + pub rows: Option>, +} + +#[derive(Debug)] +pub(crate) struct PageIndex { + pub pages: Vec, + pub logical: bool, +} + +impl PageIndex { + pub fn try_new(metadata: &PcoMetadata, rows: usize) -> VortexResult { + let logical = metadata + .logical_pages + .as_ref() + .filter(|index| index.version == 1); + let page_count = metadata.chunks.iter().try_fold(0usize, |count, chunk| { + count + .checked_add(chunk.pages.len()) + .ok_or_else(|| vortex_err!("Pco page count overflows usize")) + })?; + if let Some(index) = logical { + vortex_ensure_eq!( + index.ends.len(), + page_count, + "Pco logical page count does not match physical pages" + ); + } + let mut pages = Vec::with_capacity(page_count); + let mut value_start = 0usize; + let mut row_start = 0usize; + for (chunk, info) in metadata.chunks.iter().enumerate() { + for page in &info.pages { + let n = page.n_values as usize; + vortex_ensure!(n != 0, "Pco contains an empty physical page"); + let end = value_start + .checked_add(n) + .ok_or_else(|| vortex_err!("Pco value prefix overflows usize"))?; + vortex_ensure!(end <= rows, "Pco value prefix exceeds logical row count"); + let logical_rows = if let Some(index) = logical { + let stop = usize::try_from(index.ends[pages.len()])?; + vortex_ensure!( + stop > row_start && stop <= rows, + "Pco logical page ends are not strictly increasing within array bounds" + ); + vortex_ensure!( + stop - row_start >= n, + "Pco page has fewer logical rows than values" + ); + let range = row_start..stop; + row_start = stop; + Some(range) + } else { + None + }; + pages.push(Page { + chunk, + values: value_start..end, + rows: logical_rows, + }); + value_start = end; + } + } + if logical.is_some() && page_count != 0 { + vortex_ensure_eq!( + row_start, + rows, + "Pco logical page ends do not cover all rows" + ); + } + Ok(Self { + pages, + logical: logical.is_some(), + }) + } +} diff --git a/encodings/pco/src/probe.rs b/encodings/pco/src/probe.rs index 00aa801254a..936c55dba5c 100644 --- a/encodings/pco/src/probe.rs +++ b/encodings/pco/src/probe.rs @@ -3,51 +3,29 @@ //! PCO probes retain every decoded page and index the compacted, non-null value positions. -use std::ops::Range; - use pco::data_types::Number; -use pco::data_types::NumberType; -use pco::match_number_enum; use pco::wrapped::FileDecompressor; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::NativePType; -use vortex_array::dtype::half; use vortex_array::match_each_native_ptype; use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::BufferMut; -use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_mask::Mask; use crate::Pco; -use crate::PcoArrayExt; -use crate::array::number_type_from_ptype; +use crate::access::AccessState; use crate::array::vortex_err_from_pco; -const RANK_STRIDE: usize = 512; - /// State retained by repeated PCO probes. Construction performs no allocation. /// -/// The mask and rank index use unsliced logical rows; page boundaries use compacted value -/// positions. Every page decoded so far is retained, so the state grows with the pages touched -/// and never beyond the decompressed array. +/// Decoded pages and their local validity are retained as touched, up to the full array. +/// Legacy files additionally retain a full validity mask and its sparse rank index. #[derive(Default)] pub struct PcoProbeState { - validity: Option, - rank: Vec, - pages: Vec, - /// The page the previous probe landed in, `None` until the first probe resolves one. - last_page: Option, -} - -struct Page { - values: Range, - chunk: usize, - decoded: Option, + access: AccessState, } pub(crate) fn scalar_at( @@ -56,91 +34,19 @@ pub(crate) fn scalar_at( state: Option<&mut PcoProbeState>, ctx: &mut ExecutionCtx, ) -> VortexResult { - let Some(state) = state else { - let validity = array.unsliced_validity(); - return array - ._slice(index, index + 1) - .decompress(&validity, ctx)? - .into_array() - .execute_scalar(0, ctx); - }; - let mask = match &mut state.validity { - Some(mask) => mask, - slot @ None => slot.insert( - array - .unsliced_validity() - .execute_mask(array.unsliced_n_rows(), ctx)?, - ), - }; - let logical_index = array.slice_start() + index; - - let value_index = match mask { - Mask::AllTrue(_) => logical_index, - Mask::AllFalse(_) => return Ok(Scalar::null(array.dtype().clone())), - Mask::Values(values) => { - let bits = values.bit_buffer(); - if !bits.value(logical_index) { - return Ok(Scalar::null(array.dtype().clone())); - } - if state.rank.is_empty() { - state.rank.reserve(bits.len().div_ceil(RANK_STRIDE)); - let mut count = 0; - for start in (0..bits.len()).step_by(RANK_STRIDE) { - state.rank.push(count); - count += bits.count_range(start, (start + RANK_STRIDE).min(bits.len())); - } - } - let block = logical_index / RANK_STRIDE; - state.rank[block] + bits.count_range(block * RANK_STRIDE, logical_index) - } - }; - - if state.pages.is_empty() { - let mut start = 0; - for (chunk, metadata) in array.metadata.chunks.iter().enumerate() { - for page in &metadata.pages { - let end = start + page.n_values as usize; - state.pages.push(Page { - values: start..end, - chunk, - decoded: None, - }); - start = end; - } - } - } - let page_index = match state.last_page { - Some(last) if state.pages[last].values.contains(&value_index) => last, - _ => state - .pages - .partition_point(|page| page.values.end <= value_index), - }; - state.last_page = Some(page_index); - let page = state - .pages - .get_mut(page_index) - .vortex_expect("PCO pages cover every valid value index"); - let values = match &mut page.decoded { - Some(decoded) => decoded, - slot @ None => { - let decoded = match_number_enum!( - number_type_from_ptype(array.dtype().as_ptype()), - NumberType => { - decode_page::(array, page.chunk, page.values.len(), array.pages[page_index].as_slice(), ctx)? - } - ); - slot.insert(decoded) - } + let mut once = PcoProbeState::default(); + let state = state.unwrap_or(&mut once); + let row = array.slice_start() + index; + let Some((page, value)) = state.access.resolve(array, row, ctx)? else { + return Ok(Scalar::null(array.dtype().clone())); }; + let values = state.access.values(array, page, ctx)?; Ok(match_each_native_ptype!(values.ptype(), |T| { - Scalar::primitive( - values.as_slice::()[value_index - page.values.start], - array.dtype().nullability(), - ) + Scalar::primitive(values.as_slice::()[value], array.dtype().nullability()) })) } -fn decode_page( +pub(crate) fn decode_page( array: ArrayView<'_, Pco>, chunk: usize, n_values: usize, diff --git a/encodings/pco/src/slice.rs b/encodings/pco/src/slice.rs index 316ea835aa1..80e33972153 100644 --- a/encodings/pco/src/slice.rs +++ b/encodings/pco/src/slice.rs @@ -8,12 +8,17 @@ use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::arrays::slice::SliceReduce; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use crate::Pco; use crate::PcoArrayExt; impl SliceReduce for Pco { fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { + vortex_ensure!( + range.start <= range.end && range.end <= array.len(), + "Pco slice range exceeds array bounds" + ); let unsliced_validity = array.unsliced_validity(); Ok(Some( Pco::try_new( diff --git a/encodings/pco/src/take.rs b/encodings/pco/src/take.rs new file mode 100644 index 00000000000..4fba4b79693 --- /dev/null +++ b/encodings/pco/src/take.rs @@ -0,0 +1,187 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Batched selection using compact request descriptors and one operation-owned page cache. + +use pco::data_types::Number; +use pco::data_types::NumberType; +use pco::match_number_enum; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::PiecewiseSequence; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::dict::TakeExecute; +use vortex_array::arrays::piecewise_sequence::PiecewiseSequenceArraySlotsExt; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::half; +use vortex_array::match_each_integer_ptype; +use vortex_array::validity::Validity; +use vortex_buffer::BitBufferMut; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; + +use crate::Pco; +use crate::access::AccessState; +use crate::array::number_type_from_ptype; + +/// Compact descriptor; null indices have no start. Zero-length runs consume no source rows. +struct Run { + start: Option, + length: usize, + stride: usize, +} + +impl TakeExecute for Pco { + fn take( + array: ArrayView<'_, Pco>, + indices: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let runs = if let Some(sequence) = indices.as_opt::() { + let starts = sequence.starts().clone().execute::(ctx)?; + let lengths = sequence.lengths().clone().execute::(ctx)?; + let multipliers = sequence + .multipliers() + .clone() + .execute::(ctx)?; + let starts = primitive_usizes(&starts)?; + let lengths = primitive_usizes(&lengths)?; + let multipliers = primitive_usizes(&multipliers)?; + let mut runs = Vec::with_capacity(starts.len()); + for ((start, length), stride) in starts.into_iter().zip(lengths).zip(multipliers) { + runs.push(Run { + start: Some(start), + length, + stride, + }); + } + runs + } else { + let indices = indices.clone().execute::(ctx)?; + let mut runs = Vec::with_capacity(indices.len()); + let mask = indices.validity()?.execute_mask(indices.len(), ctx)?; + match_each_integer_ptype!(indices.ptype(), |I| { + for (&value, valid) in indices.as_slice::().iter().zip(mask.iter()) { + let start = if valid { + Some(usize::try_from(value).map_err(|_| { + vortex_err!("Pco take index is negative or exceeds usize") + })?) + } else { + None + }; + runs.push(Run { + start, + length: 1, + stride: 1, + }); + } + }); + runs + }; + validate_runs(&runs, array.len(), indices.len())?; + let nullable = array.dtype().nullability() | indices.dtype().nullability(); + let result = match_number_enum!(number_type_from_ptype(array.dtype().as_ptype()), NumberType => { + gather::(array, &runs, indices.len(), nullable, ctx)? + }); + Ok(Some(result.into_array())) + } +} + +fn primitive_usizes(array: &PrimitiveArray) -> VortexResult> { + match_each_integer_ptype!(array.ptype(), |I| { + array + .as_slice::() + .iter() + .map(|&value| { + usize::try_from(value).map_err(|_| vortex_err!("Pco take descriptor exceeds usize")) + }) + .collect() + }) +} + +fn validate_runs(runs: &[Run], source_len: usize, output_len: usize) -> VortexResult<()> { + let mut total = 0usize; + for run in runs { + total = total + .checked_add(run.length) + .ok_or_else(|| vortex_err!("Pco take output length overflows usize"))?; + if let Some(start) = run.start + && run.length != 0 + { + let last = (run.length - 1) + .checked_mul(run.stride) + .and_then(|delta| start.checked_add(delta)) + .ok_or_else(|| vortex_err!("Pco take range overflows usize"))?; + vortex_ensure!( + last < source_len, + "Pco take index {last} exceeds array length {source_len}" + ); + } + } + vortex_ensure_eq!( + total, + output_len, + "Pco take expanded length does not match declared length" + ); + Ok(()) +} + +fn gather( + array: ArrayView<'_, Pco>, + runs: &[Run], + output_len: usize, + nullability: vortex_array::dtype::Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mut state = AccessState::default(); + let mut values = BufferMut::::with_capacity_in(output_len, ctx.allocator().clone()); + let mut valid = BitBufferMut::with_capacity(output_len); + for run in runs { + let Some(start) = run.start else { + valid.append_n(false, run.length); + continue; + }; + for offset in 0..run.length { + // The complete descriptor, including its output length and last element, was + // checked before allocating output. The source slice itself is already validated. + let row = array.slice_start() + start + offset * run.stride; + if let Some((page, value)) = state.resolve(array, row, ctx)? { + values.push(state.values(array, page, ctx)?.as_slice::()[value]); + valid.append_true(); + } else { + valid.append_false(); + } + } + } + let validity = if nullability.is_nullable() { + Validity::from(valid.freeze()) + } else { + Validity::NonNullable + }; + Ok(PrimitiveArray::from_values_byte_buffer( + values.freeze().into_byte_buffer(), + array.dtype().as_ptype(), + validity, + output_len, + ctx, + )) +} + +pub(crate) fn slice( + array: ArrayView<'_, Pco>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let runs = [Run { + start: Some(0), + length: array.len(), + stride: 1, + }]; + match_number_enum!(number_type_from_ptype(array.dtype().as_ptype()), NumberType => { + gather::(array, &runs, array.len(), array.dtype().nullability(), ctx) + }) +} diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs index 1c309112078..6d589461afa 100644 --- a/vortex-file/src/lib.rs +++ b/vortex-file/src/lib.rs @@ -129,7 +129,6 @@ pub use strategy::*; use vortex_array::arrays::Patched; use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::session::ArraySessionExt; -use vortex_pco::Pco; use vortex_session::VortexSession; pub use writer::*; @@ -169,6 +168,7 @@ mod forever_constant { /// Registration covers reading: a session can decode every encoding registered here. The /// writer is gated separately by the editions enabled on its session. pub fn register_default_encodings(session: &VortexSession) { + vortex_pco::initialize(session); vortex_bytebool::initialize(session); vortex_fsst::initialize(session); vortex_onpair::initialize(session); @@ -178,7 +178,6 @@ pub fn register_default_encodings(session: &VortexSession) { { let arrays = session.arrays(); - arrays.register(Pco); if use_experimental_patches() { arrays.register(Patched); } From 77bc480e3dd7a1242711486aab5dc333bda5353e Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 12:34:21 -0400 Subject: [PATCH 02/36] Copy contiguous PCO runs and add public compatibility coverage Signed-off-by: Daniel King --- .github/workflows/pco-compat.yml | 54 ++++ encodings/pco/src/access.rs | 68 +++- encodings/pco/src/array.rs | 9 +- encodings/pco/src/lib.rs | 5 +- encodings/pco/src/select_tests.rs | 292 ++++++++++++++++++ encodings/pco/src/take.rs | 55 +++- vortex-file/tests/pco_logical_pages_compat.rs | 79 +++++ 7 files changed, 550 insertions(+), 12 deletions(-) create mode 100644 .github/workflows/pco-compat.yml create mode 100644 encodings/pco/src/select_tests.rs create mode 100644 vortex-file/tests/pco_logical_pages_compat.rs diff --git a/.github/workflows/pco-compat.yml b/.github/workflows/pco-compat.yml new file mode 100644 index 00000000000..aa870ab3612 --- /dev/null +++ b/.github/workflows/pco-compat.yml @@ -0,0 +1,54 @@ +name: PCO selective access compatibility + +on: + pull_request: + paths: + - 'encodings/pco/**' + - 'vortex-file/src/lib.rs' + - 'vortex-file/tests/pco_logical_pages_compat.rs' + - '.github/workflows/pco-compat.yml' + workflow_dispatch: {} + +permissions: + contents: read + +jobs: + pco-compat: + runs-on: >- + ${{ github.repository == 'vortex-data/vortex' + && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=pco-compat', github.run_id) + || 'ubuntu-latest' }} + timeout-minutes: 45 + steps: + - uses: runs-on/action@v2 + if: github.repository == 'vortex-data/vortex' + with: + sccache: s3 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7 + with: + fetch-depth: 0 + - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: ${{ github.repository == 'vortex-data/vortex' && 'true' || 'false' }} + - name: PCO public semantics + run: cargo test --profile ci --locked -p vortex-pco + - name: Generate new files with current public writer + env: + PCO_COMPAT_MODE: write + PCO_COMPAT_DIR: ${{ runner.temp }}/pco-new-files + run: cargo test --profile ci --locked -p vortex-file --test pco_logical_pages_compat + - name: Read new files using unchanged pre-change reader; generate legacy files + env: + PCO_COMPAT_DIR: ${{ runner.temp }}/pco-new-files + run: | + git worktree add --detach "$RUNNER_TEMP/pco-old-reader" 910f9933e8e2318c592c1eddcc2dfee8d24c6cd7 + # Only the public test harness is added. The old reader/writer/library sources remain + # exactly the pinned pre-change revision and do not know the new metadata field. + cp vortex-file/tests/pco_logical_pages_compat.rs "$RUNNER_TEMP/pco-old-reader/vortex-file/tests/" + PCO_COMPAT_MODE=read cargo test --manifest-path "$RUNNER_TEMP/pco-old-reader/Cargo.toml" --profile ci --locked -p vortex-file --test pco_logical_pages_compat + PCO_COMPAT_MODE=write PCO_COMPAT_DIR="$RUNNER_TEMP/pco-old-files" cargo test --manifest-path "$RUNNER_TEMP/pco-old-reader/Cargo.toml" --profile ci --locked -p vortex-file --test pco_logical_pages_compat + - name: Read legacy files using current reader + env: + PCO_COMPAT_MODE: read + PCO_COMPAT_DIR: ${{ runner.temp }}/pco-old-files + run: cargo test --profile ci --locked -p vortex-file --test pco_logical_pages_compat diff --git a/encodings/pco/src/access.rs b/encodings/pco/src/access.rs index 5a54388ad1a..58f05d6d729 100644 --- a/encodings/pco/src/access.rs +++ b/encodings/pco/src/access.rs @@ -57,16 +57,27 @@ impl RankMask { Self { mask, ranks } } + 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 block = row / RANK_STRIDE; + self.ranks[block] + values.bit_buffer().count_range(block * RANK_STRIDE, row) + } + } + } + fn value_index(&self, row: usize) -> Option { match &self.mask { Mask::AllTrue(_) => Some(row), Mask::AllFalse(_) => None, Mask::Values(values) => { let bits = values.bit_buffer(); - bits.value(row).then(|| { - let block = row / RANK_STRIDE; - self.ranks[block] + bits.count_range(block * RANK_STRIDE, row) - }) + bits.value(row).then(|| self.rank(row)) } } } @@ -145,6 +156,48 @@ impl AccessState { Ok(Some((page_idx, value - page.values.start))) } + /// Select a contiguous logical span within one indexed page. Validity is sliced only after + /// the complete page has been checked, so copies use one typed contiguous compact-value run. + pub fn span( + &mut self, + array: ArrayView<'_, Pco>, + row: usize, + stop: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let index = array.page_index()?; + if !index.logical || index.pages.is_empty() { + return Ok(None); + } + let page_idx = index + .pages + .partition_point(|page| page.rows.as_ref().is_some_and(|range| range.end <= row)); + let page = index + .pages + .get(page_idx) + .ok_or_else(|| vortex_err!("Pco logical page index does not cover span"))?; + let rows = page + .rows + .as_ref() + .ok_or_else(|| vortex_err!("Pco logical page range is missing"))?; + // Resolve performs whole-page validity validation, even when the first row is null. + self.resolve(array, row, ctx)?; + let cached = self + .pages + .get(&page_idx) + .and_then(|page| page.mask.as_ref()) + .ok_or_else(|| vortex_err!("Pco page validity is missing"))?; + let end = stop.min(rows.end); + let start_in_page = row - rows.start; + let end_in_page = end - rows.start; + Ok(Some(PageSpan { + page: page_idx, + stop: end, + values: cached.rank(start_in_page)..cached.rank(end_in_page), + mask: cached.mask.slice(start_in_page..end_in_page), + })) + } + pub fn values( &mut self, array: ArrayView<'_, Pco>, @@ -179,3 +232,10 @@ impl AccessState { .ok_or_else(|| vortex_err!("Pco page decoding did not produce values")) } } + +pub(crate) struct PageSpan { + pub page: usize, + pub stop: usize, + pub values: std::ops::Range, + pub mask: Mask, +} diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array.rs index 07186672181..a8b13952359 100644 --- a/encodings/pco/src/array.rs +++ b/encodings/pco/src/array.rs @@ -326,16 +326,16 @@ pub(crate) fn number_type_from_ptype(ptype: PType) -> NumberType { fn collect_valid( parray: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, -) -> VortexResult { +) -> VortexResult<(PrimitiveArray, vortex_mask::Mask)> { let mask = parray .array() .validity()? .execute_mask(parray.array().len(), ctx)?; let result = parray .array() - .filter(mask)? + .filter(mask.clone())? .execute::(ctx)?; - Ok(result) + Ok((result, mask)) } pub(crate) fn vortex_err_from_pco(err: PcoError) -> VortexError { @@ -558,7 +558,7 @@ impl PcoData { .with_compression_level(level) .with_paging_spec(PagingSpec::EqualPagesUpTo(values_per_page)); - let values = collect_valid(parray, ctx)?; + let (values, mask) = collect_valid(parray, ctx)?; let n_values = values.len(); let fc = FileCompressor::default(); @@ -599,7 +599,6 @@ impl PcoData { chunk_infos.push(PcoChunkInfo { pages: page_infos }) } - let mask = parray.validity()?.execute_mask(parray.len(), ctx)?; let mut ends = Vec::with_capacity(page_buffers.len()); let mut page_counts = chunk_infos.iter().flat_map(|chunk| &chunk.pages); let mut remaining = page_counts.next().map(|page| page.n_values as usize); diff --git a/encodings/pco/src/lib.rs b/encodings/pco/src/lib.rs index dc119a1c053..832e580903e 100644 --- a/encodings/pco/src/lib.rs +++ b/encodings/pco/src/lib.rs @@ -12,11 +12,10 @@ //! To deserialize arrays manually, register the encoding in the array session: //! //! ```rust -//! use vortex_array::session::ArraySessionExt; //! use vortex_session::VortexSession; //! //! let session = VortexSession::empty(); -//! session.arrays().register(vortex_pco::Pco); +//! vortex_pco::initialize(&session); //! ``` mod access; @@ -95,5 +94,7 @@ pub fn initialize(session: &vortex_session::VortexSession) { .register_execute_parent_kernel(Dict.id(), Pco, TakeExecuteAdaptor(Pco)); } +#[cfg(test)] +mod select_tests; #[cfg(test)] mod tests; diff --git a/encodings/pco/src/select_tests.rs b/encodings/pco/src/select_tests.rs new file mode 100644 index 00000000000..8d85064a772 --- /dev/null +++ b/encodings/pco/src/select_tests.rs @@ -0,0 +1,292 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::LazyLock; + +use rstest::rstest; +use vortex_array::ArrayContext; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::PiecewiseSequenceArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; +use vortex_array::serde::SerializeOptions; +use vortex_array::session::ArraySessionExt; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +use crate::Pco; +use crate::PcoArray; +use crate::PcoArrayExt; +use crate::PcoData; +use crate::PcoLogicalPages; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = VortexSession::empty(); + crate::initialize(&session); + session +}); + +fn legacy(array: &PcoArray, unknown: bool) -> VortexResult { + let mut metadata = array.metadata.as_ref().clone(); + metadata.logical_pages = unknown.then_some(PcoLogicalPages { + version: 999, + ends: vec![u64::MAX], + }); + // SAFETY: The unchanged compressed buffers retain the historical metadata contract. The + // optional unknown version is deliberately ignored and Pco::try_new validates the result. + let data = unsafe { + PcoData::new_unchecked( + array.chunk_metas.to_vec(), + array.pages.to_vec(), + array.dtype().as_ptype(), + metadata, + array.len(), + ) + }; + Pco::try_new(array.dtype().clone(), data, array.unsliced_validity()) +} + +#[rstest] +fn reordered_nullable_take(#[values(0, 1, 2)] format: usize) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = + PrimitiveArray::from_option_iter((0..4096i32).map(|i| (i % 7 != 0).then_some(i * 19))); + let encoded = Pco::from_primitive(input.as_view(), 3, 31, &mut ctx)?; + let encoded = match format { + 0 => encoded, + 1 => legacy(&encoded, false)?, + _ => legacy(&encoded, true)?, + } + .into_array(); + let source = encoded.slice(17..4001)?.slice(13..3900)?; + let input = input.slice(30..3917)?; + let indices = PrimitiveArray::from_option_iter([ + Some(3000u32), + Some(0), + None, + Some(127), + Some(1), + Some(3000), + Some(128), + None, + Some(1000), + ]) + .into_array(); + assert_arrays_eq!( + source.take(indices.clone())?, + input.take(indices)?, + &mut ctx + ); + Ok(()) +} + +#[rstest] +#[case(vec![0u64, 40, 0, 70], vec![15u64, 20, 3, 0], vec![1u64, 1, 0, 1])] +#[case(vec![100u64, 0, 20], vec![20u64, 20, 20], vec![1u64, 1, 2])] +#[case(vec![99u64, 0], vec![0u64, 0], vec![u64::MAX, 0])] +fn piecewise_ranges( + #[case] starts: Vec, + #[case] lengths: Vec, + #[case] strides: Vec, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = + PrimitiveArray::from_option_iter((0..256i64).map(|i| (i % 5 != 0).then_some(i * i))); + let encoded = Pco::from_primitive(input.as_view(), 3, 7, &mut ctx)?.into_array(); + let mut expected_indices = Vec::new(); + for ((&start, &length), &stride) in starts.iter().zip(&lengths).zip(&strides) { + expected_indices.extend((0..length).map(|i| start + i * stride)); + } + let indices = PiecewiseSequenceArray::try_new( + PrimitiveArray::from_iter(starts).into_array(), + PrimitiveArray::from_iter(lengths).into_array(), + PrimitiveArray::from_iter(strides).into_array(), + expected_indices.len(), + )? + .into_array(); + assert_arrays_eq!( + encoded.take(indices)?, + input.take(PrimitiveArray::from_iter(expected_indices).into_array())?, + &mut ctx + ); + Ok(()) +} + +#[test] +fn all_null_and_long_null_runs() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + for input in [ + PrimitiveArray::new(vec![0i32; 1024], Validity::AllInvalid), + PrimitiveArray::from_option_iter( + (0..100_000).map(|i| [1, 90_000, 99_990].contains(&i).then_some(i)), + ), + ] { + let encoded = Pco::from_primitive(input.as_view(), 3, 1, &mut ctx)?.into_array(); + let indices = PrimitiveArray::from_iter([0u32, 1, 100, 1000, 2, 1]).into_array(); + assert_arrays_eq!( + encoded.take(indices.clone())?, + input.take(indices)?, + &mut ctx + ); + assert_arrays_eq!(encoded.slice(1..1000)?, input.slice(1..1000)?, &mut ctx); + } + Ok(()) +} + +#[test] +fn rejects_invalid_indices_and_descriptors() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = PrimitiveArray::from_iter(0..64i32); + let encoded = Pco::from_primitive(input.as_view(), 3, 4, &mut ctx)?.into_array(); + for indices in [ + PrimitiveArray::from_iter([-1i64]).into_array(), + PrimitiveArray::from_iter([64i64]).into_array(), + ] { + assert!( + encoded + .take(indices)? + .execute::(&mut ctx) + .is_err() + ); + } + for (start, length, stride, declared) in [ + (u64::MAX, 2u64, 1u64, 2usize), + (1, 2, u64::MAX, 2), + (1, 2, 1, 3), + ] { + let indices = PiecewiseSequenceArray::try_new( + PrimitiveArray::from_iter([start]).into_array(), + PrimitiveArray::from_iter([length]).into_array(), + PrimitiveArray::from_iter([stride]).into_array(), + declared, + )? + .into_array(); + assert!( + encoded + .take(indices)? + .execute::(&mut ctx) + .is_err() + ); + } + Ok(()) +} + +#[test] +fn raw_sliced_serde_requires_explicit_normalization() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = PrimitiveArray::from_option_iter((0..128i32).map(|i| (i % 3 != 0).then_some(i))); + let encoded = Pco::from_primitive(input.as_view(), 3, 7, &mut ctx)?.into_array(); + let sliced = encoded.slice(11..77)?; + assert!( + sliced + .serialize( + &ArrayContext::empty(), + &SESSION, + &SerializeOptions::default() + ) + .is_err() + ); + let normalized = sliced.execute::(&mut ctx)?; + let recompressed = Pco::from_primitive(normalized.as_view(), 3, 7, &mut ctx)?.into_array(); + assert!( + !recompressed + .serialize( + &ArrayContext::empty(), + &SESSION, + &SerializeOptions::default() + )? + .is_empty() + ); + assert_arrays_eq!(recompressed, input.slice(11..77)?, &mut ctx); + Ok(()) +} + +#[test] +fn registry_only_session_retains_legacy_take_fallback() -> VortexResult<()> { + let session = VortexSession::empty(); + session.arrays().register(Pco); + let mut ctx = session.create_execution_ctx(); + let input = PrimitiveArray::from_option_iter([Some(5i32), None, Some(9), Some(2)]); + let encoded = Pco::from_primitive(input.as_view(), 3, 1, &mut ctx)?.into_array(); + let indices: ArrayRef = PrimitiveArray::from_iter([3u32, 1, 0, 3]).into_array(); + assert_arrays_eq!( + encoded.take(indices.clone())?, + input.take(indices)?, + &mut ctx + ); + Ok(()) +} + +#[test] +fn checked_metadata_and_touched_page_validity() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = PrimitiveArray::from_option_iter((0..128i32).map(|i| (i % 2 == 0).then_some(i))); + let encoded = Pco::from_primitive(input.as_view(), 3, 2, &mut ctx)?; + for ends in [ + vec![128u64], + vec![0; encoded.pages.len()], + vec![u64::MAX; encoded.pages.len()], + ] { + let mut metadata = encoded.metadata.as_ref().clone(); + metadata.logical_pages = Some(PcoLogicalPages { version: 1, ends }); + // SAFETY: Constructor is followed immediately by fallible validation; malformed metadata + // is never published as a valid array. + let data = unsafe { + PcoData::new_unchecked( + encoded.chunk_metas.to_vec(), + encoded.pages.to_vec(), + encoded.dtype().as_ptype(), + metadata, + encoded.len(), + ) + }; + assert!(Pco::try_new(encoded.dtype().clone(), data, encoded.unsliced_validity()).is_err()); + } + let mut metadata = encoded.metadata.as_ref().clone(); + // Structurally valid but the first logical page contains three valid rows, not two. + metadata + .logical_pages + .as_mut() + .expect("writer emits logical ends") + .ends[0] = 5; + let data = unsafe { + PcoData::new_unchecked( + encoded.chunk_metas.to_vec(), + encoded.pages.to_vec(), + encoded.dtype().as_ptype(), + metadata, + encoded.len(), + ) + }; + let malformed = + Pco::try_new(encoded.dtype().clone(), data, encoded.unsliced_validity())?.into_array(); + assert!( + malformed + .take(PrimitiveArray::from_iter([0u32]).into_array())? + .execute::(&mut ctx) + .is_err() + ); + Ok(()) +} + +#[test] +fn selections_cross_compression_chunks() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = + PrimitiveArray::from_option_iter((0..8192i64).map(|i| (i % 11 != 0).then_some(i * 37))); + let data = + PcoData::from_primitive_with_values_per_chunk(input.as_view(), 3, 257, 31, &mut ctx)?; + let encoded = Pco::try_new(input.dtype().clone(), data, input.validity()?)?.into_array(); + let indices = + PrimitiveArray::from_iter([0u32, 900, 260, 3000, 259, 8191, 4095, 900]).into_array(); + assert_arrays_eq!( + encoded.take(indices.clone())?, + input.take(indices)?, + &mut ctx + ); + assert_arrays_eq!(encoded.slice(233..7901)?, input.slice(233..7901)?, &mut ctx); + Ok(()) +} diff --git a/encodings/pco/src/take.rs b/encodings/pco/src/take.rs index 4fba4b79693..287518c1e09 100644 --- a/encodings/pco/src/take.rs +++ b/encodings/pco/src/take.rs @@ -24,6 +24,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; +use vortex_mask::Mask; use crate::Pco; use crate::access::AccessState; @@ -63,7 +64,7 @@ impl TakeExecute for Pco { runs } else { let indices = indices.clone().execute::(ctx)?; - let mut runs = Vec::with_capacity(indices.len()); + let mut runs: Vec = Vec::with_capacity(indices.len()); let mask = indices.validity()?.execute_mask(indices.len(), ctx)?; match_each_integer_ptype!(indices.ptype(), |I| { for (&value, valid) in indices.as_slice::().iter().zip(mask.iter()) { @@ -74,6 +75,20 @@ impl TakeExecute for Pco { } else { None }; + if let Some(previous) = runs.last_mut() { + let contiguous = match (previous.start, start) { + (Some(first), Some(next)) => { + previous.stride == 1 + && first.checked_add(previous.length) == Some(next) + } + (None, None) => true, + _ => false, + }; + if contiguous { + previous.length += 1; + continue; + } + } runs.push(Run { start, length: 1, @@ -146,6 +161,44 @@ fn gather( valid.append_n(false, run.length); continue; }; + if run.stride == 0 && run.length != 0 { + if let Some((page, value)) = state.resolve(array, array.slice_start() + start, ctx)? { + let value = state.values(array, page, ctx)?.as_slice::()[value]; + values.extend(std::iter::repeat_n(value, run.length)); + valid.append_n(true, run.length); + } else { + valid.append_n(false, run.length); + } + continue; + } + if run.stride == 1 && run.length != 0 { + let mut row = array.slice_start() + start; + let stop = row + run.length; + while row < stop { + if let Some(span) = state.span(array, row, stop, ctx)? { + if !span.values.is_empty() { + values.extend_from_slice( + &state.values(array, span.page, ctx)?.as_slice::()[span.values], + ); + } + match &span.mask { + Mask::AllTrue(n) => valid.append_n(true, *n), + Mask::AllFalse(n) => valid.append_n(false, *n), + Mask::Values(mask) => valid.append_buffer(mask.bit_buffer()), + } + row = span.stop; + } else { + if let Some((page, value)) = state.resolve(array, row, ctx)? { + values.push(state.values(array, page, ctx)?.as_slice::()[value]); + valid.append_true(); + } else { + valid.append_false(); + } + row += 1; + } + } + continue; + } for offset in 0..run.length { // The complete descriptor, including its output length and last element, was // checked before allocating output. The source slice itself is already validated. diff --git a/vortex-file/tests/pco_logical_pages_compat.rs b/vortex-file/tests/pco_logical_pages_compat.rs new file mode 100644 index 00000000000..08833fd2af6 --- /dev/null +++ b/vortex-file/tests/pco_logical_pages_compat.rs @@ -0,0 +1,79 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Public file harness also run, unchanged, against the pre-change reader by PCO compatibility CI. + +#![expect(clippy::tests_outside_test_module)] + +use std::sync::Arc; + +use futures::TryStreamExt; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::assert_arrays_eq; +use vortex_buffer::ByteBuffer; +use vortex_error::VortexResult; +use vortex_file::OpenOptionsSessionExt; +use vortex_file::WriteOptionsSessionExt; +use vortex_io::session::RuntimeSession; +use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; +use vortex_layout::session::LayoutSession; +use vortex_pco::Pco; +use vortex_session::VortexSession; + +mod common; + +#[tokio::test] +async fn public_pco_compatibility() -> VortexResult<()> { + let session = VortexSession::empty() + .with::() + .with::(); + vortex_file::register_default_encodings(&session); + common::enable_all_registered_array_encodings(&session); + let mut ctx = session.create_execution_ctx(); + let original = PrimitiveArray::from_option_iter( + (0..131_072i64).map(|i| ((i < 1024 || i > 120_000) && i % 7 != 0).then_some(i * 19)), + ); + let mode = std::env::var("PCO_COMPAT_MODE").unwrap_or_else(|_| "roundtrip".to_owned()); + let directory = std::env::var_os("PCO_COMPAT_DIR").map(std::path::PathBuf::from); + for (name, range) in [("unsliced", 0..131_072), ("normalized_slice", 777..125_000)] { + let encoded = Pco::from_primitive(original.as_view(), 3, 127, &mut ctx)?.into_array(); + // Explicit writer-owned normalization, never a hidden serde fallback: execute selected + // pages and recompress only the chosen logical rows using the caller's configuration. + let selected = encoded + .slice(range.clone())? + .execute::(&mut ctx)?; + let normalized = Pco::from_primitive(selected.as_view(), 3, 127, &mut ctx)?.into_array(); + let expected = + StructArray::from_fields(&[("values", original.slice(range)?)])?.into_array(); + let data = StructArray::from_fields(&[("values", normalized)])?.into_array(); + let path = directory + .as_ref() + .map(|directory| directory.join(format!("{name}.vortex"))); + let bytes = if mode == "read" { + std::fs::read(path.as_ref().expect("read mode requires PCO_COMPAT_DIR"))? + } else { + let mut bytes = Vec::new(); + session + .write_options() + .with_strategy(Arc::new(FlatLayoutStrategy::default())) + .write(&mut bytes, data.to_array_stream()) + .await?; + if let Some(path) = &path { + std::fs::create_dir_all(path.parent().expect("fixture path has a parent"))?; + std::fs::write(path, &bytes)?; + } + bytes + }; + let file = session + .open_options() + .open_buffer(ByteBuffer::from(bytes))?; + let chunks = file.scan()?.into_stream()?.try_collect::>().await?; + let actual = ChunkedArray::try_new(chunks, expected.dtype().clone())?.into_array(); + assert_arrays_eq!(actual, expected, &mut ctx); + } + Ok(()) +} From 5506f66f11c5215ce73de9bb569c77c1fdd9fed9 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 12:35:48 -0400 Subject: [PATCH 03/36] Build PCO page ends with dense prefixes and bitmap select Signed-off-by: Daniel King --- encodings/pco/src/array.rs | 46 +++++++++++++++++++++++++++----------- 1 file changed, 33 insertions(+), 13 deletions(-) diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array.rs index a8b13952359..9091a629cec 100644 --- a/encodings/pco/src/array.rs +++ b/encodings/pco/src/array.rs @@ -600,21 +600,41 @@ impl PcoData { } let mut ends = Vec::with_capacity(page_buffers.len()); - let mut page_counts = chunk_infos.iter().flat_map(|chunk| &chunk.pages); - let mut remaining = page_counts.next().map(|page| page.n_values as usize); - // Compression already visits all values; build logical boundaries in the same linear - // scale, without allocating one index per valid row. - for (row, valid) in mask.iter().enumerate() { - if valid && let Some(n) = remaining.as_mut() { - *n -= 1; - if *n == 0 { - ends.push(u64::try_from(row + 1)?); - remaining = page_counts.next().map(|page| page.n_values as usize); + match &mask { + vortex_mask::Mask::AllFalse(_) => {} + vortex_mask::Mask::AllTrue(_) => { + let mut end = 0usize; + for page in chunk_infos.iter().flat_map(|chunk| &chunk.pages) { + end = end + .checked_add(page.n_values as usize) + .ok_or_else(|| vortex_err!("Pco logical page prefix overflows usize"))?; + ends.push(u64::try_from(end)?); + } + } + vortex_mask::Mask::Values(mask) => { + let bits = mask.bit_buffer(); + let mut start = 0usize; + for page in chunk_infos.iter().flat_map(|chunk| &chunk.pages) { + if ends.len() + 1 == page_buffers.len() { + // No search through trailing nulls is necessary for the final page. + ends.push(u64::try_from(parray.len())?); + break; + } + // Existing bitmap select scans words only up to the next page boundary. + // Advancing the suffix avoids rescanning prefixes or visiting every null row. + let offset = bits + .slice(start..) + .select(page.n_values as usize - 1) + .ok_or_else(|| { + vortex_err!("Pco compression validity has too few valid rows") + })?; + start = start + .checked_add(offset) + .and_then(|end| end.checked_add(1)) + .ok_or_else(|| vortex_err!("Pco logical page end overflows usize"))?; + ends.push(u64::try_from(start)?); } } - } - if let Some(last) = ends.last_mut() { - *last = u64::try_from(parray.len())?; } let metadata = PcoMetadata { header, From a7823557e9c2bc18fed06124fc952e2d296b6393 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 12:37:00 -0400 Subject: [PATCH 04/36] Use public piecewise slots and reserve only initial take runs Signed-off-by: Daniel King --- encodings/pco/src/take.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/encodings/pco/src/take.rs b/encodings/pco/src/take.rs index 287518c1e09..815f570da2d 100644 --- a/encodings/pco/src/take.rs +++ b/encodings/pco/src/take.rs @@ -13,7 +13,7 @@ use vortex_array::IntoArray; use vortex_array::arrays::PiecewiseSequence; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::dict::TakeExecute; -use vortex_array::arrays::piecewise_sequence::PiecewiseSequenceArraySlotsExt; +use vortex_array::arrays::piecewise_sequence::array::PiecewiseSequenceArraySlotsExt; use vortex_array::dtype::NativePType; use vortex_array::dtype::half; use vortex_array::match_each_integer_ptype; @@ -64,7 +64,7 @@ impl TakeExecute for Pco { runs } else { let indices = indices.clone().execute::(ctx)?; - let mut runs: Vec = Vec::with_capacity(indices.len()); + let mut runs: Vec = Vec::with_capacity(indices.len().min(64)); let mask = indices.validity()?.execute_mask(indices.len(), ctx)?; match_each_integer_ptype!(indices.ptype(), |I| { for (&value, valid) in indices.as_slice::().iter().zip(mask.iter()) { From a0c6b117a20bdfd38717a4cfa17e1c4645014c2a Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 12:38:40 -0400 Subject: [PATCH 05/36] Amortize PCO decoding and preserve bulk selection for legacy files Signed-off-by: Daniel King --- encodings/pco/src/access.rs | 73 ++++++++- encodings/pco/src/array.rs | 9 +- encodings/pco/src/decoder.rs | 63 ++++++++ encodings/pco/src/lib.rs | 1 + encodings/pco/src/probe.rs | 30 ---- encodings/pco/src/select_tests.rs | 151 ++++++++++-------- vortex-file/tests/pco_logical_pages_compat.rs | 8 +- 7 files changed, 229 insertions(+), 106 deletions(-) create mode 100644 encodings/pco/src/decoder.rs diff --git a/encodings/pco/src/access.rs b/encodings/pco/src/access.rs index 58f05d6d729..43aed590224 100644 --- a/encodings/pco/src/access.rs +++ b/encodings/pco/src/access.rs @@ -3,6 +3,7 @@ //! Operation-owned selected-page values and validity. No decoded state lives in array metadata. +use std::any::Any; use std::collections::BTreeMap; use pco::data_types::NumberType; @@ -26,6 +27,7 @@ const RANK_STRIDE: usize = 512; #[derive(Default)] pub(crate) struct AccessState { legacy: Option, + decoder: Option>, pages: BTreeMap, } @@ -167,7 +169,69 @@ impl AccessState { ) -> VortexResult> { let index = array.page_index()?; if !index.logical || index.pages.is_empty() { - return Ok(None); + if self.legacy.is_none() { + let mask = array + .unsliced_validity() + .execute_mask(array.unsliced_n_rows(), ctx)?; + let total = index.pages.last().map_or(0, |page| page.values.end); + vortex_ensure_eq!( + mask.true_count(), + total, + "Pco validity count does not match stored values" + ); + self.legacy = Some(RankMask::new(mask)); + } + let mask = self + .legacy + .as_ref() + .ok_or_else(|| vortex_err!("Pco legacy validity is missing"))?; + let start_value = mask.rank(row); + let stop_value = mask.rank(stop); + if start_value == stop_value { + return Ok(Some(PageSpan { + page: 0, + stop, + values: 0..0, + mask: mask.mask.slice(row..stop), + })); + } + let page_idx = index + .pages + .partition_point(|page| page.values.end <= start_value); + let page = index + .pages + .get(page_idx) + .ok_or_else(|| vortex_err!("Pco page does not cover compact value"))?; + let value_stop = stop_value.min(page.values.end); + let end = if value_stop == stop_value { + stop + } else { + let offset = match &mask.mask { + Mask::AllTrue(_) => value_stop - start_value - 1, + Mask::Values(values) => values + .bit_buffer() + .slice(row..stop) + .select(value_stop - start_value - 1) + .ok_or_else(|| vortex_err!("Pco legacy page rank is inconsistent"))?, + Mask::AllFalse(_) => { + return Err(vortex_err!("Pco all-null validity has compact values")); + } + }; + row.checked_add(offset) + .and_then(|end| end.checked_add(1)) + .ok_or_else(|| vortex_err!("Pco legacy logical page end overflows usize"))? + }; + let selected_mask = mask.mask.slice(row..end); + self.pages.entry(page_idx).or_insert(CachedPage { + mask: None, + decoded: None, + }); + return Ok(Some(PageSpan { + page: page_idx, + stop: end, + values: (start_value - page.values.start)..(value_stop - page.values.start), + mask: selected_mask, + })); } let page_idx = index .pages @@ -221,7 +285,12 @@ impl AccessState { let decoded = match_number_enum!( number_type_from_ptype(array.dtype().as_ptype()), NumberType => { - crate::probe::decode_page::(array, page.chunk, page.values.len(), buffer.as_slice(), ctx)? + if self.decoder.is_none() { + self.decoder = Some(Box::new(crate::decoder::PageDecoder::::new(array.metadata.header.as_slice())?)); + } + let decoder = self.decoder.as_mut().and_then(|decoder| decoder.downcast_mut::>()) + .ok_or_else(|| vortex_err!("Pco operation decoder primitive type changed"))?; + decoder.read(page.chunk, array.chunk_metas[page.chunk].as_slice(), page.values.len(), buffer.as_slice(), ctx)? } ); cached.decoded = Some(decoded); diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array.rs index 9091a629cec..616c28684a6 100644 --- a/encodings/pco/src/array.rs +++ b/encodings/pco/src/array.rs @@ -51,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; @@ -807,9 +808,12 @@ impl PcoData { let index = PageIndex::try_new(&self.metadata, self.unsliced_n_rows)?; // Concurrent first readers may prepare equivalent immutable descriptors. Only one // index is retained, and no decoded values or validity are cached in the array. - let _ = self.page_index.set(index); + drop(self.page_index.set(index)); } - Ok(self.page_index.get().expect("page index was initialized")) + Ok(self + .page_index + .get() + .vortex_expect("page index was initialized")) } pub(crate) fn _slice(&self, start: usize, stop: usize) -> Self { @@ -882,6 +886,7 @@ mod tests { use vortex_array::assert_arrays_eq; use vortex_array::validity::Validity; use vortex_buffer::buffer; + use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; diff --git a/encodings/pco/src/decoder.rs b/encodings/pco/src/decoder.rs new file mode 100644 index 00000000000..419a5379793 --- /dev/null +++ b/encodings/pco/src/decoder.rs @@ -0,0 +1,63 @@ +// 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 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::VortexExpect; +use vortex_error::VortexResult; + +use crate::array::vortex_err_from_pco; + +pub(crate) struct PageDecoder { + file: FileDecompressor, + chunks: BTreeMap>, +} + +impl PageDecoder { + pub fn new(header: &[u8]) -> VortexResult { + 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 { + if !self.chunks.contains_key(&chunk) { + let (decoder, _) = self + .file + .chunk_decompressor::(chunk_meta) + .map_err(vortex_err_from_pco)?; + self.chunks.insert(chunk, decoder); + } + let chunk = self + .chunks + .get_mut(&chunk) + .vortex_expect("chunk decoder was inserted"); + let mut decoder = chunk + .page_decompressor(buffer, n_values) + .map_err(vortex_err_from_pco)?; + // `read` takes initialized typed storage. Zero initialization also makes a failed or + // partial corrupt-page decode safe to drop; no uninitialized typed slice is published. + let mut values = BufferMut::::zeroed_in(n_values, ctx.allocator().clone()); + decoder.read(&mut values).map_err(vortex_err_from_pco)?; + Ok(PrimitiveArray::new(values.freeze(), Validity::NonNullable)) + } +} diff --git a/encodings/pco/src/lib.rs b/encodings/pco/src/lib.rs index 832e580903e..9fb4478ee48 100644 --- a/encodings/pco/src/lib.rs +++ b/encodings/pco/src/lib.rs @@ -21,6 +21,7 @@ mod access; mod array; mod compute; +mod decoder; mod page_index; mod probe; mod rules; diff --git a/encodings/pco/src/probe.rs b/encodings/pco/src/probe.rs index 936c55dba5c..a63653a686e 100644 --- a/encodings/pco/src/probe.rs +++ b/encodings/pco/src/probe.rs @@ -3,21 +3,14 @@ //! PCO probes retain every decoded page and index the compacted, non-null value positions. -use pco::data_types::Number; -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::match_each_native_ptype; use vortex_array::scalar::Scalar; -use vortex_array::validity::Validity; -use vortex_buffer::BufferMut; use vortex_error::VortexResult; use crate::Pco; use crate::access::AccessState; -use crate::array::vortex_err_from_pco; /// State retained by repeated PCO probes. Construction performs no allocation. /// @@ -46,28 +39,5 @@ pub(crate) fn scalar_at( })) } -pub(crate) fn decode_page( - array: ArrayView<'_, Pco>, - chunk: usize, - n_values: usize, - buffer: &[u8], - ctx: &mut ExecutionCtx, -) -> VortexResult { - let (file, _) = - FileDecompressor::new(array.metadata.header.as_slice()).map_err(vortex_err_from_pco)?; - let (mut chunk, _) = file - .chunk_decompressor::(array.chunk_metas[chunk].as_ref()) - .map_err(vortex_err_from_pco)?; - let mut decoder = chunk - .page_decompressor(buffer, n_values) - .map_err(vortex_err_from_pco)?; - let mut values = BufferMut::::with_capacity_in(n_values, ctx.allocator().clone()); - // SAFETY: the buffer reserves `n_values` elements, and the page decompressor was built for - // exactly that count, so `read` writes every element before anything observes it. - unsafe { values.set_len(n_values) }; - decoder.read(&mut values).map_err(vortex_err_from_pco)?; - Ok(PrimitiveArray::new(values.freeze(), Validity::NonNullable)) -} - #[cfg(test)] mod tests; diff --git a/encodings/pco/src/select_tests.rs b/encodings/pco/src/select_tests.rs index 8d85064a772..adddbe5895a 100644 --- a/encodings/pco/src/select_tests.rs +++ b/encodings/pco/src/select_tests.rs @@ -3,24 +3,26 @@ use std::sync::LazyLock; +use prost::Message; use rstest::rstest; use vortex_array::ArrayContext; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayPlugin; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PiecewiseSequenceArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; +use vortex_array::buffer::BufferHandle; use vortex_array::serde::SerializeOptions; use vortex_array::session::ArraySessionExt; use vortex_array::validity::Validity; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; use crate::Pco; -use crate::PcoArray; -use crate::PcoArrayExt; -use crate::PcoData; use crate::PcoLogicalPages; static SESSION: LazyLock = LazyLock::new(|| { @@ -29,24 +31,44 @@ static SESSION: LazyLock = LazyLock::new(|| { session }); -fn legacy(array: &PcoArray, unknown: bool) -> VortexResult { - let mut metadata = array.metadata.as_ref().clone(); - metadata.logical_pages = unknown.then_some(PcoLogicalPages { - version: 999, - ends: vec![u64::MAX], - }); - // SAFETY: The unchanged compressed buffers retain the historical metadata contract. The - // optional unknown version is deliberately ignored and Pco::try_new validates the result. - let data = unsafe { - PcoData::new_unchecked( - array.chunk_metas.to_vec(), - array.pages.to_vec(), - array.dtype().as_ptype(), - metadata, +/// Build a wire fixture through the public serializer and deserialize its mutated protobuf +/// metadata through the public plugin boundary. No private in-memory encoding fields are used. +fn with_wire_metadata( + array: &ArrayRef, + mutate: impl FnOnce(&mut crate::PcoMetadata), +) -> VortexResult { + let mut wire = SESSION + .array_serialize(array)? + .vortex_expect("PCO is serializable"); + let mut metadata = crate::PcoMetadata::decode(wire.metadata.as_slice())?; + mutate(&mut metadata); + wire.metadata = metadata.encode_to_vec(); + let buffers = wire + .buffers + .into_iter() + .map(BufferHandle::new_host) + .collect::>(); + ArrayPlugin::deserialize( + &Pco, + ArrayDeserialization::new( + wire.serialized_id, + array.dtype(), array.len(), - ) - }; - Pco::try_new(array.dtype().clone(), data, array.unsliced_validity()) + &wire.metadata, + &buffers, + &wire.children, + ), + &SESSION, + ) +} + +fn legacy(array: &ArrayRef, unknown: bool) -> VortexResult { + with_wire_metadata(array, |metadata| { + metadata.logical_pages = unknown.then_some(PcoLogicalPages { + version: 999, + ends: vec![u64::MAX], + }); + }) } #[rstest] @@ -54,13 +76,12 @@ fn reordered_nullable_take(#[values(0, 1, 2)] format: usize) -> VortexResult<()> let mut ctx = SESSION.create_execution_ctx(); let input = PrimitiveArray::from_option_iter((0..4096i32).map(|i| (i % 7 != 0).then_some(i * 19))); - let encoded = Pco::from_primitive(input.as_view(), 3, 31, &mut ctx)?; + let encoded = Pco::from_primitive(input.as_view(), 3, 31, &mut ctx)?.into_array(); let encoded = match format { 0 => encoded, 1 => legacy(&encoded, false)?, _ => legacy(&encoded, true)?, - } - .into_array(); + }; let source = encoded.slice(17..4001)?.slice(13..3900)?; let input = input.slice(30..3917)?; let indices = PrimitiveArray::from_option_iter([ @@ -224,45 +245,23 @@ fn registry_only_session_retains_legacy_take_fallback() -> VortexResult<()> { fn checked_metadata_and_touched_page_validity() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let input = PrimitiveArray::from_option_iter((0..128i32).map(|i| (i % 2 == 0).then_some(i))); - let encoded = Pco::from_primitive(input.as_view(), 3, 2, &mut ctx)?; - for ends in [ - vec![128u64], - vec![0; encoded.pages.len()], - vec![u64::MAX; encoded.pages.len()], - ] { - let mut metadata = encoded.metadata.as_ref().clone(); - metadata.logical_pages = Some(PcoLogicalPages { version: 1, ends }); - // SAFETY: Constructor is followed immediately by fallible validation; malformed metadata - // is never published as a valid array. - let data = unsafe { - PcoData::new_unchecked( - encoded.chunk_metas.to_vec(), - encoded.pages.to_vec(), - encoded.dtype().as_ptype(), - metadata, - encoded.len(), - ) - }; - assert!(Pco::try_new(encoded.dtype().clone(), data, encoded.unsliced_validity()).is_err()); + let encoded = Pco::from_primitive(input.as_view(), 3, 2, &mut ctx)?.into_array(); + for ends in [vec![128u64], vec![0; 32], vec![u64::MAX; 32]] { + assert!( + with_wire_metadata(&encoded, |metadata| { + metadata.logical_pages = Some(PcoLogicalPages { version: 1, ends }); + }) + .is_err() + ); } - let mut metadata = encoded.metadata.as_ref().clone(); - // Structurally valid but the first logical page contains three valid rows, not two. - metadata - .logical_pages - .as_mut() - .expect("writer emits logical ends") - .ends[0] = 5; - let data = unsafe { - PcoData::new_unchecked( - encoded.chunk_metas.to_vec(), - encoded.pages.to_vec(), - encoded.dtype().as_ptype(), - metadata, - encoded.len(), - ) - }; - let malformed = - Pco::try_new(encoded.dtype().clone(), data, encoded.unsliced_validity())?.into_array(); + let malformed = with_wire_metadata(&encoded, |metadata| { + // Structurally valid but the first logical page has three valid rows, not two. + metadata + .logical_pages + .as_mut() + .vortex_expect("writer emits logical ends") + .ends[0] = 5; + })?; assert!( malformed .take(PrimitiveArray::from_iter([0u32]).into_array())? @@ -273,20 +272,32 @@ fn checked_metadata_and_touched_page_validity() -> VortexResult<()> { } #[test] -fn selections_cross_compression_chunks() -> VortexResult<()> { +fn large_public_compressor_selections() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); - let input = - PrimitiveArray::from_option_iter((0..8192i64).map(|i| (i % 11 != 0).then_some(i * 37))); - let data = - PcoData::from_primitive_with_values_per_chunk(input.as_view(), 3, 257, 31, &mut ctx)?; - let encoded = Pco::try_new(input.dtype().clone(), data, input.validity()?)?.into_array(); - let indices = - PrimitiveArray::from_iter([0u32, 900, 260, 3000, 259, 8191, 4095, 900]).into_array(); + let len = pco::DEFAULT_MAX_PAGE_N + 257; + let input = PrimitiveArray::from_iter( + (0..len).map(|i| u32::try_from(i).vortex_expect("fixture fits u32")), + ); + let encoded = Pco::from_primitive(input.as_view(), 0, 4096, &mut ctx)?.into_array(); + let indices = PrimitiveArray::from_iter([ + 0u32, + 900, + 260, + 3000, + u32::try_from(len - 1).vortex_expect("fixture fits u32"), + 4095, + 900, + ]) + .into_array(); assert_arrays_eq!( encoded.take(indices.clone())?, input.take(indices)?, &mut ctx ); - assert_arrays_eq!(encoded.slice(233..7901)?, input.slice(233..7901)?, &mut ctx); + assert_arrays_eq!( + encoded.slice(len - 1000..len - 1)?, + input.slice(len - 1000..len - 1)?, + &mut ctx + ); Ok(()) } diff --git a/vortex-file/tests/pco_logical_pages_compat.rs b/vortex-file/tests/pco_logical_pages_compat.rs index 08833fd2af6..59693ffbdf8 100644 --- a/vortex-file/tests/pco_logical_pages_compat.rs +++ b/vortex-file/tests/pco_logical_pages_compat.rs @@ -15,6 +15,7 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::assert_arrays_eq; use vortex_buffer::ByteBuffer; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_file::OpenOptionsSessionExt; use vortex_file::WriteOptionsSessionExt; @@ -54,7 +55,10 @@ async fn public_pco_compatibility() -> VortexResult<()> { .as_ref() .map(|directory| directory.join(format!("{name}.vortex"))); let bytes = if mode == "read" { - std::fs::read(path.as_ref().expect("read mode requires PCO_COMPAT_DIR"))? + std::fs::read( + path.as_ref() + .vortex_expect("read mode requires PCO_COMPAT_DIR"), + )? } else { let mut bytes = Vec::new(); session @@ -63,7 +67,7 @@ async fn public_pco_compatibility() -> VortexResult<()> { .write(&mut bytes, data.to_array_stream()) .await?; if let Some(path) = &path { - std::fs::create_dir_all(path.parent().expect("fixture path has a parent"))?; + std::fs::create_dir_all(path.parent().vortex_expect("fixture path has a parent"))?; std::fs::write(path, &bytes)?; } bytes From bfa0932136f3b4675fc5b2f0ae5c076e0553158e Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 12:40:05 -0400 Subject: [PATCH 06/36] Build nullable selections directly in their final output buffer Signed-off-by: Daniel King --- encodings/pco/src/array.rs | 5 ++-- encodings/pco/src/take.rs | 47 ++++++++++++++++++++++++++------------ 2 files changed, 35 insertions(+), 17 deletions(-) diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array.rs index 616c28684a6..3065530e8d5 100644 --- a/encodings/pco/src/array.rs +++ b/encodings/pco/src/array.rs @@ -878,6 +878,7 @@ impl OperationsVTable for Pco { #[cfg(test)] mod tests { + use std::sync::Arc; use std::sync::LazyLock; use vortex_array::IntoArray; @@ -886,7 +887,6 @@ mod tests { use vortex_array::assert_arrays_eq; use vortex_array::validity::Validity; use vortex_buffer::buffer; - use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; @@ -948,7 +948,8 @@ mod tests { let values = PrimitiveArray::from_iter([42u32]); let pco = Pco::from_primitive(values.as_view(), 0, 128, &mut ctx)?; let mut data = pco.data().clone(); - data.metadata.chunks[0].pages[0].n_values = u32::try_from(VALUES_PER_CHUNK + 1)?; + Arc::make_mut(&mut data.metadata).chunks[0].pages[0].n_values = + u32::try_from(VALUES_PER_CHUNK + 1)?; assert!( data.validate(pco.dtype(), pco.len(), &Validity::NonNullable) diff --git a/encodings/pco/src/take.rs b/encodings/pco/src/take.rs index 815f570da2d..dc8b88211ad 100644 --- a/encodings/pco/src/take.rs +++ b/encodings/pco/src/take.rs @@ -159,6 +159,7 @@ fn gather( for run in runs { let Some(start) = run.start else { valid.append_n(false, run.length); + values.extend(std::iter::repeat_n(T::default(), run.length)); continue; }; if run.stride == 0 && run.length != 0 { @@ -168,6 +169,7 @@ fn gather( valid.append_n(true, run.length); } else { valid.append_n(false, run.length); + values.extend(std::iter::repeat_n(T::default(), run.length)); } continue; } @@ -176,15 +178,34 @@ fn gather( let stop = row + run.length; while row < stop { if let Some(span) = state.span(array, row, stop, ctx)? { - if !span.values.is_empty() { - values.extend_from_slice( - &state.values(array, span.page, ctx)?.as_slice::()[span.values], - ); - } match &span.mask { - Mask::AllTrue(n) => valid.append_n(true, *n), - Mask::AllFalse(n) => valid.append_n(false, *n), - Mask::Values(mask) => valid.append_buffer(mask.bit_buffer()), + Mask::AllTrue(n) => { + values.extend_from_slice( + &state.values(array, span.page, ctx)?.as_slice::()[span.values], + ); + valid.append_n(true, *n); + } + Mask::AllFalse(n) => { + values.extend(std::iter::repeat_n(T::default(), *n)); + valid.append_n(false, *n); + } + Mask::Values(mask) => { + let source = state.values(array, span.page, ctx)?.as_slice::(); + let source = &source[span.values]; + let mut logical = 0usize; + let mut compact = 0usize; + // Copy defined runs straight into the final output. Null gaps are + // initialized once; no compact staging output is expanded afterward. + for (start, stop) in mask.bit_buffer().set_slices() { + values.extend(std::iter::repeat_n(T::default(), start - logical)); + let count = stop - start; + values.extend_from_slice(&source[compact..compact + count]); + compact += count; + logical = stop; + } + values.extend(std::iter::repeat_n(T::default(), mask.len() - logical)); + valid.append_buffer(mask.bit_buffer()); + } } row = span.stop; } else { @@ -192,6 +213,7 @@ fn gather( values.push(state.values(array, page, ctx)?.as_slice::()[value]); valid.append_true(); } else { + values.push(T::default()); valid.append_false(); } row += 1; @@ -207,6 +229,7 @@ fn gather( values.push(state.values(array, page, ctx)?.as_slice::()[value]); valid.append_true(); } else { + values.push(T::default()); valid.append_false(); } } @@ -216,13 +239,7 @@ fn gather( } else { Validity::NonNullable }; - Ok(PrimitiveArray::from_values_byte_buffer( - values.freeze().into_byte_buffer(), - array.dtype().as_ptype(), - validity, - output_len, - ctx, - )) + Ok(PrimitiveArray::new(values.freeze(), validity)) } pub(crate) fn slice( From b30c862c9afed1a20fdd057921e509a986a2c592 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 12:41:04 -0400 Subject: [PATCH 07/36] Validate corrupt page counts through the public wire boundary Signed-off-by: Daniel King --- encodings/pco/src/array.rs | 18 ------------------ encodings/pco/src/select_tests.rs | 7 +++++++ 2 files changed, 7 insertions(+), 18 deletions(-) diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array.rs index 3065530e8d5..380a559f9b6 100644 --- a/encodings/pco/src/array.rs +++ b/encodings/pco/src/array.rs @@ -878,7 +878,6 @@ impl OperationsVTable for Pco { #[cfg(test)] mod tests { - use std::sync::Arc; use std::sync::LazyLock; use vortex_array::IntoArray; @@ -890,7 +889,6 @@ mod tests { use vortex_error::VortexResult; use vortex_session::VortexSession; - use super::VALUES_PER_CHUNK; use crate::Pco; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -941,20 +939,4 @@ mod tests { assert!(result.is_err()); Ok(()) } - - #[test] - fn test_validate_rejects_oversized_chunk() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let values = PrimitiveArray::from_iter([42u32]); - let pco = Pco::from_primitive(values.as_view(), 0, 128, &mut ctx)?; - let mut data = pco.data().clone(); - Arc::make_mut(&mut data.metadata).chunks[0].pages[0].n_values = - u32::try_from(VALUES_PER_CHUNK + 1)?; - - assert!( - data.validate(pco.dtype(), pco.len(), &Validity::NonNullable) - .is_err() - ); - Ok(()) - } } diff --git a/encodings/pco/src/select_tests.rs b/encodings/pco/src/select_tests.rs index adddbe5895a..1280e97908e 100644 --- a/encodings/pco/src/select_tests.rs +++ b/encodings/pco/src/select_tests.rs @@ -254,6 +254,13 @@ fn checked_metadata_and_touched_page_validity() -> VortexResult<()> { .is_err() ); } + assert!( + with_wire_metadata(&encoded, |metadata| { + metadata.chunks[0].pages[0].n_values = u32::try_from(pco::DEFAULT_MAX_PAGE_N + 1) + .vortex_expect("invalid fixture count fits u32"); + }) + .is_err() + ); let malformed = with_wire_metadata(&encoded, |metadata| { // Structurally valid but the first logical page has three valid rows, not two. metadata From 0371a2b619d171c8252517965cff67950a67eb88 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 12:41:52 -0400 Subject: [PATCH 08/36] Elide gather validity allocation for non-nullable selections Signed-off-by: Daniel King --- encodings/pco/src/take.rs | 58 +++++++++++++++++++++++++++++++-------- 1 file changed, 46 insertions(+), 12 deletions(-) diff --git a/encodings/pco/src/take.rs b/encodings/pco/src/take.rs index dc8b88211ad..4598f927ae7 100644 --- a/encodings/pco/src/take.rs +++ b/encodings/pco/src/take.rs @@ -152,13 +152,29 @@ fn gather( output_len: usize, nullability: vortex_array::dtype::Nullability, ctx: &mut ExecutionCtx, +) -> VortexResult { + if nullability.is_nullable() { + gather_typed::(array, runs, output_len, ctx) + } else { + gather_typed::(array, runs, output_len, ctx) + } +} + +fn gather_typed( + array: ArrayView<'_, Pco>, + runs: &[Run], + output_len: usize, + ctx: &mut ExecutionCtx, ) -> VortexResult { let mut state = AccessState::default(); let mut values = BufferMut::::with_capacity_in(output_len, ctx.allocator().clone()); - let mut valid = BitBufferMut::with_capacity(output_len); + // The non-nullable monomorphization has no bitmap allocation or append work. + let mut valid = BitBufferMut::with_capacity(if NULLABLE { output_len } else { 0 }); for run in runs { let Some(start) = run.start else { - valid.append_n(false, run.length); + if NULLABLE { + valid.append_n(false, run.length); + } values.extend(std::iter::repeat_n(T::default(), run.length)); continue; }; @@ -166,9 +182,13 @@ fn gather( if let Some((page, value)) = state.resolve(array, array.slice_start() + start, ctx)? { let value = state.values(array, page, ctx)?.as_slice::()[value]; values.extend(std::iter::repeat_n(value, run.length)); - valid.append_n(true, run.length); + if NULLABLE { + valid.append_n(true, run.length); + } } else { - valid.append_n(false, run.length); + if NULLABLE { + valid.append_n(false, run.length); + } values.extend(std::iter::repeat_n(T::default(), run.length)); } continue; @@ -183,11 +203,15 @@ fn gather( values.extend_from_slice( &state.values(array, span.page, ctx)?.as_slice::()[span.values], ); - valid.append_n(true, *n); + if NULLABLE { + valid.append_n(true, *n); + } } Mask::AllFalse(n) => { values.extend(std::iter::repeat_n(T::default(), *n)); - valid.append_n(false, *n); + if NULLABLE { + valid.append_n(false, *n); + } } Mask::Values(mask) => { let source = state.values(array, span.page, ctx)?.as_slice::(); @@ -204,17 +228,23 @@ fn gather( logical = stop; } values.extend(std::iter::repeat_n(T::default(), mask.len() - logical)); - valid.append_buffer(mask.bit_buffer()); + if NULLABLE { + valid.append_buffer(mask.bit_buffer()); + } } } row = span.stop; } else { if let Some((page, value)) = state.resolve(array, row, ctx)? { values.push(state.values(array, page, ctx)?.as_slice::()[value]); - valid.append_true(); + if NULLABLE { + valid.append_true(); + } } else { values.push(T::default()); - valid.append_false(); + if NULLABLE { + valid.append_false(); + } } row += 1; } @@ -227,14 +257,18 @@ fn gather( let row = array.slice_start() + start + offset * run.stride; if let Some((page, value)) = state.resolve(array, row, ctx)? { values.push(state.values(array, page, ctx)?.as_slice::()[value]); - valid.append_true(); + if NULLABLE { + valid.append_true(); + } } else { values.push(T::default()); - valid.append_false(); + if NULLABLE { + valid.append_false(); + } } } } - let validity = if nullability.is_nullable() { + let validity = if NULLABLE { Validity::from(valid.freeze()) } else { Validity::NonNullable From d26ad59859357faca204305d3c9d5639b77836f8 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 12:49:12 -0400 Subject: [PATCH 09/36] Add matched public PCO resource controls and satisfy upstream lint Signed-off-by: Daniel King --- .github/workflows/pco-compat.yml | 51 ++++- encodings/pco/src/access.rs | 13 +- encodings/pco/src/decoder.rs | 22 +-- encodings/pco/src/lib.rs | 8 +- encodings/pco/src/select_tests.rs | 8 +- encodings/pco/src/take.rs | 77 ++++---- vortex-file/examples/pco_access_profile.rs | 214 +++++++++++++++++++++ 7 files changed, 325 insertions(+), 68 deletions(-) create mode 100644 vortex-file/examples/pco_access_profile.rs diff --git a/.github/workflows/pco-compat.yml b/.github/workflows/pco-compat.yml index aa870ab3612..8810a5dde8d 100644 --- a/.github/workflows/pco-compat.yml +++ b/.github/workflows/pco-compat.yml @@ -3,11 +3,15 @@ name: PCO selective access compatibility on: pull_request: paths: - - 'encodings/pco/**' - - 'vortex-file/src/lib.rs' - - 'vortex-file/tests/pco_logical_pages_compat.rs' - - '.github/workflows/pco-compat.yml' - workflow_dispatch: {} + - "encodings/pco/**" + - "vortex-file/src/lib.rs" + - "vortex-file/tests/pco_logical_pages_compat.rs" + - ".github/workflows/pco-compat.yml" + workflow_dispatch: { } + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true permissions: contents: read @@ -24,7 +28,7 @@ jobs: if: github.repository == 'vortex-data/vortex' with: sccache: s3 - - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7 with: fetch-depth: 0 - uses: ./.github/actions/setup-prebuild @@ -52,3 +56,38 @@ jobs: PCO_COMPAT_MODE: read PCO_COMPAT_DIR: ${{ runner.temp }}/pco-old-files run: cargo test --profile ci --locked -p vortex-file --test pco_logical_pages_compat + - name: Build matching optimized public resource harnesses + run: | + cp vortex-file/examples/pco_access_profile.rs "$RUNNER_TEMP/pco-old-reader/vortex-file/examples/" + cargo build --profile bench --locked -p vortex-file --example pco_access_profile --message-format=json-render-diagnostics > "$RUNNER_TEMP/pco-candidate-build.jsonl" + cargo build --manifest-path "$RUNNER_TEMP/pco-old-reader/Cargo.toml" --profile bench --locked -p vortex-file --example pco_access_profile --message-format=json-render-diagnostics > "$RUNNER_TEMP/pco-baseline-build.jsonl" + - name: Matched first and warmed resource controls + shell: bash + run: | + mkdir -p "$RUNNER_TEMP/pco-resource-results" + getconf CLK_TCK > "$RUNNER_TEMP/pco-resource-results/clock-ticks-per-second.txt" + git rev-parse HEAD > "$RUNNER_TEMP/pco-resource-results/candidate-source.txt" + git -C "$RUNNER_TEMP/pco-old-reader" rev-parse HEAD > "$RUNNER_TEMP/pco-resource-results/baseline-source.txt" + candidate=$(jq -r 'select(.reason == "compiler-artifact" and .target.name == "pco_access_profile") | .executable // empty' "$RUNNER_TEMP/pco-candidate-build.jsonl") + baseline=$(jq -r 'select(.reason == "compiler-artifact" and .target.name == "pco_access_profile") | .executable // empty' "$RUNNER_TEMP/pco-baseline-build.jsonl") + for mode in write full slice take ranges probe-first probe; do + for pattern in dense nullable all-null long-null compressed nullable-legacy; do + # ABBA process order; every process retains all eight iterations including cold. + order=0 + for engine in baseline candidate candidate baseline; do + order=$((order + 1)) + executable="$candidate" + if [ "$engine" = baseline ]; then + executable="$baseline" + fi + stem="$RUNNER_TEMP/pco-resource-results/$mode-$pattern-$order-$engine" + /usr/bin/time -v "$executable" "$mode" "$pattern" 8388608 8 > "$stem.jsonl" 2> "$stem.time" + done + done + done + - uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7 + if: always() + with: + name: pco-resource-results + path: ${{ runner.temp }}/pco-resource-results + if-no-files-found: ignore diff --git a/encodings/pco/src/access.rs b/encodings/pco/src/access.rs index 43aed590224..ff4f15f467b 100644 --- a/encodings/pco/src/access.rs +++ b/encodings/pco/src/access.rs @@ -110,7 +110,7 @@ impl AccessState { .rows .as_ref() .ok_or_else(|| vortex_err!("Pco logical page range is missing"))?; - if !self.pages.contains_key(&page_idx) { + if let std::collections::btree_map::Entry::Vacant(entry) = self.pages.entry(page_idx) { let mask = array .unsliced_validity() .slice(rows.clone())? @@ -120,13 +120,10 @@ impl AccessState { page.values.len(), "Pco logical page validity count does not match stored values" ); - self.pages.insert( - page_idx, - CachedPage { - mask: Some(RankMask::new(mask)), - decoded: None, - }, - ); + entry.insert(CachedPage { + mask: Some(RankMask::new(mask)), + decoded: None, + }); } return Ok(self.pages[&page_idx] .mask diff --git a/encodings/pco/src/decoder.rs b/encodings/pco/src/decoder.rs index 419a5379793..4e0f8ab2f4d 100644 --- a/encodings/pco/src/decoder.rs +++ b/encodings/pco/src/decoder.rs @@ -13,7 +13,6 @@ 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 crate::array::vortex_err_from_pco; @@ -40,17 +39,16 @@ impl PageDecoder { buffer: &[u8], ctx: &mut ExecutionCtx, ) -> VortexResult { - if !self.chunks.contains_key(&chunk) { - let (decoder, _) = self - .file - .chunk_decompressor::(chunk_meta) - .map_err(vortex_err_from_pco)?; - self.chunks.insert(chunk, decoder); - } - let chunk = self - .chunks - .get_mut(&chunk) - .vortex_expect("chunk decoder was inserted"); + let chunk = match self.chunks.entry(chunk) { + std::collections::btree_map::Entry::Occupied(entry) => entry.into_mut(), + std::collections::btree_map::Entry::Vacant(entry) => { + let (decoder, _) = self + .file + .chunk_decompressor::(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)?; diff --git a/encodings/pco/src/lib.rs b/encodings/pco/src/lib.rs index 9fb4478ee48..179f12ed078 100644 --- a/encodings/pco/src/lib.rs +++ b/encodings/pco/src/lib.rs @@ -30,7 +30,7 @@ mod take; pub use array::*; -#[derive(Clone, PartialEq, prost::Message)] +#[derive(Clone, PartialEq, Eq, prost::Message)] /// Metadata for one Pco page. pub struct PcoPageInfo { // Since pco limits to 2^24 values per chunk, u32 is sufficient for the @@ -42,7 +42,7 @@ pub struct PcoPageInfo { // We're calling this Info instead of Metadata because ChunkMeta refers to a specific // component of a Pco file. -#[derive(Clone, PartialEq, prost::Message)] +#[derive(Clone, PartialEq, Eq, prost::Message)] /// Metadata for one Pco chunk. pub struct PcoChunkInfo { /// Pages contained in this chunk. @@ -50,7 +50,7 @@ pub struct PcoChunkInfo { pub pages: Vec, } -#[derive(Clone, PartialEq, prost::Message)] +#[derive(Clone, PartialEq, Eq, prost::Message)] /// Serialized metadata for a [`PcoArray`]. pub struct PcoMetadata { // would be nice to reuse one header per vortex file, but it's really only 1 byte, so @@ -67,7 +67,7 @@ pub struct PcoMetadata { } /// Versioned optional logical-row page index. Unknown versions retain legacy decoding. -#[derive(Clone, PartialEq, prost::Message)] +#[derive(Clone, PartialEq, Eq, prost::Message)] pub struct PcoLogicalPages { /// Index interpretation version; currently one. #[prost(uint32, tag = "1")] diff --git a/encodings/pco/src/select_tests.rs b/encodings/pco/src/select_tests.rs index 1280e97908e..b4ffc49470d 100644 --- a/encodings/pco/src/select_tests.rs +++ b/encodings/pco/src/select_tests.rs @@ -168,8 +168,8 @@ fn rejects_invalid_indices_and_descriptors() -> VortexResult<()> { ] { assert!( encoded - .take(indices)? - .execute::(&mut ctx) + .take(indices) + .and_then(|array| array.execute::(&mut ctx)) .is_err() ); } @@ -187,8 +187,8 @@ fn rejects_invalid_indices_and_descriptors() -> VortexResult<()> { .into_array(); assert!( encoded - .take(indices)? - .execute::(&mut ctx) + .take(indices) + .and_then(|array| array.execute::(&mut ctx)) .is_err() ); } diff --git a/encodings/pco/src/take.rs b/encodings/pco/src/take.rs index 4598f927ae7..d482061d617 100644 --- a/encodings/pco/src/take.rs +++ b/encodings/pco/src/take.rs @@ -64,39 +64,10 @@ impl TakeExecute for Pco { runs } else { let indices = indices.clone().execute::(ctx)?; - let mut runs: Vec = Vec::with_capacity(indices.len().min(64)); let mask = indices.validity()?.execute_mask(indices.len(), ctx)?; match_each_integer_ptype!(indices.ptype(), |I| { - for (&value, valid) in indices.as_slice::().iter().zip(mask.iter()) { - let start = if valid { - Some(usize::try_from(value).map_err(|_| { - vortex_err!("Pco take index is negative or exceeds usize") - })?) - } else { - None - }; - if let Some(previous) = runs.last_mut() { - let contiguous = match (previous.start, start) { - (Some(first), Some(next)) => { - previous.stride == 1 - && first.checked_add(previous.length) == Some(next) - } - (None, None) => true, - _ => false, - }; - if contiguous { - previous.length += 1; - continue; - } - } - runs.push(Run { - start, - length: 1, - stride: 1, - }); - } - }); - runs + primitive_runs(indices.as_slice::(), &mask)? + }) }; validate_runs(&runs, array.len(), indices.len())?; let nullable = array.dtype().nullability() | indices.dtype().nullability(); @@ -107,14 +78,52 @@ impl TakeExecute for Pco { } } +fn checked_index(value: I) -> VortexResult +where + usize: TryFrom, +{ + usize::try_from(value).map_err(|_| vortex_err!("Pco take index is negative or exceeds usize")) +} + +fn primitive_runs(indices: &[I], mask: &Mask) -> VortexResult> +where + usize: TryFrom, +{ + let mut runs: Vec = Vec::with_capacity(indices.len().min(64)); + for (&value, valid) in indices.iter().zip(mask.iter()) { + let start = if valid { + Some(checked_index(value)?) + } else { + None + }; + if let Some(previous) = runs.last_mut() { + let contiguous = match (previous.start, start) { + (Some(first), Some(next)) => { + previous.stride == 1 && first.checked_add(previous.length) == Some(next) + } + (None, None) => true, + _ => false, + }; + if contiguous { + previous.length += 1; + continue; + } + } + runs.push(Run { + start, + length: 1, + stride: 1, + }); + } + Ok(runs) +} + fn primitive_usizes(array: &PrimitiveArray) -> VortexResult> { match_each_integer_ptype!(array.ptype(), |I| { array .as_slice::() .iter() - .map(|&value| { - usize::try_from(value).map_err(|_| vortex_err!("Pco take descriptor exceeds usize")) - }) + .map(|&value| checked_index(value)) .collect() }) } diff --git a/vortex-file/examples/pco_access_profile.rs b/vortex-file/examples/pco_access_profile.rs new file mode 100644 index 00000000000..084838593ee --- /dev/null +++ b/vortex-file/examples/pco_access_profile.rs @@ -0,0 +1,214 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Public API resource harness. Run identical code on a candidate and a historical checkout. +//! Arguments: operation (write/full/slice/take/ranges/probe), pattern, rows, repetitions. +//! JSON lines retain cold and warmed iterations; Linux process accounting is read outside timing. + +use std::hint::black_box; +use std::time::Instant; + +use prost::Message; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::PiecewiseSequenceArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::buffer::BufferHandle; +use vortex_array::session::ArraySessionExt; +use vortex_array::validity::Validity; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_pco::Pco; +use vortex_runend::RunEnd; +use vortex_session::VortexSession; + +#[derive(Clone, prost::Message)] +struct LegacyMetadata { + #[prost(bytes, tag = "1")] + header: Vec, + #[prost(message, repeated, tag = "2")] + chunks: Vec, +} + +fn legacy(array: &ArrayRef, session: &VortexSession) -> VortexResult { + let wire = session + .array_serialize(array)? + .vortex_expect("PCO serialization"); + let metadata = LegacyMetadata::decode(wire.metadata.as_slice())?.encode_to_vec(); + let buffers = wire + .buffers + .into_iter() + .map(BufferHandle::new_host) + .collect::>(); + ArrayPlugin::deserialize( + &Pco, + ArrayDeserialization::new( + wire.serialized_id, + array.dtype(), + array.len(), + &metadata, + &buffers, + &wire.children, + ), + session, + ) +} + +fn linux_accounting() -> (Option, Option) { + let ticks = std::fs::read_to_string("/proc/self/stat") + .ok() + .and_then(|stat| { + let fields = stat + .rsplit_once(')')? + .1 + .split_whitespace() + .collect::>(); + Some(fields.get(11)?.parse::().ok()? + fields.get(12)?.parse::().ok()?) + }); + let rss = std::fs::read_to_string("/proc/self/status") + .ok() + .and_then(|status| { + status + .lines() + .find(|line| line.starts_with("VmRSS:"))? + .split_whitespace() + .nth(1)? + .parse() + .ok() + }); + (ticks, rss) +} + +fn main() -> VortexResult<()> { + let args = std::env::args().collect::>(); + let mode = args.get(1).map(String::as_str).unwrap_or("slice"); + let pattern = args.get(2).map(String::as_str).unwrap_or("nullable"); + let rows = args + .get(3) + .map_or(Ok(8_388_608), |value| value.parse::()) + .map_err(|error| vortex_err!("invalid row count: {error}"))?; + let repetitions = args + .get(4) + .map_or(Ok(8), |value| value.parse::()) + .map_err(|error| vortex_err!("invalid repetition count: {error}"))?; + vortex_ensure!(rows >= 8192, "profile requires at least 8192 rows"); + let session = VortexSession::empty(); + vortex_file::register_default_encodings(&session); + let mut ctx = session.create_execution_ctx(); + let (validity_pattern, legacy_format) = pattern + .strip_suffix("-legacy") + .map_or((pattern, false), |pattern| (pattern, true)); + let validity = match validity_pattern { + "dense" => Validity::NonNullable, + "all-null" => Validity::AllInvalid, + "long-null" => Validity::from_iter((0..rows).map(|i| i < 2048 || i > rows - 4096)), + "compressed" => Validity::Array( + RunEnd::encode( + BoolArray::from_iter((0..rows).map(|i| i % 65536 >= 32768)).into_array(), + &mut ctx, + )? + .into_array(), + ), + "nullable" => Validity::from_iter((0..rows).map(|i| i % 7 != 0)), + _ => return Err(vortex_err!("unknown pattern: {pattern}")), + }; + let input = PrimitiveArray::new( + (0..rows) + .map(|i| i64::try_from(i / 16).vortex_expect("fixture fits i64")) + .collect::>(), + validity, + ); + let encoded = if mode == "write" { + input.clone().into_array() + } else { + Pco::from_primitive(input.as_view(), 3, 1024, &mut ctx)?.into_array() + }; + let encoded = if legacy_format && mode != "write" { + legacy(&encoded, &session)? + } else { + encoded + }; + let take_indices = PrimitiveArray::from_iter( + (0..128usize).map(|i| u64::try_from((i * 104729) % rows).vortex_expect("index fits u64")), + ) + .into_array(); + let range_starts = (0..16usize) + .map(|i| u64::try_from((i * (rows / 17)) % (rows - 1024)).vortex_expect("range fits u64")) + .collect::>(); + let range_indices = PiecewiseSequenceArray::try_new( + PrimitiveArray::from_iter(range_starts).into_array(), + PrimitiveArray::from_iter(vec![512u64; 16]).into_array(), + PrimitiveArray::from_iter(vec![1u64; 16]).into_array(), + 8192, + )? + .into_array(); + let mut probe = encoded.repeated_probe(); + for iteration in 0..repetitions { + let before = linux_accounting(); + let start = Instant::now(); + let output: Option = match mode { + "write" => Some(Pco::from_primitive(input.as_view(), 3, 1024, &mut ctx)?.into_array()), + "full" => Some( + encoded + .clone() + .execute::(&mut ctx)? + .into_array(), + ), + "slice" => Some( + encoded + .slice(rows - 4096..rows - 2048)? + .execute::(&mut ctx)? + .into_array(), + ), + "take" => Some( + encoded + .take(take_indices.clone())? + .execute::(&mut ctx)? + .into_array(), + ), + "ranges" => Some( + encoded + .take(range_indices.clone())? + .execute::(&mut ctx)? + .into_array(), + ), + "probe-first" => { + black_box(encoded.probe().execute_scalar(rows - 2048, &mut ctx)?); + None + } + "probe" => { + for i in 0..1024 { + black_box(probe.execute_scalar(rows - 2048 + i % 256, &mut ctx)?); + } + None + } + _ => return Err(vortex_err!("unknown operation: {mode}")), + }; + black_box(&output); + let seconds = start.elapsed().as_secs_f64(); + let after = linux_accounting(); + // Keep results live and inspect a public scalar after the measurement. Independent value + // equivalence is tested by the public semantics/file suites, not by timing assertions. + if let Some(output) = output { + black_box(output.execute_scalar(0, &mut ctx)?); + } + let ticks = before + .0 + .zip(after.0) + .map(|(before, after)| after.saturating_sub(before)); + let ticks = ticks.map_or_else(|| "null".to_owned(), |value| value.to_string()); + let rss = after + .1 + .map_or_else(|| "null".to_owned(), |value| value.to_string()); + println!( + "{{\"operation\":\"{mode}\",\"pattern\":\"{pattern}\",\"rows\":{rows},\"iteration\":{iteration},\"wall_seconds\":{seconds},\"cpu_ticks\":{ticks},\"rss_kib\":{rss}}}" + ); + } + Ok(()) +} From 73ee941eaab61b261bede4c5396953cc9b2138ab Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 12:52:44 -0400 Subject: [PATCH 10/36] Measure complete array preparation and large PCO selection controls Signed-off-by: Daniel King --- .github/workflows/pco-compat.yml | 3 +- vortex-file/examples/pco_access_profile.rs | 57 ++++++++++++++++++++++ 2 files changed, 59 insertions(+), 1 deletion(-) diff --git a/.github/workflows/pco-compat.yml b/.github/workflows/pco-compat.yml index 8810a5dde8d..c7d85e50d78 100644 --- a/.github/workflows/pco-compat.yml +++ b/.github/workflows/pco-compat.yml @@ -70,7 +70,8 @@ jobs: git -C "$RUNNER_TEMP/pco-old-reader" rev-parse HEAD > "$RUNNER_TEMP/pco-resource-results/baseline-source.txt" candidate=$(jq -r 'select(.reason == "compiler-artifact" and .target.name == "pco_access_profile") | .executable // empty' "$RUNNER_TEMP/pco-candidate-build.jsonl") baseline=$(jq -r 'select(.reason == "compiler-artifact" and .target.name == "pco_access_profile") | .executable // empty' "$RUNNER_TEMP/pco-baseline-build.jsonl") - for mode in write full slice take ranges probe-first probe; do + sha256sum "$candidate" "$baseline" > "$RUNNER_TEMP/pco-resource-results/executables.sha256" + for mode in write full slice slice-near-full take take-large take-repeated ranges open-only open-first probe-first probe; do for pattern in dense nullable all-null long-null compressed nullable-legacy; do # ABBA process order; every process retains all eight iterations including cold. order=0 diff --git a/vortex-file/examples/pco_access_profile.rs b/vortex-file/examples/pco_access_profile.rs index 084838593ee..43ed0b949a1 100644 --- a/vortex-file/examples/pco_access_profile.rs +++ b/vortex-file/examples/pco_access_profile.rs @@ -148,6 +148,18 @@ fn main() -> VortexResult<()> { 8192, )? .into_array(); + let large_indices = PrimitiveArray::from_iter( + (0..rows / 8).map(|i| u64::try_from((i * 104729) % rows).vortex_expect("index fits u64")), + ) + .into_array(); + let repeated_indices = + PrimitiveArray::from_iter((0..rows / 8).map(|i| { + u64::try_from(((i % 16) * (rows / 17)) % rows).vortex_expect("index fits u64") + })) + .into_array(); + let wire = session + .array_serialize(&encoded)? + .vortex_expect("prepared array serialization"); let mut probe = encoded.repeated_probe(); for iteration in 0..repetitions { let before = linux_accounting(); @@ -178,6 +190,51 @@ fn main() -> VortexResult<()> { .execute::(&mut ctx)? .into_array(), ), + "slice-near-full" => Some( + encoded + .slice(1..rows - 1)? + .execute::(&mut ctx)? + .into_array(), + ), + "take-large" => Some( + encoded + .take(large_indices.clone())? + .execute::(&mut ctx)? + .into_array(), + ), + "take-repeated" => Some( + encoded + .take(repeated_indices.clone())? + .execute::(&mut ctx)? + .into_array(), + ), + "open-only" | "open-first" => { + // Array-plugin open with compressed buffers and the validity child already + // available: includes protobuf decode/validation/index preparation, no file IO. + let buffers = wire + .buffers + .iter() + .cloned() + .map(BufferHandle::new_host) + .collect::>(); + let opened = ArrayPlugin::deserialize( + &Pco, + ArrayDeserialization::new( + wire.serialized_id, + encoded.dtype(), + encoded.len(), + &wire.metadata, + &buffers, + &wire.children, + ), + &session, + )?; + if mode == "open-first" { + black_box(opened.execute_scalar(rows - 2048, &mut ctx)?); + } + black_box(opened); + None + } "probe-first" => { black_box(encoded.probe().execute_scalar(rows - 2048, &mut ctx)?); None From 902580c04366c2f22d493d02d58771aec02f25cb Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 12:54:21 -0400 Subject: [PATCH 11/36] Prepare historical harness directory and document resource boundaries Signed-off-by: Daniel King --- .github/workflows/pco-compat.yml | 2 ++ vortex-file/examples/pco_access_profile.rs | 9 +++++---- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/.github/workflows/pco-compat.yml b/.github/workflows/pco-compat.yml index c7d85e50d78..7584d162490 100644 --- a/.github/workflows/pco-compat.yml +++ b/.github/workflows/pco-compat.yml @@ -58,6 +58,7 @@ jobs: run: cargo test --profile ci --locked -p vortex-file --test pco_logical_pages_compat - name: Build matching optimized public resource harnesses run: | + mkdir -p "$RUNNER_TEMP/pco-old-reader/vortex-file/examples" cp vortex-file/examples/pco_access_profile.rs "$RUNNER_TEMP/pco-old-reader/vortex-file/examples/" cargo build --profile bench --locked -p vortex-file --example pco_access_profile --message-format=json-render-diagnostics > "$RUNNER_TEMP/pco-candidate-build.jsonl" cargo build --manifest-path "$RUNNER_TEMP/pco-old-reader/Cargo.toml" --profile bench --locked -p vortex-file --example pco_access_profile --message-format=json-render-diagnostics > "$RUNNER_TEMP/pco-baseline-build.jsonl" @@ -65,6 +66,7 @@ jobs: shell: bash run: | mkdir -p "$RUNNER_TEMP/pco-resource-results" + printf '%s\n' 'Wall timings exclude accounting and post-timing scalar inspection.' 'open-only/open-first include public array-plugin metadata decode/validation/index preparation with compressed buffers and validity children already available; no file IO.' 'probe-first starts with prepared metadata/index; probe retains touched pages across iterations.' 'CPU ticks are OS process user+system ticks; sub-tick readings are resolution-limited, not zero CPU.' 'RSS snapshots may miss transient peaks. GNU time maximum RSS is whole-process peak including fixtures/compression and post-timing inspection, not operation-only peak.' > "$RUNNER_TEMP/pco-resource-results/README.txt" getconf CLK_TCK > "$RUNNER_TEMP/pco-resource-results/clock-ticks-per-second.txt" git rev-parse HEAD > "$RUNNER_TEMP/pco-resource-results/candidate-source.txt" git -C "$RUNNER_TEMP/pco-old-reader" rev-parse HEAD > "$RUNNER_TEMP/pco-resource-results/baseline-source.txt" diff --git a/vortex-file/examples/pco_access_profile.rs b/vortex-file/examples/pco_access_profile.rs index 43ed0b949a1..133785cf989 100644 --- a/vortex-file/examples/pco_access_profile.rs +++ b/vortex-file/examples/pco_access_profile.rs @@ -149,13 +149,14 @@ fn main() -> VortexResult<()> { )? .into_array(); let large_indices = PrimitiveArray::from_iter( - (0..rows / 8).map(|i| u64::try_from((i * 104729) % rows).vortex_expect("index fits u64")), + (0..if mode == "take-large" { rows / 8 } else { 0 }) + .map(|i| u64::try_from((i * 104729) % rows).vortex_expect("index fits u64")), ) .into_array(); let repeated_indices = - PrimitiveArray::from_iter((0..rows / 8).map(|i| { - u64::try_from(((i % 16) * (rows / 17)) % rows).vortex_expect("index fits u64") - })) + PrimitiveArray::from_iter((0..if mode == "take-repeated" { rows / 8 } else { 0 }).map( + |i| u64::try_from(((i % 16) * (rows / 17)) % rows).vortex_expect("index fits u64"), + )) .into_array(); let wire = session .array_serialize(&encoded)? From 72cbb011a4ac165e71a325be1e34200ea1c099dd Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 13:06:07 -0400 Subject: [PATCH 12/36] Retain the last selected PCO page as a constant-time probe fast path Signed-off-by: Daniel King --- encodings/pco/src/access.rs | 131 +++++++++++++----- encodings/pco/src/decoder.rs | 5 +- encodings/pco/src/select_tests.rs | 7 + vortex-file/examples/pco_access_profile.rs | 16 ++- vortex-file/tests/pco_logical_pages_compat.rs | 43 ++++++ 5 files changed, 160 insertions(+), 42 deletions(-) diff --git a/encodings/pco/src/access.rs b/encodings/pco/src/access.rs index ff4f15f467b..9e2f956b8a2 100644 --- a/encodings/pco/src/access.rs +++ b/encodings/pco/src/access.rs @@ -21,6 +21,7 @@ use vortex_mask::Mask; use crate::Pco; use crate::PcoArrayExt; use crate::array::number_type_from_ptype; +use crate::page_index::PageIndex; const RANK_STRIDE: usize = 512; @@ -29,6 +30,7 @@ pub(crate) struct AccessState { legacy: Option, decoder: Option>, pages: BTreeMap, + last_page: Option<(usize, CachedPage)>, } struct CachedPage { @@ -86,6 +88,57 @@ impl RankMask { } impl AccessState { + fn cached(&self, page: usize) -> Option<&CachedPage> { + if let Some((last, cached)) = &self.last_page + && *last == page + { + Some(cached) + } else { + self.pages.get(&page) + } + } + + fn cached_mut(&mut self, page: usize) -> Option<&mut CachedPage> { + if self.last_page.as_ref().map(|(last, _)| *last) != Some(page) { + let cached = self.pages.remove(&page)?; + self.insert(page, cached); + } + self.last_page.as_mut().map(|(_, cached)| cached) + } + + fn insert(&mut self, page: usize, cached: CachedPage) { + if let Some((previous, cached)) = self.last_page.replace((page, cached)) { + self.pages.insert(previous, cached); + } + } + + fn logical_page(&self, index: &PageIndex, row: usize) -> usize { + if let Some((last, _)) = &self.last_page + && index + .pages + .get(*last) + .and_then(|page| page.rows.as_ref()) + .is_some_and(|range| range.contains(&row)) + { + return *last; + } + index + .pages + .partition_point(|page| page.rows.as_ref().is_some_and(|range| range.end <= row)) + } + + fn compact_page(&self, index: &PageIndex, value: usize) -> usize { + if let Some((last, _)) = &self.last_page + && index + .pages + .get(*last) + .is_some_and(|page| page.values.contains(&value)) + { + return *last; + } + index.pages.partition_point(|page| page.values.end <= value) + } + /// Resolve an unsliced logical row to a physical page and compact within-page index. pub fn resolve( &mut self, @@ -99,9 +152,7 @@ impl AccessState { ); let index = array.page_index()?; if index.logical && !index.pages.is_empty() { - let page_idx = index - .pages - .partition_point(|page| page.rows.as_ref().is_some_and(|range| range.end <= row)); + let page_idx = self.logical_page(index, row); let page = index .pages .get(page_idx) @@ -110,7 +161,7 @@ impl AccessState { .rows .as_ref() .ok_or_else(|| vortex_err!("Pco logical page range is missing"))?; - if let std::collections::btree_map::Entry::Vacant(entry) = self.pages.entry(page_idx) { + if self.cached(page_idx).is_none() { let mask = array .unsliced_validity() .slice(rows.clone())? @@ -120,12 +171,17 @@ impl AccessState { page.values.len(), "Pco logical page validity count does not match stored values" ); - entry.insert(CachedPage { - mask: Some(RankMask::new(mask)), - decoded: None, - }); + self.insert( + page_idx, + CachedPage { + mask: Some(RankMask::new(mask)), + decoded: None, + }, + ); } - return Ok(self.pages[&page_idx] + return Ok(self + .cached_mut(page_idx) + .ok_or_else(|| vortex_err!("Pco selected page cache is missing"))? .mask .as_ref() .and_then(|mask| mask.value_index(row - rows.start)) @@ -144,14 +200,19 @@ impl AccessState { let Some(value) = self.legacy.as_ref().and_then(|mask| mask.value_index(row)) else { return Ok(None); }; - let page_idx = index.pages.partition_point(|page| page.values.end <= value); + let page_idx = self.compact_page(index, value); let page = index.pages.get(page_idx).ok_or_else(|| { vortex_err!("Pco validity contains more valid rows than stored values") })?; - self.pages.entry(page_idx).or_insert(CachedPage { - mask: None, - decoded: None, - }); + if self.cached(page_idx).is_none() { + self.insert( + page_idx, + CachedPage { + mask: None, + decoded: None, + }, + ); + } Ok(Some((page_idx, value - page.values.start))) } @@ -192,9 +253,7 @@ impl AccessState { mask: mask.mask.slice(row..stop), })); } - let page_idx = index - .pages - .partition_point(|page| page.values.end <= start_value); + let page_idx = self.compact_page(index, start_value); let page = index .pages .get(page_idx) @@ -219,10 +278,15 @@ impl AccessState { .ok_or_else(|| vortex_err!("Pco legacy logical page end overflows usize"))? }; let selected_mask = mask.mask.slice(row..end); - self.pages.entry(page_idx).or_insert(CachedPage { - mask: None, - decoded: None, - }); + if self.cached(page_idx).is_none() { + self.insert( + page_idx, + CachedPage { + mask: None, + decoded: None, + }, + ); + } return Ok(Some(PageSpan { page: page_idx, stop: end, @@ -230,9 +294,7 @@ impl AccessState { mask: selected_mask, })); } - let page_idx = index - .pages - .partition_point(|page| page.rows.as_ref().is_some_and(|range| range.end <= row)); + let page_idx = self.logical_page(index, row); let page = index .pages .get(page_idx) @@ -244,8 +306,7 @@ impl AccessState { // Resolve performs whole-page validity validation, even when the first row is null. self.resolve(array, row, ctx)?; let cached = self - .pages - .get(&page_idx) + .cached(page_idx) .and_then(|page| page.mask.as_ref()) .ok_or_else(|| vortex_err!("Pco page validity is missing"))?; let end = stop.min(rows.end); @@ -265,11 +326,12 @@ impl AccessState { page_idx: usize, ctx: &mut ExecutionCtx, ) -> VortexResult<&PrimitiveArray> { - let cached = self - .pages - .get_mut(&page_idx) - .ok_or_else(|| vortex_err!("Pco page must be resolved before decoding"))?; - if cached.decoded.is_none() { + let needs_decode = self + .cached(page_idx) + .ok_or_else(|| vortex_err!("Pco page must be resolved before decoding"))? + .decoded + .is_none(); + if needs_decode { let index = array.page_index()?; let page = index .pages @@ -290,9 +352,12 @@ impl AccessState { decoder.read(page.chunk, array.chunk_metas[page.chunk].as_slice(), page.values.len(), buffer.as_slice(), ctx)? } ); - cached.decoded = Some(decoded); + self.cached_mut(page_idx) + .ok_or_else(|| vortex_err!("Pco selected page cache is missing"))? + .decoded = Some(decoded); } - cached + self.cached_mut(page_idx) + .ok_or_else(|| vortex_err!("Pco selected page cache is missing"))? .decoded .as_ref() .ok_or_else(|| vortex_err!("Pco page decoding did not produce values")) diff --git a/encodings/pco/src/decoder.rs b/encodings/pco/src/decoder.rs index 4e0f8ab2f4d..fcd0cab9840 100644 --- a/encodings/pco/src/decoder.rs +++ b/encodings/pco/src/decoder.rs @@ -4,6 +4,7 @@ //! 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; @@ -40,8 +41,8 @@ impl PageDecoder { ctx: &mut ExecutionCtx, ) -> VortexResult { let chunk = match self.chunks.entry(chunk) { - std::collections::btree_map::Entry::Occupied(entry) => entry.into_mut(), - std::collections::btree_map::Entry::Vacant(entry) => { + Entry::Occupied(entry) => entry.into_mut(), + Entry::Vacant(entry) => { let (decoder, _) = self .file .chunk_decompressor::(chunk_meta) diff --git a/encodings/pco/src/select_tests.rs b/encodings/pco/src/select_tests.rs index b4ffc49470d..1b34875776d 100644 --- a/encodings/pco/src/select_tests.rs +++ b/encodings/pco/src/select_tests.rs @@ -84,6 +84,13 @@ fn reordered_nullable_take(#[values(0, 1, 2)] format: usize) -> VortexResult<()> }; let source = encoded.slice(17..4001)?.slice(13..3900)?; let input = input.slice(30..3917)?; + let mut probe = source.repeated_probe(); + for index in [3000, 0, 127, 1, 3000, 128, 1000, 0] { + assert_eq!( + probe.execute_scalar(index, &mut ctx)?, + input.execute_scalar(index, &mut ctx)? + ); + } let indices = PrimitiveArray::from_option_iter([ Some(3000u32), Some(0), diff --git a/vortex-file/examples/pco_access_profile.rs b/vortex-file/examples/pco_access_profile.rs index 133785cf989..7349567c55a 100644 --- a/vortex-file/examples/pco_access_profile.rs +++ b/vortex-file/examples/pco_access_profile.rs @@ -108,13 +108,15 @@ fn main() -> VortexResult<()> { "dense" => Validity::NonNullable, "all-null" => Validity::AllInvalid, "long-null" => Validity::from_iter((0..rows).map(|i| i < 2048 || i > rows - 4096)), - "compressed" => Validity::Array( - RunEnd::encode( - BoolArray::from_iter((0..rows).map(|i| i % 65536 >= 32768)).into_array(), - &mut ctx, - )? - .into_array(), - ), + "compressed" => { + let runs = rows.div_ceil(32768); + let ends = PrimitiveArray::from_iter((0..runs).map(|run| { + u64::try_from(((run + 1) * 32768).min(rows)).vortex_expect("run end fits u64") + })) + .into_array(); + let values = BoolArray::from_iter((0..runs).map(|run| run % 2 != 0)).into_array(); + Validity::Array(RunEnd::try_new(ends, values, &mut ctx)?.into_array()) + } "nullable" => Validity::from_iter((0..rows).map(|i| i % 7 != 0)), _ => return Err(vortex_err!("unknown pattern: {pattern}")), }; diff --git a/vortex-file/tests/pco_logical_pages_compat.rs b/vortex-file/tests/pco_logical_pages_compat.rs index 59693ffbdf8..a878c903208 100644 --- a/vortex-file/tests/pco_logical_pages_compat.rs +++ b/vortex-file/tests/pco_logical_pages_compat.rs @@ -10,10 +10,12 @@ use std::sync::Arc; use futures::TryStreamExt; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::assert_arrays_eq; +use vortex_array::validity::Validity; use vortex_buffer::ByteBuffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -23,6 +25,7 @@ use vortex_io::session::RuntimeSession; use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; use vortex_layout::session::LayoutSession; use vortex_pco::Pco; +use vortex_runend::RunEnd; use vortex_session::VortexSession; mod common; @@ -81,3 +84,43 @@ async fn public_pco_compatibility() -> VortexResult<()> { } Ok(()) } + +#[test] +fn compressed_validity_selection_oracle() -> VortexResult<()> { + let session = VortexSession::empty(); + vortex_file::register_default_encodings(&session); + let mut ctx = session.create_execution_ctx(); + let validity = Validity::Array( + RunEnd::try_new( + PrimitiveArray::from_iter([32768u64, 65536, 98304, 131072]).into_array(), + BoolArray::from_iter([false, true, false, true]).into_array(), + &mut ctx, + )? + .into_array(), + ); + let input = PrimitiveArray::new((0..131072i64).collect::>(), validity); + let expected = + PrimitiveArray::from_option_iter((0..131072i64).map(|i| (i / 32768 % 2 == 1).then_some(i))); + let encoded = Pco::from_primitive(input.as_view(), 3, 127, &mut ctx)?.into_array(); + let indices = + PrimitiveArray::from_iter([0u32, 32767, 32768, 65535, 65536, 98304, 131071, 32768]) + .into_array(); + assert_arrays_eq!( + encoded.take(indices.clone())?, + expected.take(indices)?, + &mut ctx + ); + assert_arrays_eq!( + encoded.slice(30001..130001)?, + expected.slice(30001..130001)?, + &mut ctx + ); + let mut probe = encoded.repeated_probe(); + for index in [0, 32767, 32768, 65535, 65536, 98304, 131071, 32768] { + assert_eq!( + probe.execute_scalar(index, &mut ctx)?, + expected.execute_scalar(index, &mut ctx)? + ); + } + Ok(()) +} From b2ee3f9ed61d877ab37f60dceaabb75e79066ecf Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 13:08:55 -0400 Subject: [PATCH 13/36] Account for optional PCO page metadata in the compact file size golden Signed-off-by: Daniel King --- vortex-file/src/tests.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 7e15c104226..c4312825f92 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -174,7 +174,8 @@ async fn test_eof_values() { feature = "zstd", case::compact( BtrBlocksCompressorBuilder::from_session(&SESSION).with_compact(), - 55_344 + // PCO's optional logical-page metadata adds 64 bytes to this compact fixture. + 55_408 ) )] #[tokio::test] From 9da3bb7f8fbfaad2dac0e2cd4541776bc9d31c4a Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 13:10:56 -0400 Subject: [PATCH 14/36] Record old and new public file sizes alongside resource controls Signed-off-by: Daniel King --- .github/workflows/pco-compat.yml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.github/workflows/pco-compat.yml b/.github/workflows/pco-compat.yml index 7584d162490..c0d5a62f34e 100644 --- a/.github/workflows/pco-compat.yml +++ b/.github/workflows/pco-compat.yml @@ -56,6 +56,10 @@ jobs: PCO_COMPAT_MODE: read PCO_COMPAT_DIR: ${{ runner.temp }}/pco-old-files run: cargo test --profile ci --locked -p vortex-file --test pco_logical_pages_compat + - name: Record actual public file sizes + run: | + mkdir -p "$RUNNER_TEMP/pco-resource-results" + wc -c "$RUNNER_TEMP"/pco-{new,old}-files/*.vortex > "$RUNNER_TEMP/pco-resource-results/public-file-sizes.txt" - name: Build matching optimized public resource harnesses run: | mkdir -p "$RUNNER_TEMP/pco-old-reader/vortex-file/examples" From bd0da5034cb45d2f4a91d8dc3e65a96878293584 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 13:25:35 -0400 Subject: [PATCH 15/36] Profile frozen adverse PCO writer and selection paths Signed-off-by: Daniel King --- .github/workflows/pco-attribution.yml | 65 +++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) create mode 100644 .github/workflows/pco-attribution.yml diff --git a/.github/workflows/pco-attribution.yml b/.github/workflows/pco-attribution.yml new file mode 100644 index 00000000000..fb52feb7f35 --- /dev/null +++ b/.github/workflows/pco-attribution.yml @@ -0,0 +1,65 @@ +name: PCO attribution + +on: + pull_request: + paths: + - ".github/workflows/pco-attribution.yml" + workflow_dispatch: { } + +permissions: + contents: read + +jobs: + profile: + runs-on: >- + ${{ github.repository == 'vortex-data/vortex' + && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=pco-profile', github.run_id) + || 'ubuntu-latest' }} + timeout-minutes: 25 + steps: + - uses: runs-on/action@v2 + if: github.repository == 'vortex-data/vortex' + with: + sccache: s3 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7 + with: + fetch-depth: 0 + - uses: ./.github/actions/setup-prebuild + with: + enable-sccache: ${{ github.repository == 'vortex-data/vortex' && 'true' || 'false' }} + - name: Build frozen adverse candidate and baseline + run: | + git worktree add --detach "$RUNNER_TEMP/pco-profile-candidate" 9da3bb7f8fbfaad2dac0e2cd4541776bc9d31c4a + git worktree add --detach "$RUNNER_TEMP/pco-profile-base" 910f9933e8e2318c592c1eddcc2dfee8d24c6cd7 + mkdir -p "$RUNNER_TEMP/pco-profile-base/vortex-file/examples" + cp "$RUNNER_TEMP/pco-profile-candidate/vortex-file/examples/pco_access_profile.rs" "$RUNNER_TEMP/pco-profile-base/vortex-file/examples/" + for engine in candidate base; do + cargo build --manifest-path "$RUNNER_TEMP/pco-profile-$engine/Cargo.toml" --profile bench --locked -p vortex-file --example pco_access_profile --message-format=json-render-diagnostics > "$RUNNER_TEMP/$engine-build.jsonl" + done + - name: Bounded software attribution, separate from timings + shell: bash + run: | + mkdir -p "$RUNNER_TEMP/attribution" + sudo apt-get update -qq + sudo apt-get install -y -qq linux-tools-generic + perf=$(find /usr/lib/linux-tools -name perf -type f | sort | tail -1) + test -n "$perf" + for engine in base candidate; do + executable=$(jq -r 'select(.reason == "compiler-artifact" and .target.name == "pco_access_profile") | .executable // empty' "$RUNNER_TEMP/$engine-build.jsonl") + sha256sum "$executable" > "$RUNNER_TEMP/attribution/$engine-executable.sha256" + for mode in write probe take-large; do + repeats=64 + if [ "$mode" = probe ]; then repeats=100000; fi + sudo "$perf" record -F 499 --call-graph dwarf,8192 -o "$RUNNER_TEMP/attribution/$engine-$mode.data" -- "$executable" "$mode" nullable 8388608 "$repeats" > "$RUNNER_TEMP/$engine-$mode.jsonl" + sudo "$perf" report -i "$RUNNER_TEMP/attribution/$engine-$mode.data" --stdio --no-children --percent-limit 0.5 > "$RUNNER_TEMP/attribution/$engine-$mode-flat.txt" + sudo "$perf" report -i "$RUNNER_TEMP/attribution/$engine-$mode.data" --stdio --children --percent-limit 1 > "$RUNNER_TEMP/attribution/$engine-$mode-callgraph.txt" + sudo rm "$RUNNER_TEMP/attribution/$engine-$mode.data" + done + done + printf '%s\n' 'Frozen 9da adverse candidate and develop910f. Software samples only; profiled times are not uninstrumented latency controls. Startup/fixture compression included: identify query loops by stack.' > "$RUNNER_TEMP/attribution/README.txt" + - uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7 + if: always() + with: + name: pco-attribution + path: ${{ runner.temp }}/attribution + if-no-files-found: ignore From 52eb02776aeff2ed14776354a53980ab49bd5d75 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 13:26:37 -0400 Subject: [PATCH 16/36] Avoid touching nullable output gaps and defer page rank indexes Signed-off-by: Daniel King --- encodings/pco/src/access.rs | 71 ++++++++++++++++++++++++--------- encodings/pco/src/take.rs | 78 ++++++++++++++++++++++++++++++++----- 2 files changed, 121 insertions(+), 28 deletions(-) diff --git a/encodings/pco/src/access.rs b/encodings/pco/src/access.rs index 9e2f956b8a2..45dcb62d065 100644 --- a/encodings/pco/src/access.rs +++ b/encodings/pco/src/access.rs @@ -5,6 +5,7 @@ use std::any::Any; use std::collections::BTreeMap; +use std::sync::OnceLock; use pco::data_types::NumberType; use pco::match_number_enum; @@ -40,25 +41,30 @@ struct CachedPage { struct RankMask { mask: Mask, - ranks: Vec, + ranks: OnceLock>, } impl RankMask { fn new(mask: Mask) -> Self { - let ranks = match &mask { + 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(); - 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())); + if row <= bits.len() - row { + bits.count_range(0, row) + } else { + self.mask.true_count() - bits.count_range(row, bits.len()) } - ranks } - _ => Vec::new(), - }; - Self { mask, ranks } + } } fn rank(&self, row: usize) -> usize { @@ -69,19 +75,35 @@ impl RankMask { 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; - self.ranks[block] + values.bit_buffer().count_range(block * RANK_STRIDE, row) + ranks[block] + values.bit_buffer().count_range(block * RANK_STRIDE, row) } } } - fn value_index(&self, row: usize) -> Option { + fn value_index(&self, row: usize) -> Option { match &self.mask { Mask::AllTrue(_) => Some(row), Mask::AllFalse(_) => None, Mask::Values(values) => { let bits = values.bit_buffer(); - bits.value(row).then(|| self.rank(row)) + bits.value(row).then(|| { + if INDEXED { + self.rank(row) + } else { + self.boundary_rank(row) + } + }) } } } @@ -145,6 +167,15 @@ impl AccessState { array: ArrayView<'_, Pco>, row: usize, ctx: &mut ExecutionCtx, + ) -> VortexResult> { + self.resolve_impl::(array, row, ctx) + } + + fn resolve_impl( + &mut self, + array: ArrayView<'_, Pco>, + row: usize, + ctx: &mut ExecutionCtx, ) -> VortexResult> { vortex_ensure!( row < array.unsliced_n_rows(), @@ -184,7 +215,7 @@ impl AccessState { .ok_or_else(|| vortex_err!("Pco selected page cache is missing"))? .mask .as_ref() - .and_then(|mask| mask.value_index(row - rows.start)) + .and_then(|mask| mask.value_index::(row - rows.start)) .map(|value| (page_idx, value))); } @@ -197,7 +228,11 @@ impl AccessState { .execute_mask(array.unsliced_n_rows(), ctx)?, )); } - let Some(value) = self.legacy.as_ref().and_then(|mask| mask.value_index(row)) else { + let Some(value) = self + .legacy + .as_ref() + .and_then(|mask| mask.value_index::(row)) + else { return Ok(None); }; let page_idx = self.compact_page(index, value); @@ -304,7 +339,7 @@ impl AccessState { .as_ref() .ok_or_else(|| vortex_err!("Pco logical page range is missing"))?; // Resolve performs whole-page validity validation, even when the first row is null. - self.resolve(array, row, ctx)?; + self.resolve_impl::(array, row, ctx)?; let cached = self .cached(page_idx) .and_then(|page| page.mask.as_ref()) @@ -315,7 +350,7 @@ impl AccessState { Ok(Some(PageSpan { page: page_idx, stop: end, - values: cached.rank(start_in_page)..cached.rank(end_in_page), + values: cached.boundary_rank(start_in_page)..cached.boundary_rank(end_in_page), mask: cached.mask.slice(start_in_page..end_in_page), })) } diff --git a/encodings/pco/src/take.rs b/encodings/pco/src/take.rs index d482061d617..617e9acfa0c 100644 --- a/encodings/pco/src/take.rs +++ b/encodings/pco/src/take.rs @@ -169,6 +169,59 @@ fn gather( } } +/// One initialized final output buffer. Nullable outputs leave null positions at the +/// allocator's safe zero initialization instead of eagerly touching every null gap. +struct Output { + values: BufferMut, + len: usize, +} + +impl Output { + fn new(len: usize, ctx: &ExecutionCtx) -> Self { + let values = if NULLABLE { + BufferMut::zeroed_in(len, ctx.allocator().clone()) + } else { + BufferMut::with_capacity_in(len, ctx.allocator().clone()) + }; + Self { values, len: 0 } + } + + fn push(&mut self, value: T) { + if NULLABLE { + self.values[self.len] = value; + } else { + self.values.push(value); + } + self.len += 1; + } + + fn extend_from_slice(&mut self, values: &[T]) { + if NULLABLE { + self.values[self.len..self.len + values.len()].copy_from_slice(values); + } else { + self.values.extend_from_slice(values); + } + self.len += values.len(); + } + + fn repeat(&mut self, value: T, len: usize) { + if NULLABLE { + self.values[self.len..self.len + len].fill(value); + } else { + self.values.extend(std::iter::repeat_n(value, len)); + } + self.len += len; + } + + fn nulls(&mut self, len: usize) { + if NULLABLE { + self.len += len; + } else { + self.repeat(T::default(), len); + } + } +} + fn gather_typed( array: ArrayView<'_, Pco>, runs: &[Run], @@ -176,7 +229,7 @@ fn gather_typed( ctx: &mut ExecutionCtx, ) -> VortexResult { let mut state = AccessState::default(); - let mut values = BufferMut::::with_capacity_in(output_len, ctx.allocator().clone()); + let mut values = Output::::new(output_len, ctx); // The non-nullable monomorphization has no bitmap allocation or append work. let mut valid = BitBufferMut::with_capacity(if NULLABLE { output_len } else { 0 }); for run in runs { @@ -184,13 +237,13 @@ fn gather_typed( if NULLABLE { valid.append_n(false, run.length); } - values.extend(std::iter::repeat_n(T::default(), run.length)); + values.nulls(run.length); continue; }; if run.stride == 0 && run.length != 0 { if let Some((page, value)) = state.resolve(array, array.slice_start() + start, ctx)? { let value = state.values(array, page, ctx)?.as_slice::()[value]; - values.extend(std::iter::repeat_n(value, run.length)); + values.repeat(value, run.length); if NULLABLE { valid.append_n(true, run.length); } @@ -198,7 +251,7 @@ fn gather_typed( if NULLABLE { valid.append_n(false, run.length); } - values.extend(std::iter::repeat_n(T::default(), run.length)); + values.nulls(run.length); } continue; } @@ -217,7 +270,7 @@ fn gather_typed( } } Mask::AllFalse(n) => { - values.extend(std::iter::repeat_n(T::default(), *n)); + values.nulls(*n); if NULLABLE { valid.append_n(false, *n); } @@ -230,13 +283,13 @@ fn gather_typed( // Copy defined runs straight into the final output. Null gaps are // initialized once; no compact staging output is expanded afterward. for (start, stop) in mask.bit_buffer().set_slices() { - values.extend(std::iter::repeat_n(T::default(), start - logical)); + values.nulls(start - logical); let count = stop - start; values.extend_from_slice(&source[compact..compact + count]); compact += count; logical = stop; } - values.extend(std::iter::repeat_n(T::default(), mask.len() - logical)); + values.nulls(mask.len() - logical); if NULLABLE { valid.append_buffer(mask.bit_buffer()); } @@ -250,7 +303,7 @@ fn gather_typed( valid.append_true(); } } else { - values.push(T::default()); + values.nulls(1); if NULLABLE { valid.append_false(); } @@ -270,7 +323,7 @@ fn gather_typed( valid.append_true(); } } else { - values.push(T::default()); + values.nulls(1); if NULLABLE { valid.append_false(); } @@ -282,7 +335,12 @@ fn gather_typed( } else { Validity::NonNullable }; - Ok(PrimitiveArray::new(values.freeze(), validity)) + vortex_ensure_eq!( + values.len, + output_len, + "Pco selected output length mismatch" + ); + Ok(PrimitiveArray::new(values.values.freeze(), validity)) } pub(crate) fn slice( From 0dbc4ce40e9405ce54a79fbff0c772750a5bd000 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 13:30:07 -0400 Subject: [PATCH 17/36] Group scattered requests and preserve direct cached scalar hits Signed-off-by: Daniel King --- encodings/pco/src/access.rs | 43 ++++++++++ encodings/pco/src/probe.rs | 3 + encodings/pco/src/select_tests.rs | 38 +++++++++ encodings/pco/src/take.rs | 131 +++++++++++++++++++++++++++++- 4 files changed, 214 insertions(+), 1 deletion(-) diff --git a/encodings/pco/src/access.rs b/encodings/pco/src/access.rs index 45dcb62d065..f816096c062 100644 --- a/encodings/pco/src/access.rs +++ b/encodings/pco/src/access.rs @@ -5,6 +5,7 @@ use std::any::Any; use std::collections::BTreeMap; +use std::ops::Range; use std::sync::OnceLock; use pco::data_types::NumberType; @@ -13,6 +14,8 @@ use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::half; +use vortex_array::match_each_native_ptype; +use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; @@ -35,6 +38,8 @@ pub(crate) struct AccessState { } struct CachedPage { + rows: Option>, + values: Range, mask: Option, decoded: Option, } @@ -110,6 +115,38 @@ impl RankMask { } impl AccessState { + /// Read the last fully validated decoded page directly. A miss preserves all + /// operation-owned cache state and follows the ordinary fallible resolver. + pub fn cached_scalar(&self, array: ArrayView<'_, Pco>, row: usize) -> Option { + if row >= array.unsliced_n_rows() { + return None; + } + let (_, cached) = self.last_page.as_ref()?; + let decoded = cached.decoded.as_ref()?; + let value = if let Some(rows) = &cached.rows { + if !rows.contains(&row) { + return None; + } + cached.mask.as_ref()?.value_index::(row - rows.start) + } else { + let value = self.legacy.as_ref()?.value_index::(row); + if let Some(value) = value { + if !cached.values.contains(&value) { + return None; + } + Some(value - cached.values.start) + } else { + None + } + }; + Some(match value { + None => Scalar::null(array.dtype().clone()), + Some(value) => match_each_native_ptype!(decoded.ptype(), |T| { + Scalar::primitive(decoded.as_slice::()[value], array.dtype().nullability()) + }), + }) + } + fn cached(&self, page: usize) -> Option<&CachedPage> { if let Some((last, cached)) = &self.last_page && *last == page @@ -205,6 +242,8 @@ impl AccessState { self.insert( page_idx, CachedPage { + rows: Some(rows.clone()), + values: page.values.clone(), mask: Some(RankMask::new(mask)), decoded: None, }, @@ -243,6 +282,8 @@ impl AccessState { self.insert( page_idx, CachedPage { + rows: None, + values: page.values.clone(), mask: None, decoded: None, }, @@ -317,6 +358,8 @@ impl AccessState { self.insert( page_idx, CachedPage { + rows: None, + values: page.values.clone(), mask: None, decoded: None, }, diff --git a/encodings/pco/src/probe.rs b/encodings/pco/src/probe.rs index a63653a686e..398b988402c 100644 --- a/encodings/pco/src/probe.rs +++ b/encodings/pco/src/probe.rs @@ -30,6 +30,9 @@ pub(crate) fn scalar_at( let mut once = PcoProbeState::default(); let state = state.unwrap_or(&mut once); let row = array.slice_start() + index; + if let Some(value) = state.access.cached_scalar(array, row) { + return Ok(value); + } let Some((page, value)) = state.access.resolve(array, row, ctx)? else { return Ok(Scalar::null(array.dtype().clone())); }; diff --git a/encodings/pco/src/select_tests.rs b/encodings/pco/src/select_tests.rs index 1b34875776d..672f34dcd84 100644 --- a/encodings/pco/src/select_tests.rs +++ b/encodings/pco/src/select_tests.rs @@ -315,3 +315,41 @@ fn large_public_compressor_selections() -> VortexResult<()> { ); Ok(()) } + +#[rstest] +fn repeated_scattered_requests_preserve_order_and_nulls( + #[values(false, true)] legacy_format: bool, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = (0..32768i64) + .map(|row| (row % 7 != 0).then_some(row * 13)) + .collect::>(); + let input = PrimitiveArray::from_option_iter(values.iter().copied()); + let encoded = Pco::from_primitive(input.as_view(), 3, 1024, &mut ctx)?.into_array(); + let encoded = if legacy_format { + legacy(&encoded, false)? + } else { + encoded + }; + let requests = (0..1024usize) + .map(|i| (i % 11 != 0).then_some((i * 97) % 16)) + .collect::>(); + let indices = PrimitiveArray::from_option_iter( + requests + .iter() + .map(|row| row.map(|row| u64::try_from(row).vortex_expect("small index"))), + ) + .into_array(); + let expected = PrimitiveArray::from_option_iter( + requests.iter().map(|row| row.and_then(|row| values[row])), + ); + assert_arrays_eq!(encoded.take(indices)?, expected, &mut ctx); + let invalid = PrimitiveArray::from_iter([0i64, -1, 32768]).into_array(); + assert!( + encoded + .take(invalid)? + .execute::(&mut ctx) + .is_err() + ); + Ok(()) +} diff --git a/encodings/pco/src/take.rs b/encodings/pco/src/take.rs index 617e9acfa0c..871f2f18460 100644 --- a/encodings/pco/src/take.rs +++ b/encodings/pco/src/take.rs @@ -3,6 +3,8 @@ //! Batched selection using compact request descriptors and one operation-owned page cache. +use std::collections::BTreeSet; + use pco::data_types::Number; use pco::data_types::NumberType; use pco::match_number_enum; @@ -70,6 +72,11 @@ impl TakeExecute for Pco { }) }; validate_runs(&runs, array.len(), indices.len())?; + // For broad scattered requests the ordinary canonical decoder plus native primitive + // take has less setup/decoded work than retaining almost every page separately. + if scattered(&runs) && prefer_canonical(array, &runs)? { + return Ok(None); + } let nullable = array.dtype().nullability() | indices.dtype().nullability(); let result = match_number_enum!(number_type_from_ptype(array.dtype().as_ptype()), NumberType => { gather::(array, &runs, indices.len(), nullable, ctx)? @@ -78,6 +85,52 @@ impl TakeExecute for Pco { } } +fn scattered(runs: &[Run]) -> bool { + runs.len() > 1 && runs.iter().all(|run| run.length <= 1 || run.stride != 1) +} + +/// Conservative work estimate in bytes: selected physical values plus request descriptors +/// versus one full canonical values buffer. Stop as soon as canonical work wins, so broad +/// requests do not first build a complete selected-page set. No source-sized scratch bitmap. +fn prefer_canonical(array: ArrayView<'_, Pco>, runs: &[Run]) -> VortexResult { + let index = array.page_index()?; + let total = index.pages.last().map_or(0, |page| page.values.end); + let width = array.dtype().as_ptype().byte_width(); + let full_bytes = total.saturating_mul(width); + let requests = runs.iter().map(|run| run.length).sum::(); + let setup_bytes = requests.saturating_mul(std::mem::size_of::<(usize, usize)>()); + if setup_bytes >= full_bytes { + return Ok(true); + } + if !index.logical { + // With no logical page index, page coverage itself requires global validity rank. + // Use the page-count upper bound rather than spending arbitrary-request setup first. + return Ok(requests >= index.pages.len()); + } + let mut touched = BTreeSet::new(); + let mut selected = 0usize; + for run in runs { + let Some(start) = run.start else { continue }; + for offset in 0..run.length { + let row = array.slice_start() + start + offset * run.stride; + let page = index + .pages + .partition_point(|page| page.rows.as_ref().is_some_and(|rows| rows.end <= row)); + if touched.insert(page) { + let selected_page = index + .pages + .get(page) + .ok_or_else(|| vortex_err!("Pco selected row has no logical page"))?; + selected = selected.saturating_add(selected_page.values.len()); + if setup_bytes.saturating_add(selected.saturating_mul(width)) >= full_bytes { + return Ok(true); + } + } + } + } + Ok(false) +} + fn checked_index(value: I) -> VortexResult where usize: TryFrom, @@ -163,12 +216,88 @@ fn gather( ctx: &mut ExecutionCtx, ) -> VortexResult { if nullability.is_nullable() { - gather_typed::(array, runs, output_len, ctx) + if scattered(runs) { + gather_grouped::(array, runs, output_len, ctx) + } else { + gather_typed::(array, runs, output_len, ctx) + } + } else if scattered(runs) { + gather_grouped::(array, runs, output_len, ctx) } else { gather_typed::(array, runs, output_len, ctx) } } +/// Sort requests once, resolve each repeated logical position once, then invoke the native +/// primitive gather once per physical page. Only the final scatter visits output positions; +/// it does no page-index search, validity rank, decoder lookup, or cache lookup. +fn gather_grouped( + array: ArrayView<'_, Pco>, + runs: &[Run], + output_len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mut requests = Vec::with_capacity(output_len); + let mut output = 0usize; + for run in runs { + if let Some(start) = run.start { + for offset in 0..run.length { + requests.push(( + array.slice_start() + start + offset * run.stride, + output + offset, + )); + } + } + output += run.length; + } + requests.sort_unstable(); + let mut state = AccessState::default(); + let mut resolved = Vec::with_capacity(requests.len().min(64)); + let mut cursor = 0; + while cursor < requests.len() { + let row = requests[cursor].0; + let end = cursor + requests[cursor..].partition_point(|request| request.0 == row); + if let Some((page, value)) = state.resolve(array, row, ctx)? { + resolved.push((page, value, cursor..end)); + } + cursor = end; + } + let mut values = BufferMut::::zeroed_in(output_len, ctx.allocator().clone()); + let mut valid = BitBufferMut::new_unset(if NULLABLE { output_len } else { 0 }); + let mut cursor = 0; + while cursor < resolved.len() { + let page = resolved[cursor].0; + let end = cursor + resolved[cursor..].partition_point(|request| request.0 == page); + let compact = PrimitiveArray::from_iter( + resolved[cursor..end] + .iter() + .map(|request| { + u64::try_from(request.1) + .map_err(|_| vortex_err!("Pco compact index exceeds u64")) + }) + .collect::>>()?, + ) + .into_array(); + let decoded = state.values(array, page, ctx)?.clone(); + let selected = decoded.take(compact)?.execute::(ctx)?; + for (request, &value) in resolved[cursor..end].iter().zip(selected.as_slice::()) { + for &(_, output) in &requests[request.2.clone()] { + values[output] = value; + if NULLABLE { + valid.set(output); + } + } + } + cursor = end; + } + let validity = if NULLABLE { + Validity::from(valid.freeze()) + } else { + Validity::NonNullable + }; + Ok(PrimitiveArray::new(values.freeze(), validity)) +} + /// One initialized final output buffer. Nullable outputs leave null positions at the /// allocator's safe zero initialization instead of eagerly touching every null gap. struct Output { From 2195e6662d875bfa7a6d7ecba8ca6cb3c82e7ee3 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 13:30:48 -0400 Subject: [PATCH 18/36] Avoid output validity materialization for validated zero-page slices Signed-off-by: Daniel King --- encodings/pco/src/take.rs | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/encodings/pco/src/take.rs b/encodings/pco/src/take.rs index 871f2f18460..b13a741f399 100644 --- a/encodings/pco/src/take.rs +++ b/encodings/pco/src/take.rs @@ -215,6 +215,20 @@ fn gather( nullability: vortex_array::dtype::Nullability, ctx: &mut ExecutionCtx, ) -> VortexResult { + if array.page_index()?.pages.is_empty() { + let mask = array + .unsliced_validity() + .execute_mask(array.unsliced_n_rows(), ctx)?; + vortex_ensure_eq!( + mask.true_count(), + 0, + "Pco zero-page source has defined rows" + ); + return Ok(PrimitiveArray::new( + BufferMut::::zeroed_in(output_len, ctx.allocator().clone()).freeze(), + Validity::AllInvalid, + )); + } if nullability.is_nullable() { if scattered(runs) { gather_grouped::(array, runs, output_len, ctx) From 6f10d72303a940e80f8a31b06d3304017bfb858c Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 13:34:53 -0400 Subject: [PATCH 19/36] Plan visible-domain take before sparse allocations and retain compact spans Signed-off-by: Daniel King --- encodings/pco/src/select_tests.rs | 44 +++++++ encodings/pco/src/take.rs | 193 ++++++++++++++++++++++++------ 2 files changed, 199 insertions(+), 38 deletions(-) diff --git a/encodings/pco/src/select_tests.rs b/encodings/pco/src/select_tests.rs index 672f34dcd84..209fd43cfaa 100644 --- a/encodings/pco/src/select_tests.rs +++ b/encodings/pco/src/select_tests.rs @@ -353,3 +353,47 @@ fn repeated_scattered_requests_preserve_order_and_nulls( ); Ok(()) } + +#[test] +fn mixed_fragments_narrow_slices_and_compact_repeats() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = PrimitiveArray::from_option_iter( + (0..32768i64).map(|row| (row % 5 != 0).then_some(row * 11)), + ); + let encoded = Pco::from_primitive(input.as_view(), 3, 1024, &mut ctx)?.into_array(); + let mixed = PrimitiveArray::from_iter( + (0..4096u64).map(|i| if i < 2 { i } else { (i * 104729) % 32768 }), + ) + .into_array(); + assert_arrays_eq!(encoded.take(mixed.clone())?, input.take(mixed)?, &mut ctx); + let repeated = + PrimitiveArray::from_iter((0..100000u64).map(|i| if i % 2 == 0 { 0 } else { 127 })) + .into_array(); + assert_arrays_eq!( + encoded.slice(30000..30128)?.take(repeated.clone())?, + input.slice(30000..30128)?.take(repeated)?, + &mut ctx + ); + let compact = PiecewiseSequenceArray::try_new( + PrimitiveArray::from_iter([1u64, 5]).into_array(), + PrimitiveArray::from_iter([500000u64, 500000]).into_array(), + PrimitiveArray::from_iter([0u64, 0]).into_array(), + 1000000, + )? + .into_array(); + let expected = PrimitiveArray::from_option_iter( + std::iter::repeat_n(Some(11i64), 500000).chain(std::iter::repeat_n(None, 500000)), + ); + assert_arrays_eq!(encoded.take(compact)?, expected, &mut ctx); + let nulls = PrimitiveArray::new(vec![0i64; 100000], Validity::AllInvalid); + let encoded_nulls = Pco::from_primitive(nulls.as_view(), 3, 1024, &mut ctx)?.into_array(); + let two = PrimitiveArray::from_iter([0u64, 99999]).into_array(); + assert_arrays_eq!(encoded_nulls.take(two.clone())?, nulls.take(two)?, &mut ctx); + assert!( + encoded_nulls + .take(PrimitiveArray::from_iter([100000u64]).into_array())? + .execute::(&mut ctx) + .is_err() + ); + Ok(()) +} diff --git a/encodings/pco/src/take.rs b/encodings/pco/src/take.rs index b13a741f399..c5fc170f8a5 100644 --- a/encodings/pco/src/take.rs +++ b/encodings/pco/src/take.rs @@ -4,6 +4,7 @@ //! Batched selection using compact request descriptors and one operation-owned page cache. use std::collections::BTreeSet; +use std::ops::Range; use pco::data_types::Number; use pco::data_types::NumberType; @@ -12,6 +13,7 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PiecewiseSequence; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::dict::TakeExecute; @@ -19,6 +21,7 @@ use vortex_array::arrays::piecewise_sequence::array::PiecewiseSequenceArraySlots use vortex_array::dtype::NativePType; use vortex_array::dtype::half; use vortex_array::match_each_integer_ptype; +use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::BitBufferMut; use vortex_buffer::BufferMut; @@ -45,6 +48,8 @@ impl TakeExecute for Pco { indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { + let compact = indices.as_opt::().is_some(); + let mut grouped = false; let runs = if let Some(sequence) = indices.as_opt::() { let starts = sequence.starts().clone().execute::(ctx)?; let lengths = sequence.lengths().clone().execute::(ctx)?; @@ -67,64 +72,175 @@ impl TakeExecute for Pco { } else { let indices = indices.clone().execute::(ctx)?; let mask = indices.validity()?.execute_mask(indices.len(), ctx)?; + let summary = match_each_integer_ptype!(indices.ptype(), |I| { + validate_primitive(indices.as_slice::(), &mask, array.len())? + }); + if summary.defined == 0 { + return Ok(Some( + ConstantArray::new(Scalar::null(array.dtype().as_nullable()), indices.len()) + .into_array(), + )); + } + if let Some(result) = zero_page_take(array, indices.len(), ctx)? { + return Ok(Some(result)); + } + grouped = !summary.contiguous; + if grouped && descriptor_bytes(summary.defined) >= canonical_bytes(array) { + return Ok(None); + } + if grouped + && match_each_integer_ptype!(indices.ptype(), |I| { + prefer_canonical_primitive( + array, + indices.as_slice::(), + &mask, + summary.defined, + )? + }) + { + return Ok(None); + } match_each_integer_ptype!(indices.ptype(), |I| { primitive_runs(indices.as_slice::(), &mask)? }) }; validate_runs(&runs, array.len(), indices.len())?; - // For broad scattered requests the ordinary canonical decoder plus native primitive - // take has less setup/decoded work than retaining almost every page separately. - if scattered(&runs) && prefer_canonical(array, &runs)? { - return Ok(None); + if compact { + if let Some(result) = zero_page_take(array, indices.len(), ctx)? { + return Ok(Some(result)); + } + // Compact ranges/repeats remain compact. Only genuine strided requests carry + // per-position selection work; output cardinality alone is not setup work. + let random = runs + .iter() + .filter(|run| run.stride > 1) + .map(|run| run.length) + .sum(); + if descriptor_bytes(random) >= canonical_bytes(array) && random != 0 { + return Ok(None); + } } let nullable = array.dtype().nullability() | indices.dtype().nullability(); let result = match_number_enum!(number_type_from_ptype(array.dtype().as_ptype()), NumberType => { - gather::(array, &runs, indices.len(), nullable, ctx)? + gather::(array, &runs, indices.len(), nullable, grouped, ctx)? }); Ok(Some(result.into_array())) } } -fn scattered(runs: &[Run]) -> bool { - runs.len() > 1 && runs.iter().all(|run| run.length <= 1 || run.stride != 1) +struct PlainSummary { + defined: usize, + contiguous: bool, } -/// Conservative work estimate in bytes: selected physical values plus request descriptors -/// versus one full canonical values buffer. Stop as soon as canonical work wins, so broad -/// requests do not first build a complete selected-page set. No source-sized scratch bitmap. -fn prefer_canonical(array: ArrayView<'_, Pco>, runs: &[Run]) -> VortexResult { - let index = array.page_index()?; - let total = index.pages.last().map_or(0, |page| page.values.end); - let width = array.dtype().as_ptype().byte_width(); - let full_bytes = total.saturating_mul(width); - let requests = runs.iter().map(|run| run.length).sum::(); - let setup_bytes = requests.saturating_mul(std::mem::size_of::<(usize, usize)>()); - if setup_bytes >= full_bytes { - return Ok(true); +fn validate_primitive( + indices: &[I], + mask: &Mask, + source_len: usize, +) -> VortexResult +where + usize: TryFrom, +{ + let mut previous: Option = None; + let mut defined = 0usize; + let mut contiguous = true; + for (&index, valid) in indices.iter().zip(mask.iter()) { + if !valid { + contiguous = false; + continue; + } + let row = checked_index(index)?; + vortex_ensure!( + row < source_len, + "Pco take index {row} exceeds array length {source_len}" + ); + if let Some(previous) = previous { + contiguous &= previous.checked_add(1) == Some(row); + } + previous = Some(row); + defined += 1; } + Ok(PlainSummary { + defined, + contiguous, + }) +} + +fn zero_page_take( + array: ArrayView<'_, Pco>, + len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + if !array.page_index()?.pages.is_empty() { + return Ok(None); + } + let mask = array + .unsliced_validity() + .execute_mask(array.unsliced_n_rows(), ctx)?; + vortex_ensure_eq!( + mask.true_count(), + 0, + "Pco zero-page source has defined rows" + ); + Ok(Some( + ConstantArray::new(Scalar::null(array.dtype().as_nullable()), len).into_array(), + )) +} + +fn canonical_bytes(array: ArrayView<'_, Pco>) -> usize { + // Fallback canonicalizes the visible array, not its potentially much larger owner. + array + .len() + .saturating_mul(array.dtype().as_ptype().byte_width()) +} + +fn descriptor_bytes(requests: usize) -> usize { + // Actual prospective sparse-path storage: coalesced Run upper bound, sorting pair, + // and resolved (page, compact index, repeated-request range). Output storage is common. + requests.saturating_mul( + std::mem::size_of::() + + std::mem::size_of::<(usize, usize)>() + + std::mem::size_of::<(usize, usize, Range)>(), + ) +} + +/// Borrow validated primitive inputs during planning. Broad requests stop once selected +/// physical values plus prospective descriptors exceed the visible canonical values buffer. +/// Only the chosen sparse strategy allocates Runs/sorting descriptors; no source-sized bitmap. +fn prefer_canonical_primitive( + array: ArrayView<'_, Pco>, + indices: &[I], + mask: &Mask, + requests: usize, +) -> VortexResult +where + usize: TryFrom, +{ + let index = array.page_index()?; if !index.logical { - // With no logical page index, page coverage itself requires global validity rank. - // Use the page-count upper bound rather than spending arbitrary-request setup first. return Ok(requests >= index.pages.len()); } + let full_bytes = canonical_bytes(array); + let setup = descriptor_bytes(requests); + let width = array.dtype().as_ptype().byte_width(); let mut touched = BTreeSet::new(); let mut selected = 0usize; - for run in runs { - let Some(start) = run.start else { continue }; - for offset in 0..run.length { - let row = array.slice_start() + start + offset * run.stride; + for (&row, valid) in indices.iter().zip(mask.iter()) { + if !valid { + continue; + } + let row = array.slice_start() + checked_index(row)?; + let page = index + .pages + .partition_point(|page| page.rows.as_ref().is_some_and(|rows| rows.end <= row)); + if touched.insert(page) { let page = index .pages - .partition_point(|page| page.rows.as_ref().is_some_and(|rows| rows.end <= row)); - if touched.insert(page) { - let selected_page = index - .pages - .get(page) - .ok_or_else(|| vortex_err!("Pco selected row has no logical page"))?; - selected = selected.saturating_add(selected_page.values.len()); - if setup_bytes.saturating_add(selected.saturating_mul(width)) >= full_bytes { - return Ok(true); - } + .get(page) + .ok_or_else(|| vortex_err!("Pco selected row has no logical page"))?; + selected = selected.saturating_add(page.values.len()); + if setup.saturating_add(selected.saturating_mul(width)) >= full_bytes { + return Ok(true); } } } @@ -213,6 +329,7 @@ fn gather( runs: &[Run], output_len: usize, nullability: vortex_array::dtype::Nullability, + grouped: bool, ctx: &mut ExecutionCtx, ) -> VortexResult { if array.page_index()?.pages.is_empty() { @@ -230,12 +347,12 @@ fn gather( )); } if nullability.is_nullable() { - if scattered(runs) { + if grouped { gather_grouped::(array, runs, output_len, ctx) } else { gather_typed::(array, runs, output_len, ctx) } - } else if scattered(runs) { + } else if grouped { gather_grouped::(array, runs, output_len, ctx) } else { gather_typed::(array, runs, output_len, ctx) @@ -496,6 +613,6 @@ pub(crate) fn slice( stride: 1, }]; match_number_enum!(number_type_from_ptype(array.dtype().as_ptype()), NumberType => { - gather::(array, &runs, array.len(), array.dtype().nullability(), ctx) + gather::(array, &runs, array.len(), array.dtype().nullability(), false, ctx) }) } From fa1bfe4f8ac51899d6650744116711b911c96e0a Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 13:35:26 -0400 Subject: [PATCH 20/36] Measure mixed and compact requests and attribute metadata envelope bytes Signed-off-by: Daniel King --- .github/workflows/pco-attribution.yml | 6 +- .github/workflows/pco-compat.yml | 2 +- vortex-file/examples/pco_access_profile.rs | 95 +++++++++++++++++++++- 3 files changed, 99 insertions(+), 4 deletions(-) diff --git a/.github/workflows/pco-attribution.yml b/.github/workflows/pco-attribution.yml index fb52feb7f35..1fe980a5252 100644 --- a/.github/workflows/pco-attribution.yml +++ b/.github/workflows/pco-attribution.yml @@ -32,7 +32,8 @@ jobs: git worktree add --detach "$RUNNER_TEMP/pco-profile-candidate" 9da3bb7f8fbfaad2dac0e2cd4541776bc9d31c4a git worktree add --detach "$RUNNER_TEMP/pco-profile-base" 910f9933e8e2318c592c1eddcc2dfee8d24c6cd7 mkdir -p "$RUNNER_TEMP/pco-profile-base/vortex-file/examples" - cp "$RUNNER_TEMP/pco-profile-candidate/vortex-file/examples/pco_access_profile.rs" "$RUNNER_TEMP/pco-profile-base/vortex-file/examples/" + cp vortex-file/examples/pco_access_profile.rs "$RUNNER_TEMP/pco-profile-candidate/vortex-file/examples/" + cp vortex-file/examples/pco_access_profile.rs "$RUNNER_TEMP/pco-profile-base/vortex-file/examples/" for engine in candidate base; do cargo build --manifest-path "$RUNNER_TEMP/pco-profile-$engine/Cargo.toml" --profile bench --locked -p vortex-file --example pco_access_profile --message-format=json-render-diagnostics > "$RUNNER_TEMP/$engine-build.jsonl" done @@ -42,11 +43,12 @@ jobs: mkdir -p "$RUNNER_TEMP/attribution" sudo apt-get update -qq sudo apt-get install -y -qq linux-tools-generic - perf=$(find /usr/lib/linux-tools -name perf -type f | sort | tail -1) + perf=$(find /usr/lib -path "*linux-tools*/perf" -type f | sort | tail -1) test -n "$perf" for engine in base candidate; do executable=$(jq -r 'select(.reason == "compiler-artifact" and .target.name == "pco_access_profile") | .executable // empty' "$RUNNER_TEMP/$engine-build.jsonl") sha256sum "$executable" > "$RUNNER_TEMP/attribution/$engine-executable.sha256" + "$executable" metadata dense 100000 1 > "$RUNNER_TEMP/attribution/$engine-metadata.json" for mode in write probe take-large; do repeats=64 if [ "$mode" = probe ]; then repeats=100000; fi diff --git a/.github/workflows/pco-compat.yml b/.github/workflows/pco-compat.yml index c0d5a62f34e..d3b36097964 100644 --- a/.github/workflows/pco-compat.yml +++ b/.github/workflows/pco-compat.yml @@ -77,7 +77,7 @@ jobs: candidate=$(jq -r 'select(.reason == "compiler-artifact" and .target.name == "pco_access_profile") | .executable // empty' "$RUNNER_TEMP/pco-candidate-build.jsonl") baseline=$(jq -r 'select(.reason == "compiler-artifact" and .target.name == "pco_access_profile") | .executable // empty' "$RUNNER_TEMP/pco-baseline-build.jsonl") sha256sum "$candidate" "$baseline" > "$RUNNER_TEMP/pco-resource-results/executables.sha256" - for mode in write full slice slice-near-full take take-large take-repeated ranges open-only open-first probe-first probe; do + for mode in write full slice slice-near-full take take-large take-repeated take-mixed take-narrow take-spans take-tiny ranges open-only open-first probe-first probe; do for pattern in dense nullable all-null long-null compressed nullable-legacy; do # ABBA process order; every process retains all eight iterations including cold. order=0 diff --git a/vortex-file/examples/pco_access_profile.rs b/vortex-file/examples/pco_access_profile.rs index 7349567c55a..15cc5dd9397 100644 --- a/vortex-file/examples/pco_access_profile.rs +++ b/vortex-file/examples/pco_access_profile.rs @@ -9,6 +9,7 @@ use std::hint::black_box; use std::time::Instant; use prost::Message; +use vortex_array::ArrayContext; use vortex_array::ArrayDeserialization; use vortex_array::ArrayPlugin; use vortex_array::ArrayRef; @@ -18,6 +19,7 @@ use vortex_array::arrays::BoolArray; use vortex_array::arrays::PiecewiseSequenceArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::buffer::BufferHandle; +use vortex_array::serde::SerializeOptions; use vortex_array::session::ArraySessionExt; use vortex_array::validity::Validity; use vortex_error::VortexExpect; @@ -129,13 +131,46 @@ fn main() -> VortexResult<()> { let encoded = if mode == "write" { input.clone().into_array() } else { - Pco::from_primitive(input.as_view(), 3, 1024, &mut ctx)?.into_array() + Pco::from_primitive( + input.as_view(), + 3, + if mode == "metadata" { 8192 } else { 1024 }, + &mut ctx, + )? + .into_array() }; let encoded = if legacy_format && mode != "write" { legacy(&encoded, &session)? } else { encoded }; + if mode == "metadata" { + let wire = session + .array_serialize(&encoded)? + .vortex_expect("PCO metadata"); + let stripped = LegacyMetadata::decode(wire.metadata.as_slice())?.encode_to_vec(); + let legacy = legacy(&encoded, &session)?; + let context = ArrayContext::empty(); + let options = SerializeOptions::default(); + let current = encoded.serialize(&context, &session, &options)?; + let historical = legacy.serialize(&context, &session, &options)?; + println!( + "{{\"rows\":{rows},\"page_values_limit\":8192,\"metadata_bytes\":{},\"legacy_metadata_bytes\":{},\"serialized_bytes\":{},\"legacy_serialized_bytes\":{},\"serialized_buffer_sizes\":{:?},\"legacy_serialized_buffer_sizes\":{:?}}}", + wire.metadata.len(), + stripped.len(), + current.iter().map(|buffer| buffer.len()).sum::(), + historical.iter().map(|buffer| buffer.len()).sum::(), + current + .iter() + .map(|buffer| buffer.len()) + .collect::>(), + historical + .iter() + .map(|buffer| buffer.len()) + .collect::>() + ); + return Ok(()); + } let take_indices = PrimitiveArray::from_iter( (0..128usize).map(|i| u64::try_from((i * 104729) % rows).vortex_expect("index fits u64")), ) @@ -160,6 +195,33 @@ fn main() -> VortexResult<()> { |i| u64::try_from(((i % 16) * (rows / 17)) % rows).vortex_expect("index fits u64"), )) .into_array(); + let mixed_indices = PrimitiveArray::from_iter( + (0..if mode == "take-mixed" { rows / 8 } else { 0 }).map(|i| { + u64::try_from(if i < 2 { i } else { (i * 104729) % rows }) + .vortex_expect("index fits u64") + }), + ) + .into_array(); + let narrow_indices = PrimitiveArray::from_iter( + (0..if mode == "take-narrow" { + 100000usize + } else { + 0 + }) + .map(|i| if i % 2 == 0 { 0u64 } else { 127 }), + ) + .into_array(); + let repeated_spans = PiecewiseSequenceArray::try_new( + PrimitiveArray::from_iter([ + 1u64, + u64::try_from(rows - 1).vortex_expect("index fits u64"), + ]) + .into_array(), + PrimitiveArray::from_iter([500000u64, 500000]).into_array(), + PrimitiveArray::from_iter([0u64, 0]).into_array(), + 1000000, + )? + .into_array(); let wire = session .array_serialize(&encoded)? .vortex_expect("prepared array serialization"); @@ -199,6 +261,37 @@ fn main() -> VortexResult<()> { .execute::(&mut ctx)? .into_array(), ), + "take-mixed" => Some( + encoded + .take(mixed_indices.clone())? + .execute::(&mut ctx)? + .into_array(), + ), + "take-narrow" => Some( + encoded + .slice(rows - 128..rows)? + .take(narrow_indices.clone())? + .execute::(&mut ctx)? + .into_array(), + ), + "take-spans" => Some( + encoded + .take(repeated_spans.clone())? + .execute::(&mut ctx)? + .into_array(), + ), + "take-tiny" => Some( + encoded + .take( + PrimitiveArray::from_iter([ + 0u64, + u64::try_from(rows - 1).vortex_expect("index fits u64"), + ]) + .into_array(), + )? + .execute::(&mut ctx)? + .into_array(), + ), "take-large" => Some( encoded .take(large_indices.clone())? From 364213ad7f7b775c1deeef36b87089c9ad4dc491 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 13:37:05 -0400 Subject: [PATCH 21/36] Reserve only defined requests and preserve empty take dtype Signed-off-by: Daniel King --- encodings/pco/src/select_tests.rs | 32 +++++++++++++++++++++++++++++++ encodings/pco/src/take.rs | 31 +++++++++++++++++++----------- 2 files changed, 52 insertions(+), 11 deletions(-) diff --git a/encodings/pco/src/select_tests.rs b/encodings/pco/src/select_tests.rs index 209fd43cfaa..f70a5b57a64 100644 --- a/encodings/pco/src/select_tests.rs +++ b/encodings/pco/src/select_tests.rs @@ -397,3 +397,35 @@ fn mixed_fragments_narrow_slices_and_compact_repeats() -> VortexResult<()> { ); Ok(()) } + +#[test] +fn empty_take_dtype_and_mostly_null_request_storage() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = PrimitiveArray::from_iter(0..32768i64); + let encoded = Pco::from_primitive(input.as_view(), 3, 1024, &mut ctx)?.into_array(); + for source in [encoded.clone(), encoded.slice(0..0)?] { + for indices in [ + PrimitiveArray::from_iter(std::iter::empty::()).into_array(), + PrimitiveArray::from_option_iter(std::iter::empty::>()).into_array(), + ] { + let result = source + .take(indices.clone())? + .execute::(&mut ctx)?; + let expected = input + .slice(0..0)? + .take(indices)? + .execute::(&mut ctx)?; + assert_eq!(result.dtype(), expected.dtype()); + assert_arrays_eq!(result, expected, &mut ctx); + } + } + let indices = + PrimitiveArray::from_option_iter((0..200000).map(|i| (i == 12345).then_some(1u64))) + .into_array(); + let result = encoded.take(indices)?.execute::(&mut ctx)?; + let expected = + PrimitiveArray::from_option_iter((0..200000).map(|i| (i == 12345).then_some(1i64))); + assert_eq!(result.dtype(), expected.dtype()); + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) +} diff --git a/encodings/pco/src/take.rs b/encodings/pco/src/take.rs index c5fc170f8a5..c72fcb44c24 100644 --- a/encodings/pco/src/take.rs +++ b/encodings/pco/src/take.rs @@ -11,6 +11,7 @@ use pco::data_types::NumberType; use pco::match_number_enum; use vortex_array::ArrayRef; use vortex_array::ArrayView; +use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ConstantArray; @@ -48,8 +49,15 @@ impl TakeExecute for Pco { indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { + if indices.is_empty() { + let dtype = array + .dtype() + .clone() + .union_nullability(indices.dtype().nullability()); + return Ok(Some(Canonical::empty(&dtype).into_array())); + } let compact = indices.as_opt::().is_some(); - let mut grouped = false; + let mut grouped = None; let runs = if let Some(sequence) = indices.as_opt::() { let starts = sequence.starts().clone().execute::(ctx)?; let lengths = sequence.lengths().clone().execute::(ctx)?; @@ -84,11 +92,11 @@ impl TakeExecute for Pco { if let Some(result) = zero_page_take(array, indices.len(), ctx)? { return Ok(Some(result)); } - grouped = !summary.contiguous; - if grouped && descriptor_bytes(summary.defined) >= canonical_bytes(array) { + grouped = (!summary.contiguous).then_some(summary.defined); + if grouped.is_some() && descriptor_bytes(summary.defined) >= canonical_bytes(array) { return Ok(None); } - if grouped + if grouped.is_some() && match_each_integer_ptype!(indices.ptype(), |I| { prefer_canonical_primitive( array, @@ -329,7 +337,7 @@ fn gather( runs: &[Run], output_len: usize, nullability: vortex_array::dtype::Nullability, - grouped: bool, + grouped: Option, ctx: &mut ExecutionCtx, ) -> VortexResult { if array.page_index()?.pages.is_empty() { @@ -347,13 +355,13 @@ fn gather( )); } if nullability.is_nullable() { - if grouped { - gather_grouped::(array, runs, output_len, ctx) + if let Some(defined) = grouped { + gather_grouped::(array, runs, output_len, defined, ctx) } else { gather_typed::(array, runs, output_len, ctx) } - } else if grouped { - gather_grouped::(array, runs, output_len, ctx) + } else if let Some(defined) = grouped { + gather_grouped::(array, runs, output_len, defined, ctx) } else { gather_typed::(array, runs, output_len, ctx) } @@ -366,9 +374,10 @@ fn gather_grouped( array: ArrayView<'_, Pco>, runs: &[Run], output_len: usize, + defined: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - let mut requests = Vec::with_capacity(output_len); + let mut requests = Vec::with_capacity(defined); let mut output = 0usize; for run in runs { if let Some(start) = run.start { @@ -613,6 +622,6 @@ pub(crate) fn slice( stride: 1, }]; match_number_enum!(number_type_from_ptype(array.dtype().as_ptype()), NumberType => { - gather::(array, &runs, array.len(), array.dtype().nullability(), false, ctx) + gather::(array, &runs, array.len(), array.dtype().nullability(), None, ctx) }) } From ef7eb2b7439209529a2e83b7b28c0c99bc15f980 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 13:37:35 -0400 Subject: [PATCH 22/36] Use software cpu-clock samples for bounded attribution Signed-off-by: Daniel King --- .github/workflows/pco-attribution.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/pco-attribution.yml b/.github/workflows/pco-attribution.yml index 1fe980a5252..7049373515d 100644 --- a/.github/workflows/pco-attribution.yml +++ b/.github/workflows/pco-attribution.yml @@ -52,7 +52,7 @@ jobs: for mode in write probe take-large; do repeats=64 if [ "$mode" = probe ]; then repeats=100000; fi - sudo "$perf" record -F 499 --call-graph dwarf,8192 -o "$RUNNER_TEMP/attribution/$engine-$mode.data" -- "$executable" "$mode" nullable 8388608 "$repeats" > "$RUNNER_TEMP/$engine-$mode.jsonl" + sudo "$perf" record -e cpu-clock -F 499 --call-graph dwarf,8192 -o "$RUNNER_TEMP/attribution/$engine-$mode.data" -- "$executable" "$mode" nullable 8388608 "$repeats" > "$RUNNER_TEMP/$engine-$mode.jsonl" sudo "$perf" report -i "$RUNNER_TEMP/attribution/$engine-$mode.data" --stdio --no-children --percent-limit 0.5 > "$RUNNER_TEMP/attribution/$engine-$mode-flat.txt" sudo "$perf" report -i "$RUNNER_TEMP/attribution/$engine-$mode.data" --stdio --children --percent-limit 1 > "$RUNNER_TEMP/attribution/$engine-$mode-callgraph.txt" sudo rm "$RUNNER_TEMP/attribution/$engine-$mode.data" From e80af840fbd99b5489960397fbab067820b8b4fb Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 13:39:32 -0400 Subject: [PATCH 23/36] Attribute exact compact stock fixture segment and framing bytes Signed-off-by: Daniel King --- .github/workflows/pco-attribution.yml | 4 ++ vortex-file/tests/pco_logical_pages_compat.rs | 61 +++++++++++++++++++ 2 files changed, 65 insertions(+) diff --git a/.github/workflows/pco-attribution.yml b/.github/workflows/pco-attribution.yml index 7049373515d..34fb27f3f75 100644 --- a/.github/workflows/pco-attribution.yml +++ b/.github/workflows/pco-attribution.yml @@ -35,12 +35,16 @@ jobs: cp vortex-file/examples/pco_access_profile.rs "$RUNNER_TEMP/pco-profile-candidate/vortex-file/examples/" cp vortex-file/examples/pco_access_profile.rs "$RUNNER_TEMP/pco-profile-base/vortex-file/examples/" for engine in candidate base; do + cp vortex-file/tests/pco_logical_pages_compat.rs "$RUNNER_TEMP/pco-profile-$engine/vortex-file/tests/" cargo build --manifest-path "$RUNNER_TEMP/pco-profile-$engine/Cargo.toml" --profile bench --locked -p vortex-file --example pco_access_profile --message-format=json-render-diagnostics > "$RUNNER_TEMP/$engine-build.jsonl" done - name: Bounded software attribution, separate from timings shell: bash run: | mkdir -p "$RUNNER_TEMP/attribution" + for engine in base candidate; do + PCO_SIZE_DIAGNOSTICS=1 cargo test --manifest-path "$RUNNER_TEMP/pco-profile-$engine/Cargo.toml" --profile ci --locked -p vortex-file --features zstd --test pco_logical_pages_compat public_stock_size_attribution -- --nocapture > "$RUNNER_TEMP/attribution/$engine-stock-size.txt" + done sudo apt-get update -qq sudo apt-get install -y -qq linux-tools-generic perf=$(find /usr/lib -path "*linux-tools*/perf" -type f | sort | tail -1) diff --git a/vortex-file/tests/pco_logical_pages_compat.rs b/vortex-file/tests/pco_logical_pages_compat.rs index a878c903208..11aa42fd426 100644 --- a/vortex-file/tests/pco_logical_pages_compat.rs +++ b/vortex-file/tests/pco_logical_pages_compat.rs @@ -124,3 +124,64 @@ fn compressed_validity_selection_oracle() -> VortexResult<()> { } Ok(()) } + +/// Public byte-boundary diagnostic for the exact deterministic compact stock fixture. +/// This reports format overhead; it is not a timing or topology assertion. +#[cfg(feature = "zstd")] +#[tokio::test] +async fn public_stock_size_attribution() -> VortexResult<()> { + use rand::Rng; + use rand::SeedableRng; + use rand::rngs::StdRng; + use vortex_btrblocks::BtrBlocksCompressorBuilder; + use vortex_file::strategy::WriteStrategyBuilder; + + if std::env::var_os("PCO_SIZE_DIAGNOSTICS").is_none() { + return Ok(()); + } + let session = VortexSession::empty() + .with::() + .with::(); + vortex_file::register_default_encodings(&session); + common::enable_all_registered_array_encodings(&session); + let mut rng = StdRng::seed_from_u64(0); + let input = PrimitiveArray::from_iter((0..100000i64).map(|i| rng.random_range(i..=i + 10))) + .into_array(); + let mut bytes = Vec::new(); + session + .write_options() + .with_strategy( + WriteStrategyBuilder::from_session(&session) + .with_btrblocks_builder( + BtrBlocksCompressorBuilder::from_session(&session).with_compact(), + ) + .build(), + ) + .write(&mut bytes, input.clone().to_array_stream()) + .await?; + let len = bytes.len(); + let file = session + .open_options() + .open_buffer(ByteBuffer::from(bytes))?; + let segments = file.footer().segment_map(); + let payload = segments + .iter() + .map(|segment| u64::from(segment.length)) + .sum::(); + let end = segments + .iter() + .map(|segment| segment.offset + u64::from(segment.length)) + .max() + .unwrap_or(0); + println!( + "{{\"stock_file_bytes\":{len},\"segment_payload_bytes\":{payload},\"last_segment_end\":{end},\"trailing_format_bytes\":{},\"segments\":{:?}}}", + u64::try_from(len).vortex_expect("file fits u64") - end, + segments + .iter() + .map(|segment| (segment.offset, segment.length, segment.alignment.as_usize())) + .collect::>() + ); + let actual = file.scan()?.into_array_stream()?.read_all().await?; + assert_arrays_eq!(actual, input, &mut session.create_execution_ctx()); + Ok(()) +} From acffeeff18f8f8f7ac48c3fcd7efabbc537dc9d4 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 13:39:54 -0400 Subject: [PATCH 24/36] Keep public size diagnostic output machine readable Signed-off-by: Daniel King --- vortex-file/tests/pco_logical_pages_compat.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/vortex-file/tests/pco_logical_pages_compat.rs b/vortex-file/tests/pco_logical_pages_compat.rs index 11aa42fd426..8f3ac8176ec 100644 --- a/vortex-file/tests/pco_logical_pages_compat.rs +++ b/vortex-file/tests/pco_logical_pages_compat.rs @@ -133,6 +133,7 @@ async fn public_stock_size_attribution() -> VortexResult<()> { use rand::Rng; use rand::SeedableRng; use rand::rngs::StdRng; + use vortex_array::stream::ArrayStreamExt; use vortex_btrblocks::BtrBlocksCompressorBuilder; use vortex_file::strategy::WriteStrategyBuilder; @@ -178,7 +179,11 @@ async fn public_stock_size_attribution() -> VortexResult<()> { u64::try_from(len).vortex_expect("file fits u64") - end, segments .iter() - .map(|segment| (segment.offset, segment.length, segment.alignment.as_usize())) + .map(|segment| vec![ + segment.offset, + u64::from(segment.length), + u64::try_from(segment.alignment.as_usize()).vortex_expect("alignment fits u64") + ]) .collect::>() ); let actual = file.scan()?.into_array_stream()?.read_all().await?; From e15a6e641b4dc2478291e9a17e3b1c39be6a4af2 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 13:40:35 -0400 Subject: [PATCH 25/36] Use inclusive range checks in public profiling fixtures Signed-off-by: Daniel King --- vortex-file/examples/pco_access_profile.rs | 2 +- vortex-file/tests/pco_logical_pages_compat.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/vortex-file/examples/pco_access_profile.rs b/vortex-file/examples/pco_access_profile.rs index 15cc5dd9397..3fd7c1a863e 100644 --- a/vortex-file/examples/pco_access_profile.rs +++ b/vortex-file/examples/pco_access_profile.rs @@ -109,7 +109,7 @@ fn main() -> VortexResult<()> { let validity = match validity_pattern { "dense" => Validity::NonNullable, "all-null" => Validity::AllInvalid, - "long-null" => Validity::from_iter((0..rows).map(|i| i < 2048 || i > rows - 4096)), + "long-null" => Validity::from_iter((0..rows).map(|i| !(2048..=rows - 4096).contains(&i))), "compressed" => { let runs = rows.div_ceil(32768); let ends = PrimitiveArray::from_iter((0..runs).map(|run| { diff --git a/vortex-file/tests/pco_logical_pages_compat.rs b/vortex-file/tests/pco_logical_pages_compat.rs index 8f3ac8176ec..e946303e839 100644 --- a/vortex-file/tests/pco_logical_pages_compat.rs +++ b/vortex-file/tests/pco_logical_pages_compat.rs @@ -39,7 +39,7 @@ async fn public_pco_compatibility() -> VortexResult<()> { common::enable_all_registered_array_encodings(&session); let mut ctx = session.create_execution_ctx(); let original = PrimitiveArray::from_option_iter( - (0..131_072i64).map(|i| ((i < 1024 || i > 120_000) && i % 7 != 0).then_some(i * 19)), + (0..131_072i64).map(|i| (!(1024..=120_000).contains(&i) && i % 7 != 0).then_some(i * 19)), ); let mode = std::env::var("PCO_COMPAT_MODE").unwrap_or_else(|_| "roundtrip".to_owned()); let directory = std::env::var_os("PCO_COMPAT_DIR").map(std::path::PathBuf::from); From 0e763cc69cb4304e995987eb797c6dfd57353907 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 13:42:09 -0400 Subject: [PATCH 26/36] Use public write strategy re-export in size diagnostic Signed-off-by: Daniel King --- vortex-file/tests/pco_logical_pages_compat.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-file/tests/pco_logical_pages_compat.rs b/vortex-file/tests/pco_logical_pages_compat.rs index e946303e839..31a693b0d0b 100644 --- a/vortex-file/tests/pco_logical_pages_compat.rs +++ b/vortex-file/tests/pco_logical_pages_compat.rs @@ -135,7 +135,7 @@ async fn public_stock_size_attribution() -> VortexResult<()> { use rand::rngs::StdRng; use vortex_array::stream::ArrayStreamExt; use vortex_btrblocks::BtrBlocksCompressorBuilder; - use vortex_file::strategy::WriteStrategyBuilder; + use vortex_file::WriteStrategyBuilder; if std::env::var_os("PCO_SIZE_DIAGNOSTICS").is_none() { return Ok(()); From bf5fecb6d3c42cd38495e36d5d7b3280fa8d273c Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 13:44:33 -0400 Subject: [PATCH 27/36] Fix selected-reader trait imports and warning-clean qualifications Signed-off-by: Daniel King --- encodings/pco/src/access.rs | 2 +- encodings/pco/src/take.rs | 5 ++--- vortex-file/tests/pco_logical_pages_compat.rs | 2 +- 3 files changed, 4 insertions(+), 5 deletions(-) diff --git a/encodings/pco/src/access.rs b/encodings/pco/src/access.rs index f816096c062..87777c7869d 100644 --- a/encodings/pco/src/access.rs +++ b/encodings/pco/src/access.rs @@ -445,6 +445,6 @@ impl AccessState { pub(crate) struct PageSpan { pub page: usize, pub stop: usize, - pub values: std::ops::Range, + pub values: Range, pub mask: Mask, } diff --git a/encodings/pco/src/take.rs b/encodings/pco/src/take.rs index c72fcb44c24..88538599e95 100644 --- a/encodings/pco/src/take.rs +++ b/encodings/pco/src/take.rs @@ -33,6 +33,7 @@ use vortex_error::vortex_err; use vortex_mask::Mask; use crate::Pco; +use crate::PcoArrayExt; use crate::access::AccessState; use crate::array::number_type_from_ptype; @@ -206,9 +207,7 @@ fn descriptor_bytes(requests: usize) -> usize { // Actual prospective sparse-path storage: coalesced Run upper bound, sorting pair, // and resolved (page, compact index, repeated-request range). Output storage is common. requests.saturating_mul( - std::mem::size_of::() - + std::mem::size_of::<(usize, usize)>() - + std::mem::size_of::<(usize, usize, Range)>(), + size_of::() + size_of::<(usize, usize)>() + size_of::<(usize, usize, Range)>(), ) } diff --git a/vortex-file/tests/pco_logical_pages_compat.rs b/vortex-file/tests/pco_logical_pages_compat.rs index 31a693b0d0b..841cde3fe9b 100644 --- a/vortex-file/tests/pco_logical_pages_compat.rs +++ b/vortex-file/tests/pco_logical_pages_compat.rs @@ -130,7 +130,7 @@ fn compressed_validity_selection_oracle() -> VortexResult<()> { #[cfg(feature = "zstd")] #[tokio::test] async fn public_stock_size_attribution() -> VortexResult<()> { - use rand::Rng; + use rand::RngExt; use rand::SeedableRng; use rand::rngs::StdRng; use vortex_array::stream::ArrayStreamExt; From bdae09146951cabe4739be6aa9fcf6a7ffe95e50 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 14:00:22 -0400 Subject: [PATCH 28/36] Stream contiguous pages into final output and reuse native index extrema Signed-off-by: Daniel King --- encodings/pco/src/access.rs | 8 +- encodings/pco/src/decoder.rs | 30 +++++- encodings/pco/src/select_tests.rs | 67 ++++++++++++ encodings/pco/src/take.rs | 167 ++++++++++++++++++++++++------ 4 files changed, 232 insertions(+), 40 deletions(-) diff --git a/encodings/pco/src/access.rs b/encodings/pco/src/access.rs index 87777c7869d..589e87c60e2 100644 --- a/encodings/pco/src/access.rs +++ b/encodings/pco/src/access.rs @@ -30,7 +30,7 @@ use crate::page_index::PageIndex; const RANK_STRIDE: usize = 512; #[derive(Default)] -pub(crate) struct AccessState { +pub(crate) struct AccessState { legacy: Option, decoder: Option>, pages: BTreeMap, @@ -114,7 +114,7 @@ impl RankMask { } } -impl AccessState { +impl AccessState { /// Read the last fully validated decoded page directly. A miss preserves all /// operation-owned cache state and follows the ordinary fallible resolver. pub fn cached_scalar(&self, array: ArrayView<'_, Pco>, row: usize) -> Option { @@ -166,7 +166,9 @@ impl AccessState { } fn insert(&mut self, page: usize, cached: CachedPage) { - if let Some((previous, cached)) = self.last_page.replace((page, cached)) { + if let Some((previous, cached)) = self.last_page.replace((page, cached)) + && RETAIN_PAGES + { self.pages.insert(previous, cached); } } diff --git a/encodings/pco/src/decoder.rs b/encodings/pco/src/decoder.rs index fcd0cab9840..620cb9c366f 100644 --- a/encodings/pco/src/decoder.rs +++ b/encodings/pco/src/decoder.rs @@ -15,6 +15,8 @@ use vortex_array::dtype::NativePType; use vortex_array::validity::Validity; use vortex_buffer::BufferMut; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; use crate::array::vortex_err_from_pco; @@ -40,6 +42,21 @@ impl PageDecoder { buffer: &[u8], ctx: &mut ExecutionCtx, ) -> VortexResult { + let mut values = BufferMut::::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 exactly one complete page into caller-owned initialized storage. + /// A corrupt/partial page never publishes the caller's unfinished output array. + pub fn read_into( + &mut self, + chunk: usize, + chunk_meta: &[u8], + n_values: usize, + buffer: &[u8], + output: &mut [T], + ) -> VortexResult<()> { + vortex_ensure_eq!(output.len(), n_values, "Pco page output length mismatch"); let chunk = match self.chunks.entry(chunk) { Entry::Occupied(entry) => entry.into_mut(), Entry::Vacant(entry) => { @@ -53,10 +70,13 @@ impl PageDecoder { let mut decoder = chunk .page_decompressor(buffer, n_values) .map_err(vortex_err_from_pco)?; - // `read` takes initialized typed storage. Zero initialization also makes a failed or - // partial corrupt-page decode safe to drop; no uninitialized typed slice is published. - let mut values = BufferMut::::zeroed_in(n_values, ctx.allocator().clone()); - decoder.read(&mut values).map_err(vortex_err_from_pco)?; - Ok(PrimitiveArray::new(values.freeze(), Validity::NonNullable)) + let progress = decoder.read(output).map_err(vortex_err_from_pco)?; + vortex_ensure_eq!( + progress.n_processed, + n_values, + "Pco page did not fill output" + ); + vortex_ensure!(progress.finished, "Pco page decode is incomplete"); + Ok(()) } } diff --git a/encodings/pco/src/select_tests.rs b/encodings/pco/src/select_tests.rs index f70a5b57a64..f0e0ee15ced 100644 --- a/encodings/pco/src/select_tests.rs +++ b/encodings/pco/src/select_tests.rs @@ -429,3 +429,70 @@ fn empty_take_dtype_and_mostly_null_request_storage() -> VortexResult<()> { assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } + +#[rstest] +fn streamed_slices_preserve_float_bits_and_validity( + #[values(false, true)] dense: bool, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let bits = [ + 0u64, + 1u64 << 63, + 0x7ff8000000000042, + 0x7ff0000000000000, + 0xfff0000000000000, + 0x3ff4000000000000, + ]; + let values = (0..2048) + .map(|row| f64::from_bits(bits[row % bits.len()])) + .collect::>(); + let validity = if dense { + Validity::NonNullable + } else { + Validity::from_iter((0..2048).map(|row| row % 7 != 0)) + }; + let input = PrimitiveArray::new(values.clone(), validity); + let encoded = Pco::from_primitive(input.as_view(), 3, 31, &mut ctx)?.into_array(); + for range in [1..2047, 89..129, 1000..1001] { + let actual = encoded + .slice(range.clone())? + .execute::(&mut ctx)?; + assert_eq!(actual.dtype(), input.dtype()); + let mask = actual.validity()?.execute_mask(actual.len(), &mut ctx)?; + for (offset, &value) in actual.as_slice::().iter().enumerate() { + let row = range.start + offset; + assert_eq!(mask.value(offset), dense || row % 7 != 0); + if mask.value(offset) { + assert_eq!(value.to_bits(), values[row].to_bits()); + } + } + } + Ok(()) +} + +#[test] +fn logical_null_indices_ignore_invalid_physical_values() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = PrimitiveArray::from_iter(0..128i64); + let encoded = Pco::from_primitive(input.as_view(), 3, 7, &mut ctx)?.into_array(); + let indices = PrimitiveArray::new( + vec![i64::MIN, 4, i64::MAX], + Validity::from_iter([false, true, false]), + ) + .into_array(); + let expected = PrimitiveArray::from_option_iter([None, Some(4i64), None]); + assert_arrays_eq!(encoded.take(indices.clone())?, expected, &mut ctx); + // Reusing the immutable index array also exercises cached native extrema. + assert_arrays_eq!( + encoded.take(indices)?, + PrimitiveArray::from_option_iter([None, Some(4i64), None]), + &mut ctx + ); + assert!( + encoded + .take(PrimitiveArray::from_iter([u64::MAX]).into_array())? + .execute::(&mut ctx) + .is_err() + ); + Ok(()) +} diff --git a/encodings/pco/src/take.rs b/encodings/pco/src/take.rs index 88538599e95..bc3a77e8a20 100644 --- a/encodings/pco/src/take.rs +++ b/encodings/pco/src/take.rs @@ -14,6 +14,9 @@ use vortex_array::ArrayView; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::fns::min_max::MinMaxResult; +use vortex_array::aggregate_fn::fns::min_max::min_max; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PiecewiseSequence; use vortex_array::arrays::PrimitiveArray; @@ -36,6 +39,7 @@ use crate::Pco; use crate::PcoArrayExt; use crate::access::AccessState; use crate::array::number_type_from_ptype; +use crate::decoder::PageDecoder; /// Compact descriptor; null indices have no start. Zero-length runs consume no source rows. struct Run { @@ -81,9 +85,19 @@ impl TakeExecute for Pco { } else { let indices = indices.clone().execute::(ctx)?; let mask = indices.validity()?.execute_mask(indices.len(), ctx)?; - let summary = match_each_integer_ptype!(indices.ptype(), |I| { - validate_primitive(indices.as_slice::(), &mask, array.len())? - }); + // Reuse native logical-null-aware extrema and their immutable-array stats cache. + // Numeric extrema prove every defined index is representable and in bounds. + let extrema = min_max(indices.array(), ctx, NumericalAggregateOpts::default())?; + let summary = if let Some(extrema) = extrema { + match_each_integer_ptype!(indices.ptype(), |I| { + validate_extrema(indices.as_slice::(), &mask, &extrema, array.len())? + }) + } else { + PlainSummary { + defined: 0, + contiguous: false, + } + }; if summary.defined == 0 { return Ok(Some( ConstantArray::new(Scalar::null(array.dtype().as_nullable()), indices.len()) @@ -142,35 +156,44 @@ struct PlainSummary { contiguous: bool, } -fn validate_primitive( +fn validate_extrema( indices: &[I], mask: &Mask, + extrema: &MinMaxResult, source_len: usize, ) -> VortexResult where usize: TryFrom, { - let mut previous: Option = None; - let mut defined = 0usize; - let mut contiguous = true; - for (&index, valid) in indices.iter().zip(mask.iter()) { - if !valid { - contiguous = false; - continue; - } - let row = checked_index(index)?; - vortex_ensure!( - row < source_len, - "Pco take index {row} exceeds array length {source_len}" - ); - if let Some(previous) = previous { - contiguous &= previous.checked_add(1) == Some(row); - } - previous = Some(row); - defined += 1; - } + let min = extrema + .min + .as_primitive() + .try_typed_value::()? + .ok_or_else(|| vortex_err!("Pco index minimum is null"))?; + let max = extrema + .max + .as_primitive() + .try_typed_value::()? + .ok_or_else(|| vortex_err!("Pco index maximum is null"))?; + checked_index(min)?; + let max = checked_index(max)?; + vortex_ensure!( + max < source_len, + "Pco take index {max} exceeds array length {source_len}" + ); + // Bounds/error category is guaranteed above. No contract chooses the first invalid row. + // Contiguous detection exits at the first mismatch; null physical values are never inspected. + let contiguous = match mask { + Mask::AllTrue(_) => indices.windows(2).all(|pair| { + checked_index(pair[0]) + .ok() + .and_then(|row| row.checked_add(1)) + == checked_index(pair[1]).ok() + }), + _ => false, + }; Ok(PlainSummary { - defined, + defined: mask.true_count(), contiguous, }) } @@ -390,7 +413,7 @@ fn gather_grouped( output += run.length; } requests.sort_unstable(); - let mut state = AccessState::default(); + let mut state = AccessState::::default(); let mut resolved = Vec::with_capacity(requests.len().min(64)); let mut cursor = 0; while cursor < requests.len() { @@ -496,7 +519,7 @@ fn gather_typed( output_len: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - let mut state = AccessState::default(); + let mut state = AccessState::::default(); let mut values = Output::::new(output_len, ctx); // The non-nullable monomorphization has no bitmap allocation or append work. let mut valid = BitBufferMut::with_capacity(if NULLABLE { output_len } else { 0 }); @@ -611,16 +634,96 @@ fn gather_typed( Ok(PrimitiveArray::new(values.values.freeze(), validity)) } +/// Contiguous selections stream pages into the final initialized values buffer and retain +/// exact sliced source validity. No output validity copy or retained decoded-page cache. +fn stream_slice( + array: ArrayView<'_, Pco>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mut values = BufferMut::::zeroed_in(array.len(), ctx.allocator().clone()); + let mut scratch = BufferMut::::with_capacity_in(0, ctx.allocator().clone()); + let mut state = AccessState::::default(); + let mut decoder = PageDecoder::::new(array.metadata.header.as_slice())?; + let stop = array.slice_start() + array.len(); + let mut row = array.slice_start(); + while row < stop { + let span = state + .span(array, row, stop, ctx)? + .ok_or_else(|| vortex_err!("Pco contiguous page span is missing"))?; + vortex_ensure!( + span.stop > row && span.stop <= stop, + "Pco contiguous span does not advance within bounds" + ); + vortex_ensure_eq!( + span.mask.len(), + span.stop - row, + "Pco contiguous mask length mismatch" + ); + vortex_ensure_eq!( + span.values.len(), + span.mask.true_count(), + "Pco contiguous value count mismatch" + ); + let start = row - array.slice_start(); + let output = values + .get_mut(start..start + span.mask.len()) + .ok_or_else(|| vortex_err!("Pco contiguous output exceeds visible range"))?; + if !span.mask.all_false() { + let index = array.page_index()?; + let page = index + .pages + .get(span.page) + .ok_or_else(|| vortex_err!("Pco contiguous page is missing"))?; + let buffer = array + .pages + .get(span.page) + .ok_or_else(|| vortex_err!("Pco compressed page is missing"))?; + let chunk = array + .chunk_metas + .get(page.chunk) + .ok_or_else(|| vortex_err!("Pco chunk metadata is missing"))?; + let n = page.values.len(); + if span.mask.all_true() && span.values == (0..n) { + decoder.read_into(page.chunk, chunk.as_slice(), n, buffer.as_slice(), output)?; + } else { + if scratch.len() < n { + scratch.push_n(T::default(), n - scratch.len()); + } + decoder.read_into( + page.chunk, + chunk.as_slice(), + n, + buffer.as_slice(), + &mut scratch[..n], + )?; + let source = scratch + .get(span.values) + .ok_or_else(|| vortex_err!("Pco compact range exceeds page output"))?; + match &span.mask { + Mask::AllTrue(_) => output.copy_from_slice(source), + Mask::AllFalse(_) => {} + Mask::Values(mask) => { + let mut compact = 0; + for (start, end) in mask.bit_buffer().set_slices() { + let len = end - start; + output[start..end].copy_from_slice(&source[compact..compact + len]); + compact += len; + } + } + } + } + } + row = span.stop; + } + let validity = array.unsliced_validity().slice(array.slice_start()..stop)?; + Ok(PrimitiveArray::new(values.freeze(), validity)) +} + pub(crate) fn slice( array: ArrayView<'_, Pco>, ctx: &mut ExecutionCtx, ) -> VortexResult { - let runs = [Run { - start: Some(0), - length: array.len(), - stride: 1, - }]; match_number_enum!(number_type_from_ptype(array.dtype().as_ptype()), NumberType => { - gather::(array, &runs, array.len(), array.dtype().nullability(), None, ctx) + stream_slice::(array, ctx) }) } From 459d8daf5a91019370efee3f1457ee4a45f38789 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 14:00:22 -0400 Subject: [PATCH 29/36] Format diagnostic JSON through stable Display values Signed-off-by: Daniel King --- vortex-file/examples/pco_access_profile.rs | 10 ++++++---- vortex-file/tests/pco_logical_pages_compat.rs | 12 +++++++----- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/vortex-file/examples/pco_access_profile.rs b/vortex-file/examples/pco_access_profile.rs index 3fd7c1a863e..2547fc5e4c6 100644 --- a/vortex-file/examples/pco_access_profile.rs +++ b/vortex-file/examples/pco_access_profile.rs @@ -155,19 +155,21 @@ fn main() -> VortexResult<()> { let current = encoded.serialize(&context, &session, &options)?; let historical = legacy.serialize(&context, &session, &options)?; println!( - "{{\"rows\":{rows},\"page_values_limit\":8192,\"metadata_bytes\":{},\"legacy_metadata_bytes\":{},\"serialized_bytes\":{},\"legacy_serialized_bytes\":{},\"serialized_buffer_sizes\":{:?},\"legacy_serialized_buffer_sizes\":{:?}}}", + "{{\"rows\":{rows},\"page_values_limit\":8192,\"metadata_bytes\":{},\"legacy_metadata_bytes\":{},\"serialized_bytes\":{},\"legacy_serialized_bytes\":{},\"serialized_buffer_sizes\":[{}],\"legacy_serialized_buffer_sizes\":[{}]}}", wire.metadata.len(), stripped.len(), current.iter().map(|buffer| buffer.len()).sum::(), historical.iter().map(|buffer| buffer.len()).sum::(), current .iter() - .map(|buffer| buffer.len()) - .collect::>(), + .map(|buffer| buffer.len().to_string()) + .collect::>() + .join(","), historical .iter() - .map(|buffer| buffer.len()) + .map(|buffer| buffer.len().to_string()) .collect::>() + .join(",") ); return Ok(()); } diff --git a/vortex-file/tests/pco_logical_pages_compat.rs b/vortex-file/tests/pco_logical_pages_compat.rs index 841cde3fe9b..6ff657ab51f 100644 --- a/vortex-file/tests/pco_logical_pages_compat.rs +++ b/vortex-file/tests/pco_logical_pages_compat.rs @@ -175,16 +175,18 @@ async fn public_stock_size_attribution() -> VortexResult<()> { .max() .unwrap_or(0); println!( - "{{\"stock_file_bytes\":{len},\"segment_payload_bytes\":{payload},\"last_segment_end\":{end},\"trailing_format_bytes\":{},\"segments\":{:?}}}", + "{{\"stock_file_bytes\":{len},\"segment_payload_bytes\":{payload},\"last_segment_end\":{end},\"trailing_format_bytes\":{},\"segments\":[{}]}}", u64::try_from(len).vortex_expect("file fits u64") - end, segments .iter() - .map(|segment| vec![ + .map(|segment| format!( + "[{},{},{}]", segment.offset, - u64::from(segment.length), - u64::try_from(segment.alignment.as_usize()).vortex_expect("alignment fits u64") - ]) + segment.length, + segment.alignment.as_usize() + )) .collect::>() + .join(",") ); let actual = file.scan()?.into_array_stream()?.read_all().await?; assert_arrays_eq!(actual, input, &mut session.create_execution_ctx()); From ad54a18fdce95520be5e7189205e2f68b82a49e8 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 14:03:51 -0400 Subject: [PATCH 30/36] Share streaming range decoding with contiguous take requests Signed-off-by: Daniel King --- .github/workflows/pco-compat.yml | 2 +- encodings/pco/src/select_tests.rs | 23 ++++++++++++++ encodings/pco/src/take.rs | 37 +++++++++++++++++----- vortex-file/examples/pco_access_profile.rs | 29 +++++++++++++++++ 4 files changed, 82 insertions(+), 9 deletions(-) diff --git a/.github/workflows/pco-compat.yml b/.github/workflows/pco-compat.yml index d3b36097964..1d9a16e84d8 100644 --- a/.github/workflows/pco-compat.yml +++ b/.github/workflows/pco-compat.yml @@ -77,7 +77,7 @@ jobs: candidate=$(jq -r 'select(.reason == "compiler-artifact" and .target.name == "pco_access_profile") | .executable // empty' "$RUNNER_TEMP/pco-candidate-build.jsonl") baseline=$(jq -r 'select(.reason == "compiler-artifact" and .target.name == "pco_access_profile") | .executable // empty' "$RUNNER_TEMP/pco-baseline-build.jsonl") sha256sum "$candidate" "$baseline" > "$RUNNER_TEMP/pco-resource-results/executables.sha256" - for mode in write full slice slice-near-full take take-large take-repeated take-mixed take-narrow take-spans take-tiny ranges open-only open-first probe-first probe; do + for mode in write full slice slice-near-full take take-large take-repeated take-mixed take-narrow take-spans take-tiny take-contiguous take-range ranges open-only open-first probe-first probe; do for pattern in dense nullable all-null long-null compressed nullable-legacy; do # ABBA process order; every process retains all eight iterations including cold. order=0 diff --git a/encodings/pco/src/select_tests.rs b/encodings/pco/src/select_tests.rs index f0e0ee15ced..338c5639876 100644 --- a/encodings/pco/src/select_tests.rs +++ b/encodings/pco/src/select_tests.rs @@ -496,3 +496,26 @@ fn logical_null_indices_ignore_invalid_physical_values() -> VortexResult<()> { ); Ok(()) } + +#[test] +fn contiguous_take_preserves_nullable_index_dtype() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = PrimitiveArray::from_iter(0..32768i64); + let encoded = Pco::from_primitive(input.as_view(), 3, 1024, &mut ctx)?.into_array(); + let indices = PrimitiveArray::from_option_iter((1..32767u64).map(Some)).into_array(); + let expected = input + .take(indices.clone())? + .execute::(&mut ctx)?; + let result = encoded.take(indices)?.execute::(&mut ctx)?; + assert_eq!(result.dtype(), expected.dtype()); + assert_arrays_eq!(result, expected, &mut ctx); + let range = PiecewiseSequenceArray::try_new( + PrimitiveArray::from_iter([1u64]).into_array(), + PrimitiveArray::from_iter([32766u64]).into_array(), + PrimitiveArray::from_iter([1u64]).into_array(), + 32766, + )? + .into_array(); + assert_arrays_eq!(encoded.take(range)?, input.slice(1..32767)?, &mut ctx); + Ok(()) +} diff --git a/encodings/pco/src/take.rs b/encodings/pco/src/take.rs index bc3a77e8a20..81a17c5841b 100644 --- a/encodings/pco/src/take.rs +++ b/encodings/pco/src/take.rs @@ -23,6 +23,7 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::dict::TakeExecute; use vortex_array::arrays::piecewise_sequence::array::PiecewiseSequenceArraySlotsExt; use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability; use vortex_array::dtype::half; use vortex_array::match_each_integer_ptype; use vortex_array::scalar::Scalar; @@ -358,10 +359,20 @@ fn gather( array: ArrayView<'_, Pco>, runs: &[Run], output_len: usize, - nullability: vortex_array::dtype::Nullability, + nullability: Nullability, grouped: Option, ctx: &mut ExecutionCtx, ) -> VortexResult { + if let [ + Run { + start: Some(start), + length, + stride: 1, + }, + ] = runs + { + return stream_range::(array, *start..*start + *length, nullability, ctx); + } if array.page_index()?.pages.is_empty() { let mask = array .unsliced_validity() @@ -636,16 +647,23 @@ fn gather_typed( /// Contiguous selections stream pages into the final initialized values buffer and retain /// exact sliced source validity. No output validity copy or retained decoded-page cache. -fn stream_slice( +fn stream_range( array: ArrayView<'_, Pco>, + range: Range, + nullability: Nullability, ctx: &mut ExecutionCtx, ) -> VortexResult { - let mut values = BufferMut::::zeroed_in(array.len(), ctx.allocator().clone()); + vortex_ensure!( + range.start <= range.end && range.end <= array.len(), + "Pco contiguous range exceeds visible bounds" + ); + let global_start = array.slice_start() + range.start; + let mut values = BufferMut::::zeroed_in(range.len(), ctx.allocator().clone()); let mut scratch = BufferMut::::with_capacity_in(0, ctx.allocator().clone()); let mut state = AccessState::::default(); let mut decoder = PageDecoder::::new(array.metadata.header.as_slice())?; - let stop = array.slice_start() + array.len(); - let mut row = array.slice_start(); + let stop = array.slice_start() + range.end; + let mut row = global_start; while row < stop { let span = state .span(array, row, stop, ctx)? @@ -664,7 +682,7 @@ fn stream_slice( span.mask.true_count(), "Pco contiguous value count mismatch" ); - let start = row - array.slice_start(); + let start = row - global_start; let output = values .get_mut(start..start + span.mask.len()) .ok_or_else(|| vortex_err!("Pco contiguous output exceeds visible range"))?; @@ -715,7 +733,10 @@ fn stream_slice( } row = span.stop; } - let validity = array.unsliced_validity().slice(array.slice_start()..stop)?; + let validity = array + .unsliced_validity() + .slice(global_start..stop)? + .union_nullability(nullability); Ok(PrimitiveArray::new(values.freeze(), validity)) } @@ -724,6 +745,6 @@ pub(crate) fn slice( ctx: &mut ExecutionCtx, ) -> VortexResult { match_number_enum!(number_type_from_ptype(array.dtype().as_ptype()), NumberType => { - stream_slice::(array, ctx) + stream_range::(array, 0..array.len(), array.dtype().nullability(), ctx) }) } diff --git a/vortex-file/examples/pco_access_profile.rs b/vortex-file/examples/pco_access_profile.rs index 2547fc5e4c6..0572e4f451e 100644 --- a/vortex-file/examples/pco_access_profile.rs +++ b/vortex-file/examples/pco_access_profile.rs @@ -224,6 +224,23 @@ fn main() -> VortexResult<()> { 1000000, )? .into_array(); + let contiguous_indices = PrimitiveArray::from_iter( + (0..if mode == "take-contiguous" { + rows - 2 + } else { + 0 + }) + .map(|i| u64::try_from(i + 1).vortex_expect("index fits u64")), + ) + .into_array(); + let contiguous_range = PiecewiseSequenceArray::try_new( + PrimitiveArray::from_iter([1u64]).into_array(), + PrimitiveArray::from_iter([u64::try_from(rows - 2).vortex_expect("range fits u64")]) + .into_array(), + PrimitiveArray::from_iter([1u64]).into_array(), + rows - 2, + )? + .into_array(); let wire = session .array_serialize(&encoded)? .vortex_expect("prepared array serialization"); @@ -263,6 +280,18 @@ fn main() -> VortexResult<()> { .execute::(&mut ctx)? .into_array(), ), + "take-contiguous" => Some( + encoded + .take(contiguous_indices.clone())? + .execute::(&mut ctx)? + .into_array(), + ), + "take-range" => Some( + encoded + .take(contiguous_range.clone())? + .execute::(&mut ctx)? + .into_array(), + ), "take-mixed" => Some( encoded .take(mixed_indices.clone())? From fec4aedb5048b9b7357a742d012893a8cdf7161c Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 14:05:20 -0400 Subject: [PATCH 31/36] Bound individual software profiling and report extraction commands Signed-off-by: Daniel King --- .github/workflows/pco-attribution.yml | 31 ++++++++++++++++++--------- 1 file changed, 21 insertions(+), 10 deletions(-) diff --git a/.github/workflows/pco-attribution.yml b/.github/workflows/pco-attribution.yml index 34fb27f3f75..c128116d3e0 100644 --- a/.github/workflows/pco-attribution.yml +++ b/.github/workflows/pco-attribution.yml @@ -9,6 +9,10 @@ on: permissions: contents: read +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: false + jobs: profile: runs-on: >- @@ -29,7 +33,7 @@ jobs: enable-sccache: ${{ github.repository == 'vortex-data/vortex' && 'true' || 'false' }} - name: Build frozen adverse candidate and baseline run: | - git worktree add --detach "$RUNNER_TEMP/pco-profile-candidate" 9da3bb7f8fbfaad2dac0e2cd4541776bc9d31c4a + git worktree add --detach "$RUNNER_TEMP/pco-profile-candidate" bf5fecb6d3 git worktree add --detach "$RUNNER_TEMP/pco-profile-base" 910f9933e8e2318c592c1eddcc2dfee8d24c6cd7 mkdir -p "$RUNNER_TEMP/pco-profile-base/vortex-file/examples" cp vortex-file/examples/pco_access_profile.rs "$RUNNER_TEMP/pco-profile-candidate/vortex-file/examples/" @@ -43,7 +47,8 @@ jobs: run: | mkdir -p "$RUNNER_TEMP/attribution" for engine in base candidate; do - PCO_SIZE_DIAGNOSTICS=1 cargo test --manifest-path "$RUNNER_TEMP/pco-profile-$engine/Cargo.toml" --profile ci --locked -p vortex-file --features zstd --test pco_logical_pages_compat public_stock_size_attribution -- --nocapture > "$RUNNER_TEMP/attribution/$engine-stock-size.txt" + echo "Size diagnostic $engine" + timeout --kill-after=5s 180s env PCO_SIZE_DIAGNOSTICS=1 cargo test --manifest-path "$RUNNER_TEMP/pco-profile-$engine/Cargo.toml" --profile ci --locked -p vortex-file --features zstd --test pco_logical_pages_compat public_stock_size_attribution -- --nocapture > "$RUNNER_TEMP/attribution/$engine-stock-size.txt" 2>&1 || echo "Size diagnostic exit $?" >> "$RUNNER_TEMP/attribution/$engine-stock-size.txt" done sudo apt-get update -qq sudo apt-get install -y -qq linux-tools-generic @@ -53,16 +58,22 @@ jobs: executable=$(jq -r 'select(.reason == "compiler-artifact" and .target.name == "pco_access_profile") | .executable // empty' "$RUNNER_TEMP/$engine-build.jsonl") sha256sum "$executable" > "$RUNNER_TEMP/attribution/$engine-executable.sha256" "$executable" metadata dense 100000 1 > "$RUNNER_TEMP/attribution/$engine-metadata.json" - for mode in write probe take-large; do - repeats=64 - if [ "$mode" = probe ]; then repeats=100000; fi - sudo "$perf" record -e cpu-clock -F 499 --call-graph dwarf,8192 -o "$RUNNER_TEMP/attribution/$engine-$mode.data" -- "$executable" "$mode" nullable 8388608 "$repeats" > "$RUNNER_TEMP/$engine-$mode.jsonl" - sudo "$perf" report -i "$RUNNER_TEMP/attribution/$engine-$mode.data" --stdio --no-children --percent-limit 0.5 > "$RUNNER_TEMP/attribution/$engine-$mode-flat.txt" - sudo "$perf" report -i "$RUNNER_TEMP/attribution/$engine-$mode.data" --stdio --children --percent-limit 1 > "$RUNNER_TEMP/attribution/$engine-$mode-callgraph.txt" - sudo rm "$RUNNER_TEMP/attribution/$engine-$mode.data" + for request in write:nullable probe:nullable slice-near-full:dense slice-near-full:long-null; do + IFS=: read -r mode pattern <<< "$request" + repeats=32 + if [ "$mode" = probe ]; then repeats=10000; fi + if [ "$pattern" = long-null ]; then repeats=4096; fi + stem="$RUNNER_TEMP/attribution/$engine-$mode-$pattern" + echo "Record $engine $mode $pattern" + timeout --signal=INT --kill-after=5s 25s sudo env DEBUGINFOD_URLS= "$perf" record -e cpu-clock -F 499 --call-graph dwarf,8192 -o "$stem.data" -- "$executable" "$mode" "$pattern" 8388608 "$repeats" > "$RUNNER_TEMP/$engine-$mode-$pattern.jsonl" || echo "Record exit $?" > "$stem-status.txt" + echo "Flat report $engine $mode $pattern" + timeout --kill-after=5s 20s sudo env DEBUGINFOD_URLS= "$perf" report -i "$stem.data" --stdio --no-children --call-graph none --percent-limit 0.5 > "$stem-flat.txt" || echo "Flat report exit $?" >> "$stem-status.txt" + echo "Caller report $engine $mode $pattern" + timeout --kill-after=5s 30s sudo env DEBUGINFOD_URLS= "$perf" report -i "$stem.data" --stdio --children --call-graph graph,2,caller --percent-limit 2 > "$stem-callgraph.txt" || echo "Caller report exit $?" >> "$stem-status.txt" + sudo rm -f "$stem.data" done done - printf '%s\n' 'Frozen 9da adverse candidate and develop910f. Software samples only; profiled times are not uninstrumented latency controls. Startup/fixture compression included: identify query loops by stack.' > "$RUNNER_TEMP/attribution/README.txt" + printf '%s\n' 'Frozen bf5 adverse candidate and develop910f. Software samples only; profiled times are not uninstrumented latency controls. Startup/fixture compression included: identify query loops by stack.' > "$RUNNER_TEMP/attribution/README.txt" - uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7 if: always() with: From ace2734d0a4d602bfcc2724b951918739c1c3c5a Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 14:06:52 -0400 Subject: [PATCH 32/36] Borrow typed primitive indices through their erased array handle Signed-off-by: Daniel King --- encodings/pco/src/take.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/encodings/pco/src/take.rs b/encodings/pco/src/take.rs index 81a17c5841b..678798d1ed2 100644 --- a/encodings/pco/src/take.rs +++ b/encodings/pco/src/take.rs @@ -88,7 +88,7 @@ impl TakeExecute for Pco { let mask = indices.validity()?.execute_mask(indices.len(), ctx)?; // Reuse native logical-null-aware extrema and their immutable-array stats cache. // Numeric extrema prove every defined index is representable and in bounds. - let extrema = min_max(indices.array(), ctx, NumericalAggregateOpts::default())?; + let extrema = min_max(indices.as_array(), ctx, NumericalAggregateOpts::default())?; let summary = if let Some(extrema) = extrema { match_each_integer_ptype!(indices.ptype(), |I| { validate_extrema(indices.as_slice::(), &mask, &extrema, array.len())? From a703552564a269f128c3f9a6b9c19a6be1ffa9a8 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 14:17:41 -0400 Subject: [PATCH 33/36] Inline retained scalar hits and prepare validated all-null probes Signed-off-by: Daniel King --- encodings/pco/src/access.rs | 17 ++++++++++++----- encodings/pco/src/probe.rs | 26 +++++++++++++++++++++++--- encodings/pco/src/select_tests.rs | 22 ++++++++++++++++++++++ 3 files changed, 57 insertions(+), 8 deletions(-) diff --git a/encodings/pco/src/access.rs b/encodings/pco/src/access.rs index 589e87c60e2..bf3a488d172 100644 --- a/encodings/pco/src/access.rs +++ b/encodings/pco/src/access.rs @@ -32,6 +32,7 @@ const RANK_STRIDE: usize = 512; #[derive(Default)] pub(crate) struct AccessState { legacy: Option, + validated_all_null: bool, decoder: Option>, pages: BTreeMap, last_page: Option<(usize, CachedPage)>, @@ -117,6 +118,7 @@ impl RankMask { impl AccessState { /// Read the last fully validated decoded page directly. A miss preserves all /// operation-owned cache state and follows the ordinary fallible resolver. + #[inline] pub fn cached_scalar(&self, array: ArrayView<'_, Pco>, row: usize) -> Option { if row >= array.unsliced_n_rows() { return None; @@ -147,6 +149,10 @@ impl AccessState { }) } + pub fn validated_all_null(&self) -> bool { + self.validated_all_null + } + fn cached(&self, page: usize) -> Option<&CachedPage> { if let Some((last, cached)) = &self.last_page && *last == page @@ -263,11 +269,11 @@ impl AccessState { // Legacy or unknown optional metadata: execute/rank the full validity once per operation, // never once per requested index. Empty physical arrays also take this checked path. if self.legacy.is_none() { - self.legacy = Some(RankMask::new( - array - .unsliced_validity() - .execute_mask(array.unsliced_n_rows(), ctx)?, - )); + let mask = array + .unsliced_validity() + .execute_mask(array.unsliced_n_rows(), ctx)?; + self.validated_all_null = index.pages.is_empty() && mask.all_false(); + self.legacy = Some(RankMask::new(mask)); } let Some(value) = self .legacy @@ -315,6 +321,7 @@ impl AccessState { total, "Pco validity count does not match stored values" ); + self.validated_all_null = index.pages.is_empty() && mask.all_false(); self.legacy = Some(RankMask::new(mask)); } let mask = self diff --git a/encodings/pco/src/probe.rs b/encodings/pco/src/probe.rs index 398b988402c..9dc4594f514 100644 --- a/encodings/pco/src/probe.rs +++ b/encodings/pco/src/probe.rs @@ -8,6 +8,7 @@ use vortex_array::ExecutionCtx; use vortex_array::match_each_native_ptype; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use crate::Pco; use crate::access::AccessState; @@ -19,6 +20,7 @@ use crate::access::AccessState; #[derive(Default)] pub struct PcoProbeState { access: AccessState, + all_null: Option, } pub(crate) fn scalar_at( @@ -27,14 +29,32 @@ pub(crate) fn scalar_at( state: Option<&mut PcoProbeState>, ctx: &mut ExecutionCtx, ) -> VortexResult { - let mut once = PcoProbeState::default(); - let state = state.unwrap_or(&mut once); + // Retained probes must not construct/drop a fresh fallback cache on every hit. + let mut once; + let state = match state { + Some(state) => state, + None => { + once = PcoProbeState::default(); + &mut once + } + }; + if let Some(null) = &state.all_null { + vortex_ensure!( + index < array.len(), + "Pco scalar index exceeds visible array bounds" + ); + return Ok(null.clone()); + } let row = array.slice_start() + index; if let Some(value) = state.access.cached_scalar(array, row) { return Ok(value); } let Some((page, value)) = state.access.resolve(array, row, ctx)? else { - return Ok(Scalar::null(array.dtype().clone())); + let null = Scalar::null(array.dtype().clone()); + if state.access.validated_all_null() { + state.all_null = Some(null.clone()); + } + return Ok(null); }; let values = state.access.values(array, page, ctx)?; Ok(match_each_native_ptype!(values.ptype(), |T| { diff --git a/encodings/pco/src/select_tests.rs b/encodings/pco/src/select_tests.rs index 338c5639876..e9330ff978c 100644 --- a/encodings/pco/src/select_tests.rs +++ b/encodings/pco/src/select_tests.rs @@ -519,3 +519,25 @@ fn contiguous_take_preserves_nullable_index_dtype() -> VortexResult<()> { assert_arrays_eq!(encoded.take(range)?, input.slice(1..32767)?, &mut ctx); Ok(()) } + +#[test] +fn retained_all_null_probes_preserve_slice_bounds() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = PrimitiveArray::new(vec![0i64; 100000], Validity::AllInvalid); + let encoded = Pco::from_primitive(input.as_view(), 3, 1024, &mut ctx)?.into_array(); + let sliced = encoded.slice(17..145)?; + let mut probe = sliced.repeated_probe(); + for index in [0, 127, 64, 0, 127] { + assert_eq!( + probe.execute_scalar(index, &mut ctx)?, + input.execute_scalar(17 + index, &mut ctx)? + ); + } + assert!(probe.execute_scalar(128, &mut ctx).is_err()); + assert!(probe.execute_scalar(usize::MAX, &mut ctx).is_err()); + assert_eq!( + probe.execute_scalar(0, &mut ctx)?, + input.execute_scalar(17, &mut ctx)? + ); + Ok(()) +} From 9f468abe16dd070d2d8eed5df4f45709b50ae375 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 14:19:31 -0400 Subject: [PATCH 34/36] Compare linked scalar-hit assembly for base, adverse and current sources Signed-off-by: Daniel King --- .github/workflows/pco-attribution.yml | 31 +++++++++++++-------------- 1 file changed, 15 insertions(+), 16 deletions(-) diff --git a/.github/workflows/pco-attribution.yml b/.github/workflows/pco-attribution.yml index c128116d3e0..c36dadf98b7 100644 --- a/.github/workflows/pco-attribution.yml +++ b/.github/workflows/pco-attribution.yml @@ -35,10 +35,11 @@ jobs: run: | git worktree add --detach "$RUNNER_TEMP/pco-profile-candidate" bf5fecb6d3 git worktree add --detach "$RUNNER_TEMP/pco-profile-base" 910f9933e8e2318c592c1eddcc2dfee8d24c6cd7 + git worktree add --detach "$RUNNER_TEMP/pco-profile-latest" HEAD mkdir -p "$RUNNER_TEMP/pco-profile-base/vortex-file/examples" cp vortex-file/examples/pco_access_profile.rs "$RUNNER_TEMP/pco-profile-candidate/vortex-file/examples/" cp vortex-file/examples/pco_access_profile.rs "$RUNNER_TEMP/pco-profile-base/vortex-file/examples/" - for engine in candidate base; do + for engine in candidate base latest; do cp vortex-file/tests/pco_logical_pages_compat.rs "$RUNNER_TEMP/pco-profile-$engine/vortex-file/tests/" cargo build --manifest-path "$RUNNER_TEMP/pco-profile-$engine/Cargo.toml" --profile bench --locked -p vortex-file --example pco_access_profile --message-format=json-render-diagnostics > "$RUNNER_TEMP/$engine-build.jsonl" done @@ -46,37 +47,35 @@ jobs: shell: bash run: | mkdir -p "$RUNNER_TEMP/attribution" - for engine in base candidate; do - echo "Size diagnostic $engine" - timeout --kill-after=5s 180s env PCO_SIZE_DIAGNOSTICS=1 cargo test --manifest-path "$RUNNER_TEMP/pco-profile-$engine/Cargo.toml" --profile ci --locked -p vortex-file --features zstd --test pco_logical_pages_compat public_stock_size_attribution -- --nocapture > "$RUNNER_TEMP/attribution/$engine-stock-size.txt" 2>&1 || echo "Size diagnostic exit $?" >> "$RUNNER_TEMP/attribution/$engine-stock-size.txt" - done sudo apt-get update -qq sudo apt-get install -y -qq linux-tools-generic perf=$(find /usr/lib -path "*linux-tools*/perf" -type f | sort | tail -1) test -n "$perf" - for engine in base candidate; do + for engine in base candidate latest; do executable=$(jq -r 'select(.reason == "compiler-artifact" and .target.name == "pco_access_profile") | .executable // empty' "$RUNNER_TEMP/$engine-build.jsonl") sha256sum "$executable" > "$RUNNER_TEMP/attribution/$engine-executable.sha256" - "$executable" metadata dense 100000 1 > "$RUNNER_TEMP/attribution/$engine-metadata.json" - for request in write:nullable probe:nullable slice-near-full:dense slice-near-full:long-null; do + git -C "$RUNNER_TEMP/pco-profile-$engine" rev-parse HEAD > "$RUNNER_TEMP/attribution/$engine-source.txt" + nm "$executable" | awk '/vortex_pco/ && /cached_scalar|5probe9scalar_at/ {print $NF}' > "$RUNNER_TEMP/attribution/$engine-symbols.txt" + while IFS= read -r symbol; do + timeout --kill-after=5s 20s objdump --disassemble="$symbol" --no-show-raw-insn "$executable" >> "$RUNNER_TEMP/attribution/$engine-scalar-assembly.txt" + done < "$RUNNER_TEMP/attribution/$engine-symbols.txt" + for request in probe:dense probe:nullable probe:all-null; do IFS=: read -r mode pattern <<< "$request" - repeats=32 - if [ "$mode" = probe ]; then repeats=10000; fi - if [ "$pattern" = long-null ]; then repeats=4096; fi + repeats=100000 stem="$RUNNER_TEMP/attribution/$engine-$mode-$pattern" echo "Record $engine $mode $pattern" - timeout --signal=INT --kill-after=5s 25s sudo env DEBUGINFOD_URLS= "$perf" record -e cpu-clock -F 499 --call-graph dwarf,8192 -o "$stem.data" -- "$executable" "$mode" "$pattern" 8388608 "$repeats" > "$RUNNER_TEMP/$engine-$mode-$pattern.jsonl" || echo "Record exit $?" > "$stem-status.txt" + timeout --signal=INT --kill-after=5s 25s sudo env DEBUGINFOD_URLS= "$perf" record -e cpu-clock -F 499 -o "$stem.data" -- "$executable" "$mode" "$pattern" 8388608 "$repeats" > "$RUNNER_TEMP/$engine-$mode-$pattern.jsonl" || echo "Record exit $?" > "$stem-status.txt" echo "Flat report $engine $mode $pattern" timeout --kill-after=5s 20s sudo env DEBUGINFOD_URLS= "$perf" report -i "$stem.data" --stdio --no-children --call-graph none --percent-limit 0.5 > "$stem-flat.txt" || echo "Flat report exit $?" >> "$stem-status.txt" - echo "Caller report $engine $mode $pattern" - timeout --kill-after=5s 30s sudo env DEBUGINFOD_URLS= "$perf" report -i "$stem.data" --stdio --children --call-graph graph,2,caller --percent-limit 2 > "$stem-callgraph.txt" || echo "Caller report exit $?" >> "$stem-status.txt" sudo rm -f "$stem.data" done done - printf '%s\n' 'Frozen bf5 adverse candidate and develop910f. Software samples only; profiled times are not uninstrumented latency controls. Startup/fixture compression included: identify query loops by stack.' > "$RUNNER_TEMP/attribution/README.txt" + printf '%s\n' 'Frozen bf5 adverse candidate, develop910f, and current CI merge source. Software flat samples and exact linked symbol assembly only; profiled times are not uninstrumented latency controls. Startup/fixture compression is included. Public stock/metadata byte attribution is retained in run40119098145; it is not repeated here.' > "$RUNNER_TEMP/attribution/README.txt" - uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7 if: always() with: name: pco-attribution - path: ${{ runner.temp }}/attribution + path: | + ${{ runner.temp }}/attribution + !${{ runner.temp }}/attribution/*.data if-no-files-found: ignore From 5dbe6faf05f0401c211bfe8df0c6792d1964c262 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 14:31:47 -0400 Subject: [PATCH 35/36] Separate PCO retained probe hits from temporary and miss state Signed-off-by: Daniel King --- encodings/pco/src/probe.rs | 42 +++++++++++++++++++++++--------------- 1 file changed, 26 insertions(+), 16 deletions(-) diff --git a/encodings/pco/src/probe.rs b/encodings/pco/src/probe.rs index 9dc4594f514..8df0537c2d8 100644 --- a/encodings/pco/src/probe.rs +++ b/encodings/pco/src/probe.rs @@ -20,7 +20,6 @@ use crate::access::AccessState; #[derive(Default)] pub struct PcoProbeState { access: AccessState, - all_null: Option, } pub(crate) fn scalar_at( @@ -29,32 +28,43 @@ pub(crate) fn scalar_at( state: Option<&mut PcoProbeState>, ctx: &mut ExecutionCtx, ) -> VortexResult { - // Retained probes must not construct/drop a fresh fallback cache on every hit. - let mut once; - let state = match state { - Some(state) => state, - None => { - once = PcoProbeState::default(); - &mut once - } + let Some(state) = state else { + return scalar_at_once(array, index, ctx); }; - if let Some(null) = &state.all_null { + if state.access.validated_all_null() { vortex_ensure!( index < array.len(), "Pco scalar index exceeds visible array bounds" ); - return Ok(null.clone()); + return Ok(Scalar::null(array.dtype().clone())); } let row = array.slice_start() + index; if let Some(value) = state.access.cached_scalar(array, row) { return Ok(value); } + scalar_at_uncached(array, row, state, ctx) +} + +// Keep the large temporary cache out of the retained-probe hit frame. Stateless +// probes and cache misses are ordinary workloads, so neither helper is cold. +#[inline(never)] +fn scalar_at_once( + array: ArrayView<'_, Pco>, + index: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult { + scalar_at(array, index, Some(&mut PcoProbeState::default()), ctx) +} + +#[inline(never)] +fn scalar_at_uncached( + array: ArrayView<'_, Pco>, + row: usize, + state: &mut PcoProbeState, + ctx: &mut ExecutionCtx, +) -> VortexResult { let Some((page, value)) = state.access.resolve(array, row, ctx)? else { - let null = Scalar::null(array.dtype().clone()); - if state.access.validated_all_null() { - state.all_null = Some(null.clone()); - } - return Ok(null); + return Ok(Scalar::null(array.dtype().clone())); }; let values = state.access.values(array, page, ctx)?; Ok(match_each_native_ptype!(values.ptype(), |T| { From 58fd61e34d181553a60423defaebb381e9213b74 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 24 Oct 2026 14:33:13 -0400 Subject: [PATCH 36/36] Freeze final probe comparison against the previous iteration Signed-off-by: Daniel King --- .github/workflows/pco-attribution.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/pco-attribution.yml b/.github/workflows/pco-attribution.yml index c36dadf98b7..7e1e165439f 100644 --- a/.github/workflows/pco-attribution.yml +++ b/.github/workflows/pco-attribution.yml @@ -33,7 +33,7 @@ jobs: enable-sccache: ${{ github.repository == 'vortex-data/vortex' && 'true' || 'false' }} - name: Build frozen adverse candidate and baseline run: | - git worktree add --detach "$RUNNER_TEMP/pco-profile-candidate" bf5fecb6d3 + git worktree add --detach "$RUNNER_TEMP/pco-profile-candidate" 9f468abe16dd070d2d8eed5df4f45709b50ae375 git worktree add --detach "$RUNNER_TEMP/pco-profile-base" 910f9933e8e2318c592c1eddcc2dfee8d24c6cd7 git worktree add --detach "$RUNNER_TEMP/pco-profile-latest" HEAD mkdir -p "$RUNNER_TEMP/pco-profile-base/vortex-file/examples" @@ -55,7 +55,7 @@ jobs: executable=$(jq -r 'select(.reason == "compiler-artifact" and .target.name == "pco_access_profile") | .executable // empty' "$RUNNER_TEMP/$engine-build.jsonl") sha256sum "$executable" > "$RUNNER_TEMP/attribution/$engine-executable.sha256" git -C "$RUNNER_TEMP/pco-profile-$engine" rev-parse HEAD > "$RUNNER_TEMP/attribution/$engine-source.txt" - nm "$executable" | awk '/vortex_pco/ && /cached_scalar|5probe9scalar_at/ {print $NF}' > "$RUNNER_TEMP/attribution/$engine-symbols.txt" + nm "$executable" | awk '/vortex_pco/ && /cached_scalar|5probe9scalar_at|scalar_at_once|scalar_at_uncached/ {print $NF}' > "$RUNNER_TEMP/attribution/$engine-symbols.txt" while IFS= read -r symbol; do timeout --kill-after=5s 20s objdump --disassemble="$symbol" --no-show-raw-insn "$executable" >> "$RUNNER_TEMP/attribution/$engine-scalar-assembly.txt" done < "$RUNNER_TEMP/attribution/$engine-symbols.txt" @@ -70,7 +70,7 @@ jobs: sudo rm -f "$stem.data" done done - printf '%s\n' 'Frozen bf5 adverse candidate, develop910f, and current CI merge source. Software flat samples and exact linked symbol assembly only; profiled times are not uninstrumented latency controls. Startup/fixture compression is included. Public stock/metadata byte attribution is retained in run40119098145; it is not repeated here.' > "$RUNNER_TEMP/attribution/README.txt" + printf '%s\n' 'Frozen 9f adverse candidate, develop910f, and current CI merge source. Software flat samples and exact linked symbol assembly only; profiled times are not uninstrumented latency controls. Startup/fixture compression is included. Public stock/metadata byte attribution is retained in run40119098145; it is not repeated here.' > "$RUNNER_TEMP/attribution/README.txt" - uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7 if: always() with: