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    // SAFETY: the offsets and bytes are unchanged from valid codes, and the validity was cast
27    // from the codes' own validity, so it has the same length.
28    let new_codes = unsafe {
29        VarBinArray::new_unchecked_from_handle(
30            codes.offsets().clone(),
31            codes.bytes_handle().clone(),
32            codes.dtype().with_nullability(dtype.nullability()),
33            new_codes_validity,
34        )
35    };
36
37    Ok(unsafe {
38        FSST::new_unchecked_with_symbol_table(
39            dtype.clone(),
40            array.symbol_table(),
41            new_codes,
42            array.uncompressed_lengths().clone(),
43        )
44    }
45    .into_array())
46}
47
48impl CastReduce for FSST {
49    fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
50        if !array.dtype().eq_ignore_nullability(dtype) {
51            return Ok(None);
52        }
53
54        let codes = array.codes();
55        let Some(new_codes_validity) = codes
56            .validity()?
57            .trivially_cast_nullability(dtype.nullability(), codes.len())?
58        else {
59            return Ok(None);
60        };
61
62        Ok(Some(build_with_codes_validity(
63            array,
64            dtype,
65            new_codes_validity,
66        )?))
67    }
68}
69
70impl CastKernel for FSST {
71    fn cast(
72        array: ArrayView<'_, Self>,
73        dtype: &DType,
74        ctx: &mut ExecutionCtx,
75    ) -> VortexResult<Option<ArrayRef>> {
76        if !array.dtype().eq_ignore_nullability(dtype) {
77            return Ok(None);
78        }
79
80        let codes = array.codes();
81        let new_codes_validity =
82            codes
83                .validity()?
84                .cast_nullability(dtype.nullability(), codes.len(), ctx)?;
85
86        Ok(Some(build_with_codes_validity(
87            array,
88            dtype,
89            new_codes_validity,
90        )?))
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use std::sync::LazyLock;
97
98    use rstest::rstest;
99    use vortex_array::IntoArray;
100    use vortex_array::VortexSessionExecute;
101    use vortex_array::arrays::VarBinArray;
102    use vortex_array::compute::conformance::cast::test_cast_conformance;
103    use vortex_array::dtype::DType;
104    use vortex_array::dtype::Nullability;
105    use vortex_error::VortexResult;
106    use vortex_session::VortexSession;
107
108    use crate::fsst_compress;
109    use crate::fsst_train_compressor;
110    use crate::initialize;
111
112    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
113        let session = vortex_array::array_session();
114        initialize(&session);
115        session
116    });
117
118    #[rstest]
119    #[case(VarBinArray::from_iter(
120        vec![Some("hello"), Some("world"), Some("hello world")],
121        DType::Utf8(Nullability::NonNullable)
122    ))]
123    #[case(VarBinArray::from_iter(
124        vec![Some("foo"), None, Some("bar"), Some("foobar")],
125        DType::Utf8(Nullability::Nullable)
126    ))]
127    #[case(VarBinArray::from_iter(
128        vec![Some("test")],
129        DType::Utf8(Nullability::NonNullable)
130    ))]
131    fn test_cast_fsst_conformance(#[case] array: VarBinArray) -> VortexResult<()> {
132        let mut ctx = SESSION.create_execution_ctx();
133        let array = array.into_array();
134        let compressor = fsst_train_compressor(&array, &mut ctx)?;
135        let fsst = fsst_compress(&array, &compressor, &mut ctx)?;
136        test_cast_conformance(&fsst.into_array(), &mut ctx);
137        Ok(())
138    }
139}
vortex_fsst/compute/cast.rs

vortex_fsst/compute/
cast.rs
