From 2d35497ebfcaaca3004b30471095d160c26b3bfe Mon Sep 17 00:00:00 2001 From: Will Manning Date: Thu, 3 Sep 2026 17:06:33 -0400 Subject: [PATCH 1/2] perf(compressor): Improve string dictionary selection Signed-off-by: Will Manning --- vortex-array/src/builders/dict/bytes.rs | 129 ++++- vortex-array/src/builders/dict/mod.rs | 45 +- vortex-btrblocks/Cargo.toml | 5 + .../benches/string_dict_selection.rs | 195 +++++++ .../schemes/string/scheme_selection_tests.rs | 96 ++++ vortex-compressor/src/builtins/dict/mod.rs | 1 + vortex-compressor/src/builtins/dict/string.rs | 101 ++-- .../src/builtins/dict/string_candidate.rs | 485 ++++++++++++++++++ vortex-compressor/src/compressor/mod.rs | 3 + vortex-compressor/src/compressor/sample.rs | 169 +++++- vortex-compressor/src/stats/cache.rs | 18 + 11 files changed, 1180 insertions(+), 67 deletions(-) create mode 100644 vortex-btrblocks/benches/string_dict_selection.rs create mode 100644 vortex-compressor/src/builtins/dict/string_candidate.rs diff --git a/vortex-array/src/builders/dict/bytes.rs b/vortex-array/src/builders/dict/bytes.rs index 48444a63000..02be10da250 100644 --- a/vortex-array/src/builders/dict/bytes.rs +++ b/vortex-array/src/builders/dict/bytes.rs @@ -50,19 +50,55 @@ pub struct BytesDictBuilder { dtype: DType, max_dict_bytes: usize, max_dict_len: usize, + input_byte_limit: Option, + remaining_input_bytes: Option, } pub fn bytes_dict_builder(dtype: DType, constraints: &DictConstraints) -> Box { + bytes_dict_builder_with_optional_input_byte_limit(dtype, constraints, None) +} + +pub(super) fn bytes_dict_builder_with_input_byte_limit( + dtype: DType, + constraints: &DictConstraints, + max_input_bytes: usize, +) -> Box { + bytes_dict_builder_with_optional_input_byte_limit(dtype, constraints, Some(max_input_bytes)) +} + +fn bytes_dict_builder_with_optional_input_byte_limit( + dtype: DType, + constraints: &DictConstraints, + input_byte_limit: Option, +) -> Box { match constraints.max_len as u64 { - max if max <= u8::MAX as u64 => Box::new(BytesDictBuilder::::new(dtype, constraints)), - max if max <= u16::MAX as u64 => Box::new(BytesDictBuilder::::new(dtype, constraints)), - max if max <= u32::MAX as u64 => Box::new(BytesDictBuilder::::new(dtype, constraints)), - _ => Box::new(BytesDictBuilder::::new(dtype, constraints)), + max if max <= u8::MAX as u64 + 1 => { + new_bytes_dict_builder::(dtype, constraints, input_byte_limit) + } + max if max <= u16::MAX as u64 + 1 => { + new_bytes_dict_builder::(dtype, constraints, input_byte_limit) + } + max if max <= u32::MAX as u64 + 1 => { + new_bytes_dict_builder::(dtype, constraints, input_byte_limit) + } + _ => new_bytes_dict_builder::(dtype, constraints, input_byte_limit), } } +fn new_bytes_dict_builder( + dtype: DType, + constraints: &DictConstraints, + input_byte_limit: Option, +) -> Box { + Box::new(BytesDictBuilder::::new( + dtype, + constraints, + input_byte_limit, + )) +} + impl BytesDictBuilder { - pub fn new(dtype: DType, constraints: &DictConstraints) -> Self { + fn new(dtype: DType, constraints: &DictConstraints, input_byte_limit: Option) -> Self { Self { lookup: Some(HashTable::new()), views: BufferMut::::empty(), @@ -73,6 +109,8 @@ impl BytesDictBuilder { dtype, max_dict_bytes: constraints.max_bytes.min(u32::MAX as usize), max_dict_len: constraints.max_len, + input_byte_limit, + remaining_input_bytes: input_byte_limit, } } @@ -92,6 +130,10 @@ impl BytesDictBuilder { /// Returns `None` when assigning a code would exceed the dictionary constraints, /// and callers should stop encoding after the current prefix. fn encode_value(&mut self, lookup: &mut HashTable, val: &[u8]) -> Option { + if let Some(remaining_input_bytes) = self.remaining_input_bytes { + self.remaining_input_bytes = Some(remaining_input_bytes.checked_sub(val.len())?); + } + match lookup.entry( self.hasher.hash_one(val), |idx| val == self.lookup_bytes(idx.as_()), @@ -290,6 +332,7 @@ impl DictEncoder for BytesDictBuilder { lookup.clear(); } self.null_code = OnceCell::new(); + self.remaining_input_bytes = self.input_byte_limit; let views = mem::take(&mut self.views).freeze(); let buffer = mem::take(&mut self.values).freeze(); let value_nulls = mem::take(&mut self.values_nulls).freeze(); @@ -332,8 +375,11 @@ mod test { use crate::arrays::varbinview::BinaryView; use crate::assert_arrays_eq; use crate::buffer::BufferHandle; + use crate::builders::dict::DictConstraints; + use crate::builders::dict::DictEncoder; use crate::builders::dict::UNCONSTRAINED; use crate::builders::dict::dict_encode; + use crate::builders::dict::dict_encode_with_input_byte_limit; use crate::builders::dict::dict_encoder; use crate::dtype::DType; use crate::dtype::Nullability; @@ -463,10 +509,79 @@ mod test { Ok(()) } + #[test] + fn nulls_do_not_consume_input_bytes() -> VortexResult<()> { + let array = VarBinViewArray::from_iter( + [Some("aa"), None, Some("aa"), Some("b")], + DType::Utf8(Nullability::Nullable), + ) + .into_array(); + let constraints = DictConstraints { + max_bytes: usize::MAX, + max_len: usize::MAX, + }; + let mut ctx = SESSION.create_execution_ctx(); + let limited = dict_encode_with_input_byte_limit(&array, &constraints, 4, &mut ctx)?; + let expected = array.slice(0..3)?; + + assert_eq!(limited.len(), 3); + assert_arrays_eq!(limited, expected, &mut ctx); + Ok(()) + } + + #[test] + fn limits_varbin_and_varbinview_input_bytes() -> VortexResult<()> { + let constraints = DictConstraints { + max_bytes: usize::MAX, + max_len: usize::MAX, + }; + let arrays = [ + VarBinArray::from(vec!["aa", "b", "aa"]).into_array(), + VarBinViewArray::from_iter_str(["aa", "b", "aa"]).into_array(), + ]; + let mut ctx = SESSION.create_execution_ctx(); + + for array in arrays { + let complete = dict_encode_with_input_byte_limit(&array, &constraints, 5, &mut ctx)?; + let limited = dict_encode_with_input_byte_limit(&array, &constraints, 3, &mut ctx)?; + let expected_prefix = array.slice(0..2)?; + assert_eq!(complete.len(), 3); + assert_eq!(limited.len(), 2); + assert_arrays_eq!(complete, array, &mut ctx); + assert_arrays_eq!(limited, expected_prefix, &mut ctx); + } + Ok(()) + } + + #[test] + fn input_byte_limit_accumulates_and_resets() -> VortexResult<()> { + let constraints = DictConstraints { + max_bytes: usize::MAX, + max_len: usize::MAX, + }; + let first = VarBinViewArray::from_iter_str(["aa"]).into_array(); + let second = VarBinViewArray::from_iter_str(["b"]).into_array(); + let mut encoder = BytesDictBuilder::::new( + DType::Utf8(Nullability::NonNullable), + &constraints, + Some(2), + ); + let mut ctx = SESSION.create_execution_ctx(); + + assert_eq!(encoder.encode(&first, &mut ctx)?.len(), 1); + assert_eq!(encoder.encode(&second, &mut ctx)?.len(), 0); + drop(encoder.reset()); + assert_eq!(encoder.encode(&first, &mut ctx)?.len(), 1); + Ok(()) + } + #[test] fn max_dict_bytes_cannot_exceed_the_view_offset_range() { - let builder = - BytesDictBuilder::::new(DType::Utf8(Nullability::NonNullable), &UNCONSTRAINED); + let builder = BytesDictBuilder::::new( + DType::Utf8(Nullability::NonNullable), + &UNCONSTRAINED, + None, + ); assert_eq!(builder.max_dict_bytes, u32::MAX as usize); } } diff --git a/vortex-array/src/builders/dict/mod.rs b/vortex-array/src/builders/dict/mod.rs index 94834f83cc4..69bb6e0d92f 100644 --- a/vortex-array/src/builders/dict/mod.rs +++ b/vortex-array/src/builders/dict/mod.rs @@ -24,7 +24,9 @@ mod primitive; #[derive(Clone)] pub struct DictConstraints { + /// Limits the encoded storage for dictionary values. pub max_bytes: usize, + /// Limits the number of dictionary values. pub max_len: usize, } @@ -67,7 +69,48 @@ pub fn dict_encode_with_constraints( constraints: &DictConstraints, ctx: &mut ExecutionCtx, ) -> VortexResult { - let mut encoder = dict_encoder(array, constraints); + let encoder = dict_encoder(array, constraints); + dict_encode_with_encoder(array, encoder, ctx) +} + +/// Encodes a variable-width array as a `DictArray` within storage and input-work limits. +/// +/// The input byte limit counts every valid value, including repeated values. Null values consume +/// no input bytes. The result contains the prefix completed before a limit was exceeded. +/// +/// # Errors +/// +/// Returns an error if array execution fails or the input is not variable-width. +pub fn dict_encode_with_input_byte_limit( + array: &ArrayRef, + constraints: &DictConstraints, + max_input_bytes: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let encoder = if let Some(varbinview) = array.as_opt::() { + bytes::bytes_dict_builder_with_input_byte_limit( + varbinview.dtype().clone(), + constraints, + max_input_bytes, + ) + } else if let Some(varbin) = array.as_opt::() { + bytes::bytes_dict_builder_with_input_byte_limit( + varbin.dtype().clone(), + constraints, + max_input_bytes, + ) + } else { + vortex_bail!("Input-byte limits require a variable-width array") + }; + dict_encode_with_encoder(array, encoder, ctx) +} + +/// Encodes an array with a configured dictionary encoder. +fn dict_encode_with_encoder( + array: &ArrayRef, + mut encoder: Box, + ctx: &mut ExecutionCtx, +) -> VortexResult { let codes = encoder.encode(array, ctx)?.narrow(ctx)?; // SAFETY: The encoding process will produce a value set of codes and values // All values in the dictionary are guaranteed to be referenced by at least one code diff --git a/vortex-btrblocks/Cargo.toml b/vortex-btrblocks/Cargo.toml index 4e22f042adf..1c5599ca862 100644 --- a/vortex-btrblocks/Cargo.toml +++ b/vortex-btrblocks/Cargo.toml @@ -69,3 +69,8 @@ test = false name = "compress_listview" harness = false test = false + +[[bench]] +name = "string_dict_selection" +harness = false +test = false diff --git a/vortex-btrblocks/benches/string_dict_selection.rs b/vortex-btrblocks/benches/string_dict_selection.rs new file mode 100644 index 00000000000..013c376260f --- /dev/null +++ b/vortex-btrblocks/benches/string_dict_selection.rs @@ -0,0 +1,195 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#[cfg(not(codspeed))] +mod benchmarks { + use std::sync::LazyLock; + + use divan::Bencher; + use divan::counter::ItemsCount; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_btrblocks::BtrBlocksCompressor; + use vortex_error::VortexExpect; + use vortex_session::VortexSession; + + static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + + #[derive(Clone, Copy)] + enum Distribution { + Uniform, + Skewed, + Clustered, + VariedPrefix, + } + + #[derive(Clone, Copy)] + struct Case { + name: &'static str, + rows: usize, + distinct_values: usize, + value_length: usize, + distribution: Distribution, + nullable: bool, + } + + impl std::fmt::Debug for Case { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(self.name) + } + } + + impl Case { + fn value_index(self, index: usize) -> usize { + match self.distribution { + Distribution::Uniform | Distribution::VariedPrefix => index % self.distinct_values, + Distribution::Clustered => index * self.distinct_values / self.rows, + Distribution::Skewed => { + let cold_start = self.rows * 9 / 10; + if index < cold_start { + index % 16 + } else { + 16 + (index - cold_start) % (self.distinct_values - 16) + } + } + } + } + + fn value(self, index: usize) -> String { + let value_index = self.value_index(index); + let mut value = match self.distribution { + Distribution::VariedPrefix => format!( + "{:016x}-value-{value_index:08x}", + value_index.wrapping_mul(0x9e3779b9) + ), + _ if self.value_length <= 12 => format!("{value_index:08x}"), + _ => format!("common-prefix-value-{value_index:08x}"), + }; + value.extend(std::iter::repeat_n( + 'x', + self.value_length.saturating_sub(value.len()), + )); + value + } + + fn make_array(self) -> ArrayRef { + let values = (0..self.rows) + .map(|index| self.value(index)) + .collect::>(); + let nullability = if self.nullable { + Nullability::Nullable + } else { + Nullability::NonNullable + }; + + VarBinViewArray::from_iter( + values.iter().enumerate().map(|(index, value)| { + (!self.nullable || index % 10 != 0).then_some(value.as_str()) + }), + DType::Utf8(nullability), + ) + .into_array() + } + } + + const fn case( + name: &'static str, + rows: usize, + distinct_values: usize, + value_length: usize, + distribution: Distribution, + ) -> Case { + Case { + name, + rows, + distinct_values, + value_length, + distribution, + nullable: false, + } + } + + const CASES: [Case; 17] = [ + case("Inline4096", 65_536, 4096, 8, Distribution::Uniform), + case("Outlined16", 65_536, 16, 28, Distribution::Uniform), + case("Outlined4096", 65_536, 4096, 28, Distribution::Uniform), + case("Outlined8192", 65_536, 8192, 28, Distribution::Uniform), + case( + "Outlined8192Large", + 1_048_576, + 8192, + 28, + Distribution::Uniform, + ), + Case { + nullable: true, + ..case( + "NullableOutlined4096", + 65_536, + 4096, + 28, + Distribution::Uniform, + ) + }, + case("Probe50", 8192, 4096, 28, Distribution::Uniform), + case("Probe75", 8192, 6144, 28, Distribution::Uniform), + case("Probe78", 8192, 6390, 28, Distribution::Uniform), + case( + "Probe78Large", + 1_048_576, + 817_889, + 28, + Distribution::Uniform, + ), + case("Probe79", 8192, 6471, 28, Distribution::Uniform), + case("Probe80", 8192, 6554, 28, Distribution::Uniform), + case("Skewed4096", 65_536, 4096, 28, Distribution::Skewed), + case("Clustered4096", 65_536, 4096, 28, Distribution::Clustered), + case("Long256", 65_536, 4096, 256, Distribution::Uniform), + case( + "UniqueCommonPrefix", + 65_536, + 65_536, + 28, + Distribution::Uniform, + ), + case( + "UniqueVariedPrefix", + 65_536, + 65_536, + 31, + Distribution::VariedPrefix, + ), + ]; + + #[divan::bench(args = CASES)] + fn compress(bencher: Bencher, case: Case) { + let array = case.make_array(); + let compressor = BtrBlocksCompressor::default(); + bencher + .with_inputs(|| (&array, SESSION.create_execution_ctx())) + .input_counter(|(array, _)| ItemsCount::new(array.len())) + .bench_refs(|(array, ctx)| compressor.compress(array, ctx)); + } + + #[divan::bench(args = CASES)] + fn decompress(bencher: Bencher, case: Case) { + let compressor = BtrBlocksCompressor::default(); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = compressor + .compress(&case.make_array(), &mut ctx) + .vortex_expect("benchmark input must compress"); + bencher + .with_inputs(|| (&compressed, SESSION.create_execution_ctx())) + .input_counter(|(array, _)| ItemsCount::new(array.len())) + .bench_refs(|(array, ctx)| array.clone().execute::(ctx)); + } +} + +fn main() { + divan::main(); +} diff --git a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs index ec51c5104bb..a6d2875295b 100644 --- a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs @@ -5,11 +5,13 @@ use std::sync::LazyLock; +use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::Constant; use vortex_array::arrays::Dict; use vortex_array::arrays::VarBinViewArray; +use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_error::VortexResult; @@ -17,6 +19,9 @@ use vortex_fsst::FSST; use vortex_session::VortexSession; use crate::BtrBlocksCompressor; +use crate::BtrBlocksCompressorBuilder; +use crate::schemes::string::FSSTScheme; +use crate::schemes::string::StringDictScheme; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); @@ -46,6 +51,97 @@ fn test_dict_compressed() -> VortexResult<()> { Ok(()) } +#[rstest] +#[case::outlined_utf8(4096, 28, false)] +#[case::nullable_utf8(4096, 28, true)] +#[case::outlined_utf8_8192(8192, 28, false)] +#[case::long_utf8(4096, 256, false)] +fn test_dict_compressed_with_more_values_than_sample( + #[case] distinct_count: usize, + #[case] value_length: usize, + #[case] nullable: bool, +) -> VortexResult<()> { + let distinct_values = (0..distinct_count) + .map(|value| { + let mut string = format!("common-prefix-value-{value:08x}"); + string.extend(std::iter::repeat_n( + 'x', + value_length.saturating_sub(string.len()), + )); + string + }) + .collect::>(); + let values = (0..65_536) + .map(|index| { + (!nullable || index % 10 != 0) + .then_some(distinct_values[index % distinct_values.len()].as_bytes()) + }) + .collect::>(); + let nullability = if nullable { + Nullability::Nullable + } else { + Nullability::NonNullable + }; + let array = VarBinViewArray::from_iter(values, DType::Utf8(nullability)).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let compressor = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&StringDictScheme) + .with_new_scheme(&FSSTScheme) + .build(); + let compressed = compressor.compress(&array, &mut ctx)?; + + assert!( + compressed.is::(), + "expected Dict, got {}", + compressed.encoding_id() + ); + assert_arrays_eq!(&array, &compressed, &mut ctx); + Ok(()) +} + +#[test] +fn test_unique_strings_with_common_prefix_not_dict_compressed() -> VortexResult<()> { + let values = (0usize..4096) + .map(|value| Some(format!("common-prefix-value-{value:08x}"))) + .collect::>(); + let array = + VarBinViewArray::from_iter(values, DType::Utf8(Nullability::NonNullable)).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = BtrBlocksCompressor::default().compress(&array, &mut ctx)?; + + assert!( + !compressed.is::(), + "expected a non-Dict encoding, got {}", + compressed.encoding_id() + ); + assert_arrays_eq!(&array, &compressed, &mut ctx); + Ok(()) +} + +#[test] +fn test_sample_fallback_can_select_dict() -> VortexResult<()> { + let suffix = "x".repeat((1 << 20) + 1); + let distinct_values = [format!("first-{suffix}"), format!("second-{suffix}")]; + let values = (0..4) + .map(|index| distinct_values[index % distinct_values.len()].as_str()) + .collect::>(); + let array = VarBinViewArray::from_iter_str(values).into_array(); + let compressor = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&StringDictScheme) + .with_new_scheme(&FSSTScheme) + .build(); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = compressor.compress(&array, &mut ctx)?; + + assert!( + compressed.is::(), + "expected Dict, got {}", + compressed.encoding_id() + ); + assert_arrays_eq!(&array, &compressed, &mut ctx); + Ok(()) +} + #[cfg(feature = "unstable_encodings")] #[test] fn test_unstable_all_schemes_includes_onpair() { diff --git a/vortex-compressor/src/builtins/dict/mod.rs b/vortex-compressor/src/builtins/dict/mod.rs index 4862df2b211..5b26d960f16 100644 --- a/vortex-compressor/src/builtins/dict/mod.rs +++ b/vortex-compressor/src/builtins/dict/mod.rs @@ -7,6 +7,7 @@ mod binary; mod float; mod integer; mod string; +mod string_candidate; pub use binary::BinaryDictScheme; pub use float::FloatDictScheme; diff --git a/vortex-compressor/src/builtins/dict/string.rs b/vortex-compressor/src/builtins/dict/string.rs index f5cbcd54d89..a572b957b82 100644 --- a/vortex-compressor/src/builtins/dict/string.rs +++ b/vortex-compressor/src/builtins/dict/string.rs @@ -19,21 +19,19 @@ use vortex_array::arrays::dict::DictArrayExt; use vortex_array::arrays::dict::DictArraySlotsExt; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::builders::dict::dict_encode; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::CascadingCompressor; use crate::builtins::IntDictScheme; +use crate::builtins::dict::string_candidate::CachedStringDictionary; +use crate::builtins::dict::string_candidate::string_dictionary_estimate; use crate::scheme::ChildSelection; use crate::scheme::CompressionEstimate; use crate::scheme::CompressorContext; -use crate::scheme::DeferredEstimate; use crate::scheme::DescendantExclusion; -use crate::scheme::EstimateVerdict; use crate::scheme::Scheme; use crate::scheme::SchemeExt; use crate::stats::ArrayAndStats; -use crate::stats::GenerateStatsOptions; /// Dictionary encoding for low-cardinality string values. #[derive(Debug, Copy, Clone, PartialEq, Eq)] @@ -52,12 +50,6 @@ impl Scheme for StringDictScheme { vec![Dict.id()] } - fn stats_options(&self) -> GenerateStatsOptions { - GenerateStatsOptions { - count_distinct_values: true, - } - } - /// Children: values=0, codes=1. fn num_children(&self) -> usize { 2 @@ -77,27 +69,11 @@ impl Scheme for StringDictScheme { fn expected_compression_ratio( &self, - data: &ArrayAndStats, + _data: &ArrayAndStats, _compress_ctx: CompressorContext, - exec_ctx: &mut ExecutionCtx, + _exec_ctx: &mut ExecutionCtx, ) -> CompressionEstimate { - let stats = data.varbinview_stats(exec_ctx); - - if stats.value_count() == 0 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); - } - - let estimated_distinct_values_count = stats.estimated_distinct_count().vortex_expect( - "this must be present since `DictScheme` declared that we need distinct values", - ); - - // If > 50% of the values are distinct, skip dictionary scheme. - if estimated_distinct_values_count > stats.value_count() / 2 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); - } - - // Let sampling determine the expected ratio. - CompressionEstimate::Deferred(DeferredEstimate::Sample) + string_dictionary_estimate() } fn compress( @@ -107,29 +83,52 @@ impl Scheme for StringDictScheme { compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - let dict = dict_encode(data.array(), exec_ctx)?; - - // Values = child 0. - let compressed_values = - compressor.compress_child(dict.values(), &compress_ctx, self.id(), 0, exec_ctx)?; + if let Some(candidate) = data.get::() { + return Ok(candidate.array().clone()); + } - // Codes = child 1. - let narrowed_codes = dict - .codes() - .clone() - .execute::(exec_ctx)? - .narrow(exec_ctx)? - .into_array(); - let compressed_codes = - compressor.compress_child(&narrowed_codes, &compress_ctx, self.id(), 1, exec_ctx)?; + let dict = dict_encode(data.array(), exec_ctx)?; + compress_dictionary(compressor, &dict, compress_ctx, exec_ctx) + } +} - // SAFETY: compressing codes or values does not alter the invariants. - unsafe { - Ok( - DictArray::new_unchecked(compressed_codes, compressed_values) - .set_all_values_referenced(dict.has_all_values_referenced()) - .into_array(), - ) - } +/// Compresses the value and code children of a dictionary candidate. +pub(super) fn compress_dictionary( + compressor: &CascadingCompressor, + dict: &DictArray, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + // Values = child 0. + let compressed_values = compressor.compress_child( + dict.values(), + &compress_ctx, + StringDictScheme.id(), + 0, + exec_ctx, + )?; + + // Codes = child 1. + let narrowed_codes = dict + .codes() + .clone() + .execute::(exec_ctx)? + .narrow(exec_ctx)? + .into_array(); + let compressed_codes = compressor.compress_child( + &narrowed_codes, + &compress_ctx, + StringDictScheme.id(), + 1, + exec_ctx, + )?; + + // SAFETY: compressing codes or values does not alter the invariants. + unsafe { + Ok( + DictArray::new_unchecked(compressed_codes, compressed_values) + .set_all_values_referenced(dict.has_all_values_referenced()) + .into_array(), + ) } } diff --git a/vortex-compressor/src/builtins/dict/string_candidate.rs b/vortex-compressor/src/builtins/dict/string_candidate.rs new file mode 100644 index 00000000000..e9e45368e7e --- /dev/null +++ b/vortex-compressor/src/builtins/dict/string_candidate.rs @@ -0,0 +1,485 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Bounded trial compression for UTF-8 dictionary selection. +//! +//! Ordinary samples can miss repeated strings when a column has thousands of distinct values. +//! This module builds one bounded dictionary candidate from the complete input. +//! +//! A distributed full-value probe protects high-cardinality inputs from that full pass. +//! A rejected probe uses the existing sample estimator, so the probe changes work rather than selection policy. +//! +//! If the trial dictionary wins, the compression phase reuses its completed output. + +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::Bool; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::VarBinView; +use vortex_array::arrays::bool::BoolArrayExt; +use vortex_array::arrays::varbinview::BinaryView; +use vortex_array::builders::dict::DictConstraints; +use vortex_array::builders::dict::dict_encode_with_input_byte_limit; +use vortex_error::VortexResult; +use vortex_utils::aliases::hash_set::HashSet; + +use super::StringDictScheme; +use super::string::compress_dictionary; +use crate::CascadingCompressor; +use crate::compressor::estimate_compression_ratio_with_sampling; +use crate::compressor::sample_slices; +use crate::scheme::CompressionEstimate; +use crate::scheme::CompressorContext; +use crate::scheme::DeferredEstimate; +use crate::scheme::EstimateScore; +use crate::scheme::EstimateVerdict; +use crate::stats::ArrayAndStats; + +/// Limits the number of entries so that trial codes remain two bytes wide. +const MAX_DICTIONARY_ENTRIES: usize = u16::MAX as usize + 1; + +/// Limits the value and view storage in a trial dictionary. +const MAX_DICTIONARY_BYTES: usize = 1 << 22; + +/// Limits code storage while the trial scans the complete input. +const MAX_DICTIONARY_CODE_BYTES: usize = 1 << 24; + +/// Limits complete input bytes hashed while constructing a trial dictionary. +const MAX_DICTIONARY_INPUT_BYTES: usize = 1 << 28; + +/// Limits complete value bytes hashed by the preliminary probe. +const MAX_PROBE_BYTES: usize = 1 << 20; + +/// Samples this many adjacent rows from each probe range. +const PROBE_RANGE_SIZE: u32 = 64; + +/// Distributes this many probe ranges across the input. +const PROBE_RANGE_COUNT: u32 = 128; + +/// Sets the numerator for the 80% work-admission threshold. +const PROBE_DISTINCT_NUMERATOR: usize = 4; + +/// Sets the denominator for the 80% work-admission threshold. +const PROBE_DISTINCT_DENOMINATOR: usize = 5; + +/// Stores a completed dictionary result for reuse by the compression phase. +#[derive(Debug)] +pub(super) struct CachedStringDictionary(ArrayRef); + +impl CachedStringDictionary { + /// Returns the completed dictionary result. + pub(super) fn array(&self) -> &ArrayRef { + &self.0 + } +} + +/// Selects trial dictionary construction or ordinary sample estimation. +#[derive(Debug, PartialEq, Eq)] +enum ProbeResult { + /// Permits bounded trial dictionary construction. + Candidate, + /// Uses ordinary sample estimation. + Sample, +} + +/// Reports the result of bounded trial dictionary construction. +enum CandidateResult { + /// Contains a complete trial dictionary. + Complete(DictArray), + /// Reports that a resource bound stopped trial construction. + Sample, +} + +/// Returns a deferred estimate that can reuse a completed dictionary result. +pub(super) fn string_dictionary_estimate() -> CompressionEstimate { + CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( + estimate_string_dictionary, + ))) +} + +/// Resolves the deferred dictionary estimate. +fn estimate_string_dictionary( + compressor: &CascadingCompressor, + data: &ArrayAndStats, + best_so_far: Option, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + if probe(data)? == ProbeResult::Sample { + return sample_verdict(compressor, data, best_so_far, compress_ctx, exec_ctx); + } + + let dictionary = match build_candidate(data, exec_ctx)? { + CandidateResult::Complete(dictionary) => dictionary, + CandidateResult::Sample => { + return sample_verdict(compressor, data, best_so_far, compress_ctx, exec_ctx); + } + }; + let compressed = compress_dictionary(compressor, &dictionary, compress_ctx, exec_ctx)?; + let score = EstimateScore::from_sample_sizes(data.array().nbytes(), compressed.nbytes()); + + if !score_is_best(score, best_so_far) { + return Ok(EstimateVerdict::Skip); + } + + data.get_or_insert_with(|| CachedStringDictionary(compressed)); + Ok(score_verdict(score)) +} + +/// Runs the ordinary sample estimate and applies the current threshold. +fn sample_verdict( + compressor: &CascadingCompressor, + data: &ArrayAndStats, + best_so_far: Option, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + let score = estimate_compression_ratio_with_sampling( + compressor, + &StringDictScheme, + data.array(), + compress_ctx, + exec_ctx, + )?; + if score_is_best(score, best_so_far) { + Ok(score_verdict(score)) + } else { + Ok(EstimateVerdict::Skip) + } +} + +/// Returns whether a score is valid and beats the current best score. +fn score_is_best(score: EstimateScore, best_so_far: Option) -> bool { + score.is_valid() && best_so_far.is_none_or(|best| score.beats(best)) +} + +/// Converts a score into a terminal estimate. +fn score_verdict(score: EstimateScore) -> EstimateVerdict { + match score { + EstimateScore::FiniteCompression(ratio) => EstimateVerdict::Ratio(ratio), + EstimateScore::ZeroBytes => EstimateVerdict::Skip, + } +} + +/// Checks complete values from distributed ranges before trial construction. +fn probe(data: &ArrayAndStats) -> VortexResult { + let array = data.array_as_varbinview(); + let views = array.views(); + let validity = array.validity()?; + let validity_bits = match &validity { + vortex_array::validity::Validity::NonNullable + | vortex_array::validity::Validity::AllValid => None, + vortex_array::validity::Validity::AllInvalid => return Ok(ProbeResult::Sample), + vortex_array::validity::Validity::Array(validity) => { + let Some(validity) = validity.as_opt::() else { + return Ok(ProbeResult::Sample); + }; + Some(validity.to_bit_buffer()) + } + }; + let sample_ranges = sample_slices(array.len(), PROBE_RANGE_SIZE, PROBE_RANGE_COUNT); + let sample_rows = sample_ranges + .iter() + .map(|(start, end)| end - start) + .sum::(); + let mut distinct_values = HashSet::with_capacity(sample_rows); + let mut probed_bytes = 0usize; + let mut valid_rows = 0usize; + + let maximum_range_length = sample_ranges + .iter() + .map(|(start, end)| end - start) + .max() + .unwrap_or_default(); + + 'probe: for offset in 0..maximum_range_length { + for &(start, end) in &sample_ranges { + let index = start + offset; + if index >= end { + continue; + } + if validity_bits + .as_ref() + .is_some_and(|validity| !validity.value(index)) + { + continue; + } + let value = view_bytes(&array, &views[index]); + let Some(next_probed_bytes) = probed_bytes.checked_add(value.len()) else { + break 'probe; + }; + if next_probed_bytes > MAX_PROBE_BYTES { + break 'probe; + } + probed_bytes = next_probed_bytes; + valid_rows += 1; + distinct_values.insert(value); + } + } + + if valid_rows != 0 + && distinct_values.len() * PROBE_DISTINCT_DENOMINATOR + <= valid_rows * PROBE_DISTINCT_NUMERATOR + { + Ok(ProbeResult::Candidate) + } else { + Ok(ProbeResult::Sample) + } +} + +/// Builds a dictionary within explicit storage and work bounds. +fn build_candidate( + data: &ArrayAndStats, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + let Some(code_bytes) = candidate_code_bytes(data.array_len(), MAX_DICTIONARY_ENTRIES) else { + return Ok(CandidateResult::Sample); + }; + if code_bytes > MAX_DICTIONARY_CODE_BYTES { + return Ok(CandidateResult::Sample); + } + + let constraints = DictConstraints { + max_bytes: MAX_DICTIONARY_BYTES, + max_len: MAX_DICTIONARY_ENTRIES, + }; + let dictionary = dict_encode_with_input_byte_limit( + data.array(), + &constraints, + MAX_DICTIONARY_INPUT_BYTES, + exec_ctx, + )?; + if dictionary.len() != data.array_len() { + return Ok(CandidateResult::Sample); + } + + Ok(CandidateResult::Complete(dictionary)) +} + +/// Returns the code storage required by the constrained dictionary builder. +fn candidate_code_bytes(row_count: usize, maximum_dictionary_values: usize) -> Option { + let code_width = if maximum_dictionary_values <= usize::from(u8::MAX) + 1 { + size_of::() + } else if maximum_dictionary_values <= usize::from(u16::MAX) + 1 { + size_of::() + } else { + size_of::() + }; + row_count.checked_mul(code_width) +} + +/// Returns the complete value referenced by a binary view. +fn view_bytes<'a>(array: &'a ArrayView<'a, VarBinView>, view: &'a BinaryView) -> &'a [u8] { + if view.is_inlined() { + view.as_inlined().value() + } else { + let reference = view.as_view(); + &array.buffer(reference.buffer_index as usize)[reference.as_range()] + } +} + +#[cfg(test)] +mod tests { + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::Dict; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_error::VortexResult; + + use super::CachedStringDictionary; + use super::CandidateResult; + use super::PROBE_RANGE_COUNT; + use super::PROBE_RANGE_SIZE; + use super::ProbeResult; + use super::build_candidate; + use super::probe; + use super::sample_slices; + use super::string_dictionary_estimate; + use crate::CascadingCompressor; + use crate::builtins::StringDictScheme; + use crate::scheme::CompressionEstimate; + use crate::scheme::CompressorContext; + use crate::scheme::DeferredEstimate; + use crate::scheme::EstimateVerdict; + use crate::scheme::Scheme; + use crate::stats::ArrayAndStats; + + /// Builds compression data for string candidate tests. + fn string_data(values: &[Option]) -> ArrayAndStats { + let nullability = if values.iter().any(Option::is_none) { + Nullability::Nullable + } else { + Nullability::NonNullable + }; + let array = VarBinViewArray::from_iter( + values.iter().map(|value| value.as_deref()), + DType::Utf8(nullability), + ) + .into_array(); + ArrayAndStats::new(array, Default::default()) + } + + #[test] + fn probe_distinct_ratio_boundary() -> VortexResult<()> { + for (distinct_values, expected) in + [(6553, ProbeResult::Candidate), (6554, ProbeResult::Sample)] + { + let values = (0..8192) + .map(|index| { + Some(format!( + "common-prefix-value-{:08x}", + index % distinct_values + )) + }) + .collect::>(); + let data = string_data(&values); + + assert_eq!(probe(&data)?, expected); + } + Ok(()) + } + + #[test] + fn probe_distributes_rows_before_byte_limit() -> VortexResult<()> { + let mut values = (0..65_536) + .map(|index| Some(format!("{index:0256x}"))) + .collect::>(); + for (start, end) in sample_slices(values.len(), PROBE_RANGE_SIZE, PROBE_RANGE_COUNT) + .into_iter() + .take(16) + { + values[start..end].fill(Some("x".repeat(256))); + } + let data = string_data(&values); + assert_eq!(probe(&data)?, ProbeResult::Sample); + Ok(()) + } + + #[test] + fn candidate_defers_when_dictionary_storage_exceeds_limit() -> VortexResult<()> { + let suffix = "x".repeat(70_000); + let distinct_values = (0..64) + .map(|value| format!("{value:08x}-{suffix}")) + .collect::>(); + let values = (0..128) + .map(|index| Some(distinct_values[index % distinct_values.len()].clone())) + .collect::>(); + let data = string_data(&values); + let mut exec_ctx = vortex_array::array_session().create_execution_ctx(); + + assert!(matches!( + build_candidate(&data, &mut exec_ctx)?, + CandidateResult::Sample + )); + Ok(()) + } + + #[test] + fn candidate_uses_full_u16_code_range() -> VortexResult<()> { + let values = (0..=u16::MAX) + .map(|value| Some(format!("{value:08x}"))) + .collect::>(); + let data = string_data(&values); + let mut exec_ctx = vortex_array::array_session().create_execution_ctx(); + + assert!(matches!( + build_candidate(&data, &mut exec_ctx)?, + CandidateResult::Complete(_) + )); + + let mut values_past_limit = values; + values_past_limit.push(Some("past-u16-code-range".to_owned())); + let data = string_data(&values_past_limit); + assert!(matches!( + build_candidate(&data, &mut exec_ctx)?, + CandidateResult::Sample + )); + Ok(()) + } + + #[test] + fn completed_candidate_is_reused() -> VortexResult<()> { + let values = (0..65_536) + .map(|index| Some(format!("common-prefix-value-{:08x}", index % 4096))) + .collect::>(); + let data = string_data(&values); + let compressor = CascadingCompressor::new(vec![&StringDictScheme]); + let mut exec_ctx = vortex_array::array_session().create_execution_ctx(); + let callback = match string_dictionary_estimate() { + CompressionEstimate::Deferred(DeferredEstimate::Callback(callback)) => callback, + estimate => { + return Err(vortex_error::vortex_err!( + "dictionary estimation returned {estimate:?}" + )); + } + }; + let verdict = callback( + &compressor, + &data, + None, + CompressorContext::new(), + &mut exec_ctx, + )?; + assert!(matches!(verdict, EstimateVerdict::Ratio(_))); + let cached = data + .get::() + .ok_or_else(|| vortex_error::vortex_err!("completed candidate was not cached"))?; + let compressed = StringDictScheme.compress( + &compressor, + &data, + CompressorContext::new(), + &mut exec_ctx, + )?; + + assert!(ArrayRef::ptr_eq(cached.array(), &compressed)); + Ok(()) + } + + #[test] + fn sample_fallback_can_select_dictionary() -> VortexResult<()> { + let mut values = (0..65_536) + .map(|index| Some(format!("repeated-value-{:08x}", index % 2))) + .collect::>(); + for (start, end) in sample_slices(values.len(), PROBE_RANGE_SIZE, PROBE_RANGE_COUNT) { + for index in start..end { + values[index] = Some(format!("probe-only-value-{index:08x}")); + } + } + let data = string_data(&values); + let compressor = CascadingCompressor::new(vec![&StringDictScheme]); + let mut exec_ctx = vortex_array::array_session().create_execution_ctx(); + + assert_eq!(probe(&data)?, ProbeResult::Sample); + let callback = match string_dictionary_estimate() { + CompressionEstimate::Deferred(DeferredEstimate::Callback(callback)) => callback, + estimate => { + return Err(vortex_error::vortex_err!( + "dictionary estimation returned {estimate:?}" + )); + } + }; + assert!(matches!( + callback( + &compressor, + &data, + None, + CompressorContext::new(), + &mut exec_ctx, + )?, + EstimateVerdict::Ratio(_) + )); + let compressed = StringDictScheme.compress( + &compressor, + &data, + CompressorContext::new(), + &mut exec_ctx, + )?; + + assert!(compressed.is::()); + Ok(()) + } +} diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index 219b67e2519..e228c707d4c 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -9,6 +9,9 @@ mod sample; mod select; mod structural; +pub(crate) use sample::estimate_compression_ratio_with_sampling; +pub(crate) use sample::sample_slices; + use crate::builtins::IntDictScheme; use crate::scheme::ChildSelection; use crate::scheme::DescendantExclusion; diff --git a/vortex-compressor/src/compressor/sample.rs b/vortex-compressor/src/compressor/sample.rs index ba1271d4a6d..f3700da2aa6 100644 --- a/vortex-compressor/src/compressor/sample.rs +++ b/vortex-compressor/src/compressor/sample.rs @@ -3,6 +3,8 @@ //! Sampling utilities for compression ratio estimation. +use std::mem::size_of; + use rand::RngExt; use rand::SeedableRng; use rand::prelude::StdRng; @@ -11,8 +13,12 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::VarBinView; +use vortex_array::arrays::varbinview::BinaryView; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_utils::aliases::hash_map::HashMap; use crate::CascadingCompressor; use crate::scheme::CompressorContext; @@ -42,12 +48,7 @@ pub(crate) fn sample(input: &ArrayRef, sample_size: u32, sample_count: u32) -> A return input.clone(); } - let slices = stratified_slices( - input.len(), - sample_size, - sample_count, - &mut StdRng::seed_from_u64(SAMPLE_SEED), - ); + let slices = sample_slices(input.len(), sample_size, sample_count); // For every slice, grab the relevant slice and repack into a new PrimitiveArray. let chunks: Vec<_> = slices @@ -62,6 +63,20 @@ pub(crate) fn sample(input: &ArrayRef, sample_size: u32, sample_count: u32) -> A unsafe { ChunkedArray::new_unchecked(chunks, input.dtype().clone()) }.into_array() } +/// Returns deterministic stratified sample ranges for an array length. +pub(crate) fn sample_slices( + length: usize, + sample_size: u32, + sample_count: u32, +) -> Vec<(usize, usize)> { + stratified_slices( + length, + sample_size, + sample_count, + &mut StdRng::seed_from_u64(SAMPLE_SEED), + ) +} + /// Computes the number of sample chunks to cover approximately 1% of `len` elements, /// with a minimum of `SAMPLE_SIZE * SAMPLE_COUNT` (1024) values. pub(crate) fn sample_count_approx_one_percent(len: usize) -> u32 { @@ -149,7 +164,7 @@ fn partition_indices(length: usize, num_partitions: u32) -> Vec<(usize, usize)> /// # Errors /// /// Returns an error if sample compression fails. -pub(super) fn estimate_compression_ratio_with_sampling( +pub(crate) fn estimate_compression_ratio_with_sampling( compressor: &CascadingCompressor, scheme: &S, array: &ArrayRef, @@ -178,7 +193,7 @@ pub(super) fn estimate_compression_ratio_with_sampling( }; let after = compressed.nbytes(); - let before = sample_data.array().nbytes(); + let before = canonical_visible_nbytes(sample_data.array(), exec_ctx)?; let score = EstimateScore::from_sample_sizes(before, after); @@ -189,15 +204,101 @@ pub(super) fn estimate_compression_ratio_with_sampling( Ok(score) } +/// Returns the physical canonical size without retained, unreferenced string payload bytes. +pub(crate) fn canonical_visible_nbytes( + array: &ArrayRef, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + let Some(array) = array.as_opt::() else { + return Ok(array.nbytes()); + }; + let validity = array.validity()?.execute_mask(array.len(), exec_ctx)?; + let mut referenced_ranges = HashMap::<(usize, usize), Vec<(usize, usize)>>::new(); + let mut record_view = |view: &BinaryView| { + if view.is_inlined() { + return; + } + let reference = view.as_view(); + let buffer = array.buffer(reference.buffer_index as usize); + referenced_ranges + .entry((buffer.as_ptr().addr(), buffer.len())) + .or_default() + .push(( + reference.offset as usize, + reference.offset as usize + reference.size as usize, + )); + }; + match &validity { + Mask::AllTrue(_) => array.views().iter().for_each(&mut record_view), + Mask::AllFalse(_) => {} + Mask::Values(values) => array + .views() + .iter() + .zip(values.bit_buffer().iter()) + .filter(|(_, is_valid)| *is_valid) + .for_each(|(view, _)| record_view(view)), + } + let outlined_bytes = referenced_ranges.into_values().try_fold( + 0u64, + |total, mut ranges| -> VortexResult { + ranges.sort_unstable(); + let mut range_total = 0u64; + let mut current = None::<(usize, usize)>; + for (start, end) in ranges { + match current { + Some((current_start, current_end)) if start <= current_end => { + current = Some((current_start, current_end.max(end))); + } + Some((current_start, current_end)) => { + range_total = range_total + .checked_add(u64::try_from(current_end - current_start)?) + .ok_or_else(|| { + vortex_error::vortex_err!("sample byte size overflowed u64") + })?; + current = Some((start, end)); + } + None => current = Some((start, end)), + } + } + if let Some((start, end)) = current { + range_total = range_total + .checked_add(u64::try_from(end - start)?) + .ok_or_else(|| vortex_error::vortex_err!("sample byte size overflowed u64"))?; + } + total + .checked_add(range_total) + .ok_or_else(|| vortex_error::vortex_err!("sample byte size overflowed u64")) + }, + )?; + let views_bytes = u64::try_from(array.len())? + .checked_mul(u64::try_from(size_of::())?) + .ok_or_else(|| vortex_error::vortex_err!("sample byte size overflowed u64"))?; + let validity_bytes = match validity { + Mask::Values(values) => u64::try_from(values.len().div_ceil(u8::BITS as usize))?, + Mask::AllTrue(_) | Mask::AllFalse(_) => 0, + }; + + views_bytes + .checked_add(outlined_bytes) + .and_then(|bytes| bytes.checked_add(validity_bytes)) + .ok_or_else(|| vortex_error::vortex_err!("sample byte size overflowed u64")) +} + #[cfg(test)] mod tests { + use std::sync::Arc; + use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; use vortex_array::validity::Validity; use vortex_buffer::Buffer; + use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; use super::*; @@ -219,4 +320,56 @@ mod tests { } Ok(()) } + + #[test] + fn visible_size_ignores_unreferenced_string_payload() -> VortexResult<()> { + let values = (0..128) + .map(|index| format!("outlined-string-value-{index:08x}")) + .collect::>(); + let source = VarBinViewArray::from_iter_str(&values).into_array(); + let slice = source.slice(17..19)?; + let expected = u64::try_from(2 * size_of::())? + + u64::try_from(values[17].len() + values[18].len())?; + let mut exec_ctx = array_session().create_execution_ctx(); + + assert!(slice.nbytes() > expected); + assert_eq!(canonical_visible_nbytes(&slice, &mut exec_ctx)?, expected); + Ok(()) + } + + #[test] + fn visible_size_counts_valid_outlined_values_and_validity() -> VortexResult<()> { + let outlined = "an-outlined-string-value"; + let array = VarBinViewArray::from_iter( + [Some("inline"), Some(outlined), None], + DType::Utf8(Nullability::Nullable), + ) + .into_array(); + let expected = u64::try_from(3 * size_of::() + outlined.len() + 1)?; + let mut exec_ctx = array_session().create_execution_ctx(); + + assert_eq!(canonical_visible_nbytes(&array, &mut exec_ctx)?, expected); + Ok(()) + } + + #[test] + fn visible_size_counts_shared_payload_once() -> VortexResult<()> { + let value = b"one shared outlined value"; + let view = BinaryView::make_view(value, 0, 0); + let array = VarBinViewArray::try_new( + Buffer::copy_from([view, view]), + Arc::from([ByteBuffer::copy_from(value)]), + DType::Utf8(Nullability::NonNullable), + Validity::NonNullable, + &mut array_session().create_execution_ctx(), + )? + .into_array(); + let expected = u64::try_from(2 * size_of::() + value.len())?; + + assert_eq!( + canonical_visible_nbytes(&array, &mut array_session().create_execution_ctx())?, + expected + ); + Ok(()) + } } diff --git a/vortex-compressor/src/stats/cache.rs b/vortex-compressor/src/stats/cache.rs index 2652403d3e6..07d0bf1cca1 100644 --- a/vortex-compressor/src/stats/cache.rs +++ b/vortex-compressor/src/stats/cache.rs @@ -67,6 +67,19 @@ impl StatsCache { new_arc } } + + /// Returns a cached value when one exists. + fn get(&self) -> Option> { + let type_id = TypeId::of::(); + let guard = self.entries.lock(); + let position = guard.iter().position(|(id, _)| *id == type_id)?; + Some( + Arc::clone(&guard[position].1) + .downcast::() + .ok() + .vortex_expect("we just checked the TypeID"), + ) + } } /// An array bundled with its lazily-computed statistics cache. @@ -207,4 +220,9 @@ impl ArrayAndStats { pub fn get_or_insert_with(&self, f: impl FnOnce() -> T) -> Arc { self.cache.get_or_insert_with::(f) } + + /// Returns a custom cached value when one exists. + pub(crate) fn get(&self) -> Option> { + self.cache.get::() + } } From 6ed06bd530c2dcd8d9be9df248d50d2623fc63b3 Mon Sep 17 00:00:00 2001 From: Will Manning Date: Thu, 3 Sep 2026 21:02:39 -0400 Subject: [PATCH 2/2] perf(compressor): Refine string dictionary admission Signed-off-by: Will Manning --- vortex-array/src/builders/dict/bytes.rs | 18 ++ .../benches/string_dict_selection.rs | 31 +-- .../schemes/string/scheme_selection_tests.rs | 22 ++- .../src/builtins/dict/string_candidate.rs | 187 +++++++++--------- vortex-compressor/src/compressor/sample.rs | 166 +++------------- 5 files changed, 159 insertions(+), 265 deletions(-) diff --git a/vortex-array/src/builders/dict/bytes.rs b/vortex-array/src/builders/dict/bytes.rs index 02be10da250..29aa6a478fe 100644 --- a/vortex-array/src/builders/dict/bytes.rs +++ b/vortex-array/src/builders/dict/bytes.rs @@ -360,12 +360,14 @@ mod test { use std::sync::Arc; use std::sync::LazyLock; + use rstest::rstest; use vortex_buffer::Buffer; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; use vortex_session::VortexSession; use super::BytesDictBuilder; + use super::bytes_dict_builder; use crate::IntoArray; use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; @@ -383,6 +385,7 @@ mod test { use crate::builders::dict::dict_encoder; use crate::dtype::DType; use crate::dtype::Nullability; + use crate::dtype::PType; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(crate::array_session); @@ -575,6 +578,21 @@ mod test { Ok(()) } + #[rstest] + #[case(usize::from(u8::MAX) + 1, PType::U8)] + #[case(usize::from(u8::MAX) + 2, PType::U16)] + #[case(usize::from(u16::MAX) + 1, PType::U16)] + #[case(usize::from(u16::MAX) + 2, PType::U32)] + fn selects_narrowest_code_type(#[case] max_len: usize, #[case] expected: PType) { + let constraints = DictConstraints { + max_bytes: usize::MAX, + max_len, + }; + let encoder = bytes_dict_builder(DType::Utf8(Nullability::NonNullable), &constraints); + + assert_eq!(encoder.codes_ptype(), expected); + } + #[test] fn max_dict_bytes_cannot_exceed_the_view_offset_range() { let builder = BytesDictBuilder::::new( diff --git a/vortex-btrblocks/benches/string_dict_selection.rs b/vortex-btrblocks/benches/string_dict_selection.rs index 013c376260f..8ca2c0a5a5e 100644 --- a/vortex-btrblocks/benches/string_dict_selection.rs +++ b/vortex-btrblocks/benches/string_dict_selection.rs @@ -22,7 +22,6 @@ mod benchmarks { #[derive(Clone, Copy)] enum Distribution { Uniform, - Skewed, Clustered, VariedPrefix, } @@ -48,14 +47,6 @@ mod benchmarks { match self.distribution { Distribution::Uniform | Distribution::VariedPrefix => index % self.distinct_values, Distribution::Clustered => index * self.distinct_values / self.rows, - Distribution::Skewed => { - let cold_start = self.rows * 9 / 10; - if index < cold_start { - index % 16 - } else { - 16 + (index - cold_start) % (self.distinct_values - 16) - } - } } } @@ -113,18 +104,11 @@ mod benchmarks { } } - const CASES: [Case; 17] = [ + const CASES: [Case; 9] = [ case("Inline4096", 65_536, 4096, 8, Distribution::Uniform), case("Outlined16", 65_536, 16, 28, Distribution::Uniform), case("Outlined4096", 65_536, 4096, 28, Distribution::Uniform), case("Outlined8192", 65_536, 8192, 28, Distribution::Uniform), - case( - "Outlined8192Large", - 1_048_576, - 8192, - 28, - Distribution::Uniform, - ), Case { nullable: true, ..case( @@ -135,19 +119,6 @@ mod benchmarks { Distribution::Uniform, ) }, - case("Probe50", 8192, 4096, 28, Distribution::Uniform), - case("Probe75", 8192, 6144, 28, Distribution::Uniform), - case("Probe78", 8192, 6390, 28, Distribution::Uniform), - case( - "Probe78Large", - 1_048_576, - 817_889, - 28, - Distribution::Uniform, - ), - case("Probe79", 8192, 6471, 28, Distribution::Uniform), - case("Probe80", 8192, 6554, 28, Distribution::Uniform), - case("Skewed4096", 65_536, 4096, 28, Distribution::Skewed), case("Clustered4096", 65_536, 4096, 28, Distribution::Clustered), case("Long256", 65_536, 4096, 256, Distribution::Uniform), case( diff --git a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs index a6d2875295b..e41332ab9d1 100644 --- a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs @@ -25,6 +25,12 @@ use crate::schemes::string::StringDictScheme; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +#[derive(Clone, Copy, Debug)] +enum Distribution { + Uniform, + Clustered, +} + #[test] fn test_constant_compressed() -> VortexResult<()> { let strings: Vec> = vec![Some("constant_value"); 100]; @@ -52,14 +58,15 @@ fn test_dict_compressed() -> VortexResult<()> { } #[rstest] -#[case::outlined_utf8(4096, 28, false)] -#[case::nullable_utf8(4096, 28, true)] -#[case::outlined_utf8_8192(8192, 28, false)] -#[case::long_utf8(4096, 256, false)] +#[case::outlined_utf8(4096, 28, false, Distribution::Uniform)] +#[case::nullable_utf8(4096, 28, true, Distribution::Uniform)] +#[case::outlined_utf8_8192(8192, 28, false, Distribution::Uniform)] +#[case::clustered_utf8(4096, 28, false, Distribution::Clustered)] fn test_dict_compressed_with_more_values_than_sample( #[case] distinct_count: usize, #[case] value_length: usize, #[case] nullable: bool, + #[case] distribution: Distribution, ) -> VortexResult<()> { let distinct_values = (0..distinct_count) .map(|value| { @@ -73,8 +80,11 @@ fn test_dict_compressed_with_more_values_than_sample( .collect::>(); let values = (0..65_536) .map(|index| { - (!nullable || index % 10 != 0) - .then_some(distinct_values[index % distinct_values.len()].as_bytes()) + let value_index = match distribution { + Distribution::Uniform => index % distinct_values.len(), + Distribution::Clustered => index * distinct_values.len() / 65_536, + }; + (!nullable || index % 10 != 0).then_some(distinct_values[value_index].as_bytes()) }) .collect::>(); let nullability = if nullable { diff --git a/vortex-compressor/src/builtins/dict/string_candidate.rs b/vortex-compressor/src/builtins/dict/string_candidate.rs index e9e45368e7e..2232daae067 100644 --- a/vortex-compressor/src/builtins/dict/string_candidate.rs +++ b/vortex-compressor/src/builtins/dict/string_candidate.rs @@ -7,7 +7,8 @@ //! This module builds one bounded dictionary candidate from the complete input. //! //! A distributed full-value probe protects high-cardinality inputs from that full pass. -//! A rejected probe uses the existing sample estimator, so the probe changes work rather than selection policy. +//! A high-cardinality probe rejects Dictionary. An inconclusive probe uses the existing sample +//! estimator. //! //! If the trial dictionary wins, the compression phase reuses its completed output. @@ -57,13 +58,16 @@ const PROBE_RANGE_SIZE: u32 = 64; /// Distributes this many probe ranges across the input. const PROBE_RANGE_COUNT: u32 = 128; -/// Sets the numerator for the 80% work-admission threshold. -const PROBE_DISTINCT_NUMERATOR: usize = 4; +/// Sets the numerator for the 75% work-admission threshold. +const PROBE_DISTINCT_NUMERATOR: usize = 3; -/// Sets the denominator for the 80% work-admission threshold. -const PROBE_DISTINCT_DENOMINATOR: usize = 5; +/// Sets the denominator for the 75% work-admission threshold. +const PROBE_DISTINCT_DENOMINATOR: usize = 4; /// Stores a completed dictionary result for reuse by the compression phase. +/// +/// The containing [`ArrayAndStats`] exists for one selection and compression call. The cached +/// result therefore cannot cross compressor configurations or compression contexts. #[derive(Debug)] pub(super) struct CachedStringDictionary(ArrayRef); @@ -79,16 +83,20 @@ impl CachedStringDictionary { enum ProbeResult { /// Permits bounded trial dictionary construction. Candidate, - /// Uses ordinary sample estimation. - Sample, + /// Rejects Dictionary for a high-cardinality input. + Skip, + /// Uses ordinary sample estimation because the probe did not finish. + Inconclusive, } /// Reports the result of bounded trial dictionary construction. enum CandidateResult { /// Contains a complete trial dictionary. Complete(DictArray), - /// Reports that a resource bound stopped trial construction. - Sample, + /// Reports that a preflight bound prevented trial construction. + NotAttempted, + /// Reports that trial construction exhausted a bound after work started. + Exhausted, } /// Returns a deferred estimate that can reuse a completed dictionary result. @@ -106,19 +114,23 @@ fn estimate_string_dictionary( compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - if probe(data)? == ProbeResult::Sample { - return sample_verdict(compressor, data, best_so_far, compress_ctx, exec_ctx); + match probe(data)? { + ProbeResult::Candidate => {} + ProbeResult::Skip => return Ok(EstimateVerdict::Skip), + ProbeResult::Inconclusive => { + return sample_verdict(compressor, data, best_so_far, compress_ctx, exec_ctx); + } } let dictionary = match build_candidate(data, exec_ctx)? { CandidateResult::Complete(dictionary) => dictionary, - CandidateResult::Sample => { + CandidateResult::NotAttempted => { return sample_verdict(compressor, data, best_so_far, compress_ctx, exec_ctx); } + CandidateResult::Exhausted => return Ok(EstimateVerdict::Skip), }; let compressed = compress_dictionary(compressor, &dictionary, compress_ctx, exec_ctx)?; let score = EstimateScore::from_sample_sizes(data.array().nbytes(), compressed.nbytes()); - if !score_is_best(score, best_so_far) { return Ok(EstimateVerdict::Skip); } @@ -170,10 +182,10 @@ fn probe(data: &ArrayAndStats) -> VortexResult { let validity_bits = match &validity { vortex_array::validity::Validity::NonNullable | vortex_array::validity::Validity::AllValid => None, - vortex_array::validity::Validity::AllInvalid => return Ok(ProbeResult::Sample), + vortex_array::validity::Validity::AllInvalid => return Ok(ProbeResult::Inconclusive), vortex_array::validity::Validity::Array(validity) => { let Some(validity) = validity.as_opt::() else { - return Ok(ProbeResult::Sample); + return Ok(ProbeResult::Inconclusive); }; Some(validity.to_bit_buffer()) } @@ -185,7 +197,8 @@ fn probe(data: &ArrayAndStats) -> VortexResult { .sum::(); let mut distinct_values = HashSet::with_capacity(sample_rows); let mut probed_bytes = 0usize; - let mut valid_rows = 0usize; + let mut probed_rows = 0usize; + let mut saw_null = false; let maximum_range_length = sample_ranges .iter() @@ -199,10 +212,12 @@ fn probe(data: &ArrayAndStats) -> VortexResult { if index >= end { continue; } + probed_rows += 1; if validity_bits .as_ref() .is_some_and(|validity| !validity.value(index)) { + saw_null = true; continue; } let value = view_bytes(&array, &views[index]); @@ -213,18 +228,22 @@ fn probe(data: &ArrayAndStats) -> VortexResult { break 'probe; } probed_bytes = next_probed_bytes; - valid_rows += 1; distinct_values.insert(value); } } - if valid_rows != 0 - && distinct_values.len() * PROBE_DISTINCT_DENOMINATOR - <= valid_rows * PROBE_DISTINCT_NUMERATOR - { + if probed_rows == 0 { + return Ok(ProbeResult::Inconclusive); + } + if probed_rows != sample_rows { + return Ok(ProbeResult::Inconclusive); + } + + let distinct_values = distinct_values.len() + usize::from(saw_null); + if distinct_values * PROBE_DISTINCT_DENOMINATOR <= probed_rows * PROBE_DISTINCT_NUMERATOR { Ok(ProbeResult::Candidate) } else { - Ok(ProbeResult::Sample) + Ok(ProbeResult::Skip) } } @@ -234,10 +253,10 @@ fn build_candidate( exec_ctx: &mut ExecutionCtx, ) -> VortexResult { let Some(code_bytes) = candidate_code_bytes(data.array_len(), MAX_DICTIONARY_ENTRIES) else { - return Ok(CandidateResult::Sample); + return Ok(CandidateResult::NotAttempted); }; if code_bytes > MAX_DICTIONARY_CODE_BYTES { - return Ok(CandidateResult::Sample); + return Ok(CandidateResult::NotAttempted); } let constraints = DictConstraints { @@ -251,7 +270,7 @@ fn build_candidate( exec_ctx, )?; if dictionary.len() != data.array_len() { - return Ok(CandidateResult::Sample); + return Ok(CandidateResult::Exhausted); } Ok(CandidateResult::Complete(dictionary)) @@ -284,7 +303,6 @@ mod tests { use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::arrays::Dict; use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -323,10 +341,27 @@ mod tests { ArrayAndStats::new(array, Default::default()) } + /// Resolves the Dictionary estimate without another candidate scheme. + fn estimate_dictionary( + compressor: &CascadingCompressor, + data: &ArrayAndStats, + exec_ctx: &mut vortex_array::ExecutionCtx, + ) -> VortexResult { + let callback = match string_dictionary_estimate() { + CompressionEstimate::Deferred(DeferredEstimate::Callback(callback)) => callback, + estimate => { + return Err(vortex_error::vortex_err!( + "dictionary estimation returned {estimate:?}" + )); + } + }; + callback(compressor, data, None, CompressorContext::new(), exec_ctx) + } + #[test] fn probe_distinct_ratio_boundary() -> VortexResult<()> { for (distinct_values, expected) in - [(6553, ProbeResult::Candidate), (6554, ProbeResult::Sample)] + [(6144, ProbeResult::Candidate), (6145, ProbeResult::Skip)] { let values = (0..8192) .map(|index| { @@ -344,7 +379,7 @@ mod tests { } #[test] - fn probe_distributes_rows_before_byte_limit() -> VortexResult<()> { + fn probe_reports_inconclusive_at_byte_limit() -> VortexResult<()> { let mut values = (0..65_536) .map(|index| Some(format!("{index:0256x}"))) .collect::>(); @@ -355,12 +390,40 @@ mod tests { values[start..end].fill(Some("x".repeat(256))); } let data = string_data(&values); - assert_eq!(probe(&data)?, ProbeResult::Sample); + assert_eq!(probe(&data)?, ProbeResult::Inconclusive); + Ok(()) + } + + #[test] + fn probe_counts_null_as_one_repeated_value() -> VortexResult<()> { + let values = (0..8192) + .map(|index| (index % 10 == 0).then(|| format!("unique-non-null-value-{index:08x}"))) + .collect::>(); + let data = string_data(&values); + + assert_eq!(probe(&data)?, ProbeResult::Candidate); Ok(()) } #[test] - fn candidate_defers_when_dictionary_storage_exceeds_limit() -> VortexResult<()> { + fn high_cardinality_probe_is_terminal() -> VortexResult<()> { + let values = (0..8192) + .map(|index| Some(format!("unique-value-{index:08x}"))) + .collect::>(); + let data = string_data(&values); + let compressor = CascadingCompressor::new(vec![&StringDictScheme]); + let mut exec_ctx = vortex_array::array_session().create_execution_ctx(); + + assert_eq!(probe(&data)?, ProbeResult::Skip); + assert!(matches!( + estimate_dictionary(&compressor, &data, &mut exec_ctx)?, + EstimateVerdict::Skip + )); + Ok(()) + } + + #[test] + fn candidate_stops_when_dictionary_storage_exceeds_limit() -> VortexResult<()> { let suffix = "x".repeat(70_000); let distinct_values = (0..64) .map(|value| format!("{value:08x}-{suffix}")) @@ -373,7 +436,7 @@ mod tests { assert!(matches!( build_candidate(&data, &mut exec_ctx)?, - CandidateResult::Sample + CandidateResult::Exhausted )); Ok(()) } @@ -396,7 +459,7 @@ mod tests { let data = string_data(&values_past_limit); assert!(matches!( build_candidate(&data, &mut exec_ctx)?, - CandidateResult::Sample + CandidateResult::Exhausted )); Ok(()) } @@ -409,21 +472,7 @@ mod tests { let data = string_data(&values); let compressor = CascadingCompressor::new(vec![&StringDictScheme]); let mut exec_ctx = vortex_array::array_session().create_execution_ctx(); - let callback = match string_dictionary_estimate() { - CompressionEstimate::Deferred(DeferredEstimate::Callback(callback)) => callback, - estimate => { - return Err(vortex_error::vortex_err!( - "dictionary estimation returned {estimate:?}" - )); - } - }; - let verdict = callback( - &compressor, - &data, - None, - CompressorContext::new(), - &mut exec_ctx, - )?; + let verdict = estimate_dictionary(&compressor, &data, &mut exec_ctx)?; assert!(matches!(verdict, EstimateVerdict::Ratio(_))); let cached = data .get::() @@ -438,48 +487,4 @@ mod tests { assert!(ArrayRef::ptr_eq(cached.array(), &compressed)); Ok(()) } - - #[test] - fn sample_fallback_can_select_dictionary() -> VortexResult<()> { - let mut values = (0..65_536) - .map(|index| Some(format!("repeated-value-{:08x}", index % 2))) - .collect::>(); - for (start, end) in sample_slices(values.len(), PROBE_RANGE_SIZE, PROBE_RANGE_COUNT) { - for index in start..end { - values[index] = Some(format!("probe-only-value-{index:08x}")); - } - } - let data = string_data(&values); - let compressor = CascadingCompressor::new(vec![&StringDictScheme]); - let mut exec_ctx = vortex_array::array_session().create_execution_ctx(); - - assert_eq!(probe(&data)?, ProbeResult::Sample); - let callback = match string_dictionary_estimate() { - CompressionEstimate::Deferred(DeferredEstimate::Callback(callback)) => callback, - estimate => { - return Err(vortex_error::vortex_err!( - "dictionary estimation returned {estimate:?}" - )); - } - }; - assert!(matches!( - callback( - &compressor, - &data, - None, - CompressorContext::new(), - &mut exec_ctx, - )?, - EstimateVerdict::Ratio(_) - )); - let compressed = StringDictScheme.compress( - &compressor, - &data, - CompressorContext::new(), - &mut exec_ctx, - )?; - - assert!(compressed.is::()); - Ok(()) - } } diff --git a/vortex-compressor/src/compressor/sample.rs b/vortex-compressor/src/compressor/sample.rs index f3700da2aa6..b01dc2cf0f7 100644 --- a/vortex-compressor/src/compressor/sample.rs +++ b/vortex-compressor/src/compressor/sample.rs @@ -3,8 +3,6 @@ //! Sampling utilities for compression ratio estimation. -use std::mem::size_of; - use rand::RngExt; use rand::SeedableRng; use rand::prelude::StdRng; @@ -13,12 +11,8 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ChunkedArray; -use vortex_array::arrays::VarBinView; -use vortex_array::arrays::varbinview::BinaryView; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_mask::Mask; -use vortex_utils::aliases::hash_map::HashMap; use crate::CascadingCompressor; use crate::scheme::CompressorContext; @@ -93,6 +87,21 @@ pub(crate) fn sample_count_approx_one_percent(len: usize) -> u32 { ) } +/// Materializes a compact canonical sample. +fn materialize_sample( + array: &ArrayRef, + sample_count: u32, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + let canonical: Canonical = sample(array, SAMPLE_SIZE, sample_count).execute(exec_ctx)?; + match canonical { + Canonical::VarBinView(array) => { + Ok(array.compact_with_threshold(1.0, exec_ctx)?.into_array()) + } + canonical => Ok(canonical.into_array()), + } +} + /// Divides an array into `sample_count` equal partitions and picks one random contiguous /// slice of `sample_size` elements from each partition. /// @@ -175,9 +184,7 @@ pub(crate) fn estimate_compression_ratio_with_sampling( array.clone() } else { let sample_count = sample_count_approx_one_percent(array.len()); - // `ArrayAndStats` expects a canonical array (so that it can easily compute lazy stats). - let canonical: Canonical = sample(array, SAMPLE_SIZE, sample_count).execute(exec_ctx)?; - canonical.into_array() + materialize_sample(array, sample_count, exec_ctx)? }; let sample_data = ArrayAndStats::new(sample_array, scheme.stats_options()); @@ -193,7 +200,7 @@ pub(crate) fn estimate_compression_ratio_with_sampling( }; let after = compressed.nbytes(); - let before = canonical_visible_nbytes(sample_data.array(), exec_ctx)?; + let before = sample_data.array().nbytes(); let score = EstimateScore::from_sample_sizes(before, after); @@ -204,101 +211,16 @@ pub(crate) fn estimate_compression_ratio_with_sampling( Ok(score) } -/// Returns the physical canonical size without retained, unreferenced string payload bytes. -pub(crate) fn canonical_visible_nbytes( - array: &ArrayRef, - exec_ctx: &mut ExecutionCtx, -) -> VortexResult { - let Some(array) = array.as_opt::() else { - return Ok(array.nbytes()); - }; - let validity = array.validity()?.execute_mask(array.len(), exec_ctx)?; - let mut referenced_ranges = HashMap::<(usize, usize), Vec<(usize, usize)>>::new(); - let mut record_view = |view: &BinaryView| { - if view.is_inlined() { - return; - } - let reference = view.as_view(); - let buffer = array.buffer(reference.buffer_index as usize); - referenced_ranges - .entry((buffer.as_ptr().addr(), buffer.len())) - .or_default() - .push(( - reference.offset as usize, - reference.offset as usize + reference.size as usize, - )); - }; - match &validity { - Mask::AllTrue(_) => array.views().iter().for_each(&mut record_view), - Mask::AllFalse(_) => {} - Mask::Values(values) => array - .views() - .iter() - .zip(values.bit_buffer().iter()) - .filter(|(_, is_valid)| *is_valid) - .for_each(|(view, _)| record_view(view)), - } - let outlined_bytes = referenced_ranges.into_values().try_fold( - 0u64, - |total, mut ranges| -> VortexResult { - ranges.sort_unstable(); - let mut range_total = 0u64; - let mut current = None::<(usize, usize)>; - for (start, end) in ranges { - match current { - Some((current_start, current_end)) if start <= current_end => { - current = Some((current_start, current_end.max(end))); - } - Some((current_start, current_end)) => { - range_total = range_total - .checked_add(u64::try_from(current_end - current_start)?) - .ok_or_else(|| { - vortex_error::vortex_err!("sample byte size overflowed u64") - })?; - current = Some((start, end)); - } - None => current = Some((start, end)), - } - } - if let Some((start, end)) = current { - range_total = range_total - .checked_add(u64::try_from(end - start)?) - .ok_or_else(|| vortex_error::vortex_err!("sample byte size overflowed u64"))?; - } - total - .checked_add(range_total) - .ok_or_else(|| vortex_error::vortex_err!("sample byte size overflowed u64")) - }, - )?; - let views_bytes = u64::try_from(array.len())? - .checked_mul(u64::try_from(size_of::())?) - .ok_or_else(|| vortex_error::vortex_err!("sample byte size overflowed u64"))?; - let validity_bytes = match validity { - Mask::Values(values) => u64::try_from(values.len().div_ceil(u8::BITS as usize))?, - Mask::AllTrue(_) | Mask::AllFalse(_) => 0, - }; - - views_bytes - .checked_add(outlined_bytes) - .and_then(|bytes| bytes.checked_add(validity_bytes)) - .ok_or_else(|| vortex_error::vortex_err!("sample byte size overflowed u64")) -} - #[cfg(test)] mod tests { - use std::sync::Arc; - use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; - use vortex_array::dtype::DType; - use vortex_array::dtype::Nullability; use vortex_array::validity::Validity; use vortex_buffer::Buffer; - use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; use super::*; @@ -322,53 +244,21 @@ mod tests { } #[test] - fn visible_size_ignores_unreferenced_string_payload() -> VortexResult<()> { - let values = (0..128) + fn materialized_string_sample_drops_unreferenced_payload() -> VortexResult<()> { + let values = (0..4096) .map(|index| format!("outlined-string-value-{index:08x}")) .collect::>(); let source = VarBinViewArray::from_iter_str(&values).into_array(); - let slice = source.slice(17..19)?; - let expected = u64::try_from(2 * size_of::())? - + u64::try_from(values[17].len() + values[18].len())?; - let mut exec_ctx = array_session().create_execution_ctx(); - - assert!(slice.nbytes() > expected); - assert_eq!(canonical_visible_nbytes(&slice, &mut exec_ctx)?, expected); - Ok(()) - } - - #[test] - fn visible_size_counts_valid_outlined_values_and_validity() -> VortexResult<()> { - let outlined = "an-outlined-string-value"; - let array = VarBinViewArray::from_iter( - [Some("inline"), Some(outlined), None], - DType::Utf8(Nullability::Nullable), - ) - .into_array(); - let expected = u64::try_from(3 * size_of::() + outlined.len() + 1)?; let mut exec_ctx = array_session().create_execution_ctx(); - - assert_eq!(canonical_visible_nbytes(&array, &mut exec_ctx)?, expected); - Ok(()) - } - - #[test] - fn visible_size_counts_shared_payload_once() -> VortexResult<()> { - let value = b"one shared outlined value"; - let view = BinaryView::make_view(value, 0, 0); - let array = VarBinViewArray::try_new( - Buffer::copy_from([view, view]), - Arc::from([ByteBuffer::copy_from(value)]), - DType::Utf8(Nullability::NonNullable), - Validity::NonNullable, - &mut array_session().create_execution_ctx(), - )? - .into_array(); - let expected = u64::try_from(2 * size_of::() + value.len())?; - - assert_eq!( - canonical_visible_nbytes(&array, &mut array_session().create_execution_ctx())?, - expected + let uncompacted: Canonical = + sample(&source, SAMPLE_SIZE, SAMPLE_COUNT).execute(&mut exec_ctx)?; + let compacted = materialize_sample(&source, SAMPLE_COUNT, &mut exec_ctx)?; + + assert!(compacted.nbytes() < uncompacted.into_array().nbytes()); + assert_arrays_eq!( + compacted, + sample(&source, SAMPLE_SIZE, SAMPLE_COUNT), + &mut exec_ctx ); Ok(()) }