diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index 54e88e109ef..4c92843e26a 100644 --- a/encodings/fsst/src/array.rs +++ b/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; @@ -21,11 +22,9 @@ use vortex_array::ArrayParts; use vortex_array::ArrayRef; use vortex_array::ArraySlots; 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::TypedArrayRef; use vortex_array::VortexSessionExecute; use vortex_array::array_slots; @@ -60,8 +59,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; @@ -315,13 +315,12 @@ impl VTable for FSST { 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 FSST requires a variable-binary builder") }; // Decompress the whole block of data into a new buffer, and create some views @@ -350,7 +349,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( array: ArrayView<'_, FSST>, builder: &mut VarBinBuilder, @@ -359,20 +361,26 @@ fn append_to_varbin( where usize: AsPrimitive, { - 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::

().iter().scan(0usize, |end, length| { - *end += AsPrimitive::::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]| 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::

(), + &validity, + &mut decode, + ) + } }) } diff --git a/encodings/fsst/src/canonical.rs b/encodings/fsst/src/canonical.rs index f248e123fe9..8db84b22823 100644 --- a/encodings/fsst/src/canonical.rs +++ b/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 { + let codes = fsst_array.codes().sliced_bytes(); + let lengths = fsst_array + .uncompressed_lengths() + .clone() + .execute::(ctx)?; + + #[expect(clippy::cast_possible_truncation)] + let total_size: usize = match_each_integer_ptype!(lengths.ptype(), |P| { + lengths.as_slice::

().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], + ) -> VortexResult { + 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::(ctx)?; - - #[expect(clippy::cast_possible_truncation)] - let total_size: usize = match_each_integer_ptype!(uncompressed_lens_array.ptype(), |P| { - uncompressed_lens_array - .as_slice::

() - .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(