FSST: decompress straight into the builder's byte storage

encodings/fsst/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;
 
@@ -60,8 +61,9 @@ use vortex_error::vortex_panic;
 use vortex_session::VortexSession;
 use vortex_session::registry::CachedId;
 
+use crate::canonical::FSST_DECODE_SLACK;
+use crate::canonical::FsstDecodePlan;
 use crate::canonical::canonicalize_fsst;
-use crate::canonical::fsst_decode_bytes;
 use crate::canonical::fsst_decode_views;
 use crate::rules::RULES;
 
@@ -350,7 +352,10 @@ impl VTable for FSST {
     }
 }
 
-/// Decompresses the values and appends them to `builder`.
+/// Decompresses 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 `decompress_into` is one prefix sum over them.
 fn append_to_varbin<O: IntegerPType>(
     array: ArrayView<'_, FSST>,
     builder: &mut VarBinBuilder<O>,
@@ -359,20 +364,26 @@ fn append_to_varbin<O: IntegerPType>(
 where
     usize: AsPrimitive<O>,
 {
-    let (bytes, lengths) = fsst_decode_bytes(array, ctx)?;
+    let plan = FsstDecodePlan::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,
-        )
+    let decompressor = array.decompressor();
+    // 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(&decompressor, out);
+    match_each_integer_ptype!(plan.lengths.ptype(), |P| {
+        // SAFETY: `decode_into` initializes exactly the prefix whose length it returns.
+        unsafe {
+            builder.append_decoded(
+                plan.total_size,
+                FSST_DECODE_SLACK,
+                plan.lengths.as_slice::<P>(),
+                &validity,
+                &mut decode,
+            )
+        }
     })
 }
 

encodings/fsst/src/canonical.rs

@@ -1,8 +1,10 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::mem::MaybeUninit;
 use std::sync::Arc;
 
+use fsst::Decompressor;
 use vortex_array::ArrayRef;
 use vortex_array::ArrayView;
 use vortex_array::ExecutionCtx;
@@ -42,36 +44,82 @@ pub(super) fn canonicalize_fsst(
     })
 }
 
+/// Extra headroom that keeps [`Decompressor::decompress_into`] on its fast path.
+///
+/// It never writes past the slice it is given — every store is bounded by the end of the output —
+/// but it emits whole 8-byte symbols, so it can only use the wide-store loop while at least 8
+/// bytes remain. Handing it 7 spare bytes lets that loop run through the final value instead of
+/// finishing byte at a time. This is a performance knob, not a safety requirement.
+///
+/// [`Decompressor::decompress_into`]: fsst::Decompressor::decompress_into
+pub(crate) const FSST_DECODE_SLACK: usize = 7;
+
+/// Everything needed to decode an FSST array's values in one bulk `decompress_into` call.
+pub(crate) struct FsstDecodePlan {
+    codes: ByteBuffer,
+    /// 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 FsstDecodePlan {
+    pub(crate) fn new(
+        fsst_array: ArrayView<'_, FSST>,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Self> {
+        let codes = fsst_array.codes().sliced_bytes();
+        let lengths = fsst_array
+            .uncompressed_lengths()
+            .clone()
+            .execute::<PrimitiveArray>(ctx)?;
+
+        #[expect(clippy::cast_possible_truncation)]
+        let total_size: usize = match_each_integer_ptype!(lengths.ptype(), |P| {
+            lengths.as_slice::<P>().iter().map(|x| *x as usize).sum()
+        });
+
+        Ok(Self {
+            codes,
+            lengths,
+            total_size,
+        })
+    }
+
+    /// Bulk-decompresses the whole code stream into `out`, which must hold at least
+    /// `total_size + FSST_DECODE_SLACK` bytes.
+    ///
+    /// Kept inlinable so the decoder is not called from behind an extra frame in whichever
+    /// codegen unit the caller lands in; see OnPair's equivalent for why that matters.
+    #[inline]
+    pub(crate) fn decode_into(
+        &self,
+        decompressor: &Decompressor<'_>,
+        out: &mut [MaybeUninit<u8>],
+    ) -> VortexResult<usize> {
+        let len = decompressor.decompress_into(self.codes.as_slice(), out);
+        vortex_ensure!(
+            len == self.total_size,
+            "FSST decoded {len} bytes, expected {}",
+            self.total_size
+        );
+        Ok(len)
+    }
+}
+
 pub(crate) fn fsst_decode_bytes(
     fsst_array: ArrayView<'_, FSST>,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<(ByteBufferMut, PrimitiveArray)> {
-    let bytes = fsst_array.codes().sliced_bytes();
-    let uncompressed_lens_array = fsst_array
-        .uncompressed_lengths()
-        .clone()
-        .execute::<PrimitiveArray>(ctx)?;
-
-    #[expect(clippy::cast_possible_truncation)]
-    let total_size: usize = match_each_integer_ptype!(uncompressed_lens_array.ptype(), |P| {
-        uncompressed_lens_array
-            .as_slice::<P>()
-            .iter()
-            .map(|x| *x as usize)
-            .sum()
-    });
-
-    let decompressor = fsst_array.decompressor();
-    let mut uncompressed_bytes = ByteBufferMut::with_capacity(total_size + 7);
-    let len =
-        decompressor.decompress_into(bytes.as_slice(), uncompressed_bytes.spare_capacity_mut());
-    vortex_ensure!(
-        len == total_size,
-        "FSST decoded {len} bytes, expected {total_size}"
-    );
-    // SAFETY: `decompress_into` initialized the first `len` bytes.
+    let plan = FsstDecodePlan::new(fsst_array, ctx)?;
+    let mut uncompressed_bytes = ByteBufferMut::with_capacity(plan.total_size + FSST_DECODE_SLACK);
+    let len = plan.decode_into(
+        &fsst_array.decompressor(),
+        uncompressed_bytes.spare_capacity_mut(),
+    )?;
+    // SAFETY: `decode_into` initialized the first `len` bytes.
     unsafe { uncompressed_bytes.set_len(len) };
-    Ok((uncompressed_bytes, uncompressed_lens_array))
+    Ok((uncompressed_bytes, plan.lengths))
 }
 
 pub(crate) fn fsst_decode_views(