vortex_fsst/
compress.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4// Compress a set of values into an Array.
5
6use fsst::Compressor;
7use fsst::Symbol;
8use vortex_array::ExecutionCtx;
9use vortex_array::IntoArray;
10use vortex_array::accessor::ArrayAccessor;
11use vortex_array::arrays::varbin::builder::VarBinBuilder;
12use vortex_array::dtype::DType;
13use vortex_buffer::Buffer;
14use vortex_buffer::BufferMut;
15use vortex_error::VortexExpect;
16
17/// Compress a string array using FSST.
18use crate::FSST;
19use crate::FSSTArray;
20pub fn fsst_compress<A: ArrayAccessor<[u8]>>(
21    strings: A,
22    len: usize,
23    dtype: &DType,
24    compressor: &Compressor,
25    ctx: &mut ExecutionCtx,
26) -> FSSTArray {
27    strings.with_iterator(|iter| fsst_compress_iter(iter, len, dtype.clone(), compressor, ctx))
28}
29
30/// Train a compressor from an array.
31///
32/// # Panics
33///
34/// If the provided array is not FSST compressible.
35pub fn fsst_train_compressor<A: ArrayAccessor<[u8]>>(array: &A) -> Compressor {
36    array.with_iterator(|iter| fsst_train_compressor_iter(iter))
37}
38
39/// Train a [compressor][Compressor] from an iterator of bytestrings.
40fn fsst_train_compressor_iter<'a, I>(iter: I) -> Compressor
41where
42    I: Iterator<Item = Option<&'a [u8]>>,
43{
44    let mut lines = Vec::with_capacity(8_192);
45
46    for string in iter {
47        match string {
48            None => {}
49            Some(b) => lines.push(b),
50        }
51    }
52
53    Compressor::train(&lines)
54}
55
56/// Most strings are small in practice. If we encounter a larger string, we reallocate
57/// the buffer to hold enough capacity for the worst-case compressed value.
58const DEFAULT_BUFFER_LEN: usize = 1024 * 1024;
59
60/// Compress from an iterator of bytestrings using FSST.
61pub fn fsst_compress_iter<'a, I>(
62    iter: I,
63    len: usize,
64    dtype: DType,
65    compressor: &Compressor,
66    ctx: &mut ExecutionCtx,
67) -> FSSTArray
68where
69    I: Iterator<Item = Option<&'a [u8]>>,
70{
71    let mut buffer = Vec::with_capacity(DEFAULT_BUFFER_LEN);
72
73    // Offsets are widened to i64 because the cumulative compressed bytes can exceed i32::MAX for
74    // large inputs (see issue #7833). Per-string sizes still fit in i32.
75    let mut builder = VarBinBuilder::<i64>::with_capacity(len);
76    let mut uncompressed_lengths: BufferMut<i32> = BufferMut::with_capacity(len);
77
78    for string in iter {
79        match string {
80            None => {
81                builder.append_null();
82                uncompressed_lengths.push(0);
83            }
84            Some(s) => {
85                uncompressed_lengths.push(
86                    s.len()
87                        .try_into()
88                        .vortex_expect("string length must fit in i32"),
89                );
90
91                // make sure the buffer is 2x+7 larger than the input
92                let target_size = 2 * s.len() + 7;
93                if target_size > buffer.len() {
94                    let additional_capacity = target_size - buffer.len();
95                    buffer.reserve(additional_capacity);
96                }
97
98                // SAFETY: buffer is always sized to be large enough
99                unsafe { compressor.compress_into(s, &mut buffer) };
100
101                builder.append_value(&buffer);
102            }
103        }
104    }
105
106    let codes = builder.finish(DType::Binary(dtype.nullability()));
107    let symbols: Buffer<Symbol> = Buffer::copy_from(compressor.symbol_table());
108    let symbol_lengths: Buffer<u8> = Buffer::<u8>::copy_from(compressor.symbol_lengths());
109
110    let uncompressed_lengths = uncompressed_lengths.into_array();
111
112    FSST::try_new(
113        dtype,
114        symbols,
115        symbol_lengths,
116        codes,
117        uncompressed_lengths,
118        ctx,
119    )
120    .vortex_expect("FSST parts must be valid")
121}
122
123#[cfg(test)]
124mod tests {
125    use fsst::CompressorBuilder;
126    use vortex_array::LEGACY_SESSION;
127    use vortex_array::VortexSessionExecute;
128    use vortex_array::dtype::DType;
129    use vortex_array::dtype::Nullability;
130    use vortex_array::scalar::Scalar;
131
132    use crate::compress::DEFAULT_BUFFER_LEN;
133    use crate::fsst_compress_iter;
134
135    #[test]
136    fn test_large_string() {
137        let big_string: String = "abc"
138            .chars()
139            .cycle()
140            .take(10 * DEFAULT_BUFFER_LEN)
141            .collect();
142
143        let compressor = CompressorBuilder::default().build();
144
145        let mut ctx = LEGACY_SESSION.create_execution_ctx();
146        let compressed = fsst_compress_iter(
147            [Some(big_string.as_bytes())].into_iter(),
148            1,
149            DType::Utf8(Nullability::NonNullable),
150            &compressor,
151            &mut ctx,
152        );
153
154        let decoded = compressed.execute_scalar(0, &mut ctx).unwrap();
155
156        let expected = Scalar::utf8(big_string, Nullability::NonNullable);
157
158        assert_eq!(decoded, expected);
159    }
160}
vortex_fsst/compress.rs

vortex_fsst/
compress.rs
