From d8483ad09c2dfed27ec722e9615a2413f4cab5b2 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Mon, 26 Oct 2026 17:08:05 -0400 Subject: [PATCH 1/3] Execute selective primitive Take only for arrays with logical row prefixes Signed-off-by: Daniel King --- encodings/pco/benches/selection.rs | 3 +- encodings/pco/src/access.rs | 168 ++++++ encodings/pco/src/array.rs | 18 + encodings/pco/src/decoder.rs | 74 +++ encodings/pco/src/lib.rs | 23 + encodings/pco/src/page_index.rs | 125 +++++ encodings/pco/src/range_semantics_tests.rs | 4 +- encodings/pco/src/selection_contract_tests.rs | 8 +- .../pco/src/selection_semantics_tests.rs | 6 +- encodings/pco/src/take.rs | 490 ++++++++++++++++++ vortex-file/src/lib.rs | 4 +- 11 files changed, 913 insertions(+), 10 deletions(-) create mode 100644 encodings/pco/src/access.rs create mode 100644 encodings/pco/src/decoder.rs create mode 100644 encodings/pco/src/page_index.rs create mode 100644 encodings/pco/src/take.rs diff --git a/encodings/pco/benches/selection.rs b/encodings/pco/benches/selection.rs index 6154f959594..ab34b45eba2 100644 --- a/encodings/pco/benches/selection.rs +++ b/encodings/pco/benches/selection.rs @@ -26,6 +26,7 @@ use vortex_array::session::ArraySessionExt; use vortex_array::validity::Validity; use vortex_error::VortexExpect; use vortex_pco::Pco; +use vortex_pco::initialize as initialize_pco; use vortex_session::VortexSession; const ROWS: usize = 4_194_304; @@ -35,7 +36,7 @@ static SESSION: LazyLock = LazyLock::new(|| { let session = VortexSession::empty() .with::() .with::(); - session.arrays().register(Pco); + initialize_pco(&session); session }); diff --git a/encodings/pco/src/access.rs b/encodings/pco/src/access.rs new file mode 100644 index 00000000000..b569f2420f6 --- /dev/null +++ b/encodings/pco/src/access.rs @@ -0,0 +1,168 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; +use std::sync::OnceLock; + +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use crate::Pco; +use crate::PcoArrayExt; +use crate::page_index::IndexedView; + +const RANK_STRIDE: usize = 512; +pub(crate) struct RankMask { + mask: Mask, + ranks: OnceLock>, +} + +impl RankMask { + fn new(mask: Mask) -> Self { + Self { + mask, + ranks: OnceLock::new(), + } + } + + fn boundary_rank(&self, row: usize) -> usize { + match &self.mask { + Mask::AllTrue(_) => row, + Mask::AllFalse(_) => 0, + Mask::Values(values) => { + let bits = values.bit_buffer(); + if row <= bits.len() - row { + bits.count_range(0, row) + } else { + self.mask.true_count() - bits.count_range(row, bits.len()) + } + } + } + } + + fn rank(&self, row: usize) -> usize { + match &self.mask { + Mask::AllTrue(_) => row, + Mask::AllFalse(_) => 0, + Mask::Values(values) => { + if row == self.mask.len() { + return self.mask.true_count(); + } + let ranks = self.ranks.get_or_init(|| { + let bits = values.bit_buffer(); + let mut ranks = Vec::with_capacity(bits.len().div_ceil(RANK_STRIDE)); + let mut count = 0; + for start in (0..bits.len()).step_by(RANK_STRIDE) { + ranks.push(count); + count += bits.count_range(start, (start + RANK_STRIDE).min(bits.len())); + } + ranks + }); + let block = row / RANK_STRIDE; + ranks[block] + values.bit_buffer().count_range(block * RANK_STRIDE, row) + } + } + } + + fn value_index(&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(|| { + if INDEXED { + self.rank(row) + } else { + self.boundary_rank(row) + } + }) + } + } + } +} + +/// Only new-prefix arrays use this operation-owned last-page validity preparation. +#[derive(Default)] +pub(crate) struct AccessState { + last: Option<(usize, PreparedPage)>, +} +struct PreparedPage { + rows: Range, + mask: RankMask, +} +impl AccessState { + fn page( + &mut self, + view: IndexedView<'_>, + array: ArrayView<'_, Pco>, + row: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult<(usize, &PreparedPage)> { + if !self + .last + .as_ref() + .is_some_and(|(_, cached)| cached.rows.contains(&row)) + { + let page = view.find_row(row); + let rows = view.rows(page)?; + let mask = array + .unsliced_validity() + .slice(rows.clone())? + .execute_mask(rows.len(), ctx)?; + self.last = Some(( + page, + PreparedPage { + rows, + mask: RankMask::new(mask), + }, + )); + } + let (page, cached) = self.last.as_ref().vortex_expect("page prepared above"); + Ok((*page, cached)) + } + pub fn resolve( + &mut self, + view: IndexedView<'_>, + array: ArrayView<'_, Pco>, + row: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let (page, cached) = self.page(view, array, row, ctx)?; + Ok(cached + .mask + .value_index::(row - cached.rows.start) + .map(|value| (page, value))) + } + pub fn span( + &mut self, + view: IndexedView<'_>, + array: ArrayView<'_, Pco>, + row: usize, + stop: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let (page, cached) = self.page(view, array, row, ctx)?; + let end = stop.min(cached.rows.end); + let start_in_page = row - cached.rows.start; + let end_in_page = end - cached.rows.start; + Ok(PageSpan { + page, + stop: end, + values: cached.mask.boundary_rank(start_in_page) + ..cached.mask.boundary_rank(end_in_page), + mask: cached.mask.mask.slice(start_in_page..end_in_page), + }) + } + +} +pub(crate) struct PageSpan { + pub page: usize, + pub stop: usize, + /// Dense-value positions within this physical page, never global value ordinals. + pub values: Range, + pub mask: Mask, +} diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array.rs index 71358812cee..9f911a743f3 100644 --- a/encodings/pco/src/array.rs +++ b/encodings/pco/src/array.rs @@ -8,6 +8,7 @@ use std::fmt::Formatter; use std::hash::Hash; use std::hash::Hasher; use std::sync::Arc; +use std::sync::OnceLock; use pco::ChunkConfig; use pco::PagingSpec; @@ -50,6 +51,7 @@ use vortex_array::vtable::validity_to_child; use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; use vortex_error::VortexError; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -63,6 +65,7 @@ use crate::PcoChunkInfo; use crate::PcoLogicalRowPrefix; use crate::PcoMetadata; use crate::PcoPageInfo; +use crate::page_index::PageIndex; // Overall approach here: // Chunk the array into Pco chunks (currently using the default recommended size @@ -406,6 +409,7 @@ pub struct PcoData { pub(crate) chunk_metas: Arc<[ByteBuffer]>, pub(crate) pages: Arc<[ByteBuffer]>, pub(crate) metadata: Arc, + page_index: Arc>, ptype: PType, unsliced_n_rows: usize, slice_start: usize, @@ -505,6 +509,7 @@ impl PcoData { chunk_metas: chunk_metas.into(), pages: pages.into(), metadata: Arc::new(metadata), + page_index: Arc::new(OnceLock::new()), ptype, unsliced_n_rows: len, slice_start: 0, @@ -781,6 +786,19 @@ impl PcoData { .into_byte_buffer()) } + pub(crate) fn page_index(&self) -> VortexResult<&PageIndex> { + if self.page_index.get().is_none() { + let index = PageIndex::try_new(Arc::clone(&self.metadata))?; + // Concurrent first readers may prepare equivalent immutable descriptors. Only one + // index is retained, and no decoded values or validity are cached in the array. + drop(self.page_index.set(index)); + } + Ok(self + .page_index + .get() + .vortex_expect("page index was initialized")) + } + pub(crate) fn _slice(&self, start: usize, stop: usize) -> Self { PcoData { slice_start: self.slice_start + start, diff --git a/encodings/pco/src/decoder.rs b/encodings/pco/src/decoder.rs new file mode 100644 index 00000000000..2cc836c9269 --- /dev/null +++ b/encodings/pco/src/decoder.rs @@ -0,0 +1,74 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! File and chunk decoding setup amortized over all selected pages in an operation. + +use std::collections::BTreeMap; +use std::collections::btree_map::Entry; + +use pco::data_types::Number; +use pco::wrapped::ChunkDecompressor; +use pco::wrapped::FileDecompressor; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::NativePType; +use vortex_array::validity::Validity; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; + +use crate::array::vortex_err_from_pco; + +pub(crate) struct PageDecoder { + 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 { + 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 into initialized storage and propagate codec errors. No additional semantic + /// completion check is needed to keep a partial/corrupt result memory-safe. + pub fn read_into( + &mut self, + chunk: usize, + chunk_meta: &[u8], + n_values: usize, + buffer: &[u8], + output: &mut [T], + ) -> VortexResult<()> { + let chunk = match self.chunks.entry(chunk) { + Entry::Occupied(entry) => entry.into_mut(), + Entry::Vacant(entry) => { + let (decoder, _) = self + .file + .chunk_decompressor::(chunk_meta) + .map_err(vortex_err_from_pco)?; + entry.insert(decoder) + } + }; + let mut decoder = chunk + .page_decompressor(buffer, n_values) + .map_err(vortex_err_from_pco)?; + decoder.read(output).map_err(vortex_err_from_pco)?; + Ok(()) + } +} + diff --git a/encodings/pco/src/lib.rs b/encodings/pco/src/lib.rs index c2c92eb42d1..debd67ce757 100644 --- a/encodings/pco/src/lib.rs +++ b/encodings/pco/src/lib.rs @@ -21,11 +21,22 @@ //! ``` use prost::Message; +use vortex_array::ArrayVTable; +use vortex_array::arrays::Dict; +use vortex_array::arrays::dict::TakeExecuteAdaptor; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::session::ArraySessionExt; +use vortex_session::VortexSession; + +mod access; mod array; mod compute; +mod decoder; +mod page_index; mod probe; mod rules; mod slice; +mod take; pub use array::*; @@ -75,12 +86,24 @@ pub struct PcoLogicalRowPrefix { pub logical_row_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: &VortexSession) { + session.arrays().register(Pco); + session + .kernels() + .register_execute_parent_kernel(Dict.id(), Pco, TakeExecuteAdaptor(Pco)); +} + #[cfg(test)] mod tests; #[cfg(test)] mod selection_contract_tests; + #[cfg(test)] mod selection_semantics_tests; diff --git a/encodings/pco/src/page_index.rs b/encodings/pco/src/page_index.rs new file mode 100644 index 00000000000..d9f6edf6f6a --- /dev/null +++ b/encodings/pco/src/page_index.rs @@ -0,0 +1,125 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; +use std::sync::Arc; + +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::PcoMetadata; + +#[derive(Debug)] +pub(crate) enum PageIndex { + Legacy, + IndexedFast(Arc), +} + +impl PageIndex { + pub fn try_new(metadata: Arc) -> VortexResult { + if metadata.logical_row_prefix.is_none() { + return Ok(Self::Legacy); + } + let mut chunk_ends = Vec::with_capacity(metadata.chunks.len()); + let mut end = 0usize; + for chunk in &metadata.chunks { + end = end + .checked_add(chunk.pages.len()) + .ok_or_else(|| vortex_err!("Pco page count overflows usize"))?; + chunk_ends.push(end); + } + Ok(Self::IndexedFast(Arc::new(IndexedIndex { + metadata, + chunk_ends, + }))) + } +} + +/// Shared metadata owner with only an O(chunks) page-to-chunk map; no dense-value prefixes. +#[derive(Debug)] +pub(crate) struct IndexedIndex { + metadata: Arc, + chunk_ends: Vec, +} +impl IndexedIndex { + pub fn view(&self) -> IndexedView<'_> { + IndexedView { index: self } + } +} + +#[derive(Clone, Copy)] +pub(crate) struct IndexedView<'a> { + index: &'a IndexedIndex, +} +impl<'a> IndexedView<'a> { + pub fn page_count(self) -> usize { + self.index.chunk_ends.last().copied().unwrap_or(0) + } + pub fn find_row(self, row: usize) -> usize { + self.ends().partition_point(|end| *end <= row as u64) + } + fn ends(self) -> &'a [u64] { + &self + .index + .metadata + .logical_row_prefix + .as_ref() + .vortex_expect("indexed owner has a logical-row prefix") + .logical_row_ends + } + /// Unsliced logical rows, including nulls. Inter-page nulls belong to the following page; + /// the final page includes trailing nulls. These coordinates are never slice-relative. + pub fn rows(self, page: usize) -> VortexResult> { + let ends = self.ends(); + let start = if page == 0 { + 0 + } else { + usize::try_from(ends[page - 1])? + }; + Ok(start..usize::try_from(ends[page])?) + } + pub fn physical(self, page: usize, cursor: &mut ChunkCursor) -> VortexResult { + if !cursor.pages.contains(&page) { + let chunk = self.index.chunk_ends.partition_point(|end| *end <= page); + let end = *self + .index + .chunk_ends + .get(chunk) + .ok_or_else(|| vortex_err!("Pco physical page is missing"))?; + let start = if chunk == 0 { + 0 + } else { + self.index.chunk_ends[chunk - 1] + }; + *cursor = ChunkCursor { + chunk, + pages: start..end, + }; + } + let info = self + .index + .metadata + .chunks + .get(cursor.chunk) + .and_then(|chunk| chunk.pages.get(page - cursor.pages.start)) + .ok_or_else(|| vortex_err!("Pco page metadata is missing"))?; + Ok(PhysicalPage { + chunk: cursor.chunk, + n_values: info.n_values as usize, + }) + } +} + +/// Selected physical-page decode parameters, not logical row or compressed byte offsets. +pub(crate) struct PhysicalPage { + /// Zero-based position in metadata.chunks and chunk_metas. + pub chunk: usize, + /// Number of initialized primitive output slots required by this page. + pub n_values: usize, +} +#[derive(Default)] +pub(crate) struct ChunkCursor { + chunk: usize, + pages: Range, +} diff --git a/encodings/pco/src/range_semantics_tests.rs b/encodings/pco/src/range_semantics_tests.rs index d482f2b7474..fa232e7a71d 100644 --- a/encodings/pco/src/range_semantics_tests.rs +++ b/encodings/pco/src/range_semantics_tests.rs @@ -12,17 +12,17 @@ use vortex_array::VortexSessionExecute; use vortex_array::arrays::PiecewiseSequenceArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; -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::initialize; static SESSION: LazyLock = LazyLock::new(|| { let session = VortexSession::empty(); - session.arrays().register(Pco); + initialize(&session); session }); diff --git a/encodings/pco/src/selection_contract_tests.rs b/encodings/pco/src/selection_contract_tests.rs index bba142d0672..b87a70d2af5 100644 --- a/encodings/pco/src/selection_contract_tests.rs +++ b/encodings/pco/src/selection_contract_tests.rs @@ -9,17 +9,17 @@ use vortex_array::VortexSessionExecute; use vortex_array::arrays::PiecewiseSequenceArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; -use vortex_array::session::ArraySessionExt; use vortex_array::validity::Validity; use vortex_error::VortexResult; use vortex_session::VortexSession; use crate::Pco; +use crate::initialize; #[rstest] fn selected_values_and_nulls(#[values(false, true)] nullable: bool) -> VortexResult<()> { let session = VortexSession::empty(); - session.arrays().register(Pco); + initialize(&session); let mut ctx = session.create_execution_ctx(); let input = PrimitiveArray::new( (0..4096i64).map(|i| i * 19).collect::>(), @@ -82,6 +82,7 @@ fn compact_ranges( #[case] strides: Vec, ) -> VortexResult<()> { let session = VortexSession::empty(); + initialize(&session); let mut ctx = session.create_execution_ctx(); let input = PrimitiveArray::from_option_iter((0..256i64).map(|i| (i % 5 != 0).then_some(i * i))); @@ -114,6 +115,7 @@ fn compact_ranges( #[rstest] fn null_runs_and_empty(#[values(false, true)] all_null: bool) -> VortexResult<()> { let session = VortexSession::empty(); + initialize(&session); let mut ctx = session.create_execution_ctx(); let input = PrimitiveArray::from_option_iter( (0..65536i32).map(|i| (!all_null && !(100..65000).contains(&i)).then_some(i)), @@ -147,6 +149,7 @@ fn null_runs_and_empty(#[values(false, true)] all_null: bool) -> VortexResult<() #[test] fn invalid_indices_under_nulls_are_ignored() -> VortexResult<()> { let session = VortexSession::empty(); + initialize(&session); let mut ctx = session.create_execution_ctx(); let input = PrimitiveArray::from_iter([11i32, 22, 33]); let compressed = Pco::from_primitive(input.as_view(), 3, 2, &mut ctx)?.into_array(); @@ -166,6 +169,7 @@ fn invalid_indices_under_nulls_are_ignored() -> VortexResult<()> { #[test] fn nonnullable_selection_crosses_chunks() -> VortexResult<()> { let session = VortexSession::empty(); + initialize(&session); let mut ctx = session.create_execution_ctx(); let input = PrimitiveArray::from_iter((0..300000i64).map(|i| i * 3 + 1)); let compressed = Pco::from_primitive(input.as_view(), 3, 127, &mut ctx)?.into_array(); diff --git a/encodings/pco/src/selection_semantics_tests.rs b/encodings/pco/src/selection_semantics_tests.rs index 159c52559d3..d9382352d89 100644 --- a/encodings/pco/src/selection_semantics_tests.rs +++ b/encodings/pco/src/selection_semantics_tests.rs @@ -12,17 +12,17 @@ use vortex_array::VortexSessionExecute; use vortex_array::arrays::PiecewiseSequenceArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; -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::initialize; static SESSION: LazyLock = LazyLock::new(|| { let session = VortexSession::empty(); - session.arrays().register(Pco); + initialize(&session); session }); @@ -123,7 +123,7 @@ fn all_null_and_long_null_runs() -> VortexResult<()> { #[test] fn registry_only_session_preserves_take_semantics() -> VortexResult<()> { let session = VortexSession::empty(); - session.arrays().register(Pco); + initialize(&session); 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(); diff --git a/encodings/pco/src/take.rs b/encodings/pco/src/take.rs new file mode 100644 index 00000000000..32a7a41c242 --- /dev/null +++ b/encodings/pco/src/take.rs @@ -0,0 +1,490 @@ +// 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 std::collections::BTreeSet; +use std::ops::Range; + +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::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; +use vortex_array::arrays::dict::TakeExecute; +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; +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_err; +use vortex_mask::Mask; + +use crate::Pco; +use crate::PcoArrayExt; +use crate::access::AccessState; +use crate::array::number_type_from_ptype; +use crate::decoder::PageDecoder; +use crate::page_index::ChunkCursor; +use crate::page_index::PageIndex; +use crate::page_index::IndexedView; + +/// 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> { + if array.metadata.logical_row_prefix.is_none() { + return Ok(None); + } + if indices.is_empty() { + let dtype = array + .dtype() + .clone() + .union_nullability(indices.dtype().nullability()); + return Ok(Some(Canonical::empty(&dtype).into_array())); + } + // Compact range execution is introduced separately; preserve the normal fallback. + if indices.is::() { + return Ok(None); + } + let grouped; + let runs = { + let indices = indices.clone().execute::(ctx)?; + 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.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())? + }) + } else { + PlainSummary { + defined: 0, + contiguous: false, + } + }; + 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).then_some(summary.defined); + if grouped.is_some() && descriptor_bytes(summary.defined) >= canonical_bytes(array) { + return Ok(None); + } + if grouped.is_some() + && match_each_integer_ptype!(indices.ptype(), |I| { + prefer_canonical_primitive( + array, + view, + indices.as_slice::(), + &mask, + summary.defined, + )? + }) + { + return Ok(None); + } + match_each_integer_ptype!(indices.ptype(), |I| { + primitive_runs(indices.as_slice::(), &mask)? + }) + }; + let nullable = array.dtype().nullability() | indices.dtype().nullability(); + let result = match_number_enum!(number_type_from_ptype(array.dtype().as_ptype()), NumberType => { + gather::(array, view, &runs, indices.len(), nullable, grouped, ctx)? + }); + Ok(Some(result.into_array())) + } +} + +struct PlainSummary { + defined: usize, + contiguous: bool, +} + +fn validate_extrema( + indices: &[I], + mask: &Mask, + extrema: &MinMaxResult, + source_len: usize, +) -> VortexResult +where + usize: TryFrom, +{ + 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: mask.true_count(), + contiguous, + }) +} + +fn zero_page_take( + array: ArrayView<'_, Pco>, + len: usize, + _ctx: &mut ExecutionCtx, +) -> VortexResult> { + if !array.pages.is_empty() { + return Ok(None); + } + 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 { + // Request-planning cost heuristic, not a heap bound: Run/sorting/resolved descriptors. + // Null runs, vector capacity, tree nodes, masks, ranks and decoder state are excluded. + // Output storage is common to both strategies. + requests.saturating_mul( + size_of::() + size_of::<(usize, usize)>() + 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>, + view: IndexedView<'_>, + indices: &[I], + mask: &Mask, + requests: usize, +) -> VortexResult +where + usize: TryFrom, +{ + let mut chunks = ChunkCursor::default(); + 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 (&row, valid) in indices.iter().zip(mask.iter()) { + if !valid { + continue; + } + let row = array.slice_start() + checked_index(row)?; + let page = view.find_row(row); + if touched.insert(page) { + selected = selected.saturating_add(view.physical(page, &mut chunks)?.n_values); + if setup.saturating_add(selected.saturating_mul(width)) >= full_bytes { + return Ok(true); + } + } + } + Ok(false) +} + +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 gather( + array: ArrayView<'_, Pco>, + view: IndexedView<'_>, + runs: &[Run], + output_len: usize, + nullability: Nullability, + grouped: Option, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if let [ + Run { + start: Some(start), + length, + stride: 1, + }, + ] = runs + { + return stream_range::(array, view, *start..*start + *length, nullability, ctx); + } + let defined = grouped.unwrap_or(output_len); + if nullability.is_nullable() { + gather_grouped::(array, view, runs, output_len, defined, ctx) + } else { + gather_grouped::(array, view, runs, output_len, defined, 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>, + view: IndexedView<'_>, + runs: &[Run], + output_len: usize, + defined: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mut requests = Vec::with_capacity(defined); + 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::(view, 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 }); + // Requests are already in page order. Release all validity/rank preparation before + // decoding, and keep only the current decoded page while scattering each group. + drop(state); + if !resolved.is_empty() { + let mut decoder = PageDecoder::::new(array.metadata.header.as_slice())?; + let mut chunks = ChunkCursor::default(); + 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 PageIndex::IndexedFast(index) = array.page_index()? else { return Ok(None); }; + let view = index.view(); + 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 info = view.physical(page, &mut chunks)?; + let buffer = array + .pages + .get(page) + .ok_or_else(|| vortex_err!("Pco compressed page is missing"))?; + let chunk = array + .chunk_metas + .get(info.chunk) + .ok_or_else(|| vortex_err!("Pco chunk metadata is missing"))?; + let decoded = decoder.read( + info.chunk, + chunk.as_slice(), + info.n_values, + buffer.as_slice(), + ctx, + )?; + let selected = decoded + .into_array() + .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. +/// 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_range( + array: ArrayView<'_, Pco>, + view: IndexedView<'_>, + range: Range, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + 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()); + if view.page_count() == 0 { + let validity = array + .unsliced_validity() + .slice(global_start..global_start + range.len())? + .union_nullability(nullability); + return Ok(PrimitiveArray::new(values.freeze(), validity)); + } + let mut state = AccessState::default(); + let mut decoder = PageDecoder::::new(array.metadata.header.as_slice())?; + let mut chunks = ChunkCursor::default(); + let stop = array.slice_start() + range.end; + let mut row = global_start; + // Descriptor-count iteration bounds work even when optional index contents are bad. + // Valid files visit at most one span per physical page; an all-null span needs one turn. + for _ in 0..=view.page_count() { + if row >= stop { + break; + } + let span = state.span(view, array, row, stop, ctx)?; + 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"))?; + if !span.mask.all_false() { + let page = view.physical(span.page, &mut chunks)?; + 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.n_values; + 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(global_start..stop)? + .union_nullability(nullability); + Ok(PrimitiveArray::new(values.freeze(), validity)) +} diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs index b46c766403c..7d0933cddf0 100644 --- a/vortex-file/src/lib.rs +++ b/vortex-file/src/lib.rs @@ -129,7 +129,7 @@ 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_pco::initialize as initialize_pco; use vortex_session::VortexSession; pub use writer::*; @@ -178,7 +178,7 @@ pub fn register_default_encodings(session: &VortexSession) { { let arrays = session.arrays(); - arrays.register(Pco); + initialize_pco(session); if use_experimental_patches() { arrays.register(Patched); } From c0cbdcf4724ae6966c85cfc5bea0d7c532b50508 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Mon, 26 Oct 2026 17:09:01 -0400 Subject: [PATCH 2/3] Apply pinned upstream Rust formatting Signed-off-by: Daniel King --- encodings/pco/src/access.rs | 1 - encodings/pco/src/decoder.rs | 1 - encodings/pco/src/lib.rs | 1 - encodings/pco/src/take.rs | 10 ++++++---- 4 files changed, 6 insertions(+), 7 deletions(-) diff --git a/encodings/pco/src/access.rs b/encodings/pco/src/access.rs index b569f2420f6..78fde096525 100644 --- a/encodings/pco/src/access.rs +++ b/encodings/pco/src/access.rs @@ -157,7 +157,6 @@ impl AccessState { mask: cached.mask.mask.slice(start_in_page..end_in_page), }) } - } pub(crate) struct PageSpan { pub page: usize, diff --git a/encodings/pco/src/decoder.rs b/encodings/pco/src/decoder.rs index 2cc836c9269..555a6927de4 100644 --- a/encodings/pco/src/decoder.rs +++ b/encodings/pco/src/decoder.rs @@ -71,4 +71,3 @@ impl PageDecoder { Ok(()) } } - diff --git a/encodings/pco/src/lib.rs b/encodings/pco/src/lib.rs index debd67ce757..51dfcf55223 100644 --- a/encodings/pco/src/lib.rs +++ b/encodings/pco/src/lib.rs @@ -103,7 +103,6 @@ mod tests; #[cfg(test)] mod selection_contract_tests; - #[cfg(test)] mod selection_semantics_tests; diff --git a/encodings/pco/src/take.rs b/encodings/pco/src/take.rs index 32a7a41c242..a87183db33d 100644 --- a/encodings/pco/src/take.rs +++ b/encodings/pco/src/take.rs @@ -40,8 +40,8 @@ use crate::access::AccessState; use crate::array::number_type_from_ptype; use crate::decoder::PageDecoder; use crate::page_index::ChunkCursor; -use crate::page_index::PageIndex; use crate::page_index::IndexedView; +use crate::page_index::PageIndex; /// Compact descriptor; null indices have no start. Zero-length runs consume no source rows. struct Run { @@ -351,9 +351,11 @@ fn gather_grouped( while cursor < resolved.len() { let page = resolved[cursor].0; let end = cursor + resolved[cursor..].partition_point(|request| request.0 == page); - let PageIndex::IndexedFast(index) = array.page_index()? else { return Ok(None); }; - let view = index.view(); - let compact = PrimitiveArray::from_iter( + let PageIndex::IndexedFast(index) = array.page_index()? else { + return Ok(None); + }; + let view = index.view(); + let compact = PrimitiveArray::from_iter( resolved[cursor..end] .iter() .map(|request| { From eea0ae271ec71941ed0908824c2c11f69db169e3 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Mon, 26 Oct 2026 17:14:21 -0400 Subject: [PATCH 3/3] Bind the indexed view at the primitive Take entry point Signed-off-by: Daniel King --- encodings/pco/src/take.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/encodings/pco/src/take.rs b/encodings/pco/src/take.rs index a87183db33d..e26a7f869df 100644 --- a/encodings/pco/src/take.rs +++ b/encodings/pco/src/take.rs @@ -70,6 +70,10 @@ impl TakeExecute for Pco { if indices.is::() { return Ok(None); } + let PageIndex::IndexedFast(index) = array.page_index()? else { + return Ok(None); + }; + let view = index.view(); let grouped; let runs = { let indices = indices.clone().execute::(ctx)?; @@ -351,10 +355,6 @@ fn gather_grouped( while cursor < resolved.len() { let page = resolved[cursor].0; let end = cursor + resolved[cursor..].partition_point(|request| request.0 == page); - let PageIndex::IndexedFast(index) = array.page_index()? else { - return Ok(None); - }; - let view = index.view(); let compact = PrimitiveArray::from_iter( resolved[cursor..end] .iter()