diff --git a/Cargo.lock b/Cargo.lock index 01a1fd656c7..cebb285ea86 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9472,6 +9472,7 @@ dependencies = [ "vortex-error", "vortex-fsst", "vortex-mask", + "vortex-onpair", "vortex-runend", "vortex-session", "vortex-zstd", diff --git a/encodings/onpair/src/array.rs b/encodings/onpair/src/array.rs index 9be2d3075b3..79f08f9638c 100644 --- a/encodings/onpair/src/array.rs +++ b/encodings/onpair/src/array.rs @@ -5,6 +5,7 @@ use std::fmt::Debug; use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hasher; +use std::mem::MaybeUninit; use std::sync::Arc; use std::sync::OnceLock; @@ -18,11 +19,9 @@ use vortex_array::ArrayId; use vortex_array::ArrayParts; use vortex_array::ArrayRef; use vortex_array::ArrayView; -use vortex_array::Canonical; use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; -use vortex_array::IntoArray; use vortex_array::array_slots; use vortex_array::buffer::BufferHandle; use vortex_array::builders::ArrayBuilder; @@ -50,8 +49,8 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::canonical::OnPairDecodePlan; use crate::canonical::canonicalize_onpair; -use crate::canonical::onpair_decode_bytes; use crate::canonical::onpair_decode_views; use crate::decode::collect_widened; use crate::rules::RULES; @@ -561,13 +560,12 @@ impl VTable for OnPair { return result; } + // The two arms here are every builder a `Utf8`/`Binary` dtype has: all four + // `VarBinBuilder` widths above, and `VarBinViewBuilder` below. There is deliberately no + // canonicalize-then-append fallback — it would decode to a `VarBinView` only for + // `VarBinView::append_to_builder` to reject the same remainder. let Some(builder) = builder.as_any_mut().downcast_mut::() else { - return array - .array() - .clone() - .execute::(ctx)? - .into_array() - .append_to_builder(builder, ctx); + vortex_bail!("append_to_builder for OnPair requires a variable-binary builder") }; let next_buffer_index = builder.completed_block_count() + u32::from(builder.in_progress()); @@ -592,7 +590,10 @@ impl VTable for OnPair { } } -/// Decodes the values and appends them to `builder`. +/// Decodes the code stream straight into `builder`'s byte storage. +/// +/// The offsets are the running sum of the uncompressed lengths the array already stores, so the +/// only work beyond the bulk `try_decode_into` is one prefix sum over them. fn append_to_varbin( array: ArrayView<'_, OnPair>, builder: &mut VarBinBuilder, @@ -601,20 +602,26 @@ fn append_to_varbin( where usize: AsPrimitive, { - let (bytes, lengths) = onpair_decode_bytes(array, ctx)?; + let plan = OnPairDecodePlan::new(array, ctx)?; let validity = array .array() .validity()? .execute_mask(array.array().len(), ctx)?; - match_each_integer_ptype!(lengths.ptype(), |P| { - builder.append_values( - bytes.as_slice(), - lengths.as_slice::

().iter().scan(0usize, |end, length| { - *end += AsPrimitive::::as_(*length); - Some(*end) - }), - &validity, - ) + // Built once, outside the ptype match: the decoder is `#[inline(always)]`, so creating the + // closure inside each arm would stamp out a copy of the whole decode loop per length type. + let mut decode = |out: &mut [MaybeUninit]| plan.decode_into(out); + match_each_integer_ptype!(plan.lengths.ptype(), |P| { + // SAFETY: `decode_into` initializes exactly the prefix whose length it returns. It needs + // no slack: it derives its bound from the slice it is handed and writes each value exactly. + unsafe { + builder.append_decoded( + plan.total_size, + 0, + plan.lengths.as_slice::

(), + &validity, + &mut decode, + ) + } }) } diff --git a/encodings/onpair/src/canonical.rs b/encodings/onpair/src/canonical.rs index 5179def3bf1..0069b2d074e 100644 --- a/encodings/onpair/src/canonical.rs +++ b/encodings/onpair/src/canonical.rs @@ -6,9 +6,11 @@ //! //! [`OnPairArray`]: crate::OnPairArray +use std::mem::MaybeUninit; use std::sync::Arc; use num_traits::AsPrimitive; +use onpair::CompactDictionaryView; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; @@ -44,68 +46,102 @@ pub(super) fn canonicalize_onpair( }) } -pub(crate) fn onpair_decode_bytes( - array: ArrayView<'_, OnPair>, - ctx: &mut ExecutionCtx, -) -> VortexResult<(ByteBufferMut, PrimitiveArray)> { - let lengths = array - .uncompressed_lengths() - .clone() - .execute::(ctx)?; +/// Everything needed to decode an OnPair array's values in one bulk `try_decode_into` call. +pub(crate) struct OnPairDecodePlan<'a> { + codes: Buffer, + dict: CompactDictionaryView<'a>, + /// Per-row uncompressed lengths, zero for null rows. + pub(crate) lengths: PrimitiveArray, + /// Total decoded size, i.e. the sum of `lengths`. + pub(crate) total_size: usize, +} + +impl<'a> OnPairDecodePlan<'a> { + pub(crate) fn new(array: ArrayView<'a, OnPair>, ctx: &mut ExecutionCtx) -> VortexResult { + let lengths = array + .uncompressed_lengths() + .clone() + .execute::(ctx)?; + + let total_size: usize = match_each_integer_ptype!(lengths.ptype(), |P| { + lengths + .as_slice::

() + .iter() + .map(|&l| AsPrimitive::::as_(l)) + .sum() + }); + + // `codes_offsets` holds the per-row code boundaries and may itself be a + // sliced or filtered view of the original. Its first and last entries + // bound the contiguous run of `codes` belonging to the rows present in + // this array: `slice` keeps the full `codes` child and only narrows + // `codes_offsets` (so `code_start > 0` and/or `code_end < codes.len()`), + // while `filter` rebuilds both children so the window is the whole stream. + // OnPair has no `TakeExecute`, so a reordering take is served from the + // canonical `VarBinView` and never reaches this path. We only need those + // two boundaries, so point-look them up rather than decoding every offset. + let codes_offsets = array.codes_offsets(); + let code_start = code_boundary_at(codes_offsets, 0, ctx)?; + let code_end = code_boundary_at(codes_offsets, array.len(), ctx)?; + vortex_ensure!( + code_start <= code_end, + "OnPair codes_offsets must be nondecreasing" + ); + vortex_ensure!( + code_end <= array.codes().len(), + "OnPair codes_offsets end {} exceeds codes len {}", + code_end, + array.codes().len() + ); - let total_size: usize = match_each_integer_ptype!(lengths.ptype(), |P| { - lengths - .as_slice::

() - .iter() - .map(|&l| AsPrimitive::::as_(l)) - .sum() - }); + // Slice the `codes` child to that window *before* unpacking it, so a sliced + // array materialises only its own codes rather than the whole column's. The + // contiguous decoder walks `codes` in order and never reads the per-row + // boundaries, so an empty boundary slice is sound. + let codes = collect_widened::(&array.codes().slice(code_start..code_end)?, ctx)?; + let dict = dict_view(array, ctx)?; - // `codes_offsets` holds the per-row code boundaries and may itself be a - // sliced or filtered view of the original. Its first and last entries - // bound the contiguous run of `codes` belonging to the rows present in - // this array: `slice` keeps the full `codes` child and only narrows - // `codes_offsets` (so `code_start > 0` and/or `code_end < codes.len()`), - // while `filter` rebuilds both children so the window is the whole stream. - // OnPair has no `TakeExecute`, so a reordering take is served from the - // canonical `VarBinView` and never reaches this path. We only need those - // two boundaries, so point-look them up rather than decoding every offset. - let codes_offsets = array.codes_offsets(); - let code_start = code_boundary_at(codes_offsets, 0, ctx)?; - let code_end = code_boundary_at(codes_offsets, array.len(), ctx)?; - vortex_ensure!( - code_start <= code_end, - "OnPair codes_offsets must be nondecreasing" - ); - vortex_ensure!( - code_end <= array.codes().len(), - "OnPair codes_offsets end {} exceeds codes len {}", - code_end, - array.codes().len() - ); + Ok(Self { + codes, + dict, + lengths, + total_size, + }) + } - // Slice the `codes` child to that window *before* unpacking it, so a sliced - // array materialises only its own codes rather than the whole column's. The - // contiguous decoder walks `codes` in order and never reads the per-row - // boundaries, so an empty boundary slice is sound. - let codes = collect_widened::(&array.codes().slice(code_start..code_end)?, ctx)?; - let dict = dict_view(array, ctx)?; - let mut out_bytes = ByteBufferMut::with_capacity(total_size); - let written = - match onpair::try_decode_into(codes.as_slice(), dict, out_bytes.spare_capacity_mut()) { + /// Bulk-decodes the whole code stream into `out`, which must hold at least `total_size` bytes. + /// + /// `try_decode_into` is generic over the dictionary view and only specializes its batched copy + /// loop once inlined, so this wrapper must not block that — hence `#[inline]` rather than + /// leaving it out of line, which costs ~2.5x. + #[inline] + pub(crate) fn decode_into(&self, out: &mut [MaybeUninit]) -> VortexResult { + let written = match onpair::try_decode_into(self.codes.as_slice(), self.dict, out) { Ok(written) => written, Err(_) => { vortex_panic!("OnPair codes decode to more bytes than uncompressed_lengths records") } }; - if written != total_size { - vortex_panic!( - "OnPair codes decoded to {written} bytes but uncompressed_lengths records {total_size}" - ); + if written != self.total_size { + vortex_panic!( + "OnPair codes decoded to {written} bytes but uncompressed_lengths records {}", + self.total_size + ); + } + Ok(written) } - // SAFETY: `try_decode_into` initialised exactly `written` bytes. +} + +pub(crate) fn onpair_decode_bytes( + array: ArrayView<'_, OnPair>, + ctx: &mut ExecutionCtx, +) -> VortexResult<(ByteBufferMut, PrimitiveArray)> { + let plan = OnPairDecodePlan::new(array, ctx)?; + let mut out_bytes = ByteBufferMut::with_capacity(plan.total_size); + let written = plan.decode_into(out_bytes.spare_capacity_mut())?; + // SAFETY: `decode_into` initialised exactly `written` bytes. unsafe { out_bytes.set_len(written) }; - Ok((out_bytes, lengths)) + Ok((out_bytes, plan.lengths)) } pub(crate) fn onpair_decode_views( diff --git a/vortex-arrow/Cargo.toml b/vortex-arrow/Cargo.toml index e69e6b656fa..11f2308451c 100644 --- a/vortex-arrow/Cargo.toml +++ b/vortex-arrow/Cargo.toml @@ -42,6 +42,7 @@ divan = { workspace = true } rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } vortex-fsst = { workspace = true } +vortex-onpair = { workspace = true } vortex-zstd = { workspace = true } [[bench]] diff --git a/vortex-arrow/benches/to_arrow.rs b/vortex-arrow/benches/to_arrow.rs index 3d38b3813a0..a663c71972e 100644 --- a/vortex-arrow/benches/to_arrow.rs +++ b/vortex-arrow/benches/to_arrow.rs @@ -42,6 +42,8 @@ use vortex_arrow::dtype::ToArrowType as _; use vortex_fsst::fsst_compress; use vortex_fsst::fsst_train_compressor; use vortex_mask::Mask; +use vortex_onpair::DEFAULT_DICT12_CONFIG; +use vortex_onpair::onpair_compress; use vortex_session::VortexSession; use vortex_zstd::Zstd; @@ -53,6 +55,7 @@ fn main() { static SESSION: LazyLock = LazyLock::new(|| { let session = array_session(); vortex_fsst::initialize(&session); + vortex_onpair::initialize(&session); session.arrays().register(Zstd); session }); @@ -122,6 +125,7 @@ fn ArrowExportVTable_to_arrow_field(bencher: Bencher) { enum StringEncoding { View, Fsst, + OnPair, Zstd, Dict, DictFsst, @@ -140,6 +144,7 @@ enum StringEncoding { const STRING_ENCODINGS: &[StringEncoding] = &[ StringEncoding::View, StringEncoding::Fsst, + StringEncoding::OnPair, StringEncoding::Zstd, StringEncoding::Dict, StringEncoding::DictFsst, @@ -156,12 +161,13 @@ const STRING_ENCODINGS: &[StringEncoding] = &[ /// Encodings whose `append_to_builder` the builder benchmarks reach directly. /// /// The Arrow export cannot stand in for these: `execute_until` stops at the first canonical array, -/// so a bare FSST/Zstd root is canonicalized to `VarBinView` before any builder sees it. +/// so a bare FSST/OnPair/Zstd root is canonicalized to `VarBinView` before any builder sees it. /// Only `Chunked`, `Constant` and `VarBin` roots reach an encoding's own `append_to_builder` that /// way, whereas the scan machinery appends encoded arrays into a builder directly. const BUILDER_STRING_ENCODINGS: &[StringEncoding] = &[ StringEncoding::View, StringEncoding::Fsst, + StringEncoding::OnPair, StringEncoding::Zstd, StringEncoding::Dict, StringEncoding::ChunkedFsst, @@ -250,6 +256,12 @@ fn string_array(encoding: StringEncoding) -> ArrayRef { structured_strings(OFFSET_STRING_ROWS).into_array(), &mut ctx, ), + StringEncoding::OnPair => onpair_compress( + &structured_strings(OFFSET_STRING_ROWS).into_array(), + DEFAULT_DICT12_CONFIG, + &mut ctx, + ) + .unwrap(), StringEncoding::Zstd => { let source = structured_strings(OFFSET_STRING_ROWS); Zstd::from_var_bin_view_without_dict(&source, 3, 8_192, &mut ctx)