vortex_fsst/compute/
mod.rs

1mod compare;
2mod filter;
3
4use vortex_array::arrays::VarBinVTable;
5use vortex_array::compute::{TakeKernel, TakeKernelAdapter, fill_null, take};
6use vortex_array::{Array, ArrayRef, IntoArray, register_kernel};
7use vortex_error::VortexResult;
8use vortex_scalar::{Scalar, ScalarValue};
9
10use crate::{FSSTArray, FSSTVTable};
11
12impl TakeKernel for FSSTVTable {
13    // Take on an FSSTArray is a simple take on the codes array.
14    fn take(&self, array: &FSSTArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
15        Ok(FSSTArray::try_new(
16            array
17                .dtype()
18                .clone()
19                .union_nullability(indices.dtype().nullability()),
20            array.symbols().clone(),
21            array.symbol_lengths().clone(),
22            take(array.codes().as_ref(), indices)?
23                .as_::<VarBinVTable>()
24                .clone(),
25            fill_null(
26                &take(array.uncompressed_lengths(), indices)?,
27                &Scalar::new(
28                    array.uncompressed_lengths_dtype().clone(),
29                    ScalarValue::from(0),
30                ),
31            )?,
32        )?
33        .into_array())
34    }
35}
36
37register_kernel!(TakeKernelAdapter(FSSTVTable).lift());
38
39#[cfg(test)]
40mod tests {
41    use vortex_array::arrays::{PrimitiveArray, VarBinArray};
42    use vortex_array::compute::take;
43    use vortex_dtype::{DType, Nullability};
44
45    use crate::{fsst_compress, fsst_train_compressor};
46
47    #[test]
48    fn test_take_null() {
49        let arr = VarBinArray::from_iter([Some("h")], DType::Utf8(Nullability::NonNullable));
50        let compr = fsst_train_compressor(arr.as_ref()).unwrap();
51        let fsst = fsst_compress(arr.as_ref(), &compr).unwrap();
52
53        let idx1: PrimitiveArray = (0..1).collect();
54
55        assert_eq!(
56            take(fsst.as_ref(), idx1.as_ref()).unwrap().dtype(),
57            &DType::Utf8(Nullability::NonNullable)
58        );
59
60        let idx2: PrimitiveArray = PrimitiveArray::from_option_iter(vec![Some(0)]);
61
62        assert_eq!(
63            take(fsst.as_ref(), idx2.as_ref()).unwrap().dtype(),
64            &DType::Utf8(Nullability::Nullable)
65        );
66    }
67}
vortex_fsst/compute/mod.rs

vortex_fsst/compute/
mod.rs
