diff --git a/Cargo.lock b/Cargo.lock index 159b3b020d4..f7c9dcaa4f8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7026,9 +7026,9 @@ checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" [[package]] name = "onpair" -version = "0.2.0" +version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e546e7cd983c9998a5a52e9c7eecaf38e518f2e6fc9c9ec649e8b53748a96bb0" +checksum = "34272d1ea63266148a9a5b2b96db80a45245d034862302735f2f93413a6f9172" dependencies = [ "hashbrown 0.16.1", "rand 0.9.5", diff --git a/Cargo.toml b/Cargo.toml index 3c0368e190e..7a7c94aee79 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -203,7 +203,7 @@ num_enum = { version = "0.7.3", default-features = false } object_store = { version = "0.13.2", default-features = false } object_store_opendal = "0.58.0" once_cell = "1.21.4" -onpair = "0.2.0" +onpair = "0.2.1" opendal = { version = "0.58.1", default-features = false } opentelemetry = "0.32.0" opentelemetry-otlp = "0.32.0" diff --git a/encodings/onpair/src/compress.rs b/encodings/onpair/src/compress.rs index e2cf69353ec..1b0f4c68d18 100644 --- a/encodings/onpair/src/compress.rs +++ b/encodings/onpair/src/compress.rs @@ -4,6 +4,7 @@ //! Train + compress entry points for the OnPair encoding. use onpair::Config; +use onpair::Rows; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -19,7 +20,6 @@ use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_err; use vortex_mask::AllOr; use crate::OnPair; @@ -43,14 +43,8 @@ pub fn onpair_compress( } let views = array.views(); - let flat_bytes: usize = views.iter().map(|v| v.len() as usize).sum(); - - // TODO(francesco): we flatten because onpair training needs a contiguous `(bytes, offsets)` - // pair. Allowing onpair to train on a slice-of-slices would let us skip this copy. - let mut flat: Vec = Vec::with_capacity(flat_bytes); - let mut offsets: Vec = Vec::with_capacity(len + 1); - let mut uncompressed_lengths: BufferMut = BufferMut::with_capacity(len); - offsets.push(0); + let mut uncompressed_lengths: BufferMut = BufferMut::zeroed(len); + let mut total_bytes = 0usize; let buffers = array .data_buffers() .as_ref() @@ -60,33 +54,33 @@ pub fn onpair_compress( match mask.bit_buffer() { AllOr::All => { - for view in views { - let bytes = view_bytes(view, &buffers); - flat.extend_from_slice(bytes); - offsets.push(u64::try_from(flat.len()).vortex_expect("offset must fit in u64")); - uncompressed_lengths - .push(i32::try_from(view.len()).vortex_expect("must fit in i32")); + for (view, length) in views.iter().zip(uncompressed_lengths.iter_mut()) { + *length = view.len(); + total_bytes += *length as usize; } } AllOr::None => unreachable!("all-null input handled above"), AllOr::Some(validity) => { - for (view, valid) in views.iter().zip(validity.iter()) { + for ((view, length), valid) in views + .iter() + .zip(uncompressed_lengths.iter_mut()) + .zip(validity.iter()) + { if valid { - let bytes = view_bytes(view, &buffers); - flat.extend_from_slice(bytes); - offsets.push(u64::try_from(flat.len()).vortex_expect("offset must fit in u64")); - uncompressed_lengths - .push(i32::try_from(view.len()).vortex_expect("must fit in i32")); - } else { - offsets.push(u64::try_from(flat.len()).vortex_expect("offset must fit in u64")); - uncompressed_lengths.push(0); + *length = view.len(); + total_bytes += *length as usize; } } } } - let column = onpair::compress(&flat, &offsets, config) - .map_err(|e| vortex_err!("OnPair compress failed: {e}"))?; + let rows = ViewRows { + views, + buffers: &buffers, + lengths: uncompressed_lengths.as_slice(), + total_bytes, + }; + let column = onpair::compress_rows::<_, u64>(&rows, config); let (dict, codes, row_offsets) = column.into_raw(); let (dict_bytes, dict_offsets) = dict.into_raw(); let codes_offsets = codes_offsets_array(&row_offsets); @@ -113,12 +107,34 @@ pub fn onpair_compress( Ok(encoded.into_array()) } -fn view_bytes<'a>(view: &'a BinaryView, buffers: &'a [&ByteBuffer]) -> &'a [u8] { - if view.is_inlined() { - view.as_inlined().value() - } else { - let view_ref = view.as_view(); - &buffers[view_ref.buffer_index as usize][view_ref.as_range()] +/// Reads inline and external values in place, treating null rows as empty. +struct ViewRows<'a> { + views: &'a [BinaryView], + buffers: &'a [&'a ByteBuffer], + lengths: &'a [u32], + total_bytes: usize, +} + +impl Rows for ViewRows<'_> { + fn num_rows(&self) -> usize { + self.views.len() + } + + fn total_bytes(&self) -> usize { + self.total_bytes + } + + #[inline] + fn row(&self, i: usize) -> &[u8] { + let view = &self.views[i]; + if self.lengths[i] == 0 { + &[] + } else if view.is_inlined() { + view.as_inlined().value() + } else { + let view_ref = view.as_view(); + &self.buffers[view_ref.buffer_index as usize][view_ref.as_range()] + } } }