vortex_btrblocks/schemes/string/
fsst.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! FSST (Fast Static Symbol Table) string compression.
5
6use std::sync::Arc;
7
8use vortex_array::ArrayId;
9use vortex_array::ArrayRef;
10use vortex_array::Canonical;
11use vortex_array::ExecutionCtx;
12use vortex_array::IntoArray;
13use vortex_array::VTable;
14use vortex_array::arrays::PrimitiveArray;
15use vortex_array::arrays::VarBin;
16use vortex_array::arrays::VarBinArray;
17use vortex_array::arrays::primitive::PrimitiveArrayExt;
18use vortex_array::arrays::varbin::VarBinArraySlotsExt;
19use vortex_compressor::scheme::CompressionEstimate;
20use vortex_compressor::scheme::DeferredEstimate;
21use vortex_error::VortexResult;
22use vortex_fsst::FSST;
23use vortex_fsst::FSSTArrayExt;
24use vortex_fsst::FSSTArraySlotsExt;
25use vortex_fsst::FSSTSymbolTable;
26use vortex_fsst::fsst_compress;
27use vortex_fsst::fsst_train_compressor;
28
29use crate::ArrayAndStats;
30use crate::CascadingCompressor;
31use crate::CompressorContext;
32use crate::Scheme;
33use crate::SchemeExt;
34
35/// FSST (Fast Static Symbol Table) compression.
36///
37/// One of the two string-fragmentation schemes in the default
38/// [`crate::ALL_SCHEMES`] (alongside `OnPairScheme`); the sample-based selector
39/// keeps whichever is smaller per column. FSST compresses faster, OnPair
40/// usually wins on ratio.
41#[derive(Debug, Copy, Clone, PartialEq, Eq)]
42pub struct FSSTScheme;
43
44impl Scheme for FSSTScheme {
45    fn scheme_name(&self) -> &'static str {
46        "vortex.string.fsst"
47    }
48
49    fn matches(&self, canonical: &Canonical) -> bool {
50        canonical.dtype().is_utf8()
51    }
52
53    fn produced_encodings(&self) -> Vec<ArrayId> {
54        vec![FSST.id(), VarBin.id()]
55    }
56
57    /// Children: lengths=0, code_offsets=1.
58    fn num_children(&self) -> usize {
59        2
60    }
61
62    fn expected_compression_ratio(
63        &self,
64        _data: &ArrayAndStats,
65        _compress_ctx: CompressorContext,
66        _exec_ctx: &mut ExecutionCtx,
67    ) -> CompressionEstimate {
68        CompressionEstimate::Deferred(DeferredEstimate::Sample)
69    }
70
71    fn compress(
72        &self,
73        compressor: &CascadingCompressor,
74        data: &ArrayAndStats,
75        compress_ctx: CompressorContext,
76        exec_ctx: &mut ExecutionCtx,
77    ) -> VortexResult<ArrayRef> {
78        let utf8 = data.array_as_varbinview().into_owned().into_array();
79        let compressor_fsst = fsst_train_compressor(&utf8, exec_ctx)?;
80        let fsst = fsst_compress(&utf8, &compressor_fsst, exec_ctx)?;
81
82        let uncompressed_lengths_primitive = fsst
83            .uncompressed_lengths()
84            .clone()
85            .execute::<PrimitiveArray>(exec_ctx)?
86            .narrow(exec_ctx)?;
87        let compressed_original_lengths = compressor.compress_child(
88            &uncompressed_lengths_primitive.into_array(),
89            &compress_ctx,
90            self.id(),
91            0,
92            exec_ctx,
93        )?;
94
95        let codes_offsets_primitive = fsst
96            .codes()
97            .offsets()
98            .clone()
99            .execute::<PrimitiveArray>(exec_ctx)?
100            .narrow(exec_ctx)?;
101        let compressed_codes_offsets = compressor.compress_child(
102            &codes_offsets_primitive.into_array(),
103            &compress_ctx,
104            self.id(),
105            1,
106            exec_ctx,
107        )?;
108        let compressed_codes = VarBinArray::try_new(
109            compressed_codes_offsets,
110            fsst.codes().bytes().clone(),
111            fsst.codes().dtype().clone(),
112            fsst.codes().validity()?,
113        )?;
114
115        // Reuse the padded symbol table as-is; only the codes and lengths change here.
116        let fsst = FSST::try_new_with_symbol_table(
117            fsst.dtype().clone(),
118            Arc::new(FSSTSymbolTable::new_padded(
119                fsst.padded_symbols().clone(),
120                fsst.padded_symbol_lengths().clone(),
121                fsst.n_symbols(),
122            )?),
123            compressed_codes,
124            compressed_original_lengths,
125            exec_ctx,
126        )?;
127
128        Ok(fsst.into_array())
129    }
130}
vortex_btrblocks/schemes/string/fsst.rs

vortex_btrblocks/schemes/string/
fsst.rs
