Use buffers for FSST (#2656)

TODO(ngates): compress uncompressed_lengths

Author: gatesn

encodings/fsst/goldenfiles/fsst.metadata

@@ -5,9 +5,10 @@ use vortex_array::variants::{BinaryArrayTrait, Utf8ArrayTrait};
 use vortex_array::vtable::{EncodingVTable, VTableRef};
 use vortex_array::{
     Array, ArrayImpl, ArrayRef, ArrayStatisticsImpl, ArrayValidityImpl, ArrayVariantsImpl,
-    Encoding, EncodingId, SerdeMetadata, ToCanonical,
+    Encoding, EncodingId, SerdeMetadata,
 };
-use vortex_dtype::{DType, Nullability, PType};
+use vortex_buffer::Buffer;
+use vortex_dtype::DType;
 use vortex_error::{VortexResult, vortex_bail};
 use vortex_mask::Mask;
 
@@ -16,8 +17,8 @@ use crate::serde::FSSTMetadata;
 #[derive(Clone, Debug)]
 pub struct FSSTArray {
     dtype: DType,
-    symbols: ArrayRef,
-    symbol_lengths: ArrayRef,
+    symbols: Buffer<Symbol>,
+    symbol_lengths: Buffer<u8>,
     codes: ArrayRef,
     uncompressed_lengths: ArrayRef,
     stats_set: ArrayStats,
@@ -35,9 +36,6 @@ impl EncodingVTable for FSSTEncoding {
     }
 }
 
-pub(crate) static SYMBOLS_DTYPE: DType = DType::Primitive(PType::U64, Nullability::NonNullable);
-pub(crate) static SYMBOL_LENS_DTYPE: DType = DType::Primitive(PType::U8, Nullability::NonNullable);
-
 impl FSSTArray {
     /// Build an FSST array from a set of `symbols` and `codes`.
     ///
@@ -49,20 +47,11 @@ impl FSSTArray {
     /// which tells the decoder to emit the following byte without doing a table lookup.
     pub fn try_new(
         dtype: DType,
-        symbols: ArrayRef,
-        symbol_lengths: ArrayRef,
+        symbols: Buffer<Symbol>,
+        symbol_lengths: Buffer<u8>,
         codes: ArrayRef,
         uncompressed_lengths: ArrayRef,
     ) -> VortexResult<Self> {
-        // Check: symbols must be a u64 array
-        if symbols.dtype() != &SYMBOLS_DTYPE {
-            vortex_bail!(InvalidArgument: "symbols array must be of type u64")
-        }
-
-        if symbol_lengths.dtype() != &SYMBOL_LENS_DTYPE {
-            vortex_bail!(InvalidArgument: "symbol_lengths array must be of type u8")
-        }
-
         // Check: symbols must not have length > MAX_CODE
         if symbols.len() > 255 {
             vortex_bail!(InvalidArgument: "symbols array must have length <= 255");
@@ -102,12 +91,12 @@ impl FSSTArray {
     }
 
     /// Access the symbol table array
-    pub fn symbols(&self) -> &ArrayRef {
+    pub fn symbols(&self) -> &Buffer<Symbol> {
         &self.symbols
     }
 
     /// Access the symbol table array
-    pub fn symbol_lengths(&self) -> &ArrayRef {
+    pub fn symbol_lengths(&self) -> &Buffer<u8> {
         &self.symbol_lengths
     }
 
@@ -134,27 +123,11 @@ impl FSSTArray {
     }
 
     /// Build a [`Decompressor`][fsst::Decompressor] that can be used to decompress values from
-    /// this array, and pass it to the given function.
+    /// this array.
     ///
     /// This is private to the crate to avoid leaking `fsst-rs` types as part of the public API.
-    pub(crate) fn with_decompressor<F, R>(&self, apply: F) -> VortexResult<R>
-    where
-        F: FnOnce(Decompressor) -> VortexResult<R>,
-    {
-        // canonicalize the symbols child array, so we can view it contiguously
-        let symbols_array = self.symbols().to_primitive()?;
-        let symbols = symbols_array.as_slice::<u64>();
-
-        let symbol_lengths_array = self.symbol_lengths().to_primitive()?;
-        let symbol_lengths = symbol_lengths_array.as_slice::<u8>();
-
-        // Transmute the 64-bit symbol values into fsst `Symbol`s.
-        // SAFETY: Symbol is guaranteed to be 8 bytes, guaranteed by the compiler.
-        let symbols = unsafe { std::mem::transmute::<&[u64], &[Symbol]>(symbols) };
-
-        // Build a new decompressor that uses these symbols.
-        let decompressor = Decompressor::new(symbols, symbol_lengths);
-        apply(decompressor)
+    pub(crate) fn decompressor(&self) -> Decompressor {
+        Decompressor::new(self.symbols().as_slice(), self.symbol_lengths().as_slice())
     }
 }
 

encodings/fsst/src/array.rs

@@ -13,18 +13,15 @@ use crate::FSSTArray;
 
 impl ArrayCanonicalImpl for FSSTArray {
     fn _to_canonical(&self) -> VortexResult<Canonical> {
-        self.with_decompressor(|decompressor| {
-            fsst_into_varbin_view(decompressor, self, 0).map(Canonical::VarBinView)
-        })
+        fsst_into_varbin_view(self.decompressor(), self, 0).map(Canonical::VarBinView)
     }
 
     fn _append_to_builder(&self, builder: &mut dyn ArrayBuilder) -> VortexResult<()> {
         let Some(builder) = builder.as_any_mut().downcast_mut::<VarBinViewBuilder>() else {
             return builder.extend_from_array(&self.to_canonical()?.into_array());
         };
-        let view = self.with_decompressor(|decompressor| {
-            fsst_into_varbin_view(decompressor, self, builder.completed_block_count())
-        })?;
+        let view =
+            fsst_into_varbin_view(self.decompressor(), self, builder.completed_block_count())?;
 
         builder.push_buffer_and_adjusted_views(
             view.buffers().iter().cloned(),

encodings/fsst/src/canonical.rs

@@ -5,7 +5,7 @@ use vortex_array::accessor::ArrayAccessor;
 use vortex_array::arrays::builder::VarBinBuilder;
 use vortex_array::arrays::{VarBinArray, VarBinViewArray};
 use vortex_array::{Array, IntoArray};
-use vortex_buffer::{Buffer, BufferMut, ByteBuffer};
+use vortex_buffer::{Buffer, BufferMut};
 use vortex_dtype::DType;
 use vortex_error::{VortexExpect, VortexResult, VortexUnwrap, vortex_bail};
 
@@ -113,13 +113,9 @@ where
     let codes = builder
         .finish(DType::Binary(dtype.nullability()))
         .into_array();
-    let symbols_vec: Buffer<Symbol> = Buffer::copy_from(compressor.symbol_table());
-    // SAFETY: Symbol and u64 are same size
-    let symbols_u64: Buffer<u64> = unsafe { std::mem::transmute(symbols_vec) };
-    let symbols = symbols_u64.into_array();
+    let symbols: Buffer<Symbol> = Buffer::copy_from(compressor.symbol_table());
+    let symbol_lengths: Buffer<u8> = Buffer::<u8>::copy_from(compressor.symbol_lengths());
 
-    let symbol_lengths_vec: ByteBuffer = ByteBuffer::copy_from(compressor.symbol_lengths());
-    let symbol_lengths = symbol_lengths_vec.into_array();
     let uncompressed_lengths = uncompressed_lengths.into_array();
 
     FSSTArray::try_new(dtype, symbols, symbol_lengths, codes, uncompressed_lengths)

encodings/fsst/src/compress.rs

@@ -1,4 +1,3 @@
-use fsst::Symbol;
 use vortex_array::arrays::{BoolArray, BooleanBuffer, ConstantArray};
 use vortex_array::compute::{CompareFn, Operator, compare, compare_lengths_to_empty};
 use vortex_array::validity::Validity;
@@ -69,15 +68,12 @@ fn compare_fsst_constant(
         return Ok(None);
     }
 
-    let symbols = left.symbols().to_primitive()?;
-    let symbols_u64 = symbols.as_slice::<u64>();
-
-    let symbol_lens = left.symbol_lengths().to_primitive()?;
-    let symbol_lens_u8 = symbol_lens.as_slice::<u8>();
+    let symbols = left.symbols();
+    let symbol_lens = left.symbol_lengths();
 
     let mut compressor = fsst::CompressorBuilder::new();
-    for (symbol, symbol_len) in symbols_u64.iter().zip(symbol_lens_u8.iter()) {
-        compressor.insert(Symbol::from_slice(&symbol.to_le_bytes()), *symbol_len as _);
+    for (symbol, symbol_len) in symbols.iter().zip(symbol_lens.iter()) {
+        compressor.insert(*symbol, *symbol_len as usize);
     }
     let compressor = compressor.build();
 

encodings/fsst/src/compute/compare.rs

@@ -82,10 +82,9 @@ impl ScalarAtFn<&FSSTArray> for FSSTEncoding {
             .value()
             .ok_or_else(|| vortex_err!("expected null to already be handled"))?;
 
-        array.with_decompressor(|decompressor| {
-            let decoded_buffer = ByteBuffer::from(decompressor.decompress(binary_datum.as_slice()));
-            Ok(varbin_scalar(decoded_buffer, array.dtype()))
-        })
+        let decoded_buffer =
+            ByteBuffer::from(array.decompressor().decompress(binary_datum.as_slice()));
+        Ok(varbin_scalar(decoded_buffer, array.dtype()))
     }
 }
 

encodings/fsst/src/compute/mod.rs

@@ -1,35 +1,35 @@
+use fsst::Symbol;
 use serde::{Deserialize, Serialize};
 use vortex_array::serde::ArrayParts;
 use vortex_array::vtable::SerdeVTable;
 use vortex_array::{
-    Array, ArrayChildVisitor, ArrayContext, ArrayRef, ArrayVisitorImpl, DeserializeMetadata,
-    SerdeMetadata,
+    Array, ArrayBufferVisitor, ArrayChildVisitor, ArrayContext, ArrayRef, ArrayVisitorImpl,
+    DeserializeMetadata, SerdeMetadata,
 };
+use vortex_buffer::Buffer;
 use vortex_dtype::{DType, Nullability, PType};
-use vortex_error::{VortexExpect, VortexResult};
+use vortex_error::{VortexExpect, VortexResult, vortex_bail};
 
-use crate::array::{SYMBOL_LENS_DTYPE, SYMBOLS_DTYPE};
 use crate::{FSSTArray, FSSTEncoding};
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct FSSTMetadata {
-    symbols_len: usize,
-    codes_nullability: Nullability,
     uncompressed_lengths_ptype: PType,
 }
 
 impl ArrayVisitorImpl<SerdeMetadata<FSSTMetadata>> for FSSTArray {
+    fn _buffers(&self, visitor: &mut dyn ArrayBufferVisitor) {
+        visitor.visit_buffer(&self.symbols().clone().into_byte_buffer());
+        visitor.visit_buffer(&self.symbol_lengths().clone().into_byte_buffer());
+    }
+
     fn _children(&self, visitor: &mut dyn ArrayChildVisitor) {
-        visitor.visit_child("symbols", self.symbols());
-        visitor.visit_child("symbol_lengths", self.symbol_lengths());
         visitor.visit_child("codes", self.codes());
         visitor.visit_child("uncompressed_lengths", self.uncompressed_lengths());
     }
 
     fn _metadata(&self) -> SerdeMetadata<FSSTMetadata> {
         SerdeMetadata(FSSTMetadata {
-            symbols_len: self.symbols().len(),
-            codes_nullability: self.codes().dtype().nullability(),
             uncompressed_lengths_ptype: PType::try_from(self.uncompressed_lengths().dtype())
                 .vortex_expect("Must be a valid PType"),
         })
@@ -46,17 +46,19 @@ impl SerdeVTable<&FSSTArray> for FSSTEncoding {
     ) -> VortexResult<ArrayRef> {
         let metadata = SerdeMetadata::<FSSTMetadata>::deserialize(parts.metadata())?;
 
-        let symbols = parts
-            .child(0)
-            .decode(ctx, SYMBOLS_DTYPE.clone(), metadata.symbols_len)?;
-        let symbol_lengths =
-            parts
-                .child(1)
-                .decode(ctx, SYMBOL_LENS_DTYPE.clone(), metadata.symbols_len)?;
+        if parts.nbuffers() != 2 {
+            vortex_bail!(InvalidArgument: "Expected 2 buffers, got {}", parts.nbuffers());
+        }
+        let symbols = Buffer::<Symbol>::from_byte_buffer(parts.buffer(0)?);
+        let symbol_lengths = Buffer::<u8>::from_byte_buffer(parts.buffer(1)?);
+
+        if parts.nchildren() != 2 {
+            vortex_bail!(InvalidArgument: "Expected 2 children, got {}", parts.nchildren());
+        }
         let codes = parts
-            .child(2)
-            .decode(ctx, DType::Binary(metadata.codes_nullability), len)?;
-        let uncompressed_lengths = parts.child(3).decode(
+            .child(0)
+            .decode(ctx, DType::Binary(dtype.nullability()), len)?;
+        let uncompressed_lengths = parts.child(1).decode(
             ctx,
             DType::Primitive(
                 metadata.uncompressed_lengths_ptype,
@@ -76,7 +78,7 @@ impl SerdeVTable<&FSSTArray> for FSSTEncoding {
 mod test {
     use vortex_array::SerdeMetadata;
     use vortex_array::test_harness::check_metadata;
-    use vortex_dtype::{Nullability, PType};
+    use vortex_dtype::PType;
 
     use crate::serde::FSSTMetadata;
 
@@ -86,8 +88,6 @@ mod test {
         check_metadata(
             "fsst.metadata",
             SerdeMetadata(FSSTMetadata {
-                symbols_len: usize::MAX,
-                codes_nullability: Nullability::Nullable,
                 uncompressed_lengths_ptype: PType::U64,
             }),
         );