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