OnPair: decode straight into the builder's byte storage

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;
 
@@ -50,8 +51,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;
@@ -592,7 +593,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<O: IntegerPType>(
     array: ArrayView<'_, OnPair>,
     builder: &mut VarBinBuilder<O>,
@@ -601,20 +605,26 @@ fn append_to_varbin<O: IntegerPType>(
 where
     usize: AsPrimitive<O>,
 {
-    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::<P>().iter().scan(0usize, |end, length| {
-                *end += AsPrimitive::<usize>::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<u8>]| 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::<P>(),
+                &validity,
+                &mut decode,
+            )
+        }
     })
 }
 

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,101 @@ 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::<PrimitiveArray>(ctx)?;
+/// Everything needed to decode an OnPair array's values in one bulk `try_decode_into` call.
+pub(crate) struct OnPairDecodePlan<'a> {
+    codes: Buffer<u16>,
+    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<Self> {
+        let lengths = array
+            .uncompressed_lengths()
+            .clone()
+            .execute::<PrimitiveArray>(ctx)?;
+
+        let total_size: usize = match_each_integer_ptype!(lengths.ptype(), |P| {
+            lengths
+                .as_slice::<P>()
+                .iter()
+                .map(|&l| AsPrimitive::<usize>::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::<P>()
-            .iter()
-            .map(|&l| AsPrimitive::<usize>::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::<u16>(&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::<u16>(&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 into the caller, so leaving this wrapper out of line costs ~2.5x.
+    #[inline(always)]
+    pub(crate) fn decode_into(&self, out: &mut [MaybeUninit<u8>]) -> VortexResult<usize> {
+        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(

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]]

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<VortexSession> = 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)