FSST Filter Kernel

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

Cargo.lock

@@ -18,10 +18,12 @@ workspace = true
 
 [dependencies]
 fsst-rs = { workspace = true }
+num-traits = { workspace = true }
 prost = { workspace = true }
 rand = { workspace = true, optional = true }
 vortex-array = { workspace = true }
 vortex-buffer = { workspace = true }
+vortex-compute = { workspace = true }
 vortex-dtype = { workspace = true }
 vortex-error = { workspace = true }
 vortex-mask = { workspace = true }

encodings/fsst/Cargo.toml

@@ -18,6 +18,7 @@ use vortex_array::ArrayHash;
 use vortex_array::ArrayRef;
 use vortex_array::Canonical;
 use vortex_array::DeserializeMetadata;
+use vortex_array::ExecutionCtx;
 use vortex_array::Precision;
 use vortex_array::ProstMetadata;
 use vortex_array::SerializeMetadata;
@@ -46,9 +47,11 @@ use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure;
 use vortex_error::vortex_err;
+use vortex_vector::Vector;
 
 use crate::fsst_compress;
 use crate::fsst_train_compressor;
+use crate::kernel::PARENT_KERNELS;
 
 vtable!(FSST);
 
@@ -178,6 +181,15 @@ impl VTable for FSSTVTable {
 
         Ok(())
     }
+
+    fn execute_parent(
+        array: &Self::Array,
+        parent: &ArrayRef,
+        child_idx: usize,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<Vector>> {
+        PARENT_KERNELS.execute(array, parent, child_idx, ctx)
+    }
 }
 
 #[derive(Clone)]

encodings/fsst/src/array.rs

@@ -0,0 +1,200 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::sync::Arc;
+
+use fsst::Decompressor;
+use num_traits::AsPrimitive;
+use vortex_array::Array;
+use vortex_array::ExecutionCtx;
+use vortex_array::VectorExecutor;
+use vortex_array::arrays::FilterArray;
+use vortex_array::arrays::FilterVTable;
+use vortex_array::builtins::ArrayBuiltins;
+use vortex_array::kernel::ExecuteParentKernel;
+use vortex_array::kernel::ParentKernelSet;
+use vortex_array::matchers::Exact;
+use vortex_buffer::Buffer;
+use vortex_buffer::BufferMut;
+use vortex_buffer::ByteBuffer;
+use vortex_buffer::ByteBufferMut;
+use vortex_compute::filter::Filter;
+use vortex_dtype::DType;
+use vortex_dtype::IntegerPType;
+use vortex_dtype::Nullability;
+use vortex_dtype::PType;
+use vortex_dtype::PTypeDowncastExt;
+use vortex_dtype::match_each_integer_ptype;
+use vortex_error::VortexResult;
+use vortex_mask::Mask;
+use vortex_mask::MaskValues;
+use vortex_vector::Vector;
+use vortex_vector::binaryview::BinaryVector;
+use vortex_vector::binaryview::BinaryView;
+use vortex_vector::binaryview::StringVector;
+
+use crate::FSSTArray;
+use crate::FSSTVTable;
+
+pub(super) const PARENT_KERNELS: ParentKernelSet<FSSTVTable> =
+    ParentKernelSet::new(&[ParentKernelSet::lift(&FSSTFilterKernel)]);
+
+#[derive(Debug)]
+struct FSSTFilterKernel;
+
+impl ExecuteParentKernel<FSSTVTable> for FSSTFilterKernel {
+    type Parent = Exact<FilterVTable>;
+
+    fn parent(&self) -> Self::Parent {
+        Exact::from(&FilterVTable)
+    }
+
+    fn execute_parent(
+        &self,
+        array: &FSSTArray,
+        parent: &FilterArray,
+        _child_idx: usize,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<Vector>> {
+        let mask_values = match parent.filter_mask() {
+            Mask::AllTrue(_) | Mask::AllFalse(_) => return Ok(None),
+            Mask::Values(v) => v,
+        };
+
+        // We filter the uncompressed lengths
+        let uncompressed_lens = array
+            .uncompressed_lengths()
+            .filter(parent.filter_mask().clone())?
+            .execute(ctx)?
+            .into_primitive();
+
+        // Extract the filtered validity
+        let validity = array.codes().validity_mask().filter(parent.filter_mask());
+
+        // First we unpack the codes VarBinArray to get access to the raw data.
+        let codes_data = array.codes().bytes();
+        let codes_offsets = array
+            .codes()
+            .offsets()
+            .cast(DType::Primitive(PType::U32, Nullability::NonNullable))?
+            .execute(ctx)?
+            .into_primitive()
+            .downcast::<u32>()
+            .into_nonnull_buffer();
+
+        let decompressor = array.decompressor();
+
+        let (views, buffer) = match_each_integer_ptype!(uncompressed_lens.ptype(), |S| {
+            fsst_decode::<S>(
+                decompressor,
+                codes_data,
+                &codes_offsets,
+                mask_values,
+                &validity,
+                &uncompressed_lens.downcast::<S>().into_nonnull_buffer(),
+            )
+        });
+
+        let vector = match array.dtype() {
+            DType::Binary(_) => unsafe {
+                BinaryVector::new_unchecked(views, Arc::new(vec![buffer].into()), validity)
+            }
+            .into(),
+            DType::Utf8(_) => unsafe {
+                StringVector::new_unchecked(views, Arc::new(vec![buffer].into()), validity)
+            }
+            .into(),
+            _ => unreachable!("Not a supported FSST DType"),
+        };
+
+        Ok(Some(vector))
+    }
+}
+
+fn fsst_decode<S: IntegerPType + AsPrimitive<usize> + AsPrimitive<u32>>(
+    decompressor: Decompressor,
+    codes_data: &[u8],
+    codes_offsets: &[u32],
+    filter_mask: &MaskValues,
+    filtered_validity: &Mask,
+    filtered_uncompressed_lengths: &[S],
+) -> (Buffer<BinaryView>, ByteBuffer) {
+    let total_uncompressed_size: usize = filtered_uncompressed_lengths
+        .iter()
+        .map(|x| <S as AsPrimitive<usize>>::as_(*x))
+        .sum();
+
+    // We allocate an extra 7 bytes per the FSST decompressor's requirement for padding.
+    let mut uncompressed = ByteBufferMut::with_capacity(total_uncompressed_size + 7);
+    let mut spare_capacity = uncompressed.spare_capacity_mut();
+
+    match filtered_validity {
+        Mask::AllTrue(_) => {
+            for &idx in filter_mask.indices() {
+                let start = codes_offsets[idx] as usize;
+                let end = codes_offsets[idx + 1] as usize;
+                let compressed_slice = &codes_data[start..end];
+
+                let uncompressed_len =
+                    decompressor.decompress_into(compressed_slice, spare_capacity);
+                spare_capacity = &mut spare_capacity[uncompressed_len..];
+            }
+        }
+        Mask::AllFalse(_) => {
+            // Nothing to decompress
+            unsafe { uncompressed.set_len(0) };
+            return (Buffer::empty(), uncompressed.freeze());
+        }
+        Mask::Values(values) => {
+            for (idx, is_valid) in filter_mask
+                .indices()
+                .iter()
+                .copied()
+                .zip(values.bit_buffer().iter())
+            {
+                if is_valid {
+                    let start = codes_offsets[idx] as usize;
+                    let end = codes_offsets[idx + 1] as usize;
+                    let compressed_slice = &codes_data[start..end];
+
+                    let uncompressed_len =
+                        decompressor.decompress_into(compressed_slice, spare_capacity);
+                    spare_capacity = &mut spare_capacity[uncompressed_len..];
+                } else {
+                    // We advance the output buffer to make it faster to assemble views below.
+                    spare_capacity =
+                        &mut spare_capacity[filtered_uncompressed_lengths[idx].as_()..];
+                }
+            }
+        }
+    }
+
+    for &idx in filter_mask.indices() {
+        let start = codes_offsets[idx] as usize;
+        let end = codes_offsets[idx + 1] as usize;
+        let compressed_slice = &codes_data[start..end];
+
+        let uncompressed_len = decompressor.decompress_into(compressed_slice, spare_capacity);
+        spare_capacity = &mut spare_capacity[uncompressed_len..];
+    }
+
+    unsafe { uncompressed.set_len(total_uncompressed_size) };
+    let uncompressed = uncompressed.freeze();
+    let uncompressed_slice = uncompressed.as_ref();
+
+    // Loop over the uncompressed lengths to construct the BinaryViews.
+    let mut views = BufferMut::<BinaryView>::with_capacity(filtered_uncompressed_lengths.len());
+    let mut offset = 0u32;
+    for len in filtered_uncompressed_lengths {
+        let view = BinaryView::make_view(
+            &uncompressed_slice[offset as usize..][..len.as_()],
+            0u32,
+            offset,
+        );
+        offset += <S as AsPrimitive<u32>>::as_(*len);
+        unsafe { views.push_unchecked(view) };
+    }
+    unsafe { views.set_len(filtered_uncompressed_lengths.len()) };
+
+    (views.freeze(), uncompressed)
+}

encodings/fsst/src/kernel.rs

@@ -15,6 +15,7 @@ mod array;
 mod canonical;
 mod compress;
 mod compute;
+mod kernel;
 mod ops;
 #[cfg(feature = "test-harness")]
 pub mod test_utils;