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