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..32877f5f84f --- /dev/null +++ b/vortex-btrblocks/src/schemes/binary/varbin.rs @@ -0,0 +1,98 @@ +// 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_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 { + let view = data.array_as_varbinview().into_owned(); + let len = view.len(); + // Materialize validity once; a per-element accessor here would be quadratic-ish in a + // loop this hot. + let mask = view.validity()?.execute_mask(len, exec_ctx)?; + + let varbin = VarBinArray::from_iter( + (0..len).map(|i| mask.value(i).then(|| view.bytes_at(i).as_slice().to_vec())), + view.dtype().clone(), + ); + + 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/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..64bb0df9786 --- /dev/null +++ b/vortex-btrblocks/tests/varbin_scheme.rs @@ -0,0 +1,112 @@ +// 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(()) +}