diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array.rs index 684155fa875..71358812cee 100644 --- a/encodings/pco/src/array.rs +++ b/encodings/pco/src/array.rs @@ -7,6 +7,7 @@ use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hash; use std::hash::Hasher; +use std::sync::Arc; use pco::ChunkConfig; use pco::PagingSpec; @@ -54,10 +55,12 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; +use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::PcoChunkInfo; +use crate::PcoLogicalRowPrefix; use crate::PcoMetadata; use crate::PcoPageInfo; @@ -89,11 +92,21 @@ 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_row_prefix { + index.logical_row_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 +119,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 +194,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 +214,7 @@ impl VTable for Pco { array: ArrayView<'_, Self>, _session: &VortexSession, ) -> VortexResult>> { - Ok(Some(array.metadata.clone().encode_to_vec())) + Ok(Some(array.metadata.encode_to_vec())) } fn deserialize( @@ -299,16 +315,16 @@ pub(crate) fn number_type_from_ptype(ptype: PType) -> NumberType { fn collect_valid( parray: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, -) -> VortexResult { +) -> VortexResult<(PrimitiveArray, 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 { @@ -387,9 +403,9 @@ 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, ptype: PType, unsliced_n_rows: usize, slice_start: usize, @@ -486,9 +502,9 @@ impl PcoData { len: usize, ) -> Self { Self { - chunk_metas, - pages, - metadata, + chunk_metas: chunk_metas.into(), + pages: pages.into(), + metadata: Arc::new(metadata), ptype, unsliced_n_rows: len, slice_start: 0, @@ -531,7 +547,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(); @@ -572,9 +588,49 @@ impl PcoData { chunk_infos.push(PcoChunkInfo { pages: page_infos }) } + let mut ends = Vec::with_capacity(page_buffers.len()); + match &mask { + Mask::AllFalse(_) => {} + 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)?); + } + } + 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)?); + } + } + } let metadata = PcoMetadata { header, chunks: chunk_infos, + logical_row_prefix: Some(PcoLogicalRowPrefix { + logical_row_ends: ends, + }), }; // SAFETY: The compressor produced matching chunk metadata and pages from `parray`, with // the same ptype, logical length, and valid-value count. @@ -658,7 +714,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 { @@ -787,6 +843,7 @@ impl OperationsVTable for Pco { #[cfg(test)] mod tests { + use std::sync::Arc; use std::sync::LazyLock; use vortex_array::IntoArray; @@ -849,14 +906,14 @@ 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(); - 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/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 982f6dffc34..c2c92eb42d1 100644 --- a/encodings/pco/src/lib.rs +++ b/encodings/pco/src/lib.rs @@ -12,13 +12,15 @@ //! To deserialize arrays manually, register the encoding in the array session: //! //! ```rust -//! use vortex_array::session::ArraySessionExt; //! use vortex_session::VortexSession; +//! use vortex_pco::Pco; //! //! let session = VortexSession::empty(); -//! session.arrays().register(vortex_pco::Pco); +//! use vortex_array::session::ArraySessionExt; +//! session.arrays().register(Pco); //! ``` +use prost::Message; mod array; mod compute; mod probe; @@ -27,7 +29,7 @@ mod slice; pub use array::*; -#[derive(Clone, prost::Message)] +#[derive(Clone, PartialEq, Eq, 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 +41,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, Eq, Message)] /// Metadata for one Pco chunk. pub struct PcoChunkInfo { /// Pages contained in this chunk. @@ -47,7 +49,7 @@ pub struct PcoChunkInfo { pub pages: Vec, } -#[derive(Clone, prost::Message)] +#[derive(Clone, PartialEq, Eq, 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 +60,19 @@ 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_row_prefix: Option, +} + +/// Optional cumulative logical-row boundaries in physical page order. +#[derive(Clone, PartialEq, Eq, Message)] +pub struct PcoLogicalRowPrefix { + /// 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 logical_row_ends: Vec, } #[cfg(test)] diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 7f73b2efd3e..bcf7fe1eec3 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -174,7 +174,7 @@ async fn test_eof_values() { feature = "zstd", case::compact( BtrBlocksCompressorBuilder::from_session(&SESSION).with_compact(), - 55_344 + 55_396 ) )] #[tokio::test]