vortex_fsst/compute/
cast.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_array::ArrayRef;
5use vortex_array::ArrayView;
6use vortex_array::ExecutionCtx;
7use vortex_array::IntoArray;
8use vortex_array::arrays::VarBinArray;
9use vortex_array::arrays::varbin::VarBinArrayExt;
10use vortex_array::dtype::DType;
11use vortex_array::scalar_fn::fns::cast::CastKernel;
12use vortex_array::scalar_fn::fns::cast::CastReduce;
13use vortex_array::validity::Validity;
14use vortex_error::VortexResult;
15
16use crate::FSST;
17use crate::FSSTArrayExt;
18
19fn build_with_codes_validity(
20    array: ArrayView<'_, FSST>,
21    dtype: &DType,
22    new_codes_validity: Validity,
23) -> VortexResult<ArrayRef> {
24    let codes = array.codes();
25    let new_codes = VarBinArray::try_new(
26        codes.offsets().clone(),
27        codes.bytes().clone(),
28        codes.dtype().with_nullability(dtype.nullability()),
29        new_codes_validity,
30    )?;
31
32    Ok(unsafe {
33        FSST::new_unchecked_with_symbol_table(
34            dtype.clone(),
35            array.symbol_table(),
36            new_codes,
37            array.uncompressed_lengths().clone(),
38        )
39    }
40    .into_array())
41}
42
43impl CastReduce for FSST {
44    fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
45        if !array.dtype().eq_ignore_nullability(dtype) {
46            return Ok(None);
47        }
48
49        let codes = array.codes();
50        let Some(new_codes_validity) = codes
51            .validity()?
52            .trivially_cast_nullability(dtype.nullability(), codes.len())?
53        else {
54            return Ok(None);
55        };
56
57        Ok(Some(build_with_codes_validity(
58            array,
59            dtype,
60            new_codes_validity,
61        )?))
62    }
63}
64
65impl CastKernel for FSST {
66    fn cast(
67        array: ArrayView<'_, Self>,
68        dtype: &DType,
69        ctx: &mut ExecutionCtx,
70    ) -> VortexResult<Option<ArrayRef>> {
71        if !array.dtype().eq_ignore_nullability(dtype) {
72            return Ok(None);
73        }
74
75        let codes = array.codes();
76        let new_codes_validity =
77            codes
78                .validity()?
79                .cast_nullability(dtype.nullability(), codes.len(), ctx)?;
80
81        Ok(Some(build_with_codes_validity(
82            array,
83            dtype,
84            new_codes_validity,
85        )?))
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use std::sync::LazyLock;
92
93    use rstest::rstest;
94    use vortex_array::IntoArray;
95    use vortex_array::VortexSessionExecute;
96    use vortex_array::arrays::VarBinArray;
97    use vortex_array::builtins::ArrayBuiltins;
98    use vortex_array::compute::conformance::cast::test_cast_conformance;
99    use vortex_array::dtype::DType;
100    use vortex_array::dtype::Nullability;
101    use vortex_error::VortexResult;
102    use vortex_session::VortexSession;
103
104    use crate::fsst_compress;
105    use crate::fsst_train_compressor;
106    use crate::initialize;
107
108    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
109        let session = vortex_array::array_session();
110        initialize(&session);
111        session
112    });
113
114    #[test]
115    fn test_cast_fsst_nullability() -> VortexResult<()> {
116        let mut ctx = SESSION.create_execution_ctx();
117        let strings = VarBinArray::from_iter(
118            vec![Some("hello"), Some("world"), Some("hello world")],
119            DType::Utf8(Nullability::NonNullable),
120        )
121        .into_array();
122
123        let compressor = fsst_train_compressor(&strings, &mut ctx)?;
124        let fsst = fsst_compress(&strings, &compressor, &mut ctx)?;
125
126        // Cast to nullable
127        let casted = fsst.into_array().cast(DType::Utf8(Nullability::Nullable))?;
128        assert_eq!(casted.dtype(), &DType::Utf8(Nullability::Nullable));
129        Ok(())
130    }
131
132    #[rstest]
133    #[case(VarBinArray::from_iter(
134        vec![Some("hello"), Some("world"), Some("hello world")],
135        DType::Utf8(Nullability::NonNullable)
136    ))]
137    #[case(VarBinArray::from_iter(
138        vec![Some("foo"), None, Some("bar"), Some("foobar")],
139        DType::Utf8(Nullability::Nullable)
140    ))]
141    #[case(VarBinArray::from_iter(
142        vec![Some("test")],
143        DType::Utf8(Nullability::NonNullable)
144    ))]
145    fn test_cast_fsst_conformance(#[case] array: VarBinArray) -> VortexResult<()> {
146        let mut ctx = SESSION.create_execution_ctx();
147        let array = array.into_array();
148        let compressor = fsst_train_compressor(&array, &mut ctx)?;
149        let fsst = fsst_compress(&array, &compressor, &mut ctx)?;
150        test_cast_conformance(&fsst.into_array());
151        Ok(())
152    }
153}
vortex_fsst/compute/cast.rs

vortex_fsst/compute/
cast.rs
