diff --git a/encodings/fsst/src/compute/mod.rs b/encodings/fsst/src/compute/mod.rs index 3b46768743b..f606072488c 100644 --- a/encodings/fsst/src/compute/mod.rs +++ b/encodings/fsst/src/compute/mod.rs @@ -7,24 +7,36 @@ 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::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; 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; +/// 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. @@ -36,11 +48,20 @@ impl TakeExecute for FSST { indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - if array.len() < MIN_GATHER_LEN - || indices.len().saturating_mul(SPARSE_TAKE_DENOMINATOR) >= array.len() - { + if indices.len().saturating_mul(SPARSE_TAKE_DENOMINATOR) >= array.len() { 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( FSST::try_new_with_symbol_table( @@ -63,21 +84,102 @@ 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 { + 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::ArrayRef; 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_array::validity::Validity; + 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 +230,125 @@ 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(()) + } + + #[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( + 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]