vortex_fsst/compute/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4mod cast;
5mod compare;
6mod filter;
7mod like;
8
9use vortex_array::ArrayRef;
10use vortex_array::ArrayView;
11use vortex_array::ExecutionCtx;
12use vortex_array::IntoArray;
13use vortex_array::arrays::VarBin;
14use vortex_array::arrays::dict::TakeExecute;
15use vortex_array::builtins::ArrayBuiltins;
16use vortex_array::scalar::Scalar;
17use vortex_error::VortexExpect;
18use vortex_error::VortexResult;
19use vortex_error::vortex_err;
20
21use crate::FSST;
22use crate::FSSTArrayExt;
23
24impl TakeExecute for FSST {
25    fn take(
26        array: ArrayView<'_, Self>,
27        indices: &ArrayRef,
28        _ctx: &mut ExecutionCtx,
29    ) -> VortexResult<Option<ArrayRef>> {
30        Ok(Some(
31            FSST::try_new(
32                array
33                    .dtype()
34                    .clone()
35                    .union_nullability(indices.dtype().nullability()),
36                array.symbols().clone(),
37                array.symbol_lengths().clone(),
38                {
39                    let codes = array.codes();
40                    let codes = codes.as_view();
41                    <VarBin as TakeExecute>::take(codes, indices, _ctx)?
42                        .vortex_expect("VarBin take kernel always returns Some")
43                }
44                .try_downcast::<VarBin>()
45                .map_err(|_| vortex_err!("take for codes must return varbin array"))?,
46                array
47                    .uncompressed_lengths()
48                    .take(indices.clone())?
49                    .fill_null(Scalar::zero_value(
50                        &array.uncompressed_lengths_dtype().clone(),
51                    ))?,
52            )?
53            .into_array(),
54        ))
55    }
56}
57
58#[cfg(test)]
59mod tests {
60    use rstest::rstest;
61    use vortex_array::IntoArray;
62    use vortex_array::arrays::PrimitiveArray;
63    use vortex_array::arrays::VarBinArray;
64    use vortex_array::compute::conformance::consistency::test_array_consistency;
65    use vortex_array::compute::conformance::take::test_take_conformance;
66    use vortex_array::dtype::DType;
67    use vortex_array::dtype::Nullability;
68
69    use crate::FSSTArray;
70    use crate::fsst_compress;
71    use crate::fsst_train_compressor;
72
73    #[test]
74    fn test_take_null() {
75        let arr = VarBinArray::from_iter([Some("h")], DType::Utf8(Nullability::NonNullable));
76        let compr = fsst_train_compressor(&arr);
77        let fsst = fsst_compress(&arr, arr.len(), arr.dtype(), &compr);
78
79        let idx1: PrimitiveArray = (0..1).collect();
80
81        assert_eq!(
82            fsst.take(idx1.into_array()).unwrap().dtype(),
83            &DType::Utf8(Nullability::NonNullable)
84        );
85
86        let idx2: PrimitiveArray = PrimitiveArray::from_option_iter(vec![Some(0)]);
87
88        assert_eq!(
89            fsst.take(idx2.into_array()).unwrap().dtype(),
90            &DType::Utf8(Nullability::Nullable)
91        );
92    }
93
94    #[rstest]
95    #[case(VarBinArray::from_iter(
96        ["hello world", "testing fsst", "compression test", "data array", "vortex encoding"].map(Some),
97        DType::Utf8(Nullability::NonNullable),
98    ))]
99    #[case(VarBinArray::from_iter(
100        [Some("hello"), None, Some("world"), Some("test"), None],
101        DType::Utf8(Nullability::Nullable),
102    ))]
103    #[case(VarBinArray::from_iter(
104        ["single element"].map(Some),
105        DType::Utf8(Nullability::NonNullable),
106    ))]
107    fn test_take_fsst_conformance(#[case] varbin: VarBinArray) {
108        let compressor = fsst_train_compressor(&varbin);
109        let array = fsst_compress(&varbin, varbin.len(), varbin.dtype(), &compressor);
110        test_take_conformance(&array.into_array());
111    }
112
113    #[rstest]
114    // Basic string arrays
115    #[case::fsst_simple({
116        let varbin = VarBinArray::from_iter(
117            ["hello world", "testing fsst", "compression test", "data array", "vortex encoding"].map(Some),
118            DType::Utf8(Nullability::NonNullable),
119        );
120        let compressor = fsst_train_compressor(&varbin);
121        fsst_compress(&varbin, varbin.len(), varbin.dtype(), &compressor)
122    })]
123    // Nullable strings
124    #[case::fsst_nullable({
125        let varbin = VarBinArray::from_iter(
126            [Some("hello"), None, Some("world"), Some("test"), None],
127            DType::Utf8(Nullability::Nullable),
128        );
129        let compressor = fsst_train_compressor(&varbin);
130        let len = varbin.len();
131        let dtype = varbin.dtype().clone();
132        fsst_compress(varbin, len, &dtype, &compressor)
133    })]
134    // Repetitive patterns (good for FSST compression)
135    #[case::fsst_repetitive({
136        let varbin = VarBinArray::from_iter(
137            ["http://example.com", "http://test.com", "http://vortex.dev", "http://data.org"].map(Some),
138            DType::Utf8(Nullability::NonNullable),
139        );
140        let compressor = fsst_train_compressor(&varbin);
141        fsst_compress(&varbin, varbin.len(), varbin.dtype(), &compressor)
142    })]
143    // Edge cases
144    #[case::fsst_single({
145        let varbin = VarBinArray::from_iter(
146            ["single element"].map(Some),
147            DType::Utf8(Nullability::NonNullable),
148        );
149        let compressor = fsst_train_compressor(&varbin);
150        fsst_compress(&varbin, varbin.len(), varbin.dtype(), &compressor)
151    })]
152    #[case::fsst_empty_strings({
153        let varbin = VarBinArray::from_iter(
154            ["", "test", "", "hello", ""].map(Some),
155            DType::Utf8(Nullability::NonNullable),
156        );
157        let compressor = fsst_train_compressor(&varbin);
158        let len = varbin.len();
159        let dtype = varbin.dtype().clone();
160        fsst_compress(varbin, len, &dtype, &compressor)
161    })]
162    // Large arrays
163    #[case::fsst_large({
164        let data: Vec<Option<&str>> = (0..1500)
165            .map(|i| Some(match i % 10 {
166                0 => "https://www.example.com/page",
167                1 => "https://www.test.org/data",
168                2 => "https://www.vortex.dev/docs",
169                3 => "https://www.github.com/apache/arrow",
170                4 => "https://www.rust-lang.org/learn",
171                5 => "SELECT * FROM table WHERE id = ",
172                6 => "INSERT INTO users (name, email) VALUES",
173                7 => "UPDATE records SET status = 'active'",
174                8 => "DELETE FROM logs WHERE timestamp < ",
175                _ => "CREATE TABLE data (id INT, value TEXT)",
176            }))
177            .collect();
178        let varbin = VarBinArray::from_iter(data, DType::Utf8(Nullability::NonNullable));
179        let compressor = fsst_train_compressor(&varbin);
180        let len = varbin.len();
181        let dtype = varbin.dtype().clone();
182        fsst_compress(varbin, len, &dtype, &compressor)
183    })]
184
185    fn test_fsst_consistency(#[case] array: FSSTArray) {
186        test_array_consistency(&array.into_array());
187    }
188}
vortex_fsst/compute/mod.rs

vortex_fsst/compute/
mod.rs
