diff --git a/vortex-btrblocks/src/schemes/string/fsst.rs b/vortex-btrblocks/src/schemes/string/fsst.rs index 7250be26207..43cf9304c8c 100644 --- a/vortex-btrblocks/src/schemes/string/fsst.rs +++ b/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 { + 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 { @@ -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) } diff --git a/vortex-btrblocks/tests/varbin_scheme.rs b/vortex-btrblocks/tests/varbin_scheme.rs index 64bb0df9786..1cf9196733a 100644 --- a/vortex-btrblocks/tests/varbin_scheme.rs +++ b/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 = (0..N).map(|i| format!("PREFIX_{i:09}")).collect(); + let hex_random: Vec = (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(()) +} diff --git a/vortex-compressor/src/compressor/sample.rs b/vortex-compressor/src/compressor/sample.rs index ba1271d4a6d..0bf6aed53ae 100644 --- a/vortex-compressor/src/compressor/sample.rs +++ b/vortex-compressor/src/compressor/sample.rs @@ -189,6 +189,28 @@ pub(super) fn estimate_compression_ratio_with_sampling( 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 { + estimate_compression_ratio_with_sampling(self, scheme, array, compress_ctx, exec_ctx) + } +} + #[cfg(test)] mod tests { use vortex_array::IntoArray;