diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 6f38e29cd86..4f5f580aae4 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -61,6 +61,7 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[ // Binary schemes. //////////////////////////////////////////////////////////////////////////////////////////////// &binary::BinaryDictScheme, + &binary::VarBinScheme, // Decimal schemes. &decimal::DecimalScheme, // Temporal schemes. diff --git a/vortex-btrblocks/src/schemes/binary/mod.rs b/vortex-btrblocks/src/schemes/binary/mod.rs index af66d345cc6..ee564b877c9 100644 --- a/vortex-btrblocks/src/schemes/binary/mod.rs +++ b/vortex-btrblocks/src/schemes/binary/mod.rs @@ -3,12 +3,14 @@ //! Binary compression schemes. +mod varbin; #[cfg(feature = "zstd")] mod zstd; #[cfg(all(feature = "zstd", feature = "unstable_encodings"))] mod zstd_buffers; // Re-export builtin schemes from vortex-compressor. +pub use varbin::VarBinScheme; pub use vortex_compressor::builtins::BinaryDictScheme; #[cfg(feature = "zstd")] pub use zstd::ZstdScheme; diff --git a/vortex-btrblocks/src/schemes/binary/varbin.rs b/vortex-btrblocks/src/schemes/binary/varbin.rs new file mode 100644 index 00000000000..e2b79d2be1d --- /dev/null +++ b/vortex-btrblocks/src/schemes/binary/varbin.rs @@ -0,0 +1,96 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Offset-based storage for binary arrays. +//! +//! Canonical binary arrays are [`VarBinViewArray`], which spends a fixed 16 bytes per element on +//! an opaque views buffer that no scheme can compress. Re-encoding as [`VarBinArray`] replaces +//! that buffer with an offsets child array, which the cascading compressor can then compress with +//! the ordinary integer schemes. For fixed-width values the offsets are a constant-stride +//! sequence and collapse to nothing. + +use vortex_array::ArrayId; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::VarBin; +use vortex_array::arrays::VarBinArray; +use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_array::arrays::varbin::VarBinArraySlotsExt; +use vortex_array::builders::VarBinBuilder; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::SchemeExt; +use vortex_error::VortexResult; + +use crate::ArrayAndStats; +use crate::CascadingCompressor; +use crate::CompressorContext; +use crate::Scheme; + +/// Offset-based (rather than view-based) storage for binary arrays. +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub struct VarBinScheme; + +impl Scheme for VarBinScheme { + fn scheme_name(&self) -> &'static str { + "vortex.binary.varbin" + } + + fn matches(&self, canonical: &Canonical) -> bool { + canonical.dtype().is_binary() + } + + fn produced_encodings(&self) -> Vec { + vec![VarBin.id()] + } + + fn num_children(&self) -> usize { + 1 + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Deferred(DeferredEstimate::Sample) + } + + fn compress( + &self, + compressor: &CascadingCompressor, + data: &ArrayAndStats, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + // `append_to_builder` resolves the views slice and data buffers once and appends + // borrowed slices into a single pre-sized allocation. Iterating the array per element + // instead would clone a buffer handle and allocate for every value. + let array = data.array(); + let mut builder = VarBinBuilder::::with_capacity(array.dtype().clone(), array.len()); + array.append_to_builder(&mut builder, exec_ctx)?; + let varbin = builder.finish_into_varbin(); + + let offsets = varbin + .offsets() + .clone() + .execute::(exec_ctx)? + .narrow(exec_ctx)? + .into_array(); + let compressed_offsets = + compressor.compress_child(&offsets, &compress_ctx, self.id(), 0, exec_ctx)?; + + Ok(VarBinArray::try_new( + compressed_offsets, + varbin.bytes().clone(), + varbin.dtype().clone(), + varbin.validity()?, + )? + .into_array()) + } +} diff --git a/vortex-btrblocks/src/schemes/string/fsst.rs b/vortex-btrblocks/src/schemes/string/fsst.rs index 7250be26207..fd3fd28696a 100644 --- a/vortex-btrblocks/src/schemes/string/fsst.rs +++ b/vortex-btrblocks/src/schemes/string/fsst.rs @@ -47,7 +47,7 @@ impl Scheme for FSSTScheme { } fn matches(&self, canonical: &Canonical) -> bool { - canonical.dtype().is_utf8() + canonical.dtype().is_utf8() || canonical.dtype().is_binary() } fn produced_encodings(&self) -> Vec { diff --git a/vortex-btrblocks/tests/snapshots/golden__default__binary_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__default__binary_low_cardinality.snap index 56918c0d096..788cf54ecd8 100644 --- a/vortex-btrblocks/tests/snapshots/golden__default__binary_low_cardinality.snap +++ b/vortex-btrblocks/tests/snapshots/golden__default__binary_low_cardinality.snap @@ -3,9 +3,11 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: binary, len=16384, nbytes=315856 -root: vortex.dict(binary, len=16384) nbytes=6240 +root: vortex.dict(binary, len=16384) nbytes=6196 metadata: all_values_referenced: true codes: fastlanes.bitpacked(u8, len=16384) nbytes=6144 metadata: bit_width: 3, offset: 0 - values: vortex.varbinview(binary, len=5) nbytes=96 - metadata: + values: vortex.varbin(binary, len=5) nbytes=52 + metadata: + offsets: vortex.primitive(u8, len=6) nbytes=6 + metadata: ptype: u8 diff --git a/vortex-btrblocks/tests/varbin_scheme.rs b/vortex-btrblocks/tests/varbin_scheme.rs new file mode 100644 index 00000000000..8d5dc7b5443 --- /dev/null +++ b/vortex-btrblocks/tests/varbin_scheme.rs @@ -0,0 +1,163 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Measures `VarBinScheme` against the same compressor with the scheme excluded. + +#![allow(clippy::cast_possible_truncation, clippy::tests_outside_test_module)] + +use std::sync::LazyLock; + +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_btrblocks::BtrBlocksCompressorBuilder; +use vortex_btrblocks::SchemeExt; +use vortex_btrblocks::schemes::binary::VarBinScheme; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + +const N: usize = 100_000; + +fn lcg(state: &mut u64) -> u64 { + *state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + *state +} + +fn cases() -> Vec<(&'static str, ArrayRef)> { + let mut s = 42u64; + let random: Vec> = (0..N) + .map(|_| (0..16).map(|_| (lcg(&mut s) >> 33) as u8).collect()) + .collect(); + let prefixed: Vec> = (0..N) + .map(|i| format!("PREFIX_{i:09}").into_bytes()) + .collect(); + let mut s2 = 7u64; + let wide: Vec> = (0..N) + .map(|_| (0..256).map(|_| (lcg(&mut s2) >> 33) as u8).collect()) + .collect(); + + let nullable = VarBinViewArray::from_iter( + (0..N).map(|i| (i % 7 != 0).then(|| prefixed[i].as_slice())), + DType::Binary(Nullability::Nullable), + ) + .into_array(); + + vec![ + ("nulls every 7th", nullable), + ( + "random 16B (hash)", + VarBinViewArray::from_iter_bin(random.iter().map(|v| v.as_slice())).into_array(), + ), + ( + "shared prefix", + VarBinViewArray::from_iter_bin(prefixed.iter().map(|v| v.as_slice())).into_array(), + ), + ( + "random 256B", + VarBinViewArray::from_iter_bin(wide.iter().map(|v| v.as_slice())).into_array(), + ), + ] +} + +#[test] +fn varbin_scheme_shrinks_binary() -> VortexResult<()> { + let with = BtrBlocksCompressorBuilder::default().build(); + let without = BtrBlocksCompressorBuilder::default() + .exclude_schemes([VarBinScheme.id()]) + .build(); + + println!( + "{:<20}{:>12}{:>14}{:>14}{:>9}", + "case", "input", "without", "with", "ratio" + ); + for (name, array) in cases() { + let a = { + let mut ctx = SESSION.create_execution_ctx(); + without.compress(&array, &mut ctx)?.nbytes() + }; + let b = { + let mut ctx = SESSION.create_execution_ctx(); + with.compress(&array, &mut ctx)?.nbytes() + }; + println!( + "{:<20}{:>12}{:>14}{:>14}{:>9.2}", + name, + array.nbytes(), + a, + b, + b as f64 / a as f64 + ); + + assert!( + b <= a, + "{name}: enabling VarBinScheme grew the output, {a} -> {b}" + ); + + let mut ctx = SESSION.create_execution_ctx(); + let compressed = with.compress(&array, &mut ctx)?; + let decoded = compressed + .execute::(&mut ctx)? + .into_array(); + assert_arrays_eq!(&array, &decoded, &mut ctx); + } + 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(()) +}