From 9c5df2c35a185c4f4afd44be8b89cb4b4f73f07b Mon Sep 17 00:00:00 2001 From: blaginin Date: Thu, 15 Oct 2026 11:58:32 +0100 Subject: [PATCH 1/4] perf(fsst): gather sparse takes from small arrays with large payloads Signed-off-by: blaginin --- encodings/fsst/src/compute/mod.rs | 128 +++++++++++++++++++++++++++++- 1 file changed, 126 insertions(+), 2 deletions(-) diff --git a/encodings/fsst/src/compute/mod.rs b/encodings/fsst/src/compute/mod.rs index 3b46768743b..e9c57157840 100644 --- a/encodings/fsst/src/compute/mod.rs +++ b/encodings/fsst/src/compute/mod.rs @@ -7,13 +7,18 @@ mod compare; mod filter; mod like; +use num_traits::ToPrimitive; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::arrays::Primitive; use vortex_array::arrays::dict::TakeExecute; use vortex_array::arrays::varbin::take_varbin; use vortex_array::builtins::ArrayBuiltins; +use vortex_array::expr::stats::Stat; +use vortex_array::expr::stats::StatsProviderExt; +use vortex_array::match_each_integer_ptype; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; @@ -21,8 +26,9 @@ use crate::FSST; use crate::FSSTArrayExt; use crate::FSSTArraySlotsExt; -/// Below this many rows, decoding all of them costs less than gathering compressed rows. +/// Small arrays gather only when their compressed payload is large enough to amortize setup. const MIN_GATHER_LEN: usize = 1024; +const MIN_GATHER_COMPRESSED_BYTES: u64 = 16 * 1024; /// A take with fewer indices than `1 / SPARSE_TAKE_DENOMINATOR` of the rows gathers compressed /// rows. A denser take decodes every row once and gathers the decoded strings instead: gathering @@ -36,8 +42,12 @@ impl TakeExecute for FSST { indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { + if indices.len().saturating_mul(SPARSE_TAKE_DENOMINATOR) >= array.len() { + return Ok(None); + } if array.len() < MIN_GATHER_LEN - || indices.len().saturating_mul(SPARSE_TAKE_DENOMINATOR) >= array.len() + && !compressed_byte_span(array.codes_offsets()) + .is_some_and(|bytes| bytes >= MIN_GATHER_COMPRESSED_BYTES) { return Ok(None); } @@ -63,21 +73,50 @@ impl TakeExecute for FSST { } } +// Slices retain their backing bytes, so only the logical offsets describe their decode cost. +// Inspect already-materialized offsets or cached exact bounds without executing either child. +fn compressed_byte_span(offsets: &ArrayRef) -> Option { + if let Some(offsets) = offsets.as_opt::() + && offsets.buffer_handle().as_host_opt().is_some() + { + return match_each_integer_ptype!(offsets.ptype(), |P| { + let values = offsets.as_slice::

(); + let first = values.first()?.to_u64()?; + let last = values.last()?.to_u64()?; + last.checked_sub(first) + }); + } + let stats = offsets.statistics(); + let first = stats.get_as::(Stat::Min).as_exact()?; + let last = stats.get_as::(Stat::Max).as_exact()?; + last.checked_sub(first) +} + #[cfg(test)] mod tests { + use fsst::Compressor; + use fsst::Symbol; use rstest::rstest; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; + use vortex_array::arrays::DictArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinArray; + use vortex_array::arrays::dict::TakeExecute; + use vortex_array::assert_arrays_eq; use vortex_array::compute::conformance::consistency::test_array_consistency; use vortex_array::compute::conformance::take::test_take_conformance; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; + use vortex_array::expr::stats::Precision; + use vortex_array::expr::stats::Stat; + use vortex_buffer::buffer; use vortex_error::VortexResult; + use super::compressed_byte_span; + use crate::FSST; use crate::FSSTArray; use crate::fsst_compress; use crate::fsst_train_compressor; @@ -128,6 +167,91 @@ mod tests { Ok(()) } + #[rstest] + #[case(2048, 1, false, true)] + #[case(2048, 3, true, true)] + #[case(2048, 4, false, false)] + #[case(2, 1, false, false)] + fn test_take_small_byte_gate( + #[case] repeats: usize, + #[case] selected: usize, + #[case] nullable: bool, + #[case] gathers: bool, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let input = VarBinArray::from_iter( + (0..16).map(|row| { + (!(nullable && row == 0)).then(|| format!("{}{row}", "caption ".repeat(repeats))) + }), + DType::Utf8(Nullability::Nullable), + ) + .into_array(); + let compressor = Compressor::rebuild_from([Symbol::from_slice(b"caption ")], [8u8]); + let array = fsst_compress(&input, &compressor, &mut ctx)?; + let indices = if nullable { + PrimitiveArray::from_option_iter([Some(15u64), None, Some(0)]).into_array() + } else { + PrimitiveArray::from_iter((0..selected).map(|index| (15 - index) as u64)).into_array() + }; + assert_eq!( + FSST::take(array.as_view(), &indices, &mut ctx)?.is_some(), + gathers + ); + assert_arrays_eq!( + array.into_array().take(indices.clone())?, + input.take(indices)?, + &mut ctx + ); + Ok(()) + } + + #[test] + fn test_take_small_slice_uses_logical_byte_span() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let input = VarBinArray::from_iter( + (0..16).map(|row| Some("caption ".repeat(if row < 4 { 4096 } else { 2 }))), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + let compressor = Compressor::rebuild_from([Symbol::from_slice(b"caption ")], [8u8]); + let array = fsst_compress(&input, &compressor, &mut ctx)?.into_array(); + let indices = buffer![0u64].into_array(); + for (range, gathers) in [(0..8, true), (4..16, false)] { + let sliced = array.slice(range.clone())?; + assert_eq!( + FSST::take(sliced.as_::(), &indices, &mut ctx)?.is_some(), + gathers + ); + assert_arrays_eq!( + sliced.take(indices.clone())?, + input.slice(range)?.take(indices.clone())?, + &mut ctx + ); + } + Ok(()) + } + + #[test] + fn test_take_compressed_span_requires_exact_cached_bounds() { + let offsets = DictArray::new( + buffer![0u32, 1].into_array(), + buffer![1000u64, 20000].into_array(), + ) + .into_array(); + assert_eq!(compressed_byte_span(&offsets), None); + offsets + .statistics() + .set(Stat::Min, Precision::exact(1000u64)); + offsets + .statistics() + .set(Stat::Max, Precision::exact(20000u64)); + assert_eq!(compressed_byte_span(&offsets), Some(19000)); + offsets + .statistics() + .set(Stat::Min, Precision::Inexact(1000u64.into())); + assert_eq!(compressed_byte_span(&offsets), None); + } + type FsstBuilder = fn(&mut ExecutionCtx) -> FSSTArray; #[rstest] From 7180b420fe3d5dbcdd40d25f590339d594605eea Mon Sep 17 00:00:00 2001 From: blaginin Date: Thu, 15 Oct 2026 12:28:46 +0100 Subject: [PATCH 2/4] perf(fsst): gather small-array takes from 4 KiB of codes and add a take benchmark The `fsst_take` benchmark sweeps arrays below 1024 rows from 1 to 64 KiB of codes, taking one row or the most the sparse-take gate allows. Gathering stops losing to decoding every row between 3 and 4 KiB, even for the densest takes, so the gate drops from 16 KiB to 4 KiB. Signed-off-by: blaginin Co-Authored-By: Claude Opus 5.5 --- encodings/fsst/Cargo.toml | 5 + encodings/fsst/benches/fsst_take.rs | 144 ++++++++++++++++++++++++++++ encodings/fsst/src/compute/mod.rs | 4 +- 3 files changed, 152 insertions(+), 1 deletion(-) create mode 100644 encodings/fsst/benches/fsst_take.rs diff --git a/encodings/fsst/Cargo.toml b/encodings/fsst/Cargo.toml index d08ff59db61..5b15dca85fb 100644 --- a/encodings/fsst/Cargo.toml +++ b/encodings/fsst/Cargo.toml @@ -56,5 +56,10 @@ name = "chunked_dict_fsst_builder" harness = false required-features = ["_test-harness"] +[[bench]] +name = "fsst_take" +harness = false +required-features = ["_test-harness"] + [package.metadata.cargo-shear] ignored = ["fsst-rs"] diff --git a/encodings/fsst/benches/fsst_take.rs b/encodings/fsst/benches/fsst_take.rs new file mode 100644 index 00000000000..9a91e001eea --- /dev/null +++ b/encodings/fsst/benches/fsst_take.rs @@ -0,0 +1,144 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect(clippy::unwrap_used)] +#![expect(clippy::cast_possible_truncation)] + +use std::sync::LazyLock; + +use divan::Bencher; +use mimalloc::MiMalloc; +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::VortexSessionExecute; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::VarBinArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_fsst::FSSTArray; +use vortex_fsst::fsst_compress; +use vortex_fsst::fsst_train_compressor; +use vortex_session::VortexSession; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +fn main() { + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fsst::initialize(&session); + session +}); + +// [(rows, minimum compressed KiB, taken rows)]. Every fixture has fewer than 1024 rows and takes +// one row or the most the sparse-take gate allows, so only the compressed-bytes gate chooses +// between gathering the taken rows and decoding every row. +const BENCH_ARGS: &[(usize, usize, usize)] = &[ + (16, 1, 1), + (16, 2, 1), + (16, 4, 1), + (16, 8, 1), + (16, 16, 1), + (16, 32, 1), + (16, 64, 1), + (16, 1, 3), + (16, 2, 3), + (16, 4, 3), + (16, 8, 3), + (16, 16, 3), + (16, 32, 3), + (16, 64, 3), + (256, 1, 1), + (256, 2, 1), + (256, 4, 1), + (256, 8, 1), + (256, 16, 1), + (256, 32, 1), + (256, 64, 1), + (256, 1, 63), + (256, 2, 63), + (256, 4, 63), + (256, 8, 63), + (256, 16, 63), + (256, 32, 63), + (256, 64, 63), + (1000, 1, 1), + (1000, 2, 1), + (1000, 4, 1), + (1000, 8, 1), + (1000, 16, 1), + (1000, 32, 1), + (1000, 64, 1), + (1000, 1, 249), + (1000, 2, 249), + (1000, 4, 249), + (1000, 8, 249), + (1000, 16, 249), + (1000, 32, 249), + (1000, 64, 249), +]; + +const WORDS: &[&str] = &[ + "the", "a", "photo", "of", "with", "and", "on", "in", "small", "large", "red", "blue", "green", + "white", "black", "dog", "cat", "man", "woman", "child", "table", "street", "car", "tree", + "house", "sitting", "standing", "walking", "next", "to", "near", "under", "over", "wooden", + "old", "new", "city", "park", "window", "light", +]; + +fn compressed(rows: usize, row_len: usize) -> FSSTArray { + let mut rng = StdRng::seed_from_u64(0); + let input = VarBinArray::from_iter( + (0..rows).map(|_| { + let mut row = String::with_capacity(row_len + 16); + while row.len() < row_len { + row.push_str(WORDS[rng.random_range(0..WORDS.len())]); + row.push(' '); + } + row.truncate(row_len); + Some(row) + }), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let compressor = fsst_train_compressor(&input, &mut ctx).unwrap(); + fsst_compress(&input, &compressor, &mut ctx).unwrap() +} + +fn fixture((rows, compressed_kib, taken): (usize, usize, usize)) -> (ArrayRef, ArrayRef) { + // Grow the rows until their codes reach the target, so a fixture never lands just below a + // threshold it is named after. + let target = compressed_kib * 1024; + let mut row_len = target / rows; + let mut array = compressed(rows, row_len); + while array.codes_bytes().len() < target { + row_len += row_len.div_ceil(50); + array = compressed(rows, row_len); + } + + let stride = rows / taken; + let indices = PrimitiveArray::from_iter((0..taken).map(|i| (i * stride + stride / 2) as u32)); + (array.into_array(), indices.into_array()) +} + +#[divan::bench(args = BENCH_ARGS)] +fn take(bencher: Bencher, args: (usize, usize, usize)) { + let (array, indices) = fixture(args); + + bencher + .with_inputs(|| (&array, &indices, SESSION.create_execution_ctx())) + .bench_refs(|(array, indices, ctx)| { + array + .take((*indices).clone()) + .unwrap() + .execute::(ctx) + .unwrap() + }) +} diff --git a/encodings/fsst/src/compute/mod.rs b/encodings/fsst/src/compute/mod.rs index e9c57157840..8738770f6e6 100644 --- a/encodings/fsst/src/compute/mod.rs +++ b/encodings/fsst/src/compute/mod.rs @@ -28,7 +28,9 @@ use crate::FSSTArraySlotsExt; /// Small arrays gather only when their compressed payload is large enough to amortize setup. const MIN_GATHER_LEN: usize = 1024; -const MIN_GATHER_COMPRESSED_BYTES: u64 = 16 * 1024; +/// From about this many bytes of codes, gathering beats decoding every row of a small array, even +/// for the densest take [`SPARSE_TAKE_DENOMINATOR`] allows. Measured by the `fsst_take` benchmark. +const MIN_GATHER_COMPRESSED_BYTES: u64 = 4 * 1024; /// A take with fewer indices than `1 / SPARSE_TAKE_DENOMINATOR` of the rows gathers compressed /// rows. A denser take decodes every row once and gathers the decoded strings instead: gathering From 9613a4826983ecf16cf01b2661203f6761458d0f Mon Sep 17 00:00:00 2001 From: blaginin Date: Thu, 15 Oct 2026 16:21:18 +0100 Subject: [PATCH 3/4] fix(fsst): bound compressed bytes for small sparse takes Signed-off-by: blaginin --- encodings/fsst/benches/fsst_take.rs | 89 ++++++++++++++++++++++ encodings/fsst/src/compute/mod.rs | 113 +++++++++++++++++++++++++--- 2 files changed, 193 insertions(+), 9 deletions(-) diff --git a/encodings/fsst/benches/fsst_take.rs b/encodings/fsst/benches/fsst_take.rs index 9a91e001eea..321c165d783 100644 --- a/encodings/fsst/benches/fsst_take.rs +++ b/encodings/fsst/benches/fsst_take.rs @@ -17,6 +17,7 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinArray; +use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_fsst::FSSTArray; @@ -131,7 +132,21 @@ fn fixture((rows, compressed_kib, taken): (usize, usize, usize)) -> (ArrayRef, A #[divan::bench(args = BENCH_ARGS)] fn take(bencher: Bencher, args: (usize, usize, usize)) { let (array, indices) = fixture(args); + bench_take(bencher, array, indices); +} +fn bench_take(bencher: Bencher, array: ArrayRef, indices: ArrayRef) { + let mut ctx = SESSION.create_execution_ctx(); + let canonical = array + .clone() + .execute::(&mut ctx) + .unwrap() + .into_array(); + assert_arrays_eq!( + array.take(indices.clone()).unwrap(), + canonical.take(indices.clone()).unwrap(), + &mut ctx + ); bencher .with_inputs(|| (&array, &indices, SESSION.create_execution_ctx())) .bench_refs(|(array, indices, ctx)| { @@ -142,3 +157,77 @@ fn take(bencher: Bencher, args: (usize, usize, usize)) { .unwrap() }) } + +#[derive(Clone, Copy, Debug)] +enum Selection { + Short, + AllLong, + RepeatLong, + HalfLong, + MajorityLong, + RepeatHalfLong, + RepeatMajorityLong, +} + +// (long rows among 256, long string length, selection). Short strings contain 8 bytes. +const UNEVEN_ARGS: &[(usize, usize, Selection)] = &[ + (1, 1024, Selection::Short), + (1, 1024, Selection::AllLong), + (1, 1024, Selection::RepeatLong), + (1, 16384, Selection::Short), + (1, 16384, Selection::AllLong), + (1, 16384, Selection::RepeatLong), + (1, 32768, Selection::Short), + (1, 32768, Selection::AllLong), + (1, 32768, Selection::RepeatLong), + (16, 1024, Selection::Short), + (16, 1024, Selection::AllLong), + (16, 1024, Selection::RepeatLong), + (16, 4096, Selection::Short), + (16, 4096, Selection::AllLong), + (16, 4096, Selection::RepeatLong), + (16, 16384, Selection::Short), + (16, 16384, Selection::AllLong), + (16, 16384, Selection::RepeatLong), + (16, 16384, Selection::HalfLong), + (16, 16384, Selection::MajorityLong), + (16, 16384, Selection::RepeatHalfLong), + (16, 16384, Selection::RepeatMajorityLong), +]; + +#[divan::bench(args = UNEVEN_ARGS)] +fn uneven_take(bencher: Bencher, args: (usize, usize, Selection)) { + let (long_rows, long_len, mode) = args; + let mut rng = StdRng::seed_from_u64(0); + let input = VarBinArray::from_iter( + (0..256).map(|i| { + let len = if i < long_rows { long_len } else { 8 }; + let mut value = String::new(); + while value.len() < len { + value.push_str(WORDS[rng.random_range(0..WORDS.len())]); + value.push(' '); + } + value.truncate(len); + Some(value) + }), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let compressor = fsst_train_compressor(&input, &mut ctx).unwrap(); + let array = fsst_compress(&input, &compressor, &mut ctx).unwrap(); + let selected: Vec = match mode { + Selection::Short => vec![long_rows as u32], + Selection::AllLong => (0..long_rows as u32).collect(), + Selection::RepeatLong => vec![0; 63], + Selection::HalfLong => (0..(long_rows / 2) as u32).collect(), + Selection::MajorityLong => (0..(long_rows / 2 + 1) as u32).collect(), + Selection::RepeatHalfLong => vec![0; long_rows / 2], + Selection::RepeatMajorityLong => vec![0; long_rows / 2 + 1], + }; + bench_take( + bencher, + array.into_array(), + PrimitiveArray::from_iter(selected).into_array(), + ); +} diff --git a/encodings/fsst/src/compute/mod.rs b/encodings/fsst/src/compute/mod.rs index 8738770f6e6..f606072488c 100644 --- a/encodings/fsst/src/compute/mod.rs +++ b/encodings/fsst/src/compute/mod.rs @@ -13,14 +13,19 @@ use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::dict::TakeExecute; +use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::arrays::varbin::take_varbin; use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::IntegerPType; use vortex_array::expr::stats::Stat; use vortex_array::expr::stats::StatsProviderExt; use vortex_array::match_each_integer_ptype; +use vortex_array::match_each_unsigned_integer_ptype; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; +use vortex_mask::AllOr; use crate::FSST; use crate::FSSTArrayExt; @@ -28,11 +33,10 @@ use crate::FSSTArraySlotsExt; /// Small arrays gather only when their compressed payload is large enough to amortize setup. const MIN_GATHER_LEN: usize = 1024; -/// From about this many bytes of codes, gathering beats decoding every row of a small array, even -/// for the densest take [`SPARSE_TAKE_DENOMINATOR`] allows. Measured by the `fsst_take` benchmark. -const MIN_GATHER_COMPRESSED_BYTES: u64 = 4 * 1024; +/// Minimum compressed payload for gathering from a small array, measured by `fsst_take`. +const MIN_GATHER_COMPRESSED_BYTES: u64 = 8 * 1024; -/// A take with fewer indices than `1 / SPARSE_TAKE_DENOMINATOR` of the rows gathers compressed +/// Only takes with fewer indices than `1 / SPARSE_TAKE_DENOMINATOR` of the rows gather compressed /// rows. A denser take decodes every row once and gathers the decoded strings instead: gathering /// decodes a row again for every index that repeats it, and costs more per row than decoding /// in bulk. @@ -47,11 +51,16 @@ impl TakeExecute for FSST { if indices.len().saturating_mul(SPARSE_TAKE_DENOMINATOR) >= array.len() { return Ok(None); } - if array.len() < MIN_GATHER_LEN - && !compressed_byte_span(array.codes_offsets()) - .is_some_and(|bytes| bytes >= MIN_GATHER_COMPRESSED_BYTES) - { - return Ok(None); + if array.len() < MIN_GATHER_LEN { + let Some(bytes) = compressed_byte_span(array.codes_offsets()) + .filter(|&bytes| bytes >= MIN_GATHER_COMPRESSED_BYTES) + else { + return Ok(None); + }; + // Leave at least half the codes untouched to amortize gathering and the cost check. + if !take_fits_byte_budget(array.codes_offsets(), indices, bytes / 2, ctx)? { + return Ok(None); + } } Ok(Some( @@ -75,6 +84,56 @@ impl TakeExecute for FSST { } } +fn take_fits_byte_budget( + offsets: &ArrayRef, + indices: &ArrayRef, + budget: u64, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let offsets = offsets.clone().execute::(ctx)?; + let indices = indices.clone().execute::(ctx)?; + let validity = indices.validity()?.execute_mask(indices.len(), ctx)?; + // Dispatch only on widths: offsets and valid indices are nonnegative. + let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned()); + let indices = indices.reinterpret_cast(indices.ptype().to_unsigned()); + Ok(match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { + let offsets = offsets.as_slice::(); + match_each_unsigned_integer_ptype!(indices.ptype(), |I| { + let indices = indices.as_slice::(); + match validity.indices() { + AllOr::All => selected_bytes_fit(offsets, indices.iter().copied(), budget), + AllOr::None => true, + AllOr::Some(valid) => { + selected_bytes_fit(offsets, valid.iter().map(|&i| indices[i]), budget) + } + } + }) + })) +} + +// Count repeats before copying codes: gathering decodes each selected occurrence separately. +fn selected_bytes_fit( + offsets: &[O], + mut indices: impl Iterator, + mut budget: u64, +) -> bool { + indices.all(|index| { + let Some(index) = index.to_usize().filter(|&i| i < offsets.len() - 1) else { + return false; + }; + let remaining = offsets[index + 1] + .to_u64() + .zip(offsets[index].to_u64()) + .and_then(|(end, start)| end.checked_sub(start)) + .and_then(|len| budget.checked_sub(len)); + let Some(remaining) = remaining else { + return false; + }; + budget = remaining; + true + }) +} + // Slices retain their backing bytes, so only the logical offsets describe their decode cost. // Inspect already-materialized offsets or cached exact bounds without executing either child. fn compressed_byte_span(offsets: &ArrayRef) -> Option { @@ -99,6 +158,7 @@ mod tests { use fsst::Compressor; use fsst::Symbol; use rstest::rstest; + use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; @@ -114,6 +174,7 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::expr::stats::Precision; use vortex_array::expr::stats::Stat; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; @@ -233,6 +294,40 @@ mod tests { Ok(()) } + #[rstest] + #[case::repeated_long_row(1, PrimitiveArray::from_iter([0u64; 63]).into_array(), false)] + #[case::distinct_long_rows(16, PrimitiveArray::from_iter(0i32..16).into_array(), false)] + #[case::short_row_with_null_index(1, PrimitiveArray::new( + buffer![u64::MAX, 1], Validity::from_iter([false, true]), + ).into_array(), true)] + #[case::encoded_indices(1, DictArray::new( + buffer![0u8, 0].into_array(), buffer![1u32].into_array(), + ).into_array(), true)] + fn test_take_small_uneven_lengths( + #[case] long_rows: usize, + #[case] indices: ArrayRef, + #[case] gathers: bool, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let input = VarBinArray::from_iter( + (0..256).map(|row| Some("caption ".repeat(if row < long_rows { 8192 } else { 1 }))), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + let compressor = Compressor::rebuild_from([Symbol::from_slice(b"caption ")], [8u8]); + let array = fsst_compress(&input, &compressor, &mut ctx)?; + assert_eq!( + FSST::take(array.as_view(), &indices, &mut ctx)?.is_some(), + gathers + ); + assert_arrays_eq!( + array.into_array().take(indices.clone())?, + input.take(indices)?, + &mut ctx + ); + Ok(()) + } + #[test] fn test_take_compressed_span_requires_exact_cached_bounds() { let offsets = DictArray::new( From 8c628e83e68974818f3931d6bc11803433a1340b Mon Sep 17 00:00:00 2001 From: blaginin Date: Fri, 16 Oct 2026 15:02:49 +0100 Subject: [PATCH 4/4] bench(fsst): cover sparse takes by payload size and repeated bytes Signed-off-by: blaginin --- encodings/fsst/Cargo.toml | 5 + encodings/fsst/benches/fsst_take.rs | 233 ++++++++++++++++++++++++++++ 2 files changed, 238 insertions(+) create mode 100644 encodings/fsst/benches/fsst_take.rs diff --git a/encodings/fsst/Cargo.toml b/encodings/fsst/Cargo.toml index d08ff59db61..5b15dca85fb 100644 --- a/encodings/fsst/Cargo.toml +++ b/encodings/fsst/Cargo.toml @@ -56,5 +56,10 @@ name = "chunked_dict_fsst_builder" harness = false required-features = ["_test-harness"] +[[bench]] +name = "fsst_take" +harness = false +required-features = ["_test-harness"] + [package.metadata.cargo-shear] ignored = ["fsst-rs"] diff --git a/encodings/fsst/benches/fsst_take.rs b/encodings/fsst/benches/fsst_take.rs new file mode 100644 index 00000000000..321c165d783 --- /dev/null +++ b/encodings/fsst/benches/fsst_take.rs @@ -0,0 +1,233 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect(clippy::unwrap_used)] +#![expect(clippy::cast_possible_truncation)] + +use std::sync::LazyLock; + +use divan::Bencher; +use mimalloc::MiMalloc; +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::VortexSessionExecute; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::VarBinArray; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_fsst::FSSTArray; +use vortex_fsst::fsst_compress; +use vortex_fsst::fsst_train_compressor; +use vortex_session::VortexSession; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +fn main() { + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fsst::initialize(&session); + session +}); + +// [(rows, minimum compressed KiB, taken rows)]. Every fixture has fewer than 1024 rows and takes +// one row or the most the sparse-take gate allows, so only the compressed-bytes gate chooses +// between gathering the taken rows and decoding every row. +const BENCH_ARGS: &[(usize, usize, usize)] = &[ + (16, 1, 1), + (16, 2, 1), + (16, 4, 1), + (16, 8, 1), + (16, 16, 1), + (16, 32, 1), + (16, 64, 1), + (16, 1, 3), + (16, 2, 3), + (16, 4, 3), + (16, 8, 3), + (16, 16, 3), + (16, 32, 3), + (16, 64, 3), + (256, 1, 1), + (256, 2, 1), + (256, 4, 1), + (256, 8, 1), + (256, 16, 1), + (256, 32, 1), + (256, 64, 1), + (256, 1, 63), + (256, 2, 63), + (256, 4, 63), + (256, 8, 63), + (256, 16, 63), + (256, 32, 63), + (256, 64, 63), + (1000, 1, 1), + (1000, 2, 1), + (1000, 4, 1), + (1000, 8, 1), + (1000, 16, 1), + (1000, 32, 1), + (1000, 64, 1), + (1000, 1, 249), + (1000, 2, 249), + (1000, 4, 249), + (1000, 8, 249), + (1000, 16, 249), + (1000, 32, 249), + (1000, 64, 249), +]; + +const WORDS: &[&str] = &[ + "the", "a", "photo", "of", "with", "and", "on", "in", "small", "large", "red", "blue", "green", + "white", "black", "dog", "cat", "man", "woman", "child", "table", "street", "car", "tree", + "house", "sitting", "standing", "walking", "next", "to", "near", "under", "over", "wooden", + "old", "new", "city", "park", "window", "light", +]; + +fn compressed(rows: usize, row_len: usize) -> FSSTArray { + let mut rng = StdRng::seed_from_u64(0); + let input = VarBinArray::from_iter( + (0..rows).map(|_| { + let mut row = String::with_capacity(row_len + 16); + while row.len() < row_len { + row.push_str(WORDS[rng.random_range(0..WORDS.len())]); + row.push(' '); + } + row.truncate(row_len); + Some(row) + }), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let compressor = fsst_train_compressor(&input, &mut ctx).unwrap(); + fsst_compress(&input, &compressor, &mut ctx).unwrap() +} + +fn fixture((rows, compressed_kib, taken): (usize, usize, usize)) -> (ArrayRef, ArrayRef) { + // Grow the rows until their codes reach the target, so a fixture never lands just below a + // threshold it is named after. + let target = compressed_kib * 1024; + let mut row_len = target / rows; + let mut array = compressed(rows, row_len); + while array.codes_bytes().len() < target { + row_len += row_len.div_ceil(50); + array = compressed(rows, row_len); + } + + let stride = rows / taken; + let indices = PrimitiveArray::from_iter((0..taken).map(|i| (i * stride + stride / 2) as u32)); + (array.into_array(), indices.into_array()) +} + +#[divan::bench(args = BENCH_ARGS)] +fn take(bencher: Bencher, args: (usize, usize, usize)) { + let (array, indices) = fixture(args); + bench_take(bencher, array, indices); +} + +fn bench_take(bencher: Bencher, array: ArrayRef, indices: ArrayRef) { + let mut ctx = SESSION.create_execution_ctx(); + let canonical = array + .clone() + .execute::(&mut ctx) + .unwrap() + .into_array(); + assert_arrays_eq!( + array.take(indices.clone()).unwrap(), + canonical.take(indices.clone()).unwrap(), + &mut ctx + ); + bencher + .with_inputs(|| (&array, &indices, SESSION.create_execution_ctx())) + .bench_refs(|(array, indices, ctx)| { + array + .take((*indices).clone()) + .unwrap() + .execute::(ctx) + .unwrap() + }) +} + +#[derive(Clone, Copy, Debug)] +enum Selection { + Short, + AllLong, + RepeatLong, + HalfLong, + MajorityLong, + RepeatHalfLong, + RepeatMajorityLong, +} + +// (long rows among 256, long string length, selection). Short strings contain 8 bytes. +const UNEVEN_ARGS: &[(usize, usize, Selection)] = &[ + (1, 1024, Selection::Short), + (1, 1024, Selection::AllLong), + (1, 1024, Selection::RepeatLong), + (1, 16384, Selection::Short), + (1, 16384, Selection::AllLong), + (1, 16384, Selection::RepeatLong), + (1, 32768, Selection::Short), + (1, 32768, Selection::AllLong), + (1, 32768, Selection::RepeatLong), + (16, 1024, Selection::Short), + (16, 1024, Selection::AllLong), + (16, 1024, Selection::RepeatLong), + (16, 4096, Selection::Short), + (16, 4096, Selection::AllLong), + (16, 4096, Selection::RepeatLong), + (16, 16384, Selection::Short), + (16, 16384, Selection::AllLong), + (16, 16384, Selection::RepeatLong), + (16, 16384, Selection::HalfLong), + (16, 16384, Selection::MajorityLong), + (16, 16384, Selection::RepeatHalfLong), + (16, 16384, Selection::RepeatMajorityLong), +]; + +#[divan::bench(args = UNEVEN_ARGS)] +fn uneven_take(bencher: Bencher, args: (usize, usize, Selection)) { + let (long_rows, long_len, mode) = args; + let mut rng = StdRng::seed_from_u64(0); + let input = VarBinArray::from_iter( + (0..256).map(|i| { + let len = if i < long_rows { long_len } else { 8 }; + let mut value = String::new(); + while value.len() < len { + value.push_str(WORDS[rng.random_range(0..WORDS.len())]); + value.push(' '); + } + value.truncate(len); + Some(value) + }), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let compressor = fsst_train_compressor(&input, &mut ctx).unwrap(); + let array = fsst_compress(&input, &compressor, &mut ctx).unwrap(); + let selected: Vec = match mode { + Selection::Short => vec![long_rows as u32], + Selection::AllLong => (0..long_rows as u32).collect(), + Selection::RepeatLong => vec![0; 63], + Selection::HalfLong => (0..(long_rows / 2) as u32).collect(), + Selection::MajorityLong => (0..(long_rows / 2 + 1) as u32).collect(), + Selection::RepeatHalfLong => vec![0; long_rows / 2], + Selection::RepeatMajorityLong => vec![0; long_rows / 2 + 1], + }; + bench_take( + bencher, + array.into_array(), + PrimitiveArray::from_iter(selected).into_array(), + ); +}