diff --git a/Cargo.lock b/Cargo.lock index 21e54dc1a3a..58ef8d29522 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9697,6 +9697,7 @@ dependencies = [ "vortex-array", "vortex-buffer", "vortex-error", + "vortex-fsst", "vortex-session", ] @@ -10035,6 +10036,7 @@ version = "0.1.0" dependencies = [ "codspeed-divan-compat", "fsst-rs", + "num-traits", "prost 0.14.4", "rand 0.10.1", "rstest", diff --git a/encodings/fsst/Cargo.toml b/encodings/fsst/Cargo.toml index 0d722a131d3..f749f121efa 100644 --- a/encodings/fsst/Cargo.toml +++ b/encodings/fsst/Cargo.toml @@ -18,6 +18,7 @@ workspace = true [dependencies] fsst-rs = { workspace = true } +num-traits = { workspace = true } prost = { workspace = true } rand = { workspace = true, optional = true } vortex-array = { workspace = true } diff --git a/encodings/fsst/benches/fsst_compress.rs b/encodings/fsst/benches/fsst_compress.rs index 6aad25de874..5e03d1d003d 100644 --- a/encodings/fsst/benches/fsst_compress.rs +++ b/encodings/fsst/benches/fsst_compress.rs @@ -9,6 +9,7 @@ use divan::Bencher; use rand::RngExt; use rand::SeedableRng; use rand::rngs::StdRng; +use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::RecursiveCanonical; @@ -56,60 +57,50 @@ const BENCH_ARGS: &[(usize, usize, u8)] = &[ #[divan::bench(args = BENCH_ARGS)] fn compress_fsst(bencher: Bencher, (string_count, avg_len, unique_chars): (usize, usize, u8)) { let array = generate_test_data(string_count, avg_len, unique_chars); - let compressor = fsst_train_compressor(&array); + let compressor = fsst_train_compressor(&array, &mut SESSION.create_execution_ctx()).unwrap(); bencher .with_inputs(|| (&array, &compressor, SESSION.create_execution_ctx())) - .bench_refs(|(array, compressor, ctx)| { - fsst_compress(*array, array.len(), array.dtype(), compressor, ctx) - }) + .bench_refs(|(array, compressor, ctx)| fsst_compress(array, compressor, ctx).unwrap()) } #[divan::bench(args = BENCH_ARGS)] fn decompress_fsst(bencher: Bencher, (string_count, avg_len, unique_chars): (usize, usize, u8)) { let array = generate_test_data(string_count, avg_len, unique_chars); - let compressor = fsst_train_compressor(&array); - let len = array.len(); - let dtype = array.dtype().clone(); - let encoded = fsst_compress( - array, - len, - &dtype, - &compressor, - &mut SESSION.create_execution_ctx(), - ); + let mut ctx = SESSION.create_execution_ctx(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + let encoded = fsst_compress(&array, &compressor, &mut ctx) + .unwrap() + .into_array(); bencher .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) - .bench_refs(|(encoded, ctx)| (**encoded).clone().into_array().execute::(ctx)) + .bench_refs(|(encoded, ctx)| (*encoded).clone().execute::(ctx)) } #[divan::bench(args = BENCH_ARGS)] fn train_compressor(bencher: Bencher, (string_count, avg_len, unique_chars): (usize, usize, u8)) { let array = generate_test_data(string_count, avg_len, unique_chars); bencher - .with_inputs(|| &array) - .bench_refs(|array| fsst_train_compressor(array)) + .with_inputs(|| (&array, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| fsst_train_compressor(array, ctx).unwrap()) } #[divan::bench(args = BENCH_ARGS)] fn pushdown_compare(bencher: Bencher, (string_count, avg_len, unique_chars): (usize, usize, u8)) { let array = generate_test_data(string_count, avg_len, unique_chars); - let compressor = fsst_train_compressor(&array); - let fsst_array = fsst_compress( - &array, - array.len(), - array.dtype(), - &compressor, - &mut SESSION.create_execution_ctx(), - ); - let constant = ConstantArray::new(Scalar::from(&b"const"[..]), array.len()); + let len = array.len(); + let mut ctx = SESSION.create_execution_ctx(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + let fsst_array = fsst_compress(&array, &compressor, &mut ctx) + .unwrap() + .into_array(); + let constant = ConstantArray::new(Scalar::from(&b"const"[..]), len); bencher .with_inputs(|| (&fsst_array, &constant, SESSION.create_execution_ctx())) .bench_refs(|(fsst_array, constant, ctx)| { fsst_array .clone() - .into_array() .binary(constant.clone().into_array(), Operator::Eq) .unwrap() .execute::(ctx) @@ -123,22 +114,19 @@ fn canonicalize_compare( (string_count, avg_len, unique_chars): (usize, usize, u8), ) { let array = generate_test_data(string_count, avg_len, unique_chars); - let compressor = fsst_train_compressor(&array); - let fsst_array = fsst_compress( - &array, - array.len(), - array.dtype(), - &compressor, - &mut SESSION.create_execution_ctx(), - ); - let constant = ConstantArray::new(Scalar::from(&b"const"[..]), array.len()); + let len = array.len(); + let mut ctx = SESSION.create_execution_ctx(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + let fsst_array = fsst_compress(&array, &compressor, &mut ctx) + .unwrap() + .into_array(); + let constant = ConstantArray::new(Scalar::from(&b"const"[..]), len); bencher .with_inputs(|| (&fsst_array, &constant, SESSION.create_execution_ctx())) .bench_refs(|(fsst_array, constant, ctx)| { - (*fsst_array) + fsst_array .clone() - .into_array() .execute::(ctx) .unwrap() .into_array() @@ -190,11 +178,11 @@ fn chunked_into_canonical( bencher .with_inputs(|| (&array, SESSION.create_execution_ctx())) - .bench_refs(|(array, ctx)| (**array).clone().into_array().execute::(ctx)); + .bench_refs(|(array, ctx)| array.clone().execute::(ctx)); } /// Helper function to generate random string data. -fn generate_test_data(string_count: usize, avg_len: usize, unique_chars: u8) -> VarBinArray { +fn generate_test_data(string_count: usize, avg_len: usize, unique_chars: u8) -> ArrayRef { let mut rng = StdRng::seed_from_u64(0); let mut strings = Vec::with_capacity(string_count); @@ -216,6 +204,7 @@ fn generate_test_data(string_count: usize, avg_len: usize, unique_chars: u8) -> .map(|opt_s| opt_s.map(Vec::into_boxed_slice)), DType::Binary(Nullability::NonNullable), ) + .into_array() } fn generate_chunked_test_data( @@ -223,15 +212,16 @@ fn generate_chunked_test_data( string_count: usize, avg_len: usize, unique_chars: u8, -) -> ChunkedArray { +) -> ArrayRef { let mut ctx = SESSION.create_execution_ctx(); (0..chunk_size) .map(|_| { let array = generate_test_data(string_count, avg_len, unique_chars); - let compressor = fsst_train_compressor(&array); - let len = array.len(); - let dtype = array.dtype().clone(); - fsst_compress(array, len, &dtype, &compressor, &mut ctx).into_array() + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + fsst_compress(&array, &compressor, &mut ctx) + .unwrap() + .into_array() }) .collect::() + .into_array() } diff --git a/encodings/fsst/benches/fsst_url_compare.rs b/encodings/fsst/benches/fsst_url_compare.rs index 2a18236e6be..a62dea8aa9e 100644 --- a/encodings/fsst/benches/fsst_url_compare.rs +++ b/encodings/fsst/benches/fsst_url_compare.rs @@ -12,6 +12,7 @@ use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::VarBinArray; +use vortex_array::arrays::VarBinViewArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::expr::like; use vortex_array::expr::lit; @@ -40,6 +41,16 @@ const NUM_URLS: usize = NUM_STRINGS; static URL_DATA: LazyLock = LazyLock::new(generate_url_data); +/// `URL_DATA` as a `VarBinViewArray`, used by the view-input bench variants. +static URL_VIEW_DATA: LazyLock = LazyLock::new(|| { + let mut ctx = SESSION.create_execution_ctx(); + URL_DATA + .clone() + .into_array() + .execute::(&mut ctx) + .unwrap() +}); + // --------------------------------------------------------------------------- // Eq compare benchmarks (FSST pushdown vs canonicalize) // --------------------------------------------------------------------------- @@ -58,14 +69,12 @@ fn pick_url_with_domain(data: &VarBinArray, domain: &str) -> String { #[divan::bench] fn eq_pushdown_high_match(bencher: Bencher) { let data = &*URL_DATA; - let compressor = fsst_train_compressor(data); - let fsst_array = fsst_compress( - data, - data.len(), - data.dtype(), - &compressor, - &mut SESSION.create_execution_ctx(), - ); + let mut ctx = SESSION.create_execution_ctx(); + let array = data.clone().into_array(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + let fsst_array = fsst_compress(&array, &compressor, &mut ctx) + .unwrap() + .into_array(); let match_url = pick_url_with_domain(data, HIGH_MATCH_DOMAIN); let constant = ConstantArray::new(Scalar::from(match_url.as_str()), NUM_URLS); @@ -74,7 +83,6 @@ fn eq_pushdown_high_match(bencher: Bencher) { .bench_refs(|(fsst_array, constant, ctx)| { fsst_array .clone() - .into_array() .binary(constant.clone().into_array(), Operator::Eq) .unwrap() .execute::(ctx) @@ -85,17 +93,61 @@ fn eq_pushdown_high_match(bencher: Bencher) { #[divan::bench] fn eq_pushdown_low_match(bencher: Bencher) { let data = &*URL_DATA; - let compressor = fsst_train_compressor(data); - let fsst_array = fsst_compress( - data, - data.len(), - data.dtype(), - &compressor, - &mut SESSION.create_execution_ctx(), - ); + let mut ctx = SESSION.create_execution_ctx(); + let array = data.clone().into_array(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + let fsst_array = fsst_compress(&array, &compressor, &mut ctx) + .unwrap() + .into_array(); let match_url = pick_url_with_domain(data, LOW_MATCH_DOMAIN); let constant = ConstantArray::new(Scalar::from(match_url.as_str()), NUM_URLS); + bencher + .with_inputs(|| (&fsst_array, &constant, SESSION.create_execution_ctx())) + .bench_refs(|(fsst_array, constant, ctx)| { + fsst_array + .clone() + .binary(constant.clone().into_array(), Operator::Eq) + .unwrap() + .execute::(ctx) + .unwrap() + }); +} + +#[divan::bench] +fn eq_pushdown_high_match_view(bencher: Bencher) { + let data = &*URL_VIEW_DATA; + let mut ctx = SESSION.create_execution_ctx(); + let array = data.clone().into_array(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + let fsst_array = fsst_compress(&array, &compressor, &mut ctx) + .unwrap() + .into_array(); + let match_url = pick_url_with_domain(&URL_DATA, HIGH_MATCH_DOMAIN); + let constant = ConstantArray::new(Scalar::from(match_url.as_str()), NUM_URLS); + + bencher + .with_inputs(|| (&fsst_array, &constant, SESSION.create_execution_ctx())) + .bench_refs(|(fsst_array, constant, ctx)| { + fsst_array + .clone() + .binary(constant.clone().into_array(), Operator::Eq) + .unwrap() + .execute::(ctx) + .unwrap() + }); +} + +#[divan::bench] +fn eq_pushdown_low_match_view(bencher: Bencher) { + let data = &*URL_VIEW_DATA; + let mut ctx = SESSION.create_execution_ctx(); + let array = data.clone().into_array(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + let fsst_array = fsst_compress(&array, &compressor, &mut ctx).unwrap(); + let match_url = pick_url_with_domain(&URL_DATA, LOW_MATCH_DOMAIN); + let constant = ConstantArray::new(Scalar::from(match_url.as_str()), NUM_URLS); + bencher .with_inputs(|| (&fsst_array, &constant, SESSION.create_execution_ctx())) .bench_refs(|(fsst_array, constant, ctx)| { @@ -112,23 +164,20 @@ fn eq_pushdown_low_match(bencher: Bencher) { #[divan::bench] fn eq_canonicalize_high_match(bencher: Bencher) { let data = &*URL_DATA; - let compressor = fsst_train_compressor(data); - let fsst_array = fsst_compress( - data, - data.len(), - data.dtype(), - &compressor, - &mut SESSION.create_execution_ctx(), - ); + let mut ctx = SESSION.create_execution_ctx(); + let array = data.clone().into_array(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + let fsst_array = fsst_compress(&array, &compressor, &mut ctx) + .unwrap() + .into_array(); let match_url = pick_url_with_domain(data, HIGH_MATCH_DOMAIN); let constant = ConstantArray::new(Scalar::from(match_url.as_str()), NUM_URLS); bencher .with_inputs(|| (&fsst_array, &constant, SESSION.create_execution_ctx())) .bench_refs(|(fsst_array, constant, ctx)| { - (*fsst_array) + fsst_array .clone() - .into_array() .execute::(ctx) .unwrap() .into_array() @@ -142,23 +191,20 @@ fn eq_canonicalize_high_match(bencher: Bencher) { #[divan::bench] fn eq_canonicalize_low_match(bencher: Bencher) { let data = &*URL_DATA; - let compressor = fsst_train_compressor(data); - let fsst_array = fsst_compress( - data, - data.len(), - data.dtype(), - &compressor, - &mut SESSION.create_execution_ctx(), - ); + let mut ctx = SESSION.create_execution_ctx(); + let array = data.clone().into_array(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + let fsst_array = fsst_compress(&array, &compressor, &mut ctx) + .unwrap() + .into_array(); let match_url = pick_url_with_domain(data, LOW_MATCH_DOMAIN); let constant = ConstantArray::new(Scalar::from(match_url.as_str()), NUM_URLS); bencher .with_inputs(|| (&fsst_array, &constant, SESSION.create_execution_ctx())) .bench_refs(|(fsst_array, constant, ctx)| { - (*fsst_array) + fsst_array .clone() - .into_array() .execute::(ctx) .unwrap() .into_array() @@ -176,14 +222,12 @@ fn eq_canonicalize_low_match(bencher: Bencher) { #[divan::bench] fn like_substr_high_match(bencher: Bencher) { let data = &*URL_DATA; - let compressor = fsst_train_compressor(data); - let fsst_array = fsst_compress( - data, - data.len(), - data.dtype(), - &compressor, - &mut SESSION.create_execution_ctx(), - ); + let mut ctx = SESSION.create_execution_ctx(); + let array = data.clone().into_array(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + let fsst_array = fsst_compress(&array, &compressor, &mut ctx) + .unwrap() + .into_array(); let pattern = format!("%{HIGH_MATCH_DOMAIN}%"); let expr = like(root(), lit(pattern.as_str())); @@ -192,7 +236,6 @@ fn like_substr_high_match(bencher: Bencher) { .bench_refs(|(fsst_array, ctx)| { fsst_array .clone() - .into_array() .apply(&expr) .unwrap() .execute::(ctx) @@ -203,14 +246,12 @@ fn like_substr_high_match(bencher: Bencher) { #[divan::bench] fn like_substr_low_match(bencher: Bencher) { let data = &*URL_DATA; - let compressor = fsst_train_compressor(data); - let fsst_array = fsst_compress( - data, - data.len(), - data.dtype(), - &compressor, - &mut SESSION.create_execution_ctx(), - ); + let mut ctx = SESSION.create_execution_ctx(); + let array = data.clone().into_array(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + let fsst_array = fsst_compress(&array, &compressor, &mut ctx) + .unwrap() + .into_array(); let pattern = format!("%{LOW_MATCH_DOMAIN}%"); let expr = like(root(), lit(pattern.as_str())); @@ -219,7 +260,6 @@ fn like_substr_low_match(bencher: Bencher) { .bench_refs(|(fsst_array, ctx)| { fsst_array .clone() - .into_array() .apply(&expr) .unwrap() .execute::(ctx) diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index 9a737fb3057..e3b5d74c9d7 100644 --- a/encodings/fsst/src/array.rs +++ b/encodings/fsst/src/array.rs @@ -856,7 +856,7 @@ mod test { use crate::FSST; use crate::array::FSSTArrayExt; use crate::array::FSSTMetadata; - use crate::fsst_compress_iter; + use crate::fsst_compress; #[test] fn slice_reuses_initialized_compressor() -> VortexResult<()> { @@ -868,18 +868,8 @@ mod test { let compressor = Compressor::rebuild_from(symbols.as_slice(), symbol_lengths.as_slice()); let mut ctx = LEGACY_SESSION.create_execution_ctx(); - let fsst_array = fsst_compress_iter( - [ - Some(b"abcabcab".as_ref()), - Some(b"defghijk".as_ref()), - Some(b"abcxyz".as_ref()), - ] - .into_iter(), - 3, - DType::Utf8(Nullability::NonNullable), - &compressor, - &mut ctx, - ); + let strings = VarBinViewArray::from_iter_str(["abcabcab", "defghijk", "abcxyz"]); + let fsst_array = fsst_compress(&strings.into_array(), &compressor, &mut ctx)?; let compressor_ptr = fsst_array.compressor() as *const Compressor; let sliced = fsst_array @@ -922,13 +912,8 @@ mod test { let compressor = Compressor::rebuild_from(symbols.as_slice(), symbol_lengths.as_slice()); let mut ctx = LEGACY_SESSION.create_execution_ctx(); - let fsst_array = fsst_compress_iter( - [Some(b"abcabcab".as_ref()), Some(b"defghijk".as_ref())].into_iter(), - 2, - DType::Utf8(Nullability::NonNullable), - &compressor, - &mut ctx, - ); + let input = VarBinViewArray::from_iter_str(["abcabcab", "defghijk"]); + let fsst_array = fsst_compress(&input.into_array(), &compressor, &mut ctx).unwrap(); let compressed_codes = fsst_array.codes(); diff --git a/encodings/fsst/src/canonical.rs b/encodings/fsst/src/canonical.rs index 67241ffb026..86caf98825e 100644 --- a/encodings/fsst/src/canonical.rs +++ b/encodings/fsst/src/canonical.rs @@ -151,9 +151,11 @@ mod tests { let (arr_vec, data_vec): (Vec, Vec>>>) = (0..10) .map(|_| { let (array, data) = make_data(); - let compressor = fsst_train_compressor(&array); + let array = array.into_array(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); ( - fsst_compress(&array, array.len(), array.dtype(), &compressor, &mut ctx) + fsst_compress(&array, &compressor, &mut ctx) + .unwrap() .into_array(), data, ) @@ -206,15 +208,14 @@ mod tests { let varbin = VarBinArray::from_iter( [Some(b"long enough too".to_vec().into_boxed_slice())], dtype, - ); + ) + .into_array(); let mut ctx = SESSION.create_execution_ctx(); let fsst_array = fsst_compress( &varbin, - varbin.len(), - varbin.dtype(), - &fsst_train_compressor(&varbin), + &fsst_train_compressor(&varbin, &mut ctx)?, &mut ctx, - ) + )? .into_array(); fsst_array.append_to_builder(&mut builder, &mut ctx)?; diff --git a/encodings/fsst/src/compress.rs b/encodings/fsst/src/compress.rs index 058e0b7e0e1..d0f712ec6ab 100644 --- a/encodings/fsst/src/compress.rs +++ b/encodings/fsst/src/compress.rs @@ -1,160 +1,367 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -// Compress a set of values into an Array. +//! FSST compression entry points. +//! +//! [`fsst_compress`] and [`fsst_train_compressor`] take an [`ArrayRef`] and dispatch +//! on the input encoding ([`VarBinView`] or [`VarBin`]). Callers don't need to know +//! which string encoding they hold. + +use std::sync::Arc; use fsst::Compressor; -use fsst::Symbol; +use num_traits::AsPrimitive; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::accessor::ArrayAccessor; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::VarBin; +use vortex_array::arrays::VarBinView; +use vortex_array::arrays::varbin::VarBinArrayExt; use vortex_array::arrays::varbin::builder::VarBinBuilder; +use vortex_array::arrays::varbinview::BinaryView; +use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; +use vortex_array::dtype::IntegerPType; +use vortex_array::match_each_integer_ptype; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_mask::AllOr; +use vortex_mask::Mask; -/// Compress a string array using FSST. use crate::FSST; use crate::FSSTArray; -pub fn fsst_compress>( - strings: A, - len: usize, - dtype: &DType, + +/// FSST worst case: every input byte expands to an escape + literal (2x) plus a small +/// per-string header. +const FSST_PER_BYTE_OVERHEAD: usize = 2; +const FSST_PER_ROW_OVERHEAD: usize = 7; + +/// Starting capacity for the per-row `compress_into` scratch buffer; grown monotonically. +const DEFAULT_BUFFER_LEN: usize = 1024 * 1024; + +/// Compress a string array using FSST. +/// +/// Accepts any [`VarBinView`] or [`VarBin`]-encoded array; other encodings error. +pub fn fsst_compress( + array: &ArrayRef, compressor: &Compressor, ctx: &mut ExecutionCtx, -) -> FSSTArray { - strings.with_iterator(|iter| fsst_compress_iter(iter, len, dtype.clone(), compressor, ctx)) +) -> VortexResult { + if let Some(view) = array.as_opt::() { + compress_varbinview(view, compressor, ctx) + } else if let Some(varbin) = array.as_opt::() { + compress_varbin_array(varbin, compressor, ctx) + } else { + vortex_bail!( + "fsst_compress requires VarBinView or VarBin encoding, got {}", + array.encoding_id() + ) + } } -/// Train a compressor from an array. +/// Train an FSST [`Compressor`] from a string array's non-null rows. /// -/// # Panics -/// -/// If the provided array is not FSST compressible. -pub fn fsst_train_compressor>(array: &A) -> Compressor { - array.with_iterator(|iter| fsst_train_compressor_iter(iter)) +/// Accepts any [`VarBinView`] or [`VarBin`]-encoded array; other encodings error. +pub fn fsst_train_compressor(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + if let Some(view) = array.as_opt::() { + train_varbinview(view, ctx) + } else if let Some(varbin) = array.as_opt::() { + train_varbin_array(varbin, ctx) + } else { + vortex_bail!( + "fsst_train_compressor requires VarBinView or VarBin encoding, got {}", + array.encoding_id() + ) + } } -/// Train a [compressor][Compressor] from an iterator of bytestrings. -fn fsst_train_compressor_iter<'a, I>(iter: I) -> Compressor -where - I: Iterator>, -{ - let mut lines = Vec::with_capacity(8_192); +fn compress_varbinview( + strings: ArrayView, + compressor: &Compressor, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mask = strings.validity()?.execute_mask(strings.len(), ctx)?; + let views = strings.views(); + let non_null = mask.true_count(); + + let total_input_bytes = match mask.bit_buffer() { + AllOr::All => views.iter().map(|v| v.len() as usize).sum(), + AllOr::None => 0, + AllOr::Some(bits) => views + .iter() + .zip(bits.iter()) + .filter(|&(_, b)| b) + .map(|(v, _)| v.len() as usize) + .sum(), + }; - for string in iter { - match string { - None => {} - Some(b) => lines.push(b), + if fsst_output_fits_in_i32_offsets(total_input_bytes, non_null) { + compress_views::(strings, &mask, compressor, ctx) + } else { + compress_views::(strings, &mask, compressor, ctx) + } +} + +fn compress_varbin_array( + strings: ArrayView, + compressor: &Compressor, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mask = strings.validity()?.execute_mask(strings.len(), ctx)?; + let offsets = strings.offsets().clone().execute::(ctx)?; + let total_input_bytes = match_each_integer_ptype!(offsets.ptype(), |O| { + let off = offsets.as_slice::(); + let first: usize = off[0].as_(); + let last: usize = off[off.len() - 1].as_(); + last - first + }); + let non_null = mask.true_count(); + + if fsst_output_fits_in_i32_offsets(total_input_bytes, non_null) { + compress_varbin::(strings, &offsets, &mask, compressor, ctx) + } else { + compress_varbin::(strings, &offsets, &mask, compressor, ctx) + } +} + +fn train_varbinview( + strings: ArrayView, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mask = strings.validity()?.execute_mask(strings.len(), ctx)?; + let views = strings.views(); + let buffers = strings.data_buffers(); + let mut lines: Vec<&[u8]> = Vec::with_capacity(mask.true_count()); + + match mask.bit_buffer() { + AllOr::All => { + for view in views { + lines.push(view_bytes(view, buffers)); + } + } + AllOr::None => {} + AllOr::Some(bits) => { + for (view, valid) in views.iter().zip(bits.iter()) { + if valid { + lines.push(view_bytes(view, buffers)); + } + } } } - Compressor::train(&lines) + Ok(Compressor::train(&lines)) } -/// Most strings are small in practice. If we encounter a larger string, we reallocate -/// the buffer to hold enough capacity for the worst-case compressed value. -const DEFAULT_BUFFER_LEN: usize = 1024 * 1024; +fn train_varbin_array( + strings: ArrayView, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mask = strings.validity()?.execute_mask(strings.len(), ctx)?; + let offsets = strings.offsets().clone().execute::(ctx)?; + let bytes = strings.bytes().as_slice(); + let mut lines: Vec<&[u8]> = Vec::with_capacity(mask.true_count()); + + match_each_integer_ptype!(offsets.ptype(), |I| { + let off = offsets.as_slice::(); + for_each_varbin_row(off, bytes, &mask, |row| { + if let Some(s) = row { + lines.push(s); + } + }); + }); + + Ok(Compressor::train(&lines)) +} + +#[inline] +fn fsst_output_fits_in_i32_offsets(total_input_bytes: usize, non_null: usize) -> bool { + let worst = total_input_bytes + .saturating_mul(FSST_PER_BYTE_OVERHEAD) + .saturating_add(non_null.saturating_mul(FSST_PER_ROW_OVERHEAD)); + worst <= i32::MAX as usize +} -/// Compress from an iterator of bytestrings using FSST. -pub fn fsst_compress_iter<'a, I>( - iter: I, - len: usize, - dtype: DType, +#[inline] +fn view_bytes<'a>(view: &'a BinaryView, buffers: &'a Arc<[BufferHandle]>) -> &'a [u8] { + if view.is_inlined() { + view.as_inlined().value() + } else { + let r = view.as_view(); + &buffers[r.buffer_index as usize].as_host()[r.as_range()] + } +} + +fn compress_views( + strings: ArrayView, + mask: &Mask, compressor: &Compressor, ctx: &mut ExecutionCtx, -) -> FSSTArray +) -> VortexResult where - I: Iterator>, + O: IntegerPType + 'static, { - let mut buffer = Vec::with_capacity(DEFAULT_BUFFER_LEN); - - // Offsets are widened to i64 because the cumulative compressed bytes can exceed i32::MAX for - // large inputs (see issue #7833). Per-string sizes still fit in i32. - let mut builder = VarBinBuilder::::with_capacity(len); - let mut uncompressed_lengths: BufferMut = BufferMut::with_capacity(len); - - for string in iter { - match string { - None => { - builder.append_null(); - uncompressed_lengths.push(0); + let mut sink = FsstSink::::with_capacity(strings.len(), compressor); + let views = strings.views(); + let buffers = strings.data_buffers(); + match mask.bit_buffer() { + AllOr::All => { + for view in views { + sink.emit(Some(view_bytes(view, buffers))); } - Some(s) => { - uncompressed_lengths.push( - s.len() - .try_into() - .vortex_expect("string length must fit in i32"), - ); - - // make sure the buffer is 2x+7 larger than the input - let target_size = 2 * s.len() + 7; - if target_size > buffer.len() { - let additional_capacity = target_size - buffer.len(); - buffer.reserve(additional_capacity); - } + } + AllOr::None => { + for _ in 0..mask.len() { + sink.emit(None); + } + } + AllOr::Some(bits) => { + for (view, valid) in views.iter().zip(bits.iter()) { + sink.emit(valid.then(|| view_bytes(view, buffers))); + } + } + } + sink.finish(strings.dtype().clone(), ctx) +} - // SAFETY: buffer is always sized to be large enough - unsafe { compressor.compress_into(s, &mut buffer) }; +fn compress_varbin( + strings: ArrayView, + offsets: &PrimitiveArray, + mask: &Mask, + compressor: &Compressor, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + O: IntegerPType + 'static, +{ + let mut sink = FsstSink::::with_capacity(strings.len(), compressor); + let bytes = strings.bytes().as_slice(); + match_each_integer_ptype!(offsets.ptype(), |I| { + let off = offsets.as_slice::(); + for_each_varbin_row(off, bytes, mask, |row| sink.emit(row)); + }); + sink.finish(strings.dtype().clone(), ctx) +} - builder.append_value(&buffer); +/// Call `f` once per row of a `VarBinArray` with the row bytes or `None`. +/// Validity dispatch is hoisted out of the per-row loop. +#[inline] +fn for_each_varbin_row<'a, I, F>(off: &[I], bytes: &'a [u8], mask: &Mask, mut f: F) +where + I: IntegerPType + 'static, + F: FnMut(Option<&'a [u8]>), +{ + match mask.bit_buffer() { + AllOr::All => { + for w in off.windows(2) { + f(Some(&bytes[w[0].as_()..w[1].as_()])); + } + } + AllOr::None => { + for _ in 0..mask.len() { + f(None); + } + } + AllOr::Some(bits) => { + for (w, valid) in off.windows(2).zip(bits.iter()) { + f(valid.then(|| &bytes[w[0].as_()..w[1].as_()])); } } } +} + +/// Per-row output state for an FSST compression pass. +struct FsstSink<'c, O: IntegerPType + 'static> { + buffer: Vec, + builder: VarBinBuilder, + uncompressed_lengths: BufferMut, + compressor: &'c Compressor, +} + +impl<'c, O: IntegerPType + 'static> FsstSink<'c, O> { + fn with_capacity(len: usize, compressor: &'c Compressor) -> Self { + Self { + buffer: Vec::with_capacity(DEFAULT_BUFFER_LEN), + builder: VarBinBuilder::::with_capacity(len), + uncompressed_lengths: BufferMut::with_capacity(len), + compressor, + } + } + + #[inline] + fn emit(&mut self, row: Option<&[u8]>) { + let Some(s) = row else { + self.builder.append_null(); + self.uncompressed_lengths.push(0); + return; + }; + + // A single row > i32::MAX (2 GiB) is not supported. + self.uncompressed_lengths.push( + i32::try_from(s.len()).vortex_expect("per-row uncompressed length must fit in i32"), + ); + + let target = FSST_PER_BYTE_OVERHEAD * s.len() + FSST_PER_ROW_OVERHEAD; + if target > self.buffer.len() { + self.buffer.reserve(target - self.buffer.len()); + } + + // SAFETY: `self.buffer` has capacity for the FSST worst-case output of `s`. + unsafe { self.compressor.compress_into(s, &mut self.buffer) }; - let codes = builder.finish(DType::Binary(dtype.nullability())); - let symbols: Buffer = Buffer::copy_from(compressor.symbol_table()); - let symbol_lengths: Buffer = Buffer::::copy_from(compressor.symbol_lengths()); - - let uncompressed_lengths = uncompressed_lengths.into_array(); - - FSST::try_new( - dtype, - symbols, - symbol_lengths, - codes, - uncompressed_lengths, - ctx, - ) - .vortex_expect("FSST parts must be valid") + self.builder.append_value(&self.buffer); + } + + fn finish(self, dtype: DType, ctx: &mut ExecutionCtx) -> VortexResult { + let codes = self.builder.finish(DType::Binary(dtype.nullability())); + FSST::try_new( + dtype, + Buffer::copy_from(self.compressor.symbol_table()), + Buffer::::copy_from(self.compressor.symbol_lengths()), + codes, + self.uncompressed_lengths.into_array(), + ctx, + ) + } } #[cfg(test)] mod tests { - use fsst::CompressorBuilder; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; - use vortex_array::dtype::DType; - use vortex_array::dtype::Nullability; - use vortex_array::scalar::Scalar; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::arrays::varbin::VarBinArrayExt; + use vortex_array::dtype::PType; + use vortex_error::VortexResult; - use crate::compress::DEFAULT_BUFFER_LEN; - use crate::fsst_compress_iter; + use super::fsst_compress; + use super::fsst_output_fits_in_i32_offsets; + use super::fsst_train_compressor; + use crate::array::FSSTArrayExt; + /// Regression for #7833: the i32-vs-i64 codes-offsets decision must cross at + /// `i32::MAX` against the worst-case bound `2 * total + 7 * non_null`. #[test] - fn test_large_string() { - let big_string: String = "abc" - .chars() - .cycle() - .take(10 * DEFAULT_BUFFER_LEN) - .collect(); - - let compressor = CompressorBuilder::default().build(); - - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - let compressed = fsst_compress_iter( - [Some(big_string.as_bytes())].into_iter(), - 1, - DType::Utf8(Nullability::NonNullable), - &compressor, - &mut ctx, - ); - - let decoded = compressed.execute_scalar(0, &mut ctx).unwrap(); - - let expected = Scalar::utf8(big_string, Nullability::NonNullable); + fn offset_width_boundary() { + let m = i32::MAX as usize; + assert!(fsst_output_fits_in_i32_offsets(m / 2 - 7, 1)); + assert!(!fsst_output_fits_in_i32_offsets(m / 2, 1)); + assert!(fsst_output_fits_in_i32_offsets(0, 0)); + assert!(!fsst_output_fits_in_i32_offsets(usize::MAX, 1)); + } - assert_eq!(decoded, expected); + /// Small inputs fit the i32 bound, so `fsst_compress` must pick i32 offsets. + /// The i64 branch is covered by `tests::fsst_compress_offsets_overflow_i32`. + #[test] + fn codes_offsets_dtype_small_input_is_i32() -> VortexResult<()> { + let array = VarBinViewArray::from_iter_str(["hello", "world", "fsst encoded"]); + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let compressor = fsst_train_compressor(array.as_array(), &mut ctx)?; + let fsst = fsst_compress(array.as_array(), &compressor, &mut ctx)?; + assert_eq!(fsst.codes().offsets().dtype().as_ptype(), PType::I32); + Ok(()) } } diff --git a/encodings/fsst/src/compute/cast.rs b/encodings/fsst/src/compute/cast.rs index f6e776e16d9..bdca68e841b 100644 --- a/encodings/fsst/src/compute/cast.rs +++ b/encodings/fsst/src/compute/cast.rs @@ -98,32 +98,35 @@ mod tests { use vortex_array::compute::conformance::cast::test_cast_conformance; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; + use vortex_error::VortexResult; use vortex_session::VortexSession; use crate::fsst_compress; use crate::fsst_train_compressor; + use crate::initialize; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + initialize(&session); + session + }); #[test] - fn test_cast_fsst_nullability() { + fn test_cast_fsst_nullability() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let strings = VarBinArray::from_iter( vec![Some("hello"), Some("world"), Some("hello world")], DType::Utf8(Nullability::NonNullable), - ); + ) + .into_array(); - let compressor = fsst_train_compressor(&strings); - let len = strings.len(); - let dtype = strings.dtype().clone(); - let fsst = fsst_compress(strings, len, &dtype, &compressor, &mut ctx); + let compressor = fsst_train_compressor(&strings, &mut ctx)?; + let fsst = fsst_compress(&strings, &compressor, &mut ctx)?; // Cast to nullable - let casted = fsst - .into_array() - .cast(DType::Utf8(Nullability::Nullable)) - .unwrap(); + let casted = fsst.into_array().cast(DType::Utf8(Nullability::Nullable))?; assert_eq!(casted.dtype(), &DType::Utf8(Nullability::Nullable)); + Ok(()) } #[rstest] @@ -139,10 +142,12 @@ mod tests { vec![Some("test")], DType::Utf8(Nullability::NonNullable) ))] - fn test_cast_fsst_conformance(#[case] array: VarBinArray) { + fn test_cast_fsst_conformance(#[case] array: VarBinArray) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); - let compressor = fsst_train_compressor(&array); - let fsst = fsst_compress(&array, array.len(), array.dtype(), &compressor, &mut ctx); + let array = array.into_array(); + let compressor = fsst_train_compressor(&array, &mut ctx)?; + let fsst = fsst_compress(&array, &compressor, &mut ctx)?; test_cast_conformance(&fsst.into_array()); + Ok(()) } } diff --git a/encodings/fsst/src/compute/compare.rs b/encodings/fsst/src/compute/compare.rs index 553a3608e75..732e1a446dd 100644 --- a/encodings/fsst/src/compute/compare.rs +++ b/encodings/fsst/src/compute/compare.rs @@ -134,13 +134,14 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_error::VortexResult; use crate::fsst_compress; use crate::fsst_train_compressor; #[test] #[cfg_attr(miri, ignore)] - fn test_compare_fsst() { + fn test_compare_fsst() -> VortexResult<()> { let mut ctx = LEGACY_SESSION.create_execution_ctx(); let lhs = VarBinArray::from_iter( [ @@ -151,11 +152,10 @@ mod tests { Some("this is a very long string"), ], DType::Utf8(Nullability::Nullable), - ); - let compressor = fsst_train_compressor(&lhs); - let len = lhs.len(); - let dtype = lhs.dtype().clone(); - let lhs = fsst_compress(lhs, len, &dtype, &compressor, &mut ctx); + ) + .into_array(); + let compressor = fsst_train_compressor(&lhs, &mut ctx)?; + let lhs = fsst_compress(&lhs, &compressor, &mut ctx)?; let rhs = ConstantArray::new("world", lhs.len()); @@ -163,10 +163,8 @@ mod tests { let equals = lhs .clone() .into_array() - .binary(rhs.clone().into_array(), Operator::Eq) - .unwrap() - .execute::(&mut ctx) - .unwrap(); + .binary(rhs.clone().into_array(), Operator::Eq)? + .execute::(&mut ctx)?; assert_eq!(equals.dtype(), &DType::Bool(Nullability::Nullable)); @@ -179,10 +177,8 @@ mod tests { let not_equals = lhs .clone() .into_array() - .binary(rhs.into_array(), Operator::NotEq) - .unwrap() - .execute::(&mut ctx) - .unwrap(); + .binary(rhs.into_array(), Operator::NotEq)? + .execute::(&mut ctx)?; assert_eq!(not_equals.dtype(), &DType::Bool(Nullability::Nullable)); assert_arrays_eq!( @@ -196,8 +192,7 @@ mod tests { let equals_null = lhs .clone() .into_array() - .binary(null_rhs.clone().into_array(), Operator::Eq) - .unwrap(); + .binary(null_rhs.clone().into_array(), Operator::Eq)?; assert_arrays_eq!( &equals_null, &BoolArray::from_iter([None::, None, None, None, None]) @@ -205,11 +200,11 @@ mod tests { let noteq_null = lhs .into_array() - .binary(null_rhs.into_array(), Operator::NotEq) - .unwrap(); + .binary(null_rhs.into_array(), Operator::NotEq)?; assert_arrays_eq!( ¬eq_null, &BoolArray::from_iter([None::, None, None, None, None]) ); + Ok(()) } } diff --git a/encodings/fsst/src/compute/like.rs b/encodings/fsst/src/compute/like.rs index ae2402a6e51..1b0695ecc7f 100644 --- a/encodings/fsst/src/compute/like.rs +++ b/encodings/fsst/src/compute/like.rs @@ -107,17 +107,11 @@ mod tests { static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); fn make_fsst(strings: &[Option<&str>], nullability: Nullability) -> FSSTArray { - let varbin = VarBinArray::from_iter(strings.iter().copied(), DType::Utf8(nullability)); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress( - varbin, - len, - &dtype, - &compressor, - &mut SESSION.create_execution_ctx(), - ) + let array = + VarBinArray::from_iter(strings.iter().copied(), DType::Utf8(nullability)).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + fsst_compress(&array, &compressor, &mut ctx).unwrap() } fn run_like(array: FSSTArray, pattern: &str, opts: LikeOptions) -> VortexResult { diff --git a/encodings/fsst/src/compute/mod.rs b/encodings/fsst/src/compute/mod.rs index d95725b9fbc..fafa05d293b 100644 --- a/encodings/fsst/src/compute/mod.rs +++ b/encodings/fsst/src/compute/mod.rs @@ -69,31 +69,34 @@ mod tests { use vortex_array::compute::conformance::take::test_take_conformance; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; + use vortex_error::VortexResult; use crate::FSSTArray; use crate::fsst_compress; use crate::fsst_train_compressor; #[test] - fn test_take_null() { + fn test_take_null() -> VortexResult<()> { let mut ctx = LEGACY_SESSION.create_execution_ctx(); - let arr = VarBinArray::from_iter([Some("h")], DType::Utf8(Nullability::NonNullable)); - let compr = fsst_train_compressor(&arr); - let fsst = fsst_compress(&arr, arr.len(), arr.dtype(), &compr, &mut ctx); + let arr = + VarBinArray::from_iter([Some("h")], DType::Utf8(Nullability::NonNullable)).into_array(); + let compr = fsst_train_compressor(&arr, &mut ctx)?; + let fsst = fsst_compress(&arr, &compr, &mut ctx)?; let idx1: PrimitiveArray = (0..1).collect(); assert_eq!( - fsst.take(idx1.into_array()).unwrap().dtype(), + fsst.take(idx1.into_array())?.dtype(), &DType::Utf8(Nullability::NonNullable) ); let idx2: PrimitiveArray = PrimitiveArray::from_option_iter(vec![Some(0)]); assert_eq!( - fsst.take(idx2.into_array()).unwrap().dtype(), + fsst.take(idx2.into_array())?.dtype(), &DType::Utf8(Nullability::Nullable) ); + Ok(()) } #[rstest] @@ -109,11 +112,13 @@ mod tests { ["single element"].map(Some), DType::Utf8(Nullability::NonNullable), ))] - fn test_take_fsst_conformance(#[case] varbin: VarBinArray) { + fn test_take_fsst_conformance(#[case] varbin: VarBinArray) -> VortexResult<()> { let mut ctx = LEGACY_SESSION.create_execution_ctx(); - let compressor = fsst_train_compressor(&varbin); - let array = fsst_compress(&varbin, varbin.len(), varbin.dtype(), &compressor, &mut ctx); + let varbin = varbin.into_array(); + let compressor = fsst_train_compressor(&varbin, &mut ctx)?; + let array = fsst_compress(&varbin, &compressor, &mut ctx)?; test_take_conformance(&array.into_array()); + Ok(()) } type FsstBuilder = fn(&mut ExecutionCtx) -> FSSTArray; @@ -121,51 +126,47 @@ mod tests { #[rstest] // Basic string arrays #[case::fsst_simple(|ctx: &mut ExecutionCtx| { - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( ["hello world", "testing fsst", "compression test", "data array", "vortex encoding"].map(Some), DType::Utf8(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); - fsst_compress(&varbin, varbin.len(), varbin.dtype(), &compressor, ctx) + ).into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() })] // Nullable strings #[case::fsst_nullable(|ctx: &mut ExecutionCtx| { - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( [Some("hello"), None, Some("world"), Some("test"), None], DType::Utf8(Nullability::Nullable), - ); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress(varbin, len, &dtype, &compressor, ctx) + ).into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() })] // Repetitive patterns (good for FSST compression) #[case::fsst_repetitive(|ctx: &mut ExecutionCtx| { - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( ["http://example.com", "http://test.com", "http://vortex.dev", "http://data.org"].map(Some), DType::Utf8(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); - fsst_compress(&varbin, varbin.len(), varbin.dtype(), &compressor, ctx) + ).into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() })] // Edge cases #[case::fsst_single(|ctx: &mut ExecutionCtx| { - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( ["single element"].map(Some), DType::Utf8(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); - fsst_compress(&varbin, varbin.len(), varbin.dtype(), &compressor, ctx) + ).into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() })] #[case::fsst_empty_strings(|ctx: &mut ExecutionCtx| { - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( ["", "test", "", "hello", ""].map(Some), DType::Utf8(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress(varbin, len, &dtype, &compressor, ctx) + ).into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() })] // Large arrays #[case::fsst_large(|ctx: &mut ExecutionCtx| { @@ -183,11 +184,9 @@ mod tests { _ => "CREATE TABLE data (id INT, value TEXT)", })) .collect(); - let varbin = VarBinArray::from_iter(data, DType::Utf8(Nullability::NonNullable)); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress(varbin, len, &dtype, &compressor, ctx) + let array = VarBinArray::from_iter(data, DType::Utf8(Nullability::NonNullable)).into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() })] fn test_fsst_consistency(#[case] build: FsstBuilder) { diff --git a/encodings/fsst/src/dfa/tests.rs b/encodings/fsst/src/dfa/tests.rs index af23ae581d5..101eb3989d4 100644 --- a/encodings/fsst/src/dfa/tests.rs +++ b/encodings/fsst/src/dfa/tests.rs @@ -264,20 +264,14 @@ fn test_contains_pushdown_rejects_len_255() { // --------------------------------------------------------------------------- fn make_fsst_str(strings: &[Option<&str>]) -> FSSTArray { - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( strings.iter().copied(), DType::Utf8(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress( - varbin, - len, - &dtype, - &compressor, - &mut SESSION.create_execution_ctx(), ) + .into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let compressor = fsst_train_compressor(&array, &mut ctx).unwrap(); + fsst_compress(&array, &compressor, &mut ctx).unwrap() } fn run_like(array: FSSTArray, pattern_arr: ArrayRef) -> VortexResult { diff --git a/encodings/fsst/src/kernel.rs b/encodings/fsst/src/kernel.rs index 942182097fa..f137ea95025 100644 --- a/encodings/fsst/src/kernel.rs +++ b/encodings/fsst/src/kernel.rs @@ -70,11 +70,12 @@ mod tests { builder.append_value(b"final string"); let input = builder.finish(DType::Utf8(Nullability::NonNullable)); - let compressor = fsst_train_compressor(&input); - let len = input.len(); - let dtype = input.dtype().clone(); let mut ctx = SESSION.create_execution_ctx(); - fsst_compress(input, len, &dtype, &compressor, &mut ctx).into_array() + let arr = input.into_array(); + let compressor = fsst_train_compressor(&arr, &mut ctx).unwrap(); + fsst_compress(&arr, &compressor, &mut ctx) + .unwrap() + .into_array() } #[test] @@ -144,17 +145,11 @@ mod tests { builder.append_null(); } let input = builder.finish(DType::Utf8(Nullability::Nullable)); + let array = input.clone().into_array(); - let compressor = fsst_train_compressor(&input); let mut ctx = SESSION.create_execution_ctx(); - let fsst_array: ArrayRef = fsst_compress( - input.clone(), - input.len(), - input.dtype(), - &compressor, - &mut ctx, - ) - .into_array(); + let compressor = fsst_train_compressor(&array, &mut ctx)?; + let fsst_array: ArrayRef = fsst_compress(&array, &compressor, &mut ctx)?.into_array(); // Filter: only select the last element (index 22) let mut mask = vec![false; 22]; @@ -179,17 +174,11 @@ mod tests { builder.append_null(); let input = builder.finish(DType::Utf8(Nullability::Nullable)); + let array = input.clone().into_array(); - let compressor = fsst_train_compressor(&input); let mut ctx = SESSION.create_execution_ctx(); - let fsst_array: ArrayRef = fsst_compress( - input.clone(), - input.len(), - input.dtype(), - &compressor, - &mut ctx, - ) - .into_array(); + let compressor = fsst_train_compressor(&array, &mut ctx)?; + let fsst_array: ArrayRef = fsst_compress(&array, &compressor, &mut ctx)?.into_array(); let mask = Mask::from_iter([true, false, true]); @@ -226,12 +215,12 @@ mod tests { builder.append_value("Пуховички"); // 9 characters, 18 bytes builder.append_value(b""); - let varbin = builder.finish(DType::Utf8(Nullability::NonNullable)); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); + let varbin = builder + .finish(DType::Utf8(Nullability::NonNullable)) + .into_array(); let mut ctx = SESSION.create_execution_ctx(); - let fsst = fsst_compress(varbin, len, &dtype, &compressor, &mut ctx).into_array(); + let compressor = fsst_train_compressor(&varbin, &mut ctx)?; + let fsst = fsst_compress(&varbin, &compressor, &mut ctx)?.into_array(); let result = fsst.apply(&byte_length(root()))?; let expected = PrimitiveArray::from_iter(vec![5u64, 7, 18, 0]); assert_arrays_eq!(result, expected); diff --git a/encodings/fsst/src/test_utils.rs b/encodings/fsst/src/test_utils.rs index eeca412e377..7becb34ffbc 100644 --- a/encodings/fsst/src/test_utils.rs +++ b/encodings/fsst/src/test_utils.rs @@ -41,17 +41,18 @@ pub fn gen_fsst_test_data( )); } - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( strings .into_iter() .map(|opt_s| opt_s.map(Vec::into_boxed_slice)), DType::Binary(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); + ) + .into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress(varbin, len, &dtype, &compressor, ctx).into_array() + fsst_compress(&array, &compressor, ctx) + .unwrap() + .into_array() } pub fn gen_dict_fsst_test_data( @@ -144,11 +145,9 @@ pub fn generate_url_data_n(n: usize) -> VarBinArray { } pub fn make_fsst_urls(n: usize, ctx: &mut ExecutionCtx) -> FSSTArray { - let varbin = generate_url_data_n(n); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress(varbin, len, &dtype, &compressor, ctx) + let array = generate_url_data_n(n).into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() } // --------------------------------------------------------------------------- @@ -237,14 +236,13 @@ pub fn generate_clickbench_urls(n: usize) -> Vec { pub fn make_fsst_clickbench_urls(n: usize, ctx: &mut ExecutionCtx) -> FSSTArray { let urls = generate_clickbench_urls(n); - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( urls.iter().map(|s| Some(s.as_str())), DType::Utf8(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress(varbin, len, &dtype, &compressor, ctx) + ) + .into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() } // --------------------------------------------------------------------------- @@ -305,14 +303,13 @@ pub fn generate_short_urls(n: usize) -> Vec { pub fn make_fsst_short_urls(n: usize, ctx: &mut ExecutionCtx) -> FSSTArray { let urls = generate_short_urls(n); - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( urls.iter().map(|s| Some(s.as_str())), DType::Utf8(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress(varbin, len, &dtype, &compressor, ctx) + ) + .into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() } // --------------------------------------------------------------------------- @@ -377,14 +374,13 @@ pub fn generate_log_lines(n: usize) -> Vec { pub fn make_fsst_log_lines(n: usize, ctx: &mut ExecutionCtx) -> FSSTArray { let lines = generate_log_lines(n); - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( lines.iter().map(|s| Some(s.as_str())), DType::Utf8(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress(varbin, len, &dtype, &compressor, ctx) + ) + .into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() } // --------------------------------------------------------------------------- @@ -436,14 +432,13 @@ pub fn generate_json_strings(n: usize) -> Vec { pub fn make_fsst_json_strings(n: usize, ctx: &mut ExecutionCtx) -> FSSTArray { let jsons = generate_json_strings(n); - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( jsons.iter().map(|s| Some(s.as_str())), DType::Utf8(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress(varbin, len, &dtype, &compressor, ctx) + ) + .into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() } // --------------------------------------------------------------------------- @@ -508,14 +503,13 @@ pub fn generate_file_paths(n: usize) -> Vec { pub fn make_fsst_file_paths(n: usize, ctx: &mut ExecutionCtx) -> FSSTArray { let paths = generate_file_paths(n); - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( paths.iter().map(|s| Some(s.as_str())), DType::Utf8(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress(varbin, len, &dtype, &compressor, ctx) + ) + .into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() } // --------------------------------------------------------------------------- @@ -561,14 +555,13 @@ pub fn generate_emails(n: usize) -> Vec { pub fn make_fsst_emails(n: usize, ctx: &mut ExecutionCtx) -> FSSTArray { let emails = generate_emails(n); - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( emails.iter().map(|s| Some(s.as_str())), DType::Utf8(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress(varbin, len, &dtype, &compressor, ctx) + ) + .into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() } // --------------------------------------------------------------------------- @@ -600,12 +593,11 @@ pub fn generate_rare_match_strings(n: usize, match_rate: f64) -> Vec { pub fn make_fsst_rare_match(n: usize, ctx: &mut ExecutionCtx) -> FSSTArray { let strings = generate_rare_match_strings(n, 0.00001); - let varbin = VarBinArray::from_iter( + let array = VarBinArray::from_iter( strings.iter().map(|s| Some(s.as_str())), DType::Utf8(Nullability::NonNullable), - ); - let compressor = fsst_train_compressor(&varbin); - let len = varbin.len(); - let dtype = varbin.dtype().clone(); - fsst_compress(varbin, len, &dtype, &compressor, ctx) + ) + .into_array(); + let compressor = fsst_train_compressor(&array, ctx).unwrap(); + fsst_compress(&array, &compressor, ctx).unwrap() } diff --git a/encodings/fsst/src/tests.rs b/encodings/fsst/src/tests.rs index 481b84865df..7cf7e9e9f6c 100644 --- a/encodings/fsst/src/tests.rs +++ b/encodings/fsst/src/tests.rs @@ -27,13 +27,15 @@ pub(crate) fn build_fsst_array() -> ArrayRef { b"They said it existed and that whoever dared to exceed it was mercilessly struck down", ); input_array.append_value(b"Nothing in present history can contradict them"); - let input_array = input_array.finish(DType::Utf8(Nullability::NonNullable)); + let input_array = input_array + .finish(DType::Utf8(Nullability::NonNullable)) + .into_array(); - let compressor = fsst_train_compressor(&input_array); - let len = input_array.len(); - let dtype = input_array.dtype().clone(); let mut ctx = LEGACY_SESSION.create_execution_ctx(); - fsst_compress(input_array, len, &dtype, &compressor, &mut ctx).into_array() + let compressor = fsst_train_compressor(&input_array, &mut ctx).unwrap(); + fsst_compress(&input_array, &compressor, &mut ctx) + .unwrap() + .into_array() } #[test] @@ -112,7 +114,7 @@ fn test_fsst_array_ops() { // TODO(someone): ideally CI would run this in release mode as well since debug builds make the // allocation and compression loop substantially slower. /// Regression for #7833: [`fsst_compress`] must accept inputs whose cumulative compressed -/// bytes exceed [`i32::MAX`]. Before the fix, [`fsst_compress_iter`] hardcoded +/// bytes exceed [`i32::MAX`]. Before the fix, the compress path hardcoded /// [`VarBinBuilder`] for the FSST output and panicked in /// [`VarBinBuilder::append_value`] once cumulative compressed bytes crossed the boundary. /// @@ -132,7 +134,7 @@ fn test_fsst_array_ops() { /// CI=1 cargo test --release -p vortex-fsst fsst_compress_offsets /// ``` /// -/// [`fsst_compress_iter`]: crate::compress::fsst_compress_iter +/// [`fsst_compress`]: crate::compress::fsst_compress #[test_with::env(CI)] #[test_with::no_env(VORTEX_SKIP_SLOW_TESTS)] fn fsst_compress_offsets_overflow_i32() { @@ -148,15 +150,16 @@ fn fsst_compress_offsets_overflow_i32() { for _ in 0..N { builder.append_value(&string); } - let array = builder.finish(DType::Utf8(Nullability::NonNullable)); + let array = builder + .finish(DType::Utf8(Nullability::NonNullable)) + .into_array(); let compressor = CompressorBuilder::default().build(); let len = array.len(); - let dtype = array.dtype().clone(); let mut ctx = LEGACY_SESSION.create_execution_ctx(); println!("compressing to FSST"); - let compressed = fsst_compress(array, len, &dtype, &compressor, &mut ctx); + let compressed = fsst_compress(&array, &compressor, &mut ctx).unwrap(); assert_eq!(compressed.len(), len); // Prove the regression condition was exercised: compressed bytes crossed i32::MAX. assert!(compressed.codes_bytes().len() > i32::MAX as usize); diff --git a/fuzz/Cargo.toml b/fuzz/Cargo.toml index e2d05b706f9..58c728b23a5 100644 --- a/fuzz/Cargo.toml +++ b/fuzz/Cargo.toml @@ -36,7 +36,7 @@ vortex-array = { workspace = true, features = ["arbitrary", "_test-harness"] } vortex-btrblocks = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } -vortex-fsst = { workspace = true } +vortex-fsst = { workspace = true, features = ["_test-harness"] } vortex-io = { workspace = true } vortex-mask = { workspace = true } vortex-runend = { workspace = true, features = ["arbitrary"] } diff --git a/fuzz/src/fsst_like.rs b/fuzz/src/fsst_like.rs index ec310a0b81a..9a8c9ae21ac 100644 --- a/fuzz/src/fsst_like.rs +++ b/fuzz/src/fsst_like.rs @@ -107,18 +107,15 @@ pub fn run_fsst_like_fuzz(fuzz: FuzzFsstLike) -> VortexFuzzResult { let varbin = VarBinArray::from_iter( strings.iter().map(|s| Some(s.as_str())), DType::Utf8(Nullability::NonNullable), - ); + ) + .into_array(); // Train FSST compressor and compress. - let compressor = fsst_train_compressor(&varbin); let mut ctx = SESSION.create_execution_ctx(); - let fsst_array: FSSTArray = fsst_compress( - varbin.clone(), - varbin.len(), - varbin.dtype(), - &compressor, - &mut ctx, - ); + let compressor = fsst_train_compressor(&varbin, &mut ctx) + .map_err(|err| VortexFuzzError::VortexError(err, Backtrace::capture()))?; + let fsst_array: FSSTArray = fsst_compress(&varbin, &compressor, &mut ctx) + .map_err(|err| VortexFuzzError::VortexError(err, Backtrace::capture()))?; let opts = LikeOptions { negated, @@ -126,7 +123,7 @@ pub fn run_fsst_like_fuzz(fuzz: FuzzFsstLike) -> VortexFuzzResult { }; // Run LIKE on the uncompressed array. - let expected = run_like_on_array(&varbin.into_array(), &pattern, len, opts) + let expected = run_like_on_array(&varbin, &pattern, len, opts) .map_err(|err| VortexFuzzError::VortexError(err, Backtrace::capture()))?; // Run LIKE on the FSST-compressed array. diff --git a/vortex-btrblocks/src/schemes/string/fsst.rs b/vortex-btrblocks/src/schemes/string/fsst.rs index a90fbf4f452..5e6baf73933 100644 --- a/vortex-btrblocks/src/schemes/string/fsst.rs +++ b/vortex-btrblocks/src/schemes/string/fsst.rs @@ -64,9 +64,9 @@ impl Scheme for FSSTScheme { compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - let utf8 = data.array_as_varbinview().into_owned(); - let compressor_fsst = fsst_train_compressor(&utf8); - let fsst = fsst_compress(&utf8, utf8.len(), utf8.dtype(), &compressor_fsst, exec_ctx); + let utf8 = data.array_as_varbinview().into_owned().into_array(); + let compressor_fsst = fsst_train_compressor(&utf8, exec_ctx)?; + let fsst = fsst_compress(&utf8, &compressor_fsst, exec_ctx)?; let uncompressed_lengths_primitive = fsst .uncompressed_lengths() diff --git a/vortex-cuda/src/kernel/encodings/fsst.rs b/vortex-cuda/src/kernel/encodings/fsst.rs index 56ba6001cf1..6b978ebba23 100644 --- a/vortex-cuda/src/kernel/encodings/fsst.rs +++ b/vortex-cuda/src/kernel/encodings/fsst.rs @@ -238,12 +238,10 @@ mod tests { let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); - let varbin = VarBinArray::from_iter(strings, DType::Binary(nullability)); - let compressor = fsst_train_compressor(&varbin); - let dtype = varbin.dtype().clone(); - let len = varbin.len(); + let varbin = VarBinArray::from_iter(strings, DType::Binary(nullability)).into_array(); + let compressor = fsst_train_compressor(&varbin, cuda_ctx.execution_ctx())?; let fsst_array = - fsst_compress(&varbin, len, &dtype, &compressor, cuda_ctx.execution_ctx()).into_array(); + fsst_compress(&varbin, &compressor, cuda_ctx.execution_ctx())?.into_array(); let cpu_result = crate::canonicalize_cpu(fsst_array.clone())?; let gpu_result = FSSTExecutor diff --git a/vortex-test/compat-gen/Cargo.toml b/vortex-test/compat-gen/Cargo.toml index dd8bec4703a..b22c49112b8 100644 --- a/vortex-test/compat-gen/Cargo.toml +++ b/vortex-test/compat-gen/Cargo.toml @@ -25,6 +25,7 @@ vortex = { workspace = true, features = ["files", "tokio", "zstd"] } vortex-array = { workspace = true, features = ["_test-harness"] } vortex-buffer = { workspace = true } vortex-error = { workspace = true } +vortex-fsst = { workspace = true, features = ["_test-harness"] } vortex-session = { workspace = true } # TPC-H generation diff --git a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/fsst.rs b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/fsst.rs index 17314629737..2276ebd1f74 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/fsst.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/fsst.rs @@ -45,7 +45,7 @@ impl FlatLayoutFixture for FsstFixture { .map(|i| format!("{}{}", prefixes[i % prefixes.len()], i)) .collect(); let url_refs: Vec<&str> = urls.iter().map(|s| s.as_str()).collect(); - let url_col = VarBinArray::from_strs(url_refs); + let url_col = VarBinArray::from_strs(url_refs).into_array(); let severities = ["INFO", "WARN", "ERROR", "DEBUG"]; let components = ["auth", "db", "cache", "api"]; @@ -60,36 +60,36 @@ impl FlatLayoutFixture for FsstFixture { }) .collect(); let log_refs: Vec<&str> = logs.iter().map(|s| s.as_str()).collect(); - let log_col = VarBinArray::from_strs(log_refs); + let log_col = VarBinArray::from_strs(log_refs).into_array(); let nullable_urls: Vec> = (0..N) .map(|i| (i % 7 != 0).then(|| format!("{}{}", prefixes[i % prefixes.len()], i * 3))) .collect(); let nullable_refs: Vec> = nullable_urls.iter().map(|s| s.as_deref()).collect(); - let nullable_col = VarBinArray::from_nullable_strs(nullable_refs); + let nullable_col = VarBinArray::from_nullable_strs(nullable_refs).into_array(); let short_tokens = ["a", "bb", "ccc", "dd", "e"]; let short_strs: Vec<&str> = (0..N) .map(|i| short_tokens[i % short_tokens.len()]) .collect(); - let short_col = VarBinArray::from_strs(short_strs); + let short_col = VarBinArray::from_strs(short_strs).into_array(); let empty_and_unicode_values = ["", "こんにちは", "😀", "naive", "façade", "résumé", "مرحبا"]; let empty_and_unicode: Vec<&str> = (0..N) .map(|i| empty_and_unicode_values[i % empty_and_unicode_values.len()]) .collect(); - let empty_and_unicode_col = VarBinArray::from_strs(empty_and_unicode); + let empty_and_unicode_col = VarBinArray::from_strs(empty_and_unicode).into_array(); let suffix_shared_values: Vec = (0..N) .map(|i| format!("prefix-{:04}-common-suffix", i % 64)) .collect(); let suffix_shared_refs: Vec<&str> = suffix_shared_values.iter().map(String::as_str).collect(); - let suffix_shared_col = VarBinArray::from_strs(suffix_shared_refs); + let suffix_shared_col = VarBinArray::from_strs(suffix_shared_refs).into_array(); let high_entropy_values: Vec = (0..N) .map(|i| format!("{:016x}{:016x}", i.wrapping_mul(97), i.wrapping_mul(13_579))) .collect(); let high_entropy_refs: Vec<&str> = high_entropy_values.iter().map(String::as_str).collect(); - let high_entropy_col = VarBinArray::from_strs(high_entropy_refs); + let high_entropy_col = VarBinArray::from_strs(high_entropy_refs).into_array(); let all_null_clustered = VarBinArray::from_nullable_strs( (0..N) .map(|i| { @@ -100,18 +100,19 @@ impl FlatLayoutFixture for FsstFixture { } }) .collect::>(), - ); - - let url_comp = fsst_train_compressor(&url_col); - let log_comp = fsst_train_compressor(&log_col); - let nullable_comp = fsst_train_compressor(&nullable_col); - let short_comp = fsst_train_compressor(&short_col); - let empty_and_unicode_comp = fsst_train_compressor(&empty_and_unicode_col); - let suffix_shared_comp = fsst_train_compressor(&suffix_shared_col); - let high_entropy_comp = fsst_train_compressor(&high_entropy_col); - let all_null_clustered_comp = fsst_train_compressor(&all_null_clustered); + ) + .into_array(); let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let url_comp = fsst_train_compressor(&url_col, &mut ctx)?; + let log_comp = fsst_train_compressor(&log_col, &mut ctx)?; + let nullable_comp = fsst_train_compressor(&nullable_col, &mut ctx)?; + let short_comp = fsst_train_compressor(&short_col, &mut ctx)?; + let empty_and_unicode_comp = fsst_train_compressor(&empty_and_unicode_col, &mut ctx)?; + let suffix_shared_comp = fsst_train_compressor(&suffix_shared_col, &mut ctx)?; + let high_entropy_comp = fsst_train_compressor(&high_entropy_col, &mut ctx)?; + let all_null_clustered_comp = fsst_train_compressor(&all_null_clustered, &mut ctx)?; + let arr = StructArray::try_new( FieldNames::from([ "urls", @@ -124,70 +125,16 @@ impl FlatLayoutFixture for FsstFixture { "all_null_clustered", ]), vec![ - fsst_compress( - &url_col, - url_col.len(), - url_col.dtype(), - &url_comp, - &mut ctx, - ) - .into_array(), - fsst_compress( - &log_col, - log_col.len(), - log_col.dtype(), - &log_comp, - &mut ctx, - ) - .into_array(), - fsst_compress( - &nullable_col, - nullable_col.len(), - nullable_col.dtype(), - &nullable_comp, - &mut ctx, - ) - .into_array(), - fsst_compress( - &short_col, - short_col.len(), - short_col.dtype(), - &short_comp, - &mut ctx, - ) - .into_array(), - fsst_compress( - &empty_and_unicode_col, - empty_and_unicode_col.len(), - empty_and_unicode_col.dtype(), - &empty_and_unicode_comp, - &mut ctx, - ) - .into_array(), - fsst_compress( - &suffix_shared_col, - suffix_shared_col.len(), - suffix_shared_col.dtype(), - &suffix_shared_comp, - &mut ctx, - ) - .into_array(), - fsst_compress( - &high_entropy_col, - high_entropy_col.len(), - high_entropy_col.dtype(), - &high_entropy_comp, - &mut ctx, - ) - .into_array(), - fsst_compress( - &all_null_clustered, - all_null_clustered.len(), - all_null_clustered.dtype(), - &all_null_clustered_comp, - &mut ctx, - ) - .into_array(), + fsst_compress(&url_col, &url_comp, &mut ctx)?.into_array(), + fsst_compress(&log_col, &log_comp, &mut ctx)?.into_array(), + fsst_compress(&nullable_col, &nullable_comp, &mut ctx)?.into_array(), + fsst_compress(&short_col, &short_comp, &mut ctx)?.into_array(), + fsst_compress(&empty_and_unicode_col, &empty_and_unicode_comp, &mut ctx)? + .into_array(), + fsst_compress(&suffix_shared_col, &suffix_shared_comp, &mut ctx)?.into_array(), + fsst_compress(&high_entropy_col, &high_entropy_comp, &mut ctx)?.into_array(), + fsst_compress(&all_null_clustered, &all_null_clustered_comp, &mut ctx)? + .into_array(), ], N, Validity::NonNullable, diff --git a/vortex/benches/common_encoding_tree_throughput.rs b/vortex/benches/common_encoding_tree_throughput.rs index cefc9e10723..26ba402a534 100644 --- a/vortex/benches/common_encoding_tree_throughput.rs +++ b/vortex/benches/common_encoding_tree_throughput.rs @@ -239,15 +239,9 @@ mod setup { // Train and compress unique values with FSST let mut ctx = SESSION.create_execution_ctx(); - let unique_varbinview = VarBinViewArray::from_iter_str(unique_strings); - let fsst_compressor = fsst_train_compressor(&unique_varbinview); - let fsst_values = fsst_compress( - &unique_varbinview, - unique_varbinview.len(), - unique_varbinview.dtype(), - &fsst_compressor, - &mut ctx, - ); + let unique_varbinview = VarBinViewArray::from_iter_str(unique_strings).into_array(); + let fsst_compressor = fsst_train_compressor(&unique_varbinview, &mut ctx).unwrap(); + let fsst_values = fsst_compress(&unique_varbinview, &fsst_compressor, &mut ctx).unwrap(); // Create codes array (random indices into unique values) let codes: Vec = (0..NUM_VALUES) @@ -278,15 +272,9 @@ mod setup { // Train and compress unique values with FSST let mut ctx = SESSION.create_execution_ctx(); - let unique_varbinview = VarBinViewArray::from_iter_str(unique_strings); - let fsst_compressor = fsst_train_compressor(&unique_varbinview); - let fsst = fsst_compress( - &unique_varbinview, - unique_varbinview.len(), - unique_varbinview.dtype(), - &fsst_compressor, - &mut ctx, - ); + let unique_varbinview = VarBinViewArray::from_iter_str(unique_strings).into_array(); + let fsst_compressor = fsst_train_compressor(&unique_varbinview, &mut ctx).unwrap(); + let fsst = fsst_compress(&unique_varbinview, &fsst_compressor, &mut ctx).unwrap(); // Compress the VarBin offsets with BitPacked let codes = fsst.codes(); diff --git a/vortex/benches/single_encoding_throughput.rs b/vortex/benches/single_encoding_throughput.rs index 2c2db300d82..3c58297e70b 100644 --- a/vortex/benches/single_encoding_throughput.rs +++ b/vortex/benches/single_encoding_throughput.rs @@ -408,52 +408,49 @@ fn bench_dict_decompress_string(bencher: Bencher) { #[divan::bench(name = "fsst_compress_string")] fn bench_fsst_compress_string(bencher: Bencher) { let varbinview_arr = - VarBinViewArray::from_iter_str(gen_varbin_words(NUM_VALUES as usize, 0.00005)); - let fsst_compressor = fsst_train_compressor(&varbinview_arr); + VarBinViewArray::from_iter_str(gen_varbin_words(NUM_VALUES as usize, 0.00005)).into_array(); + let fsst_compressor = + fsst_train_compressor(&varbinview_arr, &mut SESSION.create_execution_ctx()).unwrap(); let nbytes = varbinview_arr.nbytes() as u64; with_byte_counter(bencher, nbytes) .with_inputs(|| (&varbinview_arr, SESSION.create_execution_ctx())) - .bench_refs(|(a, ctx)| fsst_compress(*a, a.len(), a.dtype(), &fsst_compressor, ctx)); + .bench_refs(|(a, ctx)| fsst_compress(a, &fsst_compressor, ctx).unwrap()); } #[divan::bench(name = "fsst_decompress_string")] fn bench_fsst_decompress_string(bencher: Bencher) { let varbinview_arr = - VarBinViewArray::from_iter_str(gen_varbin_words(NUM_VALUES as usize, 0.00005)); - let fsst_compressor = fsst_train_compressor(&varbinview_arr); - let fsst_array = fsst_compress( - &varbinview_arr, - varbinview_arr.len(), - varbinview_arr.dtype(), - &fsst_compressor, - &mut SESSION.create_execution_ctx(), - ); - let nbytes = varbinview_arr.into_array().nbytes() as u64; + VarBinViewArray::from_iter_str(gen_varbin_words(NUM_VALUES as usize, 0.00005)).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let fsst_compressor = fsst_train_compressor(&varbinview_arr, &mut ctx).unwrap(); + let fsst_array = fsst_compress(&varbinview_arr, &fsst_compressor, &mut ctx) + .unwrap() + .into_array(); + let nbytes = varbinview_arr.nbytes() as u64; with_byte_counter(bencher, nbytes) .with_inputs(|| (&fsst_array, SESSION.create_execution_ctx())) - .bench_refs(|(a, ctx)| canonicalize((**a).clone(), ctx)); + .bench_refs(|(a, ctx)| canonicalize(a.clone(), ctx)); } #[cfg(feature = "zstd")] #[divan::bench(name = "zstd_compress_string")] fn bench_zstd_compress_string(bencher: Bencher) { let varbinview_arr = - VarBinViewArray::from_iter_str(gen_varbin_words(NUM_VALUES as usize, 0.00005)); + VarBinViewArray::from_iter_str(gen_varbin_words(NUM_VALUES as usize, 0.00005)).into_array(); let nbytes = varbinview_arr.nbytes() as u64; - let array = varbinview_arr.into_array(); with_byte_counter(bencher, nbytes) - .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) - .bench_values(|(a, mut ctx)| ZstdData::from_array(a, 3, 8192, &mut ctx).unwrap()); + .with_inputs(|| (&varbinview_arr, SESSION.create_execution_ctx())) + .bench_refs(|(a, ctx)| ZstdData::from_array(a.clone(), 3, 8192, ctx).unwrap()); } #[cfg(feature = "zstd")] #[divan::bench(name = "zstd_decompress_string")] fn bench_zstd_decompress_string(bencher: Bencher) { let varbinview_arr = - VarBinViewArray::from_iter_str(gen_varbin_words(NUM_VALUES as usize, 0.00005)); + VarBinViewArray::from_iter_str(gen_varbin_words(NUM_VALUES as usize, 0.00005)).into_array(); let dtype = varbinview_arr.dtype().clone(); let validity = varbinview_arr.validity().unwrap(); let compressed = Zstd::try_new( @@ -469,9 +466,9 @@ fn bench_zstd_decompress_string(bencher: Bencher) { ) .unwrap() .into_array(); - let nbytes = varbinview_arr.into_array().nbytes() as u64; + let nbytes = varbinview_arr.nbytes() as u64; with_byte_counter(bencher, nbytes) .with_inputs(|| (&compressed, SESSION.create_execution_ctx())) - .bench_refs(|(a, ctx)| canonicalize((**a).clone(), ctx)); + .bench_refs(|(a, ctx)| canonicalize(a.clone(), ctx)); }