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::VarBinArraySlotsExt;
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;
18use crate::FSSTArraySlotsExt;
19
20fn build_with_codes_validity(
21    array: ArrayView<'_, FSST>,
22    dtype: &DType,
23    new_codes_validity: Validity,
24) -> VortexResult<ArrayRef> {
25    let codes = array.codes();
26    let new_codes = VarBinArray::try_new(
27        codes.offsets().clone(),
28        codes.bytes().clone(),
29        codes.dtype().with_nullability(dtype.nullability()),
30        new_codes_validity,
31    )?;
32
33    Ok(unsafe {
34        FSST::new_unchecked_with_symbol_table(
35            dtype.clone(),
36            array.symbol_table(),
37            new_codes,
38            array.uncompressed_lengths().clone(),
39        )
40    }
41    .into_array())
42}
43
44impl CastReduce for FSST {
45    fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
46        if !array.dtype().eq_ignore_nullability(dtype) {
47            return Ok(None);
48        }
49
50        let codes = array.codes();
51        let Some(new_codes_validity) = codes
52            .validity()?
53            .trivially_cast_nullability(dtype.nullability(), codes.len())?
54        else {
55            return Ok(None);
56        };
57
58        Ok(Some(build_with_codes_validity(
59            array,
60            dtype,
61            new_codes_validity,
62        )?))
63    }
64}
65
66impl CastKernel for FSST {
67    fn cast(
68        array: ArrayView<'_, Self>,
69        dtype: &DType,
70        ctx: &mut ExecutionCtx,
71    ) -> VortexResult<Option<ArrayRef>> {
72        if !array.dtype().eq_ignore_nullability(dtype) {
73            return Ok(None);
74        }
75
76        let codes = array.codes();
77        let new_codes_validity =
78            codes
79                .validity()?
80                .cast_nullability(dtype.nullability(), codes.len(), ctx)?;
81
82        Ok(Some(build_with_codes_validity(
83            array,
84            dtype,
85            new_codes_validity,
86        )?))
87    }
88}
89
90#[cfg(test)]
91mod tests {
92    use std::sync::LazyLock;
93
94    use rstest::rstest;
95    use vortex_array::IntoArray;
96    use vortex_array::VortexSessionExecute;
97    use vortex_array::arrays::VarBinArray;
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    #[rstest]
115    #[case(VarBinArray::from_iter(
116        vec![Some("hello"), Some("world"), Some("hello world")],
117        DType::Utf8(Nullability::NonNullable)
118    ))]
119    #[case(VarBinArray::from_iter(
120        vec![Some("foo"), None, Some("bar"), Some("foobar")],
121        DType::Utf8(Nullability::Nullable)
122    ))]
123    #[case(VarBinArray::from_iter(
124        vec![Some("test")],
125        DType::Utf8(Nullability::NonNullable)
126    ))]
127    fn test_cast_fsst_conformance(#[case] array: VarBinArray) -> VortexResult<()> {
128        let mut ctx = SESSION.create_execution_ctx();
129        let array = array.into_array();
130        let compressor = fsst_train_compressor(&array, &mut ctx)?;
131        let fsst = fsst_compress(&array, &compressor, &mut ctx)?;
132        test_cast_conformance(&fsst.into_array(), &mut ctx);
133        Ok(())
134    }
135}
vortex_fsst/compute/cast.rs

vortex_fsst/compute/
cast.rs
