vortex-fsst 0.69.0

Vortex FSST string array encoding
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use vortex_array::ArrayRef;
use vortex_array::ArrayView;
use vortex_array::IntoArray;
use vortex_array::arrays::VarBin;
use vortex_array::builtins::ArrayBuiltins;
use vortex_array::dtype::DType;
use vortex_array::scalar_fn::fns::cast::CastReduce;
use vortex_error::VortexResult;

use crate::FSST;
use crate::FSSTArrayExt;
impl CastReduce for FSST {
    fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
        // FSST is a string compression encoding.
        // For nullability changes, we can cast the codes and symbols arrays
        if array.dtype().eq_ignore_nullability(dtype) {
            // Cast codes array to handle nullability
            let new_codes = array
                .codes()
                .into_array()
                .cast(array.codes_dtype().with_nullability(dtype.nullability()))?;

            Ok(Some(
                FSST::try_new(
                    dtype.clone(),
                    array.symbols().clone(),
                    array.symbol_lengths().clone(),
                    new_codes.as_::<VarBin>().into_owned(),
                    array.uncompressed_lengths().clone(),
                )?
                .into_array(),
            ))
        } else {
            Ok(None)
        }
    }
}

#[cfg(test)]
mod tests {
    use rstest::rstest;
    use vortex_array::IntoArray;
    use vortex_array::arrays::VarBinArray;
    use vortex_array::builtins::ArrayBuiltins;
    use vortex_array::compute::conformance::cast::test_cast_conformance;
    use vortex_array::dtype::DType;
    use vortex_array::dtype::Nullability;

    use crate::fsst_compress;
    use crate::fsst_train_compressor;

    #[test]
    fn test_cast_fsst_nullability() {
        let strings = VarBinArray::from_iter(
            vec![Some("hello"), Some("world"), Some("hello world")],
            DType::Utf8(Nullability::NonNullable),
        );

        let compressor = fsst_train_compressor(&strings);
        let len = strings.len();
        let dtype = strings.dtype().clone();
        let fsst = fsst_compress(strings, len, &dtype, &compressor);

        // Cast to nullable
        let casted = fsst
            .into_array()
            .cast(DType::Utf8(Nullability::Nullable))
            .unwrap();
        assert_eq!(casted.dtype(), &DType::Utf8(Nullability::Nullable));
    }

    #[rstest]
    #[case(VarBinArray::from_iter(
        vec![Some("hello"), Some("world"), Some("hello world")],
        DType::Utf8(Nullability::NonNullable)
    ))]
    #[case(VarBinArray::from_iter(
        vec![Some("foo"), None, Some("bar"), Some("foobar")],
        DType::Utf8(Nullability::Nullable)
    ))]
    #[case(VarBinArray::from_iter(
        vec![Some("test")],
        DType::Utf8(Nullability::NonNullable)
    ))]
    fn test_cast_fsst_conformance(#[case] array: VarBinArray) {
        let compressor = fsst_train_compressor(&array);
        let fsst = fsst_compress(&array, array.len(), array.dtype(), &compressor);
        test_cast_conformance(&fsst.into_array());
    }
}