Cache FSST compressor for running compare (#3343)

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

Author: gatesn

encodings/fsst/src/array.rs

@@ -1,4 +1,7 @@
-use fsst::{Decompressor, Symbol};
+use std::fmt::{Debug, Formatter};
+use std::sync::{Arc, LazyLock};
+
+use fsst::{Compressor, Decompressor, Symbol};
 use vortex_array::arrays::VarBinArray;
 use vortex_array::stats::{ArrayStats, StatsSetRef};
 use vortex_array::vtable::{
@@ -33,7 +36,7 @@ impl VTable for FSSTVTable {
     }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone)]
 pub struct FSSTArray {
     dtype: DType,
     symbols: Buffer<Symbol>,
@@ -42,6 +45,21 @@ pub struct FSSTArray {
     /// Lengths of the original values before compression, can be compressed.
     uncompressed_lengths: ArrayRef,
     stats_set: ArrayStats,
+
+    /// Memoized compressor used for push-down of compute by compressing the RHS.
+    compressor: Arc<LazyLock<Compressor, Box<dyn Fn() -> Compressor + Send>>>,
+}
+
+impl Debug for FSSTArray {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("FSSTArray")
+            .field("dtype", &self.dtype)
+            .field("symbols", &self.symbols)
+            .field("symbol_lengths", &self.symbol_lengths)
+            .field("codes", &self.codes)
+            .field("uncompressed_lengths", &self.uncompressed_lengths)
+            .finish()
+    }
 }
 
 #[derive(Clone, Debug)]
@@ -84,13 +102,21 @@ impl FSSTArray {
             vortex_bail!(InvalidArgument: "codes array must be DType::Binary type");
         }
 
+        let symbols2 = symbols.clone();
+        let symbol_lengths2 = symbol_lengths.clone();
+        let compressor = Arc::new(LazyLock::new(Box::new(move || {
+            Compressor::rebuild_from(symbols2.as_slice(), symbol_lengths2.as_slice())
+        })
+            as Box<dyn Fn() -> Compressor + Send>));
+
         Ok(Self {
             dtype,
             symbols,
             symbol_lengths,
             codes,
             uncompressed_lengths,
             stats_set: Default::default(),
+            compressor,
         })
     }
 
@@ -133,6 +159,10 @@ impl FSSTArray {
     pub(crate) fn decompressor(&self) -> Decompressor {
         Decompressor::new(self.symbols().as_slice(), self.symbol_lengths().as_slice())
     }
+
+    pub(crate) fn compressor(&self) -> &Compressor {
+        self.compressor.as_ref()
+    }
 }
 
 impl ArrayVTable<FSSTVTable> for FSSTVTable {

encodings/fsst/src/compute/compare.rs

@@ -19,9 +19,7 @@ impl CompareKernel for FSSTVTable {
         operator: Operator,
     ) -> VortexResult<Option<ArrayRef>> {
         match rhs.as_constant() {
-            Some(constant) => {
-                compare_fsst_constant(lhs, &ConstantArray::new(constant, lhs.len()), operator)
-            }
+            Some(constant) => compare_fsst_constant(lhs, &constant, operator),
             // Otherwise, fall back to the default comparison behavior.
             _ => Ok(None),
         }
@@ -33,16 +31,15 @@ register_kernel!(CompareKernelAdapter(FSSTVTable).lift());
 /// Specialized compare function implementation used when performing against a constant
 fn compare_fsst_constant(
     left: &FSSTArray,
-    right: &ConstantArray,
+    right: &Scalar,
     operator: Operator,
 ) -> VortexResult<Option<ArrayRef>> {
-    let rhs_scalar = right.scalar();
-    let is_rhs_empty = match rhs_scalar.dtype() {
-        DType::Binary(_) => rhs_scalar
+    let is_rhs_empty = match right.dtype() {
+        DType::Binary(_) => right
             .as_binary()
             .is_empty()
             .vortex_expect("RHS should not be null"),
-        DType::Utf8(_) => rhs_scalar
+        DType::Utf8(_) => right
             .as_utf8()
             .is_empty()
             .vortex_expect("RHS should not be null"),
@@ -77,20 +74,17 @@ fn compare_fsst_constant(
         return Ok(None);
     }
 
-    let compressor = fsst::Compressor::rebuild_from(left.symbols(), left.symbol_lengths());
-
+    let compressor = left.compressor();
     let encoded_buffer = match left.dtype() {
         DType::Utf8(_) => {
             let value = right
-                .scalar()
                 .as_utf8()
                 .value()
                 .vortex_expect("Expected non-null scalar");
             ByteBuffer::from(compressor.compress(value.as_bytes()))
         }
         DType::Binary(_) => {
             let value = right
-                .scalar()
                 .as_binary()
                 .value()
                 .vortex_expect("Expected non-null scalar");