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

vortex_fsst/compute/
mod.rs
