Replace DynVarBinBuilder with VarBinBuilder being an actual builder and optimise append_to_builder for string/binary compressed arrays

encodings/fsst/src/array.rs

@@ -34,13 +34,15 @@ use vortex_array::arrays::VarBinArray;
 use vortex_array::arrays::varbin::VarBinArraySlotsExt;
 use vortex_array::buffer::BufferHandle;
 use vortex_array::builders::ArrayBuilder;
-use vortex_array::builders::DynVarBinBuilder;
+use vortex_array::builders::VarBinBuilder;
 use vortex_array::builders::VarBinViewBuilder;
 use vortex_array::dtype::DType;
+use vortex_array::dtype::IntegerPType;
 use vortex_array::dtype::Nullability;
 use vortex_array::dtype::PType;
 use vortex_array::legacy_session;
 use vortex_array::match_each_integer_ptype;
+use vortex_array::match_each_varbin_builder;
 use vortex_array::serde::ArrayChildren;
 use vortex_array::validity::Validity;
 use vortex_array::vtable::VTable;
@@ -58,8 +60,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;
 
@@ -307,23 +310,10 @@ impl VTable for FSST {
         builder: &mut dyn ArrayBuilder,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<()> {
-        if let Some(builder) = builder.as_any_mut().downcast_mut::<DynVarBinBuilder>() {
-            let (bytes, lengths) = fsst_decode_bytes(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,
-                )
-            })?;
-            return Ok(());
+        if let Some(result) =
+            match_each_varbin_builder!(builder, |builder| append_to_varbin(array, builder, ctx))
+        {
+            return result;
         }
 
         let Some(builder) = builder.as_any_mut().downcast_mut::<VarBinViewBuilder>() else {
@@ -361,6 +351,35 @@ impl VTable for FSST {
     }
 }
 
+/// 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>,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<()>
+where
+    usize: AsPrimitive<O>,
+{
+    let plan = FsstDecodePlan::new(array, ctx)?;
+    let validity = array
+        .array()
+        .validity()?
+        .execute_mask(array.array().len(), ctx)?;
+    let decompressor = array.decompressor();
+    match_each_integer_ptype!(plan.lengths.ptype(), |P| {
+        builder.append_decoded(
+            plan.total_size,
+            FSST_DECODE_SLACK,
+            plan.lengths.as_slice::<P>(),
+            &validity,
+            |out| plan.decode_into(&decompressor, out),
+        )
+    })
+}
+
 #[array_slots(FSST)]
 pub struct FSSTSlots {
     /// Lengths of the original values before compression, can be compressed.

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,78 @@ pub(super) fn canonicalize_fsst(
     })
 }
 
+/// Extra writable bytes [`Decompressor::decompress_into`] may touch past the last value: it stores
+/// whole 8-byte symbols and only trims the length it reports.
+///
+/// [`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(
@@ -106,7 +150,7 @@ mod tests {
     use vortex_array::arrays::VarBinViewArray;
     use vortex_array::arrays::varbin::VarBinArrayExt;
     use vortex_array::builders::ArrayBuilder;
-    use vortex_array::builders::DynVarBinBuilder;
+    use vortex_array::builders::VarBinBuilder;
     use vortex_array::builders::VarBinViewBuilder;
     use vortex_array::dtype::DType;
     use vortex_array::dtype::Nullability;
@@ -212,7 +256,7 @@ mod tests {
 
         {
             let mut builder =
-                DynVarBinBuilder::with_capacity(chunked_arr.dtype().clone(), false, data.len());
+                VarBinBuilder::<i32>::with_capacity(chunked_arr.dtype().clone(), data.len());
             chunked_arr
                 .into_array()
                 .append_to_builder(&mut builder, &mut ctx)?;

encodings/fsst/src/compress.rs

@@ -199,7 +199,11 @@ fn compress_views<O>(
 where
     O: IntegerPType + 'static,
 {
-    let mut sink = FsstSink::<O>::with_capacity(strings.len(), compressor);
+    let mut sink = FsstSink::<O>::with_capacity(
+        DType::Binary(strings.dtype().nullability()),
+        strings.len(),
+        compressor,
+    );
     let views = strings.views();
     let buffers = strings.data_buffers();
     match mask.bit_buffer() {
@@ -232,7 +236,11 @@ fn compress_varbin<O>(
 where
     O: IntegerPType + 'static,
 {
-    let mut sink = FsstSink::<O>::with_capacity(strings.len(), compressor);
+    let mut sink = FsstSink::<O>::with_capacity(
+        DType::Binary(strings.dtype().nullability()),
+        strings.len(),
+        compressor,
+    );
     let bytes = strings.bytes().as_slice();
     match_each_integer_ptype!(offsets.ptype(), |I| {
         let off = offsets.as_slice::<I>();
@@ -277,10 +285,10 @@ struct FsstSink<'c, O: IntegerPType + 'static> {
 }
 
 impl<'c, O: IntegerPType + 'static> FsstSink<'c, O> {
-    fn with_capacity(len: usize, compressor: &'c Compressor) -> Self {
+    fn with_capacity(dtype: DType, len: usize, compressor: &'c Compressor) -> Self {
         Self {
             buffer: Vec::with_capacity(DEFAULT_BUFFER_LEN),
-            builder: VarBinBuilder::<O>::with_capacity(len),
+            builder: VarBinBuilder::<O>::with_capacity(dtype, len),
             uncompressed_lengths: BufferMut::with_capacity(len),
             compressor,
         }
@@ -289,7 +297,7 @@ impl<'c, O: IntegerPType + 'static> FsstSink<'c, O> {
     #[inline]
     fn emit(&mut self, row: Option<&[u8]>) {
         let Some(s) = row else {
-            self.builder.append_null();
+            self.builder.push_null();
             self.uncompressed_lengths.push(0);
             return;
         };
@@ -310,8 +318,8 @@ impl<'c, O: IntegerPType + 'static> FsstSink<'c, O> {
         self.builder.append_value(&self.buffer);
     }
 
-    fn finish(self, dtype: DType, ctx: &mut ExecutionCtx) -> VortexResult<FSSTArray> {
-        let codes = self.builder.finish(DType::Binary(dtype.nullability()));
+    fn finish(mut self, dtype: DType, ctx: &mut ExecutionCtx) -> VortexResult<FSSTArray> {
+        let codes = self.builder.finish_into_varbin();
         FSST::try_new(
             dtype,
             Buffer::copy_from(self.compressor.symbol_table()),

encodings/fsst/src/kernel.rs

@@ -61,7 +61,8 @@ mod tests {
     });
 
     fn build_test_fsst_array() -> ArrayRef {
-        let mut builder = VarBinBuilder::<i32>::with_capacity(10);
+        let mut builder =
+            VarBinBuilder::<i32>::with_capacity(DType::Utf8(Nullability::NonNullable), 10);
         builder.append_value(b"hello world");
         builder.append_value(b"foo bar baz");
         builder.append_value(b"testing fsst compression");
@@ -72,7 +73,7 @@ mod tests {
         builder.append_value(b"qrstuvwxyz");
         builder.append_value(b"0123456789");
         builder.append_value(b"final string");
-        let input = builder.finish(DType::Utf8(Nullability::NonNullable));
+        let input = builder.finish_into_varbin();
 
         let mut ctx = SESSION.create_execution_ctx();
         let arr = input.into_array();
@@ -131,7 +132,8 @@ mod tests {
         // Test case with special characters and nulls
         // Values: ["", "", "", "", "", "", "", "", "", "", "", ",", "A<<<<<<<", "", "", "", "", null, null, null, null, null, null]
         // Mask: only the last element is selected (true at index 22)
-        let mut builder = VarBinBuilder::<i32>::with_capacity(23);
+        let mut builder =
+            VarBinBuilder::<i32>::with_capacity(DType::Utf8(Nullability::Nullable), 23);
         // 11 empty strings
         for _ in 0..11 {
             builder.append_value(b"");
@@ -146,9 +148,9 @@ mod tests {
         }
         // 6 nulls
         for _ in 0..6 {
-            builder.append_null();
+            builder.push_null();
         }
-        let input = builder.finish(DType::Utf8(Nullability::Nullable));
+        let input = builder.finish_into_varbin();
         let array = input.clone().into_array();
 
         let mut ctx = SESSION.create_execution_ctx();
@@ -172,12 +174,13 @@ mod tests {
 
     #[test]
     fn filter_only_null() -> VortexResult<()> {
-        let mut builder = VarBinBuilder::<i32>::with_capacity(3);
-        builder.append_null();
+        let mut builder =
+            VarBinBuilder::<i32>::with_capacity(DType::Utf8(Nullability::Nullable), 3);
+        builder.push_null();
         builder.append_value(b"A");
-        builder.append_null();
+        builder.push_null();
 
-        let input = builder.finish(DType::Utf8(Nullability::Nullable));
+        let input = builder.finish_into_varbin();
         let array = input.clone().into_array();
 
         let mut ctx = SESSION.create_execution_ctx();
@@ -213,15 +216,14 @@ mod tests {
 
     #[test]
     fn test_fsst_byte_length() -> VortexResult<()> {
-        let mut builder = VarBinBuilder::<i32>::with_capacity(3);
+        let mut builder =
+            VarBinBuilder::<i32>::with_capacity(DType::Utf8(Nullability::NonNullable), 3);
         builder.append_value(b"hello");
         builder.append_value(b"world!!");
         builder.append_value("Пуховички"); // 9 characters, 18 bytes
         builder.append_value(b"");
 
-        let varbin = builder
-            .finish(DType::Utf8(Nullability::NonNullable))
-            .into_array();
+        let varbin = builder.finish_into_varbin().into_array();
         let mut ctx = SESSION.create_execution_ctx();
         let compressor = fsst_train_compressor(&varbin, &mut ctx)?;
         let fsst = fsst_compress(&varbin, &compressor, &mut ctx)?.into_array();