feat(btrblocks): let FSSTScheme compress binary arrays behind a sampled gate

vortex-btrblocks/src/schemes/string/fsst.rs

@@ -18,6 +18,7 @@ use vortex_array::arrays::primitive::PrimitiveArrayExt;
 use vortex_array::arrays::varbin::VarBinArraySlotsExt;
 use vortex_compressor::scheme::CompressionEstimate;
 use vortex_compressor::scheme::DeferredEstimate;
+use vortex_compressor::scheme::EstimateVerdict;
 use vortex_error::VortexResult;
 use vortex_fsst::FSST;
 use vortex_fsst::FSSTArrayExt;
@@ -41,13 +42,75 @@ use crate::SchemeExt;
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
 pub struct FSSTScheme;
 
+/// Number of values in the tiny sample used to gate FSST on binary columns.
+const BINARY_SAMPLE_VALUES: usize = 64;
+
+/// Minimum fraction of value bytes FSST must save over plain VarBin storage on the sample
+/// before it is allowed into scheme selection for a binary column.
+const BINARY_MIN_SAVINGS: f64 = 0.15;
+
+/// Gates FSST on binary columns with a tiny trial compression.
+///
+/// Strings almost always have intra-value structure, but arbitrary binary payloads often do
+/// not (hashes, ciphertexts, compressed blobs), and a symbol table that buys nothing still
+/// costs training on write and decoding on read. Trial-compress up to
+/// [`BINARY_SAMPLE_VALUES`] values strided across the column and admit FSST only when its
+/// code bytes undercut the raw value bytes (the VarBin baseline, views excluded) by more than
+/// [`BINARY_MIN_SAVINGS`]. The symbol table is excluded from the measurement: its size is
+/// fixed per column and amortizes away at real column lengths, while against a 64-value
+/// sample it would drown out the signal.
+///
+/// Once past the gate, the ratio entered into selection comes from the compressor's standard
+/// sampling estimator, so FSST competes with other schemes on the same measurement basis.
+fn estimate_binary_fsst(
+    compressor: &CascadingCompressor,
+    data: &ArrayAndStats,
+    compress_ctx: CompressorContext,
+    exec_ctx: &mut ExecutionCtx,
+) -> VortexResult<EstimateVerdict> {
+    let view = data.array_as_varbinview();
+    let len = view.len();
+    if len == 0 {
+        return Ok(EstimateVerdict::Skip);
+    }
+    let mask = view.validity()?.execute_mask(len, exec_ctx)?;
+
+    let step = len.div_ceil(BINARY_SAMPLE_VALUES);
+    let sample = VarBinArray::from_iter(
+        (0..len)
+            .step_by(step)
+            .map(|i| mask.value(i).then(|| view.bytes_at(i).as_slice().to_vec())),
+        view.dtype().clone(),
+    );
+    let raw_nbytes = sample.bytes().len();
+    if raw_nbytes == 0 {
+        return Ok(EstimateVerdict::Skip);
+    }
+
+    let sample = sample.into_array();
+    let trained = fsst_train_compressor(&sample, exec_ctx)?;
+    let fsst = fsst_compress(&sample, &trained, exec_ctx)?;
+    let compressed_nbytes = fsst.codes().bytes().len();
+
+    if (compressed_nbytes as f64) >= (1.0 - BINARY_MIN_SAVINGS) * raw_nbytes as f64 {
+        return Ok(EstimateVerdict::Skip);
+    }
+
+    let score =
+        compressor.estimate_by_sampling(&FSSTScheme, data.array(), compress_ctx, exec_ctx)?;
+    Ok(match score.finite_ratio() {
+        Some(ratio) => EstimateVerdict::Ratio(ratio),
+        None => EstimateVerdict::Skip,
+    })
+}
+
 impl Scheme for FSSTScheme {
     fn scheme_name(&self) -> &'static str {
         "vortex.string.fsst"
     }
 
     fn matches(&self, canonical: &Canonical) -> bool {
-        canonical.dtype().is_utf8()
+        canonical.dtype().is_utf8() || canonical.dtype().is_binary()
     }
 
     fn produced_encodings(&self) -> Vec<ArrayId> {
@@ -61,10 +124,17 @@ impl Scheme for FSSTScheme {
 
     fn expected_compression_ratio(
         &self,
-        _data: &ArrayAndStats,
+        data: &ArrayAndStats,
         _compress_ctx: CompressorContext,
         _exec_ctx: &mut ExecutionCtx,
     ) -> CompressionEstimate {
+        if data.array().dtype().is_binary() {
+            return CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new(
+                |compressor, data, _best_so_far, compress_ctx, exec_ctx| {
+                    estimate_binary_fsst(compressor, data, compress_ctx, exec_ctx)
+                },
+            )));
+        }
         CompressionEstimate::Deferred(DeferredEstimate::Sample)
     }
 

vortex-btrblocks/tests/varbin_scheme.rs

@@ -110,3 +110,78 @@ fn varbin_scheme_shrinks_binary() -> VortexResult<()> {
     }
     Ok(())
 }
+
+/// FSST on binary is gated by a tiny trial compression: structured payloads must still reach
+/// FSST somewhere in the tree, while incompressible payloads must never pay for a symbol table.
+#[test]
+fn fsst_binary_gate() -> VortexResult<()> {
+    let compressor = BtrBlocksCompressorBuilder::default().build();
+    for (name, array, expect_fsst) in [
+        ("shared prefix", &cases()[2].1, true),
+        ("random 16B (hash)", &cases()[1].1, false),
+        ("random 256B", &cases()[3].1, false),
+    ] {
+        let mut ctx = SESSION.create_execution_ctx();
+        let tree = compressor
+            .compress(array, &mut ctx)?
+            .display_tree()
+            .to_string();
+        assert_eq!(
+            tree.contains("fsst"),
+            expect_fsst,
+            "{name}: unexpected FSST selection\n{tree}"
+        );
+    }
+    Ok(())
+}
+
+/// Same bytes, two dtypes: Binary takes the `VarBinScheme` path, Utf8 takes FSST.
+#[test]
+fn fsst_versus_varbin_on_identical_bytes() -> VortexResult<()> {
+    let compressor = BtrBlocksCompressorBuilder::default().build();
+    let mut seed = 99u64;
+
+    let shared_prefix: Vec<String> = (0..N).map(|i| format!("PREFIX_{i:09}")).collect();
+    let hex_random: Vec<String> = (0..N)
+        .map(|_| {
+            (0..16)
+                .map(|_| format!("{:02x}", (lcg(&mut seed) >> 33) as u8))
+                .collect()
+        })
+        .collect();
+
+    println!(
+        "{:<20}{:>14}{:>14}{:>9}",
+        "case", "binary", "utf8(fsst)", "fsst/bin"
+    );
+    for (name, vals) in [
+        ("shared prefix", &shared_prefix),
+        ("hex random", &hex_random),
+    ] {
+        let as_bin = VarBinViewArray::from_iter_bin(vals.iter().map(|v| v.as_bytes())).into_array();
+        let as_str = VarBinViewArray::from_iter_str(vals.iter().map(|v| v.as_str())).into_array();
+
+        let bin_bytes = {
+            let mut ctx = SESSION.create_execution_ctx();
+            compressor.compress(&as_bin, &mut ctx)?.nbytes()
+        };
+        let utf8_bytes = {
+            let mut ctx = SESSION.create_execution_ctx();
+            compressor.compress(&as_str, &mut ctx)?.nbytes()
+        };
+        println!(
+            "{:<20}{:>14}{:>14}{:>9.2}",
+            name,
+            bin_bytes,
+            utf8_bytes,
+            utf8_bytes as f64 / bin_bytes as f64
+        );
+
+        // FSST compresses bytes, not codepoints, so the dtype must not change the result.
+        assert_eq!(
+            bin_bytes, utf8_bytes,
+            "{name}: binary and utf8 must compress identically"
+        );
+    }
+    Ok(())
+}

vortex-compressor/src/compressor/sample.rs

@@ -189,6 +189,28 @@ pub(super) fn estimate_compression_ratio_with_sampling<S: Scheme + ?Sized>(
     Ok(score)
 }
 
+impl CascadingCompressor {
+    /// Estimates a scheme's compression ratio exactly as
+    /// [`DeferredEstimate::Sample`](crate::scheme::DeferredEstimate::Sample) would.
+    ///
+    /// This is for [`DeferredEstimate::Callback`](crate::scheme::DeferredEstimate::Callback)
+    /// implementations that gate a scheme with a custom cheap check but, once the gate passes,
+    /// want to compete in selection on the same sampled score as every other scheme.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if sample compression fails.
+    pub fn estimate_by_sampling(
+        &self,
+        scheme: &dyn Scheme,
+        array: &ArrayRef,
+        compress_ctx: CompressorContext,
+        exec_ctx: &mut ExecutionCtx,
+    ) -> VortexResult<EstimateScore> {
+        estimate_compression_ratio_with_sampling(self, scheme, array, compress_ctx, exec_ctx)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use vortex_array::IntoArray;