vortex_fsst/compute/
mod.rs

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

vortex_fsst/compute/
mod.rs
