vortex-fsst 0.76.0

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

use vortex_array::ArrayVTable;
use vortex_array::arrays::Dict;
use vortex_array::arrays::Filter;
use vortex_array::arrays::dict::TakeExecuteAdaptor;
use vortex_array::arrays::filter::FilterExecuteAdaptor;
use vortex_array::optimizer::kernels::ArrayKernelsExt;
use vortex_array::scalar_fn::ScalarFnVTable;
use vortex_array::scalar_fn::fns::binary::Binary;
use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor;
use vortex_array::scalar_fn::fns::byte_length::ByteLength;
use vortex_array::scalar_fn::fns::byte_length::ByteLengthExecuteAdaptor;
use vortex_array::scalar_fn::fns::cast::Cast;
use vortex_array::scalar_fn::fns::cast::CastExecuteAdaptor;
use vortex_array::scalar_fn::fns::like::Like;
use vortex_array::scalar_fn::fns::like::LikeExecuteAdaptor;
use vortex_session::VortexSession;

use crate::FSST;

pub(super) fn initialize(session: &VortexSession) {
    let kernels = session.kernels();
    kernels.register_execute_parent_kernel(Cast.id(), FSST, CastExecuteAdaptor(FSST));
    kernels.register_execute_parent_kernel(Binary.id(), FSST, CompareExecuteAdaptor(FSST));
    kernels.register_execute_parent_kernel(Filter.id(), FSST, FilterExecuteAdaptor(FSST));
    kernels.register_execute_parent_kernel(Dict.id(), FSST, TakeExecuteAdaptor(FSST));
    kernels.register_execute_parent_kernel(Like.id(), FSST, LikeExecuteAdaptor(FSST));
    kernels.register_execute_parent_kernel(ByteLength.id(), FSST, ByteLengthExecuteAdaptor(FSST));
}

#[cfg(test)]
mod tests {
    use std::sync::LazyLock;

    use vortex_array::ArrayRef;
    use vortex_array::Canonical;
    use vortex_array::IntoArray;
    use vortex_array::VortexSessionExecute;
    use vortex_array::arrays::FilterArray;
    use vortex_array::arrays::PrimitiveArray;
    use vortex_array::arrays::varbin::builder::VarBinBuilder;
    use vortex_array::assert_arrays_eq;
    use vortex_array::dtype::DType;
    use vortex_array::dtype::Nullability;
    use vortex_array::expr::byte_length;
    use vortex_array::expr::root;
    use vortex_error::VortexResult;
    use vortex_mask::Mask;
    use vortex_session::VortexSession;

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

    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
        let session = vortex_array::array_session();
        crate::initialize(&session);
        session
    });

    fn build_test_fsst_array() -> ArrayRef {
        let mut builder = VarBinBuilder::<i32>::with_capacity(10);
        builder.append_value(b"hello world");
        builder.append_value(b"foo bar baz");
        builder.append_value(b"testing fsst compression");
        builder.append_value(b"another string here");
        builder.append_value(b"the quick brown fox");
        builder.append_value(b"jumps over the lazy dog");
        builder.append_value(b"abcdefghijklmnop");
        builder.append_value(b"qrstuvwxyz");
        builder.append_value(b"0123456789");
        builder.append_value(b"final string");
        let input = builder.finish(DType::Utf8(Nullability::NonNullable));

        let mut ctx = SESSION.create_execution_ctx();
        let arr = input.into_array();
        let compressor = fsst_train_compressor(&arr, &mut ctx).unwrap();
        fsst_compress(&arr, &compressor, &mut ctx)
            .unwrap()
            .into_array()
    }

    #[test]
    fn test_fsst_filter_simple() -> VortexResult<()> {
        let fsst_array = build_test_fsst_array();
        assert!(fsst_array.is::<FSST>());
        assert_eq!(fsst_array.len(), 10);

        // Filter 1/5 elements (every 5th element: indices 0 and 5)
        let mask = Mask::from_iter([
            true, false, false, false, false, true, false, false, false, false,
        ]);

        // Create FilterArray and execute
        let filter_array = FilterArray::new(fsst_array.clone(), mask.clone()).into_array();
        let mut ctx = SESSION.create_execution_ctx();
        let result = filter_array.execute::<Canonical>(&mut ctx)?;

        // Compare with filtering the canonical VarBinView.
        let expected = fsst_array.filter(mask)?;

        assert_eq!(result.len(), 2);
        assert_arrays_eq!(result.into_array(), expected, &mut ctx);
        Ok(())
    }

    #[test]
    fn test_fsst_filter_every_other() -> VortexResult<()> {
        let fsst_array = build_test_fsst_array();

        // Filter every other element
        let mask = Mask::from_iter([
            true, false, true, false, true, false, true, false, true, false,
        ]);

        let filter_array = FilterArray::new(fsst_array.clone(), mask.clone()).into_array();
        let mut ctx = SESSION.create_execution_ctx();
        let result = filter_array.execute::<Canonical>(&mut ctx)?;

        let expected = fsst_array.filter(mask)?;

        assert_eq!(result.len(), 5);
        assert_arrays_eq!(result.into_array(), expected, &mut ctx);
        Ok(())
    }

    #[test]
    fn issues_6034_test_fsst_filter_with_nulls_and_special_chars() -> VortexResult<()> {
        // Test case with special characters and nulls
        // Values: ["", "", "", "", "", "", "", "", "", "", "", ",", "A<<<<<<<", "", "", "", "", null, null, null, null, null, null]
        // Mask: only the last element is selected (true at index 22)
        let mut builder = VarBinBuilder::<i32>::with_capacity(23);
        // 11 empty strings
        for _ in 0..11 {
            builder.append_value(b"");
        }
        // ","
        builder.append_value(b",");
        // "A<<<<<<<"
        builder.append_value(b"A<<<<<<<");
        // 4 more empty strings
        for _ in 0..4 {
            builder.append_value(b"");
        }
        // 6 nulls
        for _ in 0..6 {
            builder.append_null();
        }
        let input = builder.finish(DType::Utf8(Nullability::Nullable));
        let array = input.clone().into_array();

        let mut ctx = SESSION.create_execution_ctx();
        let compressor = fsst_train_compressor(&array, &mut ctx)?;
        let fsst_array: ArrayRef = fsst_compress(&array, &compressor, &mut ctx)?.into_array();

        // Filter: only select the last element (index 22)
        let mut mask = vec![false; 22];
        mask.push(true);
        let mask = Mask::from_iter(mask);

        let filter_array = FilterArray::new(fsst_array, mask.clone()).into_array();
        let result = filter_array.execute::<Canonical>(&mut ctx)?;

        let expected = input.filter(mask)?;

        assert_eq!(result.len(), 1);
        assert_arrays_eq!(result.into_array(), expected, &mut ctx);
        Ok(())
    }

    #[test]
    fn filter_only_null() -> VortexResult<()> {
        let mut builder = VarBinBuilder::<i32>::with_capacity(3);
        builder.append_null();
        builder.append_value(b"A");
        builder.append_null();

        let input = builder.finish(DType::Utf8(Nullability::Nullable));
        let array = input.clone().into_array();

        let mut ctx = SESSION.create_execution_ctx();
        let compressor = fsst_train_compressor(&array, &mut ctx)?;
        let fsst_array: ArrayRef = fsst_compress(&array, &compressor, &mut ctx)?.into_array();

        let mask = Mask::from_iter([true, false, true]);

        let filter_array = FilterArray::new(fsst_array, mask.clone()).into_array();
        let result = filter_array.execute::<Canonical>(&mut ctx)?;

        let expected = input.filter(mask)?;

        assert_eq!(result.len(), 2);
        assert_arrays_eq!(result.into_array(), expected, &mut ctx);
        Ok(())
    }

    #[test]
    fn test_fsst_filter_all_true() -> VortexResult<()> {
        let fsst_array = build_test_fsst_array();
        assert_eq!(fsst_array.len(), 10);

        let mask = Mask::new_true(10);

        let filter_array = fsst_array.filter(mask)?;
        let mut ctx = SESSION.create_execution_ctx();
        let result = filter_array.execute::<Canonical>(&mut ctx)?.into_array();

        assert_arrays_eq!(result, fsst_array, &mut ctx);
        Ok(())
    }

    #[test]
    fn test_fsst_byte_length() -> VortexResult<()> {
        let mut builder = VarBinBuilder::<i32>::with_capacity(3);
        builder.append_value(b"hello");
        builder.append_value(b"world!!");
        builder.append_value("Пуховички"); // 9 characters, 18 bytes
        builder.append_value(b"");

        let varbin = builder
            .finish(DType::Utf8(Nullability::NonNullable))
            .into_array();
        let mut ctx = SESSION.create_execution_ctx();
        let compressor = fsst_train_compressor(&varbin, &mut ctx)?;
        let fsst = fsst_compress(&varbin, &compressor, &mut ctx)?.into_array();
        let result = fsst.apply(&byte_length(root()))?;
        let expected = PrimitiveArray::from_iter(vec![5u64, 7, 18, 0]);
        assert_arrays_eq!(result, expected, &mut ctx);
        Ok(())
    }
}