Skip to main content

vortex_fsst/
array.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Debug;
5use std::fmt::Display;
6use std::fmt::Formatter;
7use std::hash::Hasher;
8use std::mem::MaybeUninit;
9use std::sync::Arc;
10use std::sync::OnceLock;
11
12use fsst::Compressor;
13use fsst::Decompressor;
14use fsst::Symbol;
15use num_traits::AsPrimitive;
16use prost::Message as _;
17use vortex_array::Array;
18use vortex_array::ArrayEq;
19use vortex_array::ArrayHash;
20use vortex_array::ArrayId;
21use vortex_array::ArrayParts;
22use vortex_array::ArrayRef;
23use vortex_array::ArraySlots;
24use vortex_array::ArrayView;
25use vortex_array::EqMode;
26use vortex_array::ExecutionCtx;
27use vortex_array::ExecutionResult;
28use vortex_array::TypedArrayRef;
29use vortex_array::VortexSessionExecute;
30use vortex_array::array_slots;
31use vortex_array::arrays::Primitive;
32use vortex_array::arrays::VarBin;
33use vortex_array::arrays::VarBinArray;
34use vortex_array::arrays::varbin::VarBinArraySlotsExt;
35use vortex_array::buffer::BufferHandle;
36use vortex_array::builders::ArrayBuilder;
37use vortex_array::builders::VarBinBuilder;
38use vortex_array::builders::VarBinViewBuilder;
39use vortex_array::dtype::DType;
40use vortex_array::dtype::Nullability;
41use vortex_array::dtype::OffsetBuilderPType;
42use vortex_array::dtype::PType;
43use vortex_array::legacy_session;
44use vortex_array::match_each_integer_ptype;
45use vortex_array::match_each_varbin_builder;
46use vortex_array::require_child;
47use vortex_array::require_validity;
48use vortex_array::serde::ArrayChildren;
49use vortex_array::validity::Validity;
50use vortex_array::vtable::VTable;
51use vortex_array::vtable::ValidityVTable;
52use vortex_array::vtable::child_to_validity;
53use vortex_array::vtable::validity_to_child;
54use vortex_buffer::Buffer;
55use vortex_buffer::BufferMut;
56use vortex_buffer::ByteBuffer;
57use vortex_error::VortexExpect;
58use vortex_error::VortexResult;
59use vortex_error::vortex_bail;
60use vortex_error::vortex_ensure;
61use vortex_error::vortex_ensure_eq;
62use vortex_error::vortex_err;
63use vortex_error::vortex_panic;
64use vortex_session::VortexSession;
65use vortex_session::registry::CachedId;
66
67use crate::canonical::FSST_DECODE_SLACK;
68use crate::canonical::FsstDecodePlan;
69use crate::canonical::canonicalize_fsst;
70use crate::canonical::fsst_decode_bytes;
71use crate::rules::RULES;
72
73/// A [`FSST`]-encoded Vortex array.
74pub type FSSTArray = Array<FSST>;
75
76#[derive(Clone, prost::Message)]
77pub struct FSSTMetadata {
78    #[prost(enumeration = "PType", tag = "1")]
79    uncompressed_lengths_ptype: i32,
80
81    #[prost(enumeration = "PType", tag = "2")]
82    codes_offsets_ptype: i32,
83}
84
85impl FSSTMetadata {
86    pub fn get_uncompressed_lengths_ptype(&self) -> VortexResult<PType> {
87        PType::try_from(self.uncompressed_lengths_ptype)
88            .map_err(|_| vortex_err!("Invalid PType {}", self.uncompressed_lengths_ptype))
89    }
90}
91
92/// The number of entries in a fully-populated FSST symbol table.
93///
94/// Code 255 is reserved as the escape code, leaving 255 usable codes. [`Decompressor`] borrows
95/// the symbol table as fixed-size arrays of exactly this length, so Vortex pads the symbols and
96/// symbol lengths buffers out to it on construction.
97pub const FSST_SYMBOL_TABLE_LEN: usize = 255;
98
99impl ArrayHash for FSSTData {
100    fn array_hash<H: Hasher>(&self, state: &mut H, precision: EqMode) {
101        self.padded_symbols().array_hash(state, precision);
102        self.padded_symbol_lengths().array_hash(state, precision);
103        self.codes_bytes.as_host().array_hash(state, precision);
104    }
105}
106
107impl ArrayEq for FSSTData {
108    fn array_eq(&self, other: &Self, precision: EqMode) -> bool {
109        self.padded_symbols()
110            .array_eq(other.padded_symbols(), precision)
111            && self
112                .padded_symbol_lengths()
113                .array_eq(other.padded_symbol_lengths(), precision)
114            && self
115                .codes_bytes
116                .as_host()
117                .array_eq(other.codes_bytes.as_host(), precision)
118    }
119}
120
121impl VTable for FSST {
122    type TypedArrayData = FSSTData;
123    type OperationsVTable = Self;
124    type ValidityVTable = Self;
125
126    fn id(&self) -> ArrayId {
127        static ID: CachedId = CachedId::new("vortex.fsst");
128        *ID
129    }
130
131    #[allow(clippy::disallowed_methods)]
132    fn validate(
133        &self,
134        data: &Self::TypedArrayData,
135        dtype: &DType,
136        len: usize,
137        slots: &[Option<ArrayRef>],
138    ) -> VortexResult<()> {
139        // TODO(ctx): trait fixes - VTable::validate has a fixed signature.
140        let mut ctx = legacy_session().create_execution_ctx();
141        data.validate(dtype, len, slots, &mut ctx)
142    }
143
144    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
145        3
146    }
147
148    fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
149        match idx {
150            0 => BufferHandle::new_host(
151                array
152                    .padded_symbols()
153                    .slice(0..array.n_symbols())
154                    .into_byte_buffer(),
155            ),
156            1 => BufferHandle::new_host(
157                array
158                    .padded_symbol_lengths()
159                    .slice(0..array.n_symbols())
160                    .into_byte_buffer(),
161            ),
162            2 => array.codes_bytes_handle().clone(),
163            _ => vortex_panic!("FSSTArray buffer index {idx} out of bounds"),
164        }
165    }
166
167    fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option<String> {
168        match idx {
169            0 => Some("symbols".to_string()),
170            1 => Some("symbol_lengths".to_string()),
171            2 => Some("compressed_codes".to_string()),
172            _ => vortex_panic!("FSSTArray buffer_name index {idx} out of bounds"),
173        }
174    }
175
176    fn with_buffers(
177        &self,
178        array: ArrayView<'_, Self>,
179        buffers: &[BufferHandle],
180    ) -> VortexResult<ArrayParts<Self>> {
181        vortex_ensure_eq!(buffers.len(), 3);
182        let symbols = Buffer::<Symbol>::from_byte_buffer(buffers[0].clone().try_to_host_sync()?);
183        let symbol_lengths = Buffer::<u8>::from_byte_buffer(buffers[1].clone().try_to_host_sync()?);
184        let data = FSSTData::try_new(symbols, symbol_lengths, buffers[2].clone(), array.len())?;
185        Ok(ArrayParts::new(
186            self.clone(),
187            array.dtype().clone(),
188            array.len(),
189            data,
190            array.slots().iter().cloned().collect(),
191        ))
192    }
193
194    fn serialize(
195        array: ArrayView<'_, Self>,
196        _session: &VortexSession,
197    ) -> VortexResult<Option<Vec<u8>>> {
198        let codes_offsets = array.codes_offsets();
199        Ok(Some(
200            FSSTMetadata {
201                uncompressed_lengths_ptype: array.uncompressed_lengths().dtype().as_ptype().into(),
202                codes_offsets_ptype: codes_offsets.dtype().as_ptype().into(),
203            }
204            .encode_to_vec(),
205        ))
206    }
207
208    /// Deserializes an FSST array from its serialized components.
209    ///
210    /// Supports two serialization formats:
211    ///
212    /// ## Legacy format (2 buffers, 2 children)
213    ///
214    /// The original FSST layout stored the compressed codes as a full `VarBinArray` child.
215    /// - **Buffers**: `[symbols, symbol_lengths]`
216    /// - **Children**: `[codes (VarBinArray), uncompressed_lengths (Primitive)]`
217    ///
218    /// The codes VarBinArray child is decomposed: its bytes become the `codes_bytes` buffer,
219    /// and its offsets/validity are extracted into slots.
220    /// See `FSST::deserialize_legacy`.
221    ///
222    /// ## Current format (3 buffers, 2-3 children)
223    ///
224    /// The current layout stores the compressed bytes as a raw buffer alongside the symbol
225    /// table, with offsets and validity as separate children.
226    /// - **Buffers**: `[symbols, symbol_lengths, compressed_codes_bytes]`
227    /// - **Children**: `[uncompressed_lengths, codes_offsets, (optional) codes_validity]`
228    ///
229    /// The `codes_bytes` buffer is stored directly in `FSSTData`. A `VarBinArray` for the
230    /// codes can be reconstructed on demand via [`FSSTArrayExt::codes()`] using the bytes
231    /// from `FSSTData` combined with offsets and validity from the array's slots.
232    fn deserialize(
233        &self,
234        dtype: &DType,
235        len: usize,
236        metadata: &[u8],
237        buffers: &[BufferHandle],
238        children: &dyn ArrayChildren,
239        session: &VortexSession,
240    ) -> VortexResult<ArrayParts<Self>> {
241        let metadata = FSSTMetadata::decode(metadata)?;
242        let symbols = Buffer::<Symbol>::from_byte_buffer(buffers[0].clone().try_to_host_sync()?);
243        let symbol_lengths = Buffer::<u8>::from_byte_buffer(buffers[1].clone().try_to_host_sync()?);
244
245        let mut ctx = session.create_execution_ctx();
246        if buffers.len() == 2 {
247            return Self::deserialize_legacy(
248                self,
249                dtype,
250                len,
251                &metadata,
252                &symbols,
253                &symbol_lengths,
254                children,
255                &mut ctx,
256            );
257        }
258
259        if buffers.len() == 3 {
260            let uncompressed_lengths = children.get(
261                0,
262                &DType::Primitive(
263                    metadata.get_uncompressed_lengths_ptype()?,
264                    Nullability::NonNullable,
265                ),
266                len,
267            )?;
268
269            let codes_bytes = buffers[2].clone();
270            let codes_offsets = children.get(
271                1,
272                &DType::Primitive(
273                    PType::try_from(metadata.codes_offsets_ptype)?,
274                    Nullability::NonNullable,
275                ),
276                // VarBin offsets are len + 1
277                len + 1,
278            )?;
279
280            let codes_validity = if children.len() == 2 {
281                Validity::from(dtype.nullability())
282            } else if children.len() == 3 {
283                let validity = children.get(2, &Validity::DTYPE, len)?;
284                Validity::Array(validity)
285            } else {
286                vortex_bail!("Expected 2 or 3 children, got {}", children.len());
287            };
288
289            FSSTData::validate_parts(
290                symbols.as_slice(),
291                symbol_lengths.as_slice(),
292                &codes_bytes,
293                &codes_offsets,
294                dtype.nullability(),
295                &uncompressed_lengths,
296                dtype,
297                len,
298                &mut ctx,
299            )?;
300            let slots = FSSTSlots {
301                uncompressed_lengths,
302                codes_offsets,
303                codes_validity: validity_to_child(&codes_validity, len),
304            }
305            .into_slots();
306            let data = FSSTData::try_new(symbols, symbol_lengths, codes_bytes, len)?;
307            return Ok(ArrayParts::new(
308                self.clone(),
309                dtype.clone(),
310                len,
311                data,
312                slots,
313            ));
314        }
315
316        vortex_bail!(
317            "InvalidArgument: Expected 2 or 3 buffers, got {}",
318            buffers.len()
319        );
320    }
321
322    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
323        FSSTSlots::NAMES[idx].to_string()
324    }
325
326    fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
327        let array = require_child!(
328            array,
329            array.uncompressed_lengths(),
330            FSSTSlots::UNCOMPRESSED_LENGTHS => Primitive
331        );
332        let array = require_child!(
333            array,
334            array.codes_offsets(),
335            FSSTSlots::CODES_OFFSETS => Primitive
336        );
337        require_validity!(array, FSSTSlots::CODES_VALIDITY);
338        canonicalize_fsst(array.as_view(), ctx).map(ExecutionResult::done)
339    }
340
341    fn append_to_builder(
342        array: ArrayView<'_, Self>,
343        builder: &mut dyn ArrayBuilder,
344        ctx: &mut ExecutionCtx,
345    ) -> VortexResult<()> {
346        if let Some(result) =
347            match_each_varbin_builder!(builder, |builder| append_to_varbin(array, builder, ctx))
348        {
349            return result;
350        }
351
352        // The two arms here are every builder a `Utf8`/`Binary` dtype has: all four
353        // `VarBinBuilder` widths above, and `VarBinViewBuilder` below. There is deliberately no
354        // canonicalize-then-append fallback — it would decode to a `VarBinView` only for
355        // `VarBinView::append_to_builder` to reject the same remainder.
356        let Some(builder) = builder.as_any_mut().downcast_mut::<VarBinViewBuilder>() else {
357            vortex_bail!("append_to_builder for FSST requires a variable-binary builder")
358        };
359
360        // Decompress the whole block of data into a new buffer, which the builder adopts as a
361        // data buffer with views built over it in place.
362        let validity = array
363            .array()
364            .validity()?
365            .execute_mask(array.array().len(), ctx)?;
366        let (uncompressed_bytes, uncompressed_lens) = fsst_decode_bytes(array, ctx)?;
367        match_each_integer_ptype!(uncompressed_lens.ptype(), |P| {
368            builder.append_buffer_with_lengths(
369                uncompressed_bytes.freeze(),
370                uncompressed_lens.as_slice::<P>(),
371                &validity,
372            )
373        });
374        Ok(())
375    }
376
377    fn reduce_parent(
378        array: ArrayView<'_, Self>,
379        parent: &ArrayRef,
380        child_idx: usize,
381    ) -> VortexResult<Option<ArrayRef>> {
382        RULES.evaluate(array, parent, child_idx)
383    }
384}
385
386/// Decompresses the code stream straight into `builder`'s byte storage.
387///
388/// The offsets are the running sum of the uncompressed lengths the array already stores, so the
389/// only work beyond the bulk `decompress_into` is one prefix sum over them.
390fn append_to_varbin<O: OffsetBuilderPType>(
391    array: ArrayView<'_, FSST>,
392    builder: &mut VarBinBuilder<O>,
393    ctx: &mut ExecutionCtx,
394) -> VortexResult<()>
395where
396    usize: AsPrimitive<O>,
397{
398    let plan = FsstDecodePlan::new(array, ctx)?;
399    let validity = array
400        .array()
401        .validity()?
402        .execute_mask(array.array().len(), ctx)?;
403    let decompressor = array.decompressor();
404    // Built once, outside the ptype match: the decoder is `#[inline(always)]`, so creating the
405    // closure inside each arm would stamp out a copy of the whole decode loop per length type.
406    let mut decode = |out: &mut [MaybeUninit<u8>]| plan.decode_into(&decompressor, out);
407    match_each_integer_ptype!(plan.lengths.ptype(), |P| {
408        // SAFETY: `decode_into` initializes exactly the prefix whose length it returns.
409        unsafe {
410            builder.append_decoded(
411                plan.total_size,
412                FSST_DECODE_SLACK,
413                plan.lengths.as_slice::<P>(),
414                &validity,
415                &mut decode,
416            )
417        }
418    })
419}
420
421#[array_slots(FSST)]
422pub struct FSSTSlots {
423    /// Lengths of the original values before compression, can be compressed.
424    #[slot(0)]
425    pub uncompressed_lengths: ArrayRef,
426    /// The offsets array for the FSST-compressed codes.
427    #[slot(1)]
428    pub codes_offsets: ArrayRef,
429    /// The validity bitmap for the compressed codes.
430    #[slot(2)]
431    pub codes_validity: Option<ArrayRef>,
432}
433
434/// The inner data for an FSST-compressed array.
435///
436/// Holds the FSST symbol table (`symbols` + `symbol_lengths`) and the raw compressed
437/// codes bytes buffer. The codes offsets and validity live in the outer array's slots
438/// (slots 1 and 2 respectively).
439///
440/// A full [`VarBinArray`] representing the codes can be reconstructed on demand via
441/// [`FSSTArrayExt::codes()`], combining this buffer with the offsets/validity from slots.
442#[derive(Clone)]
443pub struct FSSTData {
444    symbol_table: Arc<FSSTSymbolTable>,
445    /// The raw compressed codes bytes, equivalent to `VarBinData::bytes`.
446    codes_bytes: BufferHandle,
447    /// Cached length (number of elements).
448    len: usize,
449}
450
451impl Display for FSSTData {
452    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
453        write!(
454            f,
455            "len: {}, nsymbols: {}",
456            self.len, self.symbol_table.n_symbols
457        )
458    }
459}
460
461impl Debug for FSSTData {
462    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
463        f.debug_struct("FSSTArray")
464            .field("symbols", &self.symbols())
465            .field("symbol_lengths", &self.symbol_lengths())
466            .field("codes_bytes_len", &self.codes_bytes.len())
467            .field("len", &self.len)
468            .field("uncompressed_lengths", &"<outer slot>")
469            .field("codes_offsets", &"<outer slot>")
470            .field("codes_validity", &"<outer slot>")
471            .finish()
472    }
473}
474
475pub struct FSSTSymbolTable {
476    /// Symbols padded out to [`FSST_SYMBOL_TABLE_LEN`] entries, zero-filled past `n_symbols`,
477    /// so that a [`Decompressor`] can borrow them without copying.
478    padded_symbols: Buffer<Symbol>,
479    /// Symbol lengths padded out to [`FSST_SYMBOL_TABLE_LEN`] entries, zero-filled past
480    /// `n_symbols`.
481    padded_symbol_lengths: Buffer<u8>,
482    /// The number of populated symbols. Entries at or past this index are padding.
483    n_symbols: usize,
484    /// Memoized compressor used for push-down of compute by compressing the RHS.
485    compressor: OnceLock<Compressor>,
486}
487
488impl FSSTSymbolTable {
489    /// Builds a symbol table, padding `symbols` and `symbol_lengths` out to
490    /// [`FSST_SYMBOL_TABLE_LEN`] entries if they are shorter.
491    ///
492    /// `n_symbols` is the number of populated entries; everything past it is padding. Buffers
493    /// longer than [`FSST_SYMBOL_TABLE_LEN`] are truncated; callers are expected to have rejected
494    /// them already (see [`FSSTData::try_new`]).
495    pub fn new(symbols: Buffer<Symbol>, symbol_lengths: Buffer<u8>, n_symbols: usize) -> Self {
496        Self {
497            padded_symbols: pad_symbol_table(symbols, Symbol::ZERO),
498            padded_symbol_lengths: pad_symbol_table(symbol_lengths, 0),
499            n_symbols: n_symbols.min(FSST_SYMBOL_TABLE_LEN),
500            compressor: OnceLock::new(),
501        }
502    }
503
504    /// Builds a symbol table, padding `symbols` and `symbol_lengths` out to
505    /// [`FSST_SYMBOL_TABLE_LEN`] entries if they are shorter.
506    ///
507    /// `n_symbols` is the number of populated entries; everything past it is padding. Buffers
508    /// longer than [`FSST_SYMBOL_TABLE_LEN`] are truncated; callers are expected to have rejected
509    /// them already (see [`FSSTData::try_new`]).
510    pub fn new_padded(
511        padded_symbols: Buffer<Symbol>,
512        padded_symbol_lengths: Buffer<u8>,
513        n_symbols: usize,
514    ) -> VortexResult<Self> {
515        vortex_ensure!(
516            padded_symbols.len() == FSST_SYMBOL_TABLE_LEN
517                && padded_symbol_lengths.len() == FSST_SYMBOL_TABLE_LEN,
518            InvalidArgument: "padded symbol table must have exactly {FSST_SYMBOL_TABLE_LEN} entries, found {} symbols and {} symbol lengths",
519            padded_symbols.len(),
520            padded_symbol_lengths.len()
521        );
522        vortex_ensure!(
523            n_symbols <= FSST_SYMBOL_TABLE_LEN,
524            InvalidArgument: "n_symbols must be <= {FSST_SYMBOL_TABLE_LEN}, found {n_symbols}"
525        );
526        Ok(Self {
527            padded_symbols,
528            padded_symbol_lengths,
529            n_symbols,
530            compressor: OnceLock::new(),
531        })
532    }
533
534    /// The populated symbols, excluding the padding added by [`Self::new`].
535    fn symbols(&self) -> &[Symbol] {
536        &self.padded_symbols.as_slice()[..self.n_symbols]
537    }
538
539    /// The populated symbol lengths, excluding the padding added by [`Self::new`].
540    fn symbol_lengths(&self) -> &[u8] {
541        &self.padded_symbol_lengths.as_slice()[..self.n_symbols]
542    }
543
544    /// The symbols buffer, padded to exactly [`FSST_SYMBOL_TABLE_LEN`] entries with
545    /// [`Symbol::ZERO`].
546    fn padded_symbols(&self) -> &Buffer<Symbol> {
547        &self.padded_symbols
548    }
549
550    /// The symbol lengths buffer, padded to exactly [`FSST_SYMBOL_TABLE_LEN`] entries with zeros.
551    fn padded_symbol_lengths(&self) -> &Buffer<u8> {
552        &self.padded_symbol_lengths
553    }
554
555    /// Borrow the padded symbol table as the fixed-size arrays expected by [`Decompressor`].
556    ///
557    /// Both buffers are padded to [`FSST_SYMBOL_TABLE_LEN`] on construction, so this is a
558    /// length check and a pointer cast rather than a copy.
559    fn decompressor(&self) -> Decompressor<'_> {
560        const PADDED: &str = "FSST symbol table is padded to FSST_SYMBOL_TABLE_LEN entries";
561        let symbols = self
562            .padded_symbols
563            .as_slice()
564            .first_chunk::<FSST_SYMBOL_TABLE_LEN>()
565            .vortex_expect(PADDED);
566        let symbol_lengths = self
567            .padded_symbol_lengths
568            .as_slice()
569            .first_chunk::<FSST_SYMBOL_TABLE_LEN>()
570            .vortex_expect(PADDED);
571        Decompressor::new(symbols, symbol_lengths)
572    }
573
574    fn compressor(&self) -> &Compressor {
575        self.compressor
576            .get_or_init(|| Compressor::rebuild_from(self.symbols(), self.symbol_lengths()))
577    }
578}
579
580/// Returns `buffer` resized to exactly [`FSST_SYMBOL_TABLE_LEN`] entries, filling any tail with
581/// `pad`. Buffers that are already the right length are returned untouched.
582fn pad_symbol_table<T: Copy>(buffer: Buffer<T>, pad: T) -> Buffer<T> {
583    if buffer.len() == FSST_SYMBOL_TABLE_LEN {
584        return buffer;
585    }
586    padded_symbol_table(buffer.as_slice(), pad)
587}
588
589/// Copies `values` into a buffer of exactly [`FSST_SYMBOL_TABLE_LEN`] entries, filling the tail
590/// with `pad`.
591///
592/// FSST symbol tables are stored padded so that [`FSSTData::decompressor`] can borrow them as the
593/// fixed-size arrays [`Decompressor`] requires, without copying.
594pub(crate) fn padded_symbol_table<T: Copy>(values: &[T], pad: T) -> Buffer<T> {
595    let populated = values.len().min(FSST_SYMBOL_TABLE_LEN);
596    let mut padded = BufferMut::with_capacity(FSST_SYMBOL_TABLE_LEN);
597    padded.extend_from_slice(&values[..populated]);
598    padded.push_n(pad, FSST_SYMBOL_TABLE_LEN - populated);
599    padded.freeze()
600}
601
602#[derive(Clone, Debug)]
603pub struct FSST;
604
605impl FSST {
606    /// Build an FSST array from a set of `symbols` and `codes`.
607    ///
608    /// The `codes` VarBinArray is decomposed: its bytes are stored in [`FSSTData`], while
609    /// its offsets and validity become array slots. The codes VarBinArray can be
610    /// reconstructed on demand via [`FSSTArrayExt::codes()`].
611    pub fn try_new(
612        dtype: DType,
613        symbols: Buffer<Symbol>,
614        symbol_lengths: Buffer<u8>,
615        codes: VarBinArray,
616        uncompressed_lengths: ArrayRef,
617        ctx: &mut ExecutionCtx,
618    ) -> VortexResult<FSSTArray> {
619        let len = codes.len();
620        FSSTData::validate_parts_from_codes(
621            symbols.as_slice(),
622            symbol_lengths.as_slice(),
623            &codes,
624            &uncompressed_lengths,
625            &dtype,
626            len,
627            ctx,
628        )?;
629        let slots = FSSTData::make_slots(&codes, &uncompressed_lengths);
630        let codes_bytes = codes.bytes_handle().clone();
631        let data = FSSTData::try_new(symbols, symbol_lengths, codes_bytes, len)?;
632        Ok(unsafe { Array::from_parts_unchecked(ArrayParts::new(FSST, dtype, len, data, slots)) })
633    }
634
635    pub fn try_new_with_symbol_table(
636        dtype: DType,
637        symbol_table: Arc<FSSTSymbolTable>,
638        codes: VarBinArray,
639        uncompressed_lengths: ArrayRef,
640        ctx: &mut ExecutionCtx,
641    ) -> VortexResult<FSSTArray> {
642        let len = codes.len();
643        FSSTData::validate_parts_from_codes(
644            symbol_table.symbols(),
645            symbol_table.symbol_lengths(),
646            &codes,
647            &uncompressed_lengths,
648            &dtype,
649            len,
650            ctx,
651        )?;
652        let slots = FSSTData::make_slots(&codes, &uncompressed_lengths);
653        let codes_bytes = codes.bytes_handle().clone();
654        let data =
655            unsafe { FSSTData::new_unchecked_with_symbol_table(symbol_table, codes_bytes, len) };
656        Ok(unsafe { Array::from_parts_unchecked(ArrayParts::new(FSST, dtype, len, data, slots)) })
657    }
658
659    /// Legacy deserialization path (2 buffers): the codes were stored as a full
660    /// `VarBinArray` child. We decompose the VarBinArray into its bytes (stored in
661    /// FSSTData) and offsets/validity (stored in slots).
662    #[allow(clippy::too_many_arguments)]
663    fn deserialize_legacy(
664        &self,
665        dtype: &DType,
666        len: usize,
667        metadata: &FSSTMetadata,
668        symbols: &Buffer<Symbol>,
669        symbol_lengths: &Buffer<u8>,
670        children: &dyn ArrayChildren,
671        ctx: &mut ExecutionCtx,
672    ) -> VortexResult<ArrayParts<Self>> {
673        if children.len() != 2 {
674            vortex_bail!(InvalidArgument: "Expected 2 children, got {}", children.len());
675        }
676        let codes = children.get(0, &DType::Binary(dtype.nullability()), len)?;
677        let codes: VarBinArray = codes
678            .as_opt::<VarBin>()
679            .ok_or_else(|| {
680                vortex_err!(
681                    "Expected VarBinArray for codes, got {}",
682                    codes.encoding_id()
683                )
684            })?
685            .into_owned();
686        let uncompressed_lengths = children.get(
687            1,
688            &DType::Primitive(
689                metadata.get_uncompressed_lengths_ptype()?,
690                Nullability::NonNullable,
691            ),
692            len,
693        )?;
694
695        FSSTData::validate_parts_from_codes(
696            symbols.as_slice(),
697            symbol_lengths.as_slice(),
698            &codes,
699            &uncompressed_lengths,
700            dtype,
701            len,
702            ctx,
703        )?;
704        let slots = FSSTData::make_slots(&codes, &uncompressed_lengths);
705        let codes_bytes = codes.bytes_handle().clone();
706        let data = FSSTData::try_new(symbols.clone(), symbol_lengths.clone(), codes_bytes, len)?;
707        Ok(ArrayParts::new(
708            self.clone(),
709            dtype.clone(),
710            len,
711            data,
712            slots,
713        ))
714    }
715
716    pub(crate) unsafe fn new_unchecked_with_symbol_table(
717        dtype: DType,
718        symbol_table: Arc<FSSTSymbolTable>,
719        codes: VarBinArray,
720        uncompressed_lengths: ArrayRef,
721    ) -> FSSTArray {
722        let len = codes.len();
723        let slots = FSSTData::make_slots(&codes, &uncompressed_lengths);
724        let codes_bytes = codes.bytes_handle().clone();
725        let data =
726            unsafe { FSSTData::new_unchecked_with_symbol_table(symbol_table, codes_bytes, len) };
727        unsafe { Array::from_parts_unchecked(ArrayParts::new(FSST, dtype, len, data, slots)) }
728    }
729}
730
731impl FSSTData {
732    fn make_slots(codes: &VarBinArray, uncompressed_lengths: &ArrayRef) -> ArraySlots {
733        FSSTSlots {
734            uncompressed_lengths: uncompressed_lengths.clone(),
735            codes_offsets: codes.offsets().clone(),
736            codes_validity: validity_to_child(
737                &codes
738                    .validity()
739                    .vortex_expect("FSST codes validity should be derivable"),
740                codes.len(),
741            ),
742        }
743        .into_slots()
744    }
745
746    /// Build FSST data from a set of `symbols`, `symbol_lengths`, and compressed codes bytes.
747    ///
748    /// Symbols are 8-bytes and can represent short strings, each of which is assigned
749    /// a code.
750    ///
751    /// The `codes_bytes` buffer contains the concatenated compressed bytecodes for all elements.
752    /// Each element's compressed bytecodes are a sequence of 8-bit codes, where each code
753    /// corresponds either to a symbol or to the "escape code" (which tells the decoder to
754    /// emit the following byte without doing a table lookup).
755    ///
756    /// `symbols` and `symbol_lengths` hold only the populated entries; they are padded out to
757    /// [`FSST_SYMBOL_TABLE_LEN`] entries so that [`Self::decompressor`] can borrow them without
758    /// copying.
759    ///
760    /// The offsets and validity for the codes are stored in the array's slots, not here.
761    /// Use [`FSSTArrayExt::codes()`] to reconstruct a full `VarBinArray`.
762    pub fn try_new(
763        symbols: Buffer<Symbol>,
764        symbol_lengths: Buffer<u8>,
765        codes_bytes: BufferHandle,
766        len: usize,
767    ) -> VortexResult<Self> {
768        vortex_ensure_eq!(
769            symbols.len(),
770            symbol_lengths.len(),
771            InvalidArgument: "symbols and symbol_lengths arrays must have same length"
772        );
773        vortex_ensure!(
774            symbols.len() <= FSST_SYMBOL_TABLE_LEN,
775            InvalidArgument: "symbols array must have length <= {FSST_SYMBOL_TABLE_LEN}, found {}",
776            symbols.len()
777        );
778        let n_symbols = symbols.len();
779        // SAFETY: the symbol table shape is validated above.
780        let symbol_table = Arc::new(FSSTSymbolTable::new(symbols, symbol_lengths, n_symbols));
781        unsafe {
782            Ok(Self::new_unchecked_with_symbol_table(
783                symbol_table,
784                codes_bytes,
785                len,
786            ))
787        }
788    }
789
790    pub fn validate(
791        &self,
792        dtype: &DType,
793        len: usize,
794        slots: &[Option<ArrayRef>],
795        ctx: &mut ExecutionCtx,
796    ) -> VortexResult<()> {
797        let fsst_slots = FSSTSlotsView::from_slots(slots);
798        Self::validate_parts(
799            self.symbol_table.symbols(),
800            self.symbol_table.symbol_lengths(),
801            &self.codes_bytes,
802            fsst_slots.codes_offsets,
803            dtype.nullability(),
804            fsst_slots.uncompressed_lengths,
805            dtype,
806            len,
807            ctx,
808        )
809    }
810
811    /// Validate using the decomposed components (codes bytes + offsets + nullability).
812    #[expect(clippy::too_many_arguments)]
813    fn validate_parts(
814        symbols: &[Symbol],
815        symbol_lengths: &[u8],
816        codes_bytes: &BufferHandle,
817        codes_offsets: &ArrayRef,
818        codes_nullability: Nullability,
819        uncompressed_lengths: &ArrayRef,
820        dtype: &DType,
821        len: usize,
822        ctx: &mut ExecutionCtx,
823    ) -> VortexResult<()> {
824        vortex_ensure!(
825            matches!(dtype, DType::Binary(_) | DType::Utf8(_)),
826            "FSST arrays must be Binary or Utf8, found {dtype}"
827        );
828
829        if symbols.len() > FSST_SYMBOL_TABLE_LEN {
830            vortex_bail!(InvalidArgument: "symbols array must have length <= {FSST_SYMBOL_TABLE_LEN}");
831        }
832
833        if symbols.len() != symbol_lengths.len() {
834            vortex_bail!(InvalidArgument: "symbols and symbol_lengths arrays must have same length");
835        }
836
837        Self::validate_symbol_lengths(symbol_lengths)?;
838
839        // codes_offsets.len() - 1 == number of elements
840        let codes_len = codes_offsets.len().saturating_sub(1);
841        if codes_len != len {
842            vortex_bail!(InvalidArgument: "codes must have same len as outer array");
843        }
844
845        if uncompressed_lengths.len() != len {
846            vortex_bail!(InvalidArgument: "uncompressed_lengths must be same len as codes");
847        }
848
849        if !uncompressed_lengths.dtype().is_int() || uncompressed_lengths.dtype().is_nullable() {
850            vortex_bail!(InvalidArgument: "uncompressed_lengths must have integer type and cannot be nullable, found {}", uncompressed_lengths.dtype());
851        }
852
853        // Offsets must be non-nullable integer.
854        if !codes_offsets.dtype().is_int() || codes_offsets.dtype().is_nullable() {
855            vortex_bail!(InvalidArgument: "codes offsets must be non-nullable integer type, found {}", codes_offsets.dtype());
856        }
857
858        if codes_nullability != dtype.nullability() {
859            vortex_bail!(InvalidArgument: "codes nullability must match outer dtype nullability");
860        }
861
862        // Validate that last offset doesn't exceed bytes length (when host-resident). A
863        // compressed offsets child would have to be decoded to read its last offset, so it is
864        // checked by the decode plan instead, once `execute` has made it primitive.
865        if codes_bytes.is_on_host()
866            && codes_offsets.is_host()
867            && codes_offsets.is::<Primitive>()
868            && !codes_offsets.is_empty()
869        {
870            let last_offset: usize = (&codes_offsets
871                .execute_scalar(codes_offsets.len() - 1, ctx)
872                .vortex_expect("offsets must support scalar_at"))
873                .try_into()
874                .vortex_expect("Failed to convert offset to usize");
875            vortex_ensure!(
876                last_offset <= codes_bytes.len(),
877                InvalidArgument: "Last codes offset {} exceeds codes bytes length {}",
878                last_offset,
879                codes_bytes.len()
880            );
881        }
882
883        Ok(())
884    }
885
886    fn validate_symbol_lengths(symbol_lengths: &[u8]) -> VortexResult<()> {
887        let mut expected = 2;
888        for (idx, &len) in symbol_lengths.iter().enumerate() {
889            if len > 8 || len == 0 {
890                vortex_bail!(InvalidArgument: "symbol length at index {idx} must be between 1 and 8, found {len}");
891            }
892
893            if expected == 1 {
894                if len != 1 {
895                    vortex_bail!(InvalidArgument: "symbol length at index {idx} must be 1 after one-byte symbols begin, found {len}");
896                }
897            } else {
898                if len == 1 {
899                    expected = 1;
900                }
901
902                if len < expected {
903                    vortex_bail!(InvalidArgument: "symbol length at index {idx} violates FSST symbol table ordering");
904                }
905                expected = len;
906            }
907        }
908
909        Ok(())
910    }
911
912    /// Validate using a VarBinArray for the codes (convenience for construction paths).
913    fn validate_parts_from_codes(
914        symbols: &[Symbol],
915        symbol_lengths: &[u8],
916        codes: &VarBinArray,
917        uncompressed_lengths: &ArrayRef,
918        dtype: &DType,
919        len: usize,
920        ctx: &mut ExecutionCtx,
921    ) -> VortexResult<()> {
922        Self::validate_parts(
923            symbols,
924            symbol_lengths,
925            codes.bytes_handle(),
926            codes.offsets(),
927            codes.dtype().nullability(),
928            uncompressed_lengths,
929            dtype,
930            len,
931            ctx,
932        )
933    }
934
935    pub(crate) unsafe fn new_unchecked_with_symbol_table(
936        symbol_table: Arc<FSSTSymbolTable>,
937        codes_bytes: BufferHandle,
938        len: usize,
939    ) -> Self {
940        Self {
941            symbol_table,
942            codes_bytes,
943            len,
944        }
945    }
946
947    /// Returns the number of elements in the array.
948    pub fn len(&self) -> usize {
949        self.len
950    }
951
952    /// Returns `true` if the array contains no elements.
953    pub fn is_empty(&self) -> bool {
954        self.len == 0
955    }
956
957    /// Access the symbol table array.
958    ///
959    /// This is the populated prefix of the symbol table; the padding added on construction to
960    /// reach [`FSST_SYMBOL_TABLE_LEN`] is not included. Use [`Self::padded_symbols`] to get the
961    /// whole buffer.
962    pub fn symbols(&self) -> &[Symbol] {
963        &self.symbol_table.padded_symbols().as_slice()[0..self.symbol_table.n_symbols]
964    }
965
966    /// Access the symbol lengths array.
967    ///
968    /// As with [`Self::symbols`], this excludes the padding added on construction.
969    pub fn symbol_lengths(&self) -> &[u8] {
970        &self.symbol_table.padded_symbol_lengths().as_slice()[0..self.symbol_table.n_symbols]
971    }
972
973    /// The whole symbols buffer, padded to exactly [`FSST_SYMBOL_TABLE_LEN`] entries with
974    /// [`Symbol::ZERO`]. Entries at or past [`Self::n_symbols`] are padding.
975    pub fn padded_symbols(&self) -> &Buffer<Symbol> {
976        self.symbol_table.padded_symbols()
977    }
978
979    /// The whole symbol lengths buffer, padded to exactly [`FSST_SYMBOL_TABLE_LEN`] entries with
980    /// zeros. Entries at or past [`Self::n_symbols`] are padding.
981    pub fn padded_symbol_lengths(&self) -> &Buffer<u8> {
982        self.symbol_table.padded_symbol_lengths()
983    }
984
985    /// The number of populated entries in the symbol table.
986    pub fn n_symbols(&self) -> usize {
987        self.symbol_table.n_symbols
988    }
989
990    pub fn symbol_table(&self) -> Arc<FSSTSymbolTable> {
991        Arc::clone(&self.symbol_table)
992    }
993
994    /// Access the compressed codes bytes buffer handle (may be on host or device).
995    pub fn codes_bytes_handle(&self) -> &BufferHandle {
996        &self.codes_bytes
997    }
998
999    /// Access the compressed codes bytes on the host.
1000    pub fn codes_bytes(&self) -> &ByteBuffer {
1001        self.codes_bytes.as_host()
1002    }
1003
1004    /// Build a [`Decompressor`] that can be used to decompress values from
1005    /// this array.
1006    pub fn decompressor(&self) -> Decompressor<'_> {
1007        self.symbol_table.decompressor()
1008    }
1009
1010    /// Retrieves the FSST compressor.
1011    pub fn compressor(&self) -> &Compressor {
1012        self.symbol_table.compressor()
1013    }
1014}
1015
1016pub trait FSSTArrayExt: FSSTArraySlotsExt {
1017    fn uncompressed_lengths_dtype(&self) -> &DType {
1018        self.uncompressed_lengths().dtype()
1019    }
1020
1021    /// Reconstruct a [`VarBinArray`] for the compressed codes by combining the bytes
1022    /// from [`FSSTData`] with the offsets and validity stored in the array's slots.
1023    fn codes(&self) -> VarBinArray {
1024        let offsets = self.codes_offsets().clone();
1025        let validity =
1026            child_to_validity(self.codes_validity(), self.as_ref().dtype().nullability());
1027        let codes_bytes = self.codes_bytes_handle().clone();
1028        // SAFETY: components were validated at construction time.
1029        unsafe {
1030            VarBinArray::new_unchecked_from_handle(
1031                offsets,
1032                codes_bytes,
1033                DType::Binary(self.as_ref().dtype().nullability()),
1034                validity,
1035            )
1036        }
1037    }
1038
1039    /// Get the DType of the codes array.
1040    fn codes_dtype(&self) -> DType {
1041        DType::Binary(self.as_ref().dtype().nullability())
1042    }
1043}
1044
1045impl<T: TypedArrayRef<FSST>> FSSTArrayExt for T {}
1046
1047impl ValidityVTable<FSST> for FSST {
1048    fn validity(array: ArrayView<'_, FSST>) -> VortexResult<Validity> {
1049        Ok(child_to_validity(
1050            array.codes_validity(),
1051            array.dtype().nullability(),
1052        ))
1053    }
1054}
1055
1056#[cfg(test)]
1057mod test {
1058    use fsst::Compressor;
1059    use fsst::Symbol;
1060    use prost::Message;
1061    use vortex_array::ArrayDeserialization;
1062    use vortex_array::ArrayPlugin;
1063    use vortex_array::IntoArray;
1064    use vortex_array::VortexSessionExecute;
1065    use vortex_array::array_session;
1066    use vortex_array::arrays::VarBinViewArray;
1067    use vortex_array::buffer::BufferHandle;
1068    use vortex_array::dtype::DType;
1069    use vortex_array::dtype::Nullability;
1070    use vortex_array::dtype::PType;
1071    use vortex_array::test_harness::check_metadata;
1072    use vortex_array::vtable::VTable as _;
1073    use vortex_buffer::Buffer;
1074    use vortex_error::VortexResult;
1075    use vortex_error::vortex_err;
1076
1077    use crate::FSST;
1078    use crate::array::FSST_SYMBOL_TABLE_LEN;
1079    use crate::array::FSSTArrayExt;
1080    use crate::array::FSSTArraySlotsExt;
1081    use crate::array::FSSTData;
1082    use crate::array::FSSTMetadata;
1083    use crate::array::FSSTSymbolTable;
1084    use crate::array::padded_symbol_table;
1085    use crate::fsst_compress;
1086    use crate::fsst_train_compressor;
1087
1088    #[test]
1089    fn slice_reuses_initialized_compressor() -> VortexResult<()> {
1090        let symbols = Buffer::<Symbol>::copy_from([
1091            Symbol::from_slice(b"abc00000"),
1092            Symbol::from_slice(b"defghijk"),
1093        ]);
1094        let symbol_lengths = Buffer::<u8>::copy_from([3, 8]);
1095
1096        let compressor = Compressor::rebuild_from(symbols.as_slice(), symbol_lengths.as_slice());
1097        let mut ctx = array_session().create_execution_ctx();
1098        let strings = VarBinViewArray::from_iter_str(["abcabcab", "defghijk", "abcxyz"]);
1099        let fsst_array = fsst_compress(&strings.into_array(), &compressor, &mut ctx)?;
1100
1101        let compressor_ptr = fsst_array.compressor() as *const Compressor;
1102        let sliced = fsst_array
1103            .slice(1..3)?
1104            .try_downcast::<FSST>()
1105            .map_err(|_| vortex_err!("slice must return an FSST array"))?;
1106        let sliced_compressor_ptr = sliced.compressor() as *const Compressor;
1107
1108        assert_eq!(compressor_ptr, sliced_compressor_ptr);
1109        Ok(())
1110    }
1111
1112    /// The symbol table is padded out to [`FSST_SYMBOL_TABLE_LEN`] so that `Decompressor` can
1113    /// borrow it directly, but the logical accessors and the serialized buffers must still only
1114    /// expose the populated prefix.
1115    #[test]
1116    fn symbol_table_padded_on_creation() -> VortexResult<()> {
1117        let mut ctx = array_session().create_execution_ctx();
1118        let strings = VarBinViewArray::from_iter_str(["abcabcab", "defghijk", "abcxyz"]);
1119        let compressor = Compressor::rebuild_from(
1120            [
1121                Symbol::from_slice(b"abc00000"),
1122                Symbol::from_slice(b"defghijk"),
1123            ],
1124            [3u8, 8],
1125        );
1126        let fsst_array = fsst_compress(&strings.into_array(), &compressor, &mut ctx)?;
1127
1128        assert_eq!(fsst_array.padded_symbols().len(), FSST_SYMBOL_TABLE_LEN);
1129        assert_eq!(
1130            fsst_array.padded_symbol_lengths().len(),
1131            FSST_SYMBOL_TABLE_LEN
1132        );
1133        assert_eq!(fsst_array.padded_symbol_lengths().as_slice()[2..], [0; 253]);
1134
1135        // Accessors and serialized buffers only see the two populated symbols.
1136        assert_eq!(fsst_array.n_symbols(), 2);
1137        assert_eq!(fsst_array.symbols().len(), 2);
1138        assert_eq!(fsst_array.symbol_lengths(), &[3, 8]);
1139        assert_eq!(
1140            FSST::buffer(fsst_array.as_view(), 0).len(),
1141            2 * size_of::<Symbol>()
1142        );
1143        assert_eq!(FSST::buffer(fsst_array.as_view(), 1).len(), 2);
1144
1145        let decompressed = fsst_array
1146            .into_array()
1147            .execute::<VarBinViewArray>(&mut ctx)?;
1148        assert_eq!(decompressed.bytes_at(0).as_slice(), b"abcabcab".as_ref());
1149        assert_eq!(decompressed.bytes_at(1).as_slice(), b"defghijk".as_ref());
1150        assert_eq!(decompressed.bytes_at(2).as_slice(), b"abcxyz".as_ref());
1151        Ok(())
1152    }
1153
1154    /// Buffers arriving from a file hold the unpadded symbol table, so deserialization must pad
1155    /// them before a `Decompressor` can borrow them.
1156    #[test]
1157    fn symbol_table_padded_on_deserialize() -> VortexResult<()> {
1158        let mut ctx = array_session().create_execution_ctx();
1159        let input = VarBinViewArray::from_iter_str(["abcabcab", "defghijk"]).into_array();
1160        let compressor = fsst_train_compressor(&input, &mut ctx)?;
1161        let fsst_array = fsst_compress(&input, &compressor, &mut ctx)?;
1162
1163        let buffers = [
1164            BufferHandle::new_host(
1165                fsst_array
1166                    .padded_symbols()
1167                    .slice(0..fsst_array.n_symbols())
1168                    .into_byte_buffer(),
1169            ),
1170            BufferHandle::new_host(
1171                fsst_array
1172                    .padded_symbol_lengths()
1173                    .slice(0..fsst_array.n_symbols())
1174                    .into_byte_buffer(),
1175            ),
1176            fsst_array.codes_bytes_handle().clone(),
1177        ];
1178        assert!(buffers[1].len() < FSST_SYMBOL_TABLE_LEN);
1179
1180        let children = vec![
1181            fsst_array.uncompressed_lengths().clone(),
1182            fsst_array.codes_offsets().clone(),
1183        ];
1184
1185        let deserialized = ArrayPlugin::deserialize(
1186            &FSST,
1187            ArrayDeserialization::new(
1188                vortex_array::ArrayVTable::id(&FSST),
1189                &DType::Utf8(Nullability::NonNullable),
1190                2,
1191                &FSSTMetadata {
1192                    uncompressed_lengths_ptype: fsst_array
1193                        .uncompressed_lengths()
1194                        .dtype()
1195                        .as_ptype()
1196                        .into(),
1197                    codes_offsets_ptype: fsst_array.codes_offsets().dtype().as_ptype().into(),
1198                }
1199                .encode_to_vec(),
1200                &buffers,
1201                &children.as_slice(),
1202            ),
1203            &array_session(),
1204        )?;
1205
1206        let padded = deserialized
1207            .clone()
1208            .try_downcast::<FSST>()
1209            .map_err(|_| vortex_err!("deserialize must return an FSST array"))?;
1210        assert_eq!(padded.padded_symbols().len(), FSST_SYMBOL_TABLE_LEN);
1211        assert_eq!(padded.n_symbols(), fsst_array.symbols().len());
1212
1213        let decompressed = deserialized.execute::<VarBinViewArray>(&mut ctx)?;
1214        assert_eq!(decompressed.bytes_at(0).as_slice(), b"abcabcab".as_ref());
1215        assert_eq!(decompressed.bytes_at(1).as_slice(), b"defghijk".as_ref());
1216        Ok(())
1217    }
1218
1219    /// An already-padded table must be stored as-is rather than copied again.
1220    #[test]
1221    fn padded_constructor_does_not_repad() -> VortexResult<()> {
1222        let symbols = padded_symbol_table(&[Symbol::from_slice(b"ab000000")], Symbol::ZERO);
1223        let symbol_lengths = padded_symbol_table(&[2u8], 0);
1224        let symbols_ptr = symbols.as_slice().as_ptr();
1225        let symbol_lengths_ptr = symbol_lengths.as_slice().as_ptr();
1226
1227        let data = FSSTSymbolTable::new_padded(symbols, symbol_lengths, 1)?;
1228
1229        assert_eq!(data.padded_symbols().as_slice().as_ptr(), symbols_ptr);
1230        assert_eq!(
1231            data.padded_symbol_lengths().as_slice().as_ptr(),
1232            symbol_lengths_ptr
1233        );
1234        assert_eq!(data.symbols().len(), 1);
1235        Ok(())
1236    }
1237
1238    /// [`FSSTData::try_new_padded`] stores the buffers as given, so it must reject tables that are
1239    /// not already padded.
1240    #[test]
1241    fn rejects_unpadded_input_to_padded_constructor() {
1242        assert!(
1243            FSSTSymbolTable::new_padded(
1244                Buffer::<Symbol>::copy_from([Symbol::from_slice(b"ab000000")]),
1245                Buffer::<u8>::copy_from([2]),
1246                1,
1247            )
1248            .is_err()
1249        );
1250        assert!(
1251            FSSTSymbolTable::new_padded(
1252                Buffer::<Symbol>::full(Symbol::ZERO, FSST_SYMBOL_TABLE_LEN),
1253                Buffer::<u8>::full(0, FSST_SYMBOL_TABLE_LEN),
1254                FSST_SYMBOL_TABLE_LEN + 1,
1255            )
1256            .is_err()
1257        );
1258    }
1259
1260    #[test]
1261    fn rejects_malformed_symbol_table() {
1262        let codes_bytes = BufferHandle::new_host(Buffer::<u8>::empty());
1263        assert!(
1264            FSSTData::try_new(
1265                Buffer::<Symbol>::copy_from([Symbol::from_slice(b"ab000000")]),
1266                Buffer::<u8>::copy_from([2, 2]),
1267                codes_bytes.clone(),
1268                0,
1269            )
1270            .is_err()
1271        );
1272        assert!(
1273            FSSTData::try_new(
1274                Buffer::<Symbol>::full(Symbol::from_slice(b"ab000000"), FSST_SYMBOL_TABLE_LEN + 1,),
1275                Buffer::<u8>::full(2, FSST_SYMBOL_TABLE_LEN + 1),
1276                codes_bytes,
1277                0,
1278            )
1279            .is_err()
1280        );
1281    }
1282
1283    #[cfg_attr(miri, ignore)]
1284    #[test]
1285    fn test_fsst_metadata() {
1286        check_metadata(
1287            "fsst.metadata",
1288            &FSSTMetadata {
1289                uncompressed_lengths_ptype: PType::U64 as i32,
1290                codes_offsets_ptype: PType::I32 as i32,
1291            }
1292            .encode_to_vec(),
1293        );
1294    }
1295
1296    /// The original FSST array stored codes as a VarBinArray child and required that the child
1297    /// have this encoding. Vortex forbids this kind of introspection, therefore we had to fix
1298    /// the array to store the compressed offsets and compressed data buffer separately, and only
1299    /// use VarBinArray to delegate behavior.
1300    ///
1301    /// This test manually constructs an old-style FSST array and ensures that it can still be
1302    /// deserialized.
1303    #[test]
1304    fn test_back_compat() -> VortexResult<()> {
1305        let symbols = Buffer::<Symbol>::copy_from([
1306            Symbol::from_slice(b"abc00000"),
1307            Symbol::from_slice(b"defghijk"),
1308        ]);
1309        let symbol_lengths = Buffer::<u8>::copy_from([3, 8]);
1310
1311        let compressor = Compressor::rebuild_from(symbols.as_slice(), symbol_lengths.as_slice());
1312        let mut ctx = array_session().create_execution_ctx();
1313        let input = VarBinViewArray::from_iter_str(["abcabcab", "defghijk"]);
1314        let fsst_array = fsst_compress(&input.into_array(), &compressor, &mut ctx)?;
1315
1316        let compressed_codes = fsst_array.codes();
1317
1318        // There were two buffers:
1319        // 1. The 8 byte symbols
1320        // 2. The symbol lengths as u8.
1321        let buffers = [
1322            BufferHandle::new_host(symbols.into_byte_buffer()),
1323            BufferHandle::new_host(symbol_lengths.into_byte_buffer()),
1324        ];
1325
1326        // There were 2 children:
1327        // 1. The compressed codes, stored as a VarBinArray.
1328        // 2. The uncompressed lengths, stored as a Primitive array.
1329        let children = vec![
1330            compressed_codes.into_array(),
1331            fsst_array.uncompressed_lengths().clone(),
1332        ];
1333
1334        let fsst = ArrayPlugin::deserialize(
1335            &FSST,
1336            ArrayDeserialization::new(
1337                vortex_array::ArrayVTable::id(&FSST),
1338                &DType::Utf8(Nullability::NonNullable),
1339                2,
1340                &FSSTMetadata {
1341                    uncompressed_lengths_ptype: fsst_array
1342                        .uncompressed_lengths()
1343                        .dtype()
1344                        .as_ptype()
1345                        .into(),
1346                    // Legacy array did not store this field, use Protobuf default of 0.
1347                    codes_offsets_ptype: 0,
1348                }
1349                .encode_to_vec(),
1350                &buffers,
1351                &children.as_slice(),
1352            ),
1353            &array_session(),
1354        )?;
1355
1356        let decompressed =
1357            fsst.execute::<VarBinViewArray>(&mut array_session().create_execution_ctx())?;
1358        let mask = decompressed
1359            .validity()?
1360            .execute_mask(decompressed.len(), &mut ctx)?;
1361        assert!(mask.value(0));
1362        assert_eq!(decompressed.bytes_at(0).as_slice(), b"abcabcab".as_ref());
1363        assert!(mask.value(1));
1364        assert_eq!(decompressed.bytes_at(1).as_slice(), b"defghijk".as_ref());
1365        Ok(())
1366    }
1367}