FSST Filter Kernel (#5755)

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

Author: gatesn

Cargo.lock

@@ -18,6 +18,7 @@ workspace = true
 
 [dependencies]
 fsst-rs = { workspace = true }
+num-traits = { workspace = true }
 prost = { workspace = true }
 rand = { workspace = true, optional = true }
 vortex-array = { 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::mask::MaskExecutor;
+use vortex_array::matchers::Exact;
+use vortex_array::validity::Validity;
+use vortex_array::vtable::ValidityHelper;
+use vortex_buffer::Buffer;
+use vortex_buffer::BufferMut;
+use vortex_buffer::ByteBuffer;
+use vortex_buffer::ByteBufferMut;
+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 = match array.codes().validity().filter(parent.filter_mask())? {
+            Validity::NonNullable | Validity::AllValid => {
+                Mask::new_true(parent.filter_mask().true_count())
+            }
+            Validity::AllInvalid => Mask::new_false(parent.filter_mask().true_count()),
+            Validity::Array(a) => a.execute_mask(ctx.session())?,
+        };
+
+        // 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 (filtered_idx, (idx, is_valid)) in filter_mask
+                .indices()
+                .iter()
+                .copied()
+                .zip(values.bit_buffer().iter())
+                .enumerate()
+            {
+                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[filtered_idx].as_()..];
+                }
+            }
+        }
+    }
+
+    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;

encodings/fsst/src/lib.rs

@@ -71,9 +71,9 @@ impl ArrayReduceRule<ScalarFnVTable> for ScalarFnConstantRule {
         let result = match result {
             Datum::Scalar(s) => s,
             Datum::Vector(v) => {
-                tracing::warn!(
+                tracing::info!(
                     "Scalar function {} returned vector from execution over all scalar inputs",
-                    array.scalar_fn
+                    array.scalar_fn,
                 );
                 v.scalar_at(0)
             }

vortex-array/src/arrays/scalar_fn/rules.rs

@@ -37,6 +37,16 @@ use crate::segments::SegmentSource;
 //  actual expression? Perhaps all expressions are given a selection mask to decide for themselves?
 const EXPR_EVAL_THRESHOLD: f64 = 0.2;
 
+/// Below this mask density we will propagate filters one by one. In other words, we filter an
+/// array using a mask prior to running a filter expression, and then have to perform a more
+/// expensive rank intersection on the result. This threshold exists because filtering has a
+/// non-trivial cost, and often that cost outweighs evaluating the filter expression over a few
+/// more rows that are already known to be false.
+///
+/// TODO(ngates): this threshold should really be estimated based on the cost of the filter + the
+///  the cost of the expression itself.
+const FILTER_OF_FILTER_THRESHOLD: f64 = 0.8;
+
 pub struct FlatReader {
     layout: FlatLayout,
     name: Arc<str>,
@@ -147,11 +157,20 @@ impl LayoutReader for FlatReader {
             }
 
             let array_mask = if *USE_VORTEX_OPERATORS {
-                // Apply the expression to the array.
-                let array = array.apply(&expr)?;
-                // Evaluate the array into a mask.
-                let array_mask = array.execute_mask(&session)?;
-                mask.bitand(&array_mask)
+                if mask.density() < FILTER_OF_FILTER_THRESHOLD {
+                    // Run only over the pre-filtered rows.
+                    let array = array.filter(mask.clone())?;
+                    let array = array.apply(&expr)?;
+                    let array_mask = array.execute_mask(&session)?;
+
+                    mask.intersect_by_rank(&array_mask)
+                } else {
+                    // Run over the full array, with a simpler bitand at the end.
+                    let array = array.apply(&expr)?;
+                    let array_mask = array.execute_mask(&session)?;
+
+                    mask.bitand(&array_mask)
+                }
             } else {
                 // TODO(ngates): the mask may actually be dense within a range, as is often the case when
                 //  we have approximate mask results from a zone map. In which case we could look at