diff --git a/Cargo.lock b/Cargo.lock index 2448cb9aff8..a913204d928 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4042,9 +4042,9 @@ dependencies = [ [[package]] name = "fastlanes" -version = "0.7.0" +version = "0.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34f6c951d711d8a10f08524071f6171bc95c5f857e8b39d91e36ad6353fc4f4f" +checksum = "9e218082bba8aee4ba5704355ea1ce1f742facbcf748eadbca90de56aa551cd9" dependencies = [ "const_for", "num-traits", @@ -10426,7 +10426,7 @@ dependencies = [ "anyhow", "arrow-array 59.2.0", "codspeed-divan-compat", - "fastlanes 0.7.0", + "fastlanes 0.7.1", "mimalloc", "parquet 59.2.0", "rand 0.10.2", @@ -10826,7 +10826,7 @@ dependencies = [ "bindgen", "codspeed-criterion-compat-walltime", "cudarc", - "fastlanes 0.7.0", + "fastlanes 0.7.1", "futures", "itertools 0.14.0", "kanal", @@ -11000,7 +11000,7 @@ name = "vortex-fastlanes" version = "0.1.0" dependencies = [ "codspeed-divan-compat", - "fastlanes 0.7.0", + "fastlanes 0.7.1", "itertools 0.14.0", "lending-iterator", "num-traits", @@ -11009,6 +11009,7 @@ dependencies = [ "rstest", "vortex-alp", "vortex-array", + "vortex-bench-support", "vortex-buffer", "vortex-error", "vortex-fastlanes", diff --git a/Cargo.toml b/Cargo.toml index ae164db13c2..906dcc203ac 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -160,7 +160,7 @@ datafusion-sqllogictest = { version = "55.0.0" } divan = { package = "codspeed-divan-compat", version = "5.0.0" } enum-iterator = "2.0.0" env_logger = "0.11" -fastlanes = "0.7.0" +fastlanes = "0.7.1" flatbuffers = "25.2.10" fsst-rs = "0.6.0" futures = { version = "0.3.31", default-features = false } diff --git a/encodings/fastlanes/Cargo.toml b/encodings/fastlanes/Cargo.toml index 9085390b67b..dfe2fb7ea70 100644 --- a/encodings/fastlanes/Cargo.toml +++ b/encodings/fastlanes/Cargo.toml @@ -39,6 +39,7 @@ rand = { workspace = true } rstest = { workspace = true } vortex-alp = { path = "../alp" } vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-bench-support = { workspace = true } vortex-fastlanes = { path = ".", features = ["_test-harness"] } [features] @@ -48,6 +49,10 @@ _test-harness = ["dep:rand"] name = "bitpacking_take" harness = false +[[bench]] +name = "bitpacking_filter" +harness = false + [[bench]] name = "canonicalize_bench" harness = false diff --git a/encodings/fastlanes/benches/bitpacking_filter.rs b/encodings/fastlanes/benches/bitpacking_filter.rs new file mode 100644 index 00000000000..51ac0828bf9 --- /dev/null +++ b/encodings/fastlanes/benches/bitpacking_filter.rs @@ -0,0 +1,131 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Measures selective filtering around the sparse extraction thresholds. + +#![expect(clippy::cast_possible_truncation)] +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use divan::counter::ItemsCount; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::NativePType; +use vortex_array::validity::Validity; +use vortex_buffer::BufferMut; +use vortex_fastlanes::BitPackedData; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +fn main() { + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fastlanes::initialize(&session); + session +}); + +const NUM_ARRAY_CHUNKS: usize = 64; +// Keep the array density below the outer full-decode policy. +const NUM_SELECTED_CHUNKS: usize = 8; +const CHUNK_SIZE: usize = 1_024; +const LEN: usize = NUM_ARRAY_CHUNKS * CHUNK_SIZE; + +trait BenchInt: NativePType { + fn from_counter(value: u64) -> Self; +} + +macro_rules! impl_bench_int { + ($($T:ty),+) => { + $(impl BenchInt for $T { + fn from_counter(value: u64) -> Self { + value as $T + } + })+ + }; +} + +impl_bench_int!(u8, u16, u32, u64); + +fn fixture(bit_width: usize, selected_per_chunk: usize) -> (ArrayRef, Mask) { + let limit = if bit_width == 64 { + u64::MAX + } else { + 1_u64 << bit_width + }; + let values: BufferMut = (0..LEN) + .map(|index| T::from_counter(index as u64 % limit)) + .collect(); + let packed = BitPackedData::encode( + &PrimitiveArray::new(values.freeze(), Validity::NonNullable).into_array(), + bit_width as u8, + &mut SESSION.create_execution_ctx(), + ) + .unwrap() + .into_array(); + let indices = (0..NUM_SELECTED_CHUNKS).flat_map(|chunk| { + (0..selected_per_chunk) + .map(move |index| chunk * CHUNK_SIZE + index * CHUNK_SIZE / selected_per_chunk) + }); + (packed, Mask::from_indices(LEN, indices)) +} + +macro_rules! bench_width { + ($module:ident, $T:ty, $bit_width:expr, [$($selected:expr),+ $(,)?]) => { + mod $module { + use super::*; + + #[vortex_bench_support::cpu_features] + #[divan::bench(args = [$($selected),+])] + fn filter(bencher: Bencher, selected_per_chunk: usize) { + let (packed, mask) = fixture::<$T>($bit_width, selected_per_chunk); + bencher + .counter(ItemsCount::new(LEN)) + .with_inputs(|| (mask.clone(), SESSION.create_execution_ctx())) + .bench_refs(|(mask, ctx)| { + packed + .filter(mask.clone()) + .unwrap() + .execute::(ctx) + .unwrap() + }); + } + } + }; +} + +macro_rules! bench_type { + ($module:ident, $T:ty, [$(($width_module:ident, $bit_width:expr)),+ $(,)?], $selected:tt) => { + mod $module { + use super::*; + + $(bench_width!($width_module, $T, $bit_width, $selected);)+ + } + }; +} + +bench_type!(u8, u8, [(width1, 1), (width4, 4), (width7, 7)], [8, 16, 24]); +bench_type!( + u16, + u16, + [(width1, 1), (width8, 8), (width15, 15)], + [8, 32, 48] +); +bench_type!( + u32, + u32, + [(width1, 1), (width16, 16), (width31, 31)], + [8, 64, 80, 96] +); +bench_type!( + u64, + u64, + [(width1, 1), (width32, 32), (width63, 63)], + [8, 128, 160, 192] +); diff --git a/encodings/fastlanes/benches/bitpacking_take.rs b/encodings/fastlanes/benches/bitpacking_take.rs index eb072017ae3..b4d43ac0411 100644 --- a/encodings/fastlanes/benches/bitpacking_take.rs +++ b/encodings/fastlanes/benches/bitpacking_take.rs @@ -7,18 +7,23 @@ use std::sync::LazyLock; use divan::Bencher; +use divan::counter::ItemsCount; use rand::RngExt; use rand::SeedableRng; use rand::distr::Uniform; use rand::prelude::StdRng; +use vortex_array::ArrayRef; use vortex_array::IntoArray as _; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::NativePType; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; use vortex_buffer::buffer; use vortex_fastlanes::BitPackedArrayExt; +use vortex_fastlanes::BitPackedData; use vortex_fastlanes::bitpack_compress::bitpack_to_best_bit_width; use vortex_session::VortexSession; @@ -32,6 +37,109 @@ static SESSION: LazyLock = LazyLock::new(|| { session }); +const NUM_ARRAY_CHUNKS: usize = 64; +// Keep the selected count below the outer full-decode policy. +const NUM_SELECTED_CHUNKS: usize = 8; +const CHUNK_SIZE: usize = 1_024; +const THRESHOLD_FIXTURE_LEN: usize = NUM_ARRAY_CHUNKS * CHUNK_SIZE; + +trait BenchInt: NativePType { + fn from_counter(value: u64) -> Self; +} + +macro_rules! impl_bench_int { + ($($T:ty),+) => { + $(impl BenchInt for $T { + fn from_counter(value: u64) -> Self { + value as $T + } + })+ + }; +} + +impl_bench_int!(u8, u16, u32, u64); + +fn threshold_fixture( + bit_width: usize, + selected_per_chunk: usize, +) -> (ArrayRef, ArrayRef) { + let limit = if bit_width == 64 { + u64::MAX + } else { + 1_u64 << bit_width + }; + let values: BufferMut = (0..THRESHOLD_FIXTURE_LEN) + .map(|index| T::from_counter(index as u64 % limit)) + .collect(); + let packed = BitPackedData::encode( + &PrimitiveArray::new(values.freeze(), Validity::NonNullable).into_array(), + bit_width as u8, + &mut SESSION.create_execution_ctx(), + ) + .unwrap() + .into_array(); + let indices = PrimitiveArray::from_iter((0..NUM_SELECTED_CHUNKS).flat_map(|chunk| { + (0..selected_per_chunk) + .map(move |index| (chunk * CHUNK_SIZE + index * CHUNK_SIZE / selected_per_chunk) as u32) + })) + .into_array(); + (packed, indices) +} + +macro_rules! bench_width { + ($module:ident, $T:ty, $bit_width:expr, [$($selected:expr),+ $(,)?]) => { + mod $module { + use super::*; + + #[vortex_bench_support::cpu_features] + #[divan::bench(args = [$($selected),+])] + fn threshold(bencher: Bencher, selected_per_chunk: usize) { + let (packed, indices) = threshold_fixture::<$T>($bit_width, selected_per_chunk); + bencher + .counter(ItemsCount::new(indices.len())) + .with_inputs(|| (indices.clone(), SESSION.create_execution_ctx())) + .bench_refs(|(indices, ctx)| { + packed + .take(indices.clone()) + .unwrap() + .execute::(ctx) + .unwrap() + }); + } + } + }; +} + +macro_rules! bench_type { + ($module:ident, $T:ty, [$(($width_module:ident, $bit_width:expr)),+ $(,)?], $selected:tt) => { + mod $module { + use super::*; + + $(bench_width!($width_module, $T, $bit_width, $selected);)+ + } + }; +} + +bench_type!(u8, u8, [(width1, 1), (width4, 4), (width7, 7)], [8, 16, 24]); +bench_type!( + u16, + u16, + [(width1, 1), (width8, 8), (width15, 15)], + [8, 32, 48] +); +bench_type!( + u32, + u32, + [(width1, 1), (width16, 16), (width31, 31)], + [8, 64, 80, 96, 112] +); +bench_type!( + u64, + u64, + [(width1, 1), (width32, 32), (width63, 63)], + [8, 128, 160, 192] +); + #[divan::bench] fn take_10_stratified(bencher: Bencher) { let values = fixture(65_536, 8); diff --git a/encodings/fastlanes/src/bitpacking/compute/filter.rs b/encodings/fastlanes/src/bitpacking/compute/filter.rs index 0b1b9422f86..5544e460a68 100644 --- a/encodings/fastlanes/src/bitpacking/compute/filter.rs +++ b/encodings/fastlanes/src/bitpacking/compute/filter.rs @@ -22,7 +22,8 @@ use vortex_mask::Mask; use vortex_mask::MaskValuesRef; use super::chunked_indices; -use super::take::UNPACK_CHUNK_THRESHOLD; +use super::unpack_chunk_threshold; +use super::unpack_indices_into; use crate::BitPacked; use crate::BitPackedArrayExt; use crate::BitPackedData; @@ -150,8 +151,10 @@ fn filter_with_indices( &mut values.as_mut_slice()[values_len..], ); } - } else if indices_within_chunk.len() > UNPACK_CHUNK_THRESHOLD { + } else if indices_within_chunk.len() > unpack_chunk_threshold::() { // Unpack into a temporary chunk and then copy the values. + // SAFETY: The validated bit width fits `T`. The source and destination contain + // one complete FastLanes block. The call initializes every destination value. unsafe { let dst: &mut [MaybeUninit] = &mut unpacked; let dst: &mut [T] = std::mem::transmute(dst); @@ -160,13 +163,11 @@ fn filter_with_indices( values.extend_trusted( indices_within_chunk .iter() + // SAFETY: The preceding unpack initialized the complete temporary block. .map(|&idx| unsafe { unpacked.get_unchecked(idx).assume_init() }), ); } else { - // Otherwise, unpack each element individually. - values.extend_trusted(indices_within_chunk.iter().map(|&idx| unsafe { - BitPacking::unchecked_unpack_single(bit_width, packed, idx) - })); + unpack_indices_into(&mut values, bit_width, packed, indices_within_chunk); } }, ); @@ -186,9 +187,12 @@ mod tests { use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; + use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; + use super::filter_with_indices; + use super::unpack_chunk_threshold; use crate::BitPackedData; use crate::bitpacking::array::BitPackedArrayExt; @@ -198,6 +202,56 @@ mod tests { session }); + #[test] + fn sparse_extraction_covers_batch_and_full_chunk_paths() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + + macro_rules! check_type { + ($T:ty, $bit_width:expr) => {{ + let values = (0..2_048) + .map(|index| (index % 127) as $T) + .collect::>(); + let packed = BitPackedData::encode( + &PrimitiveArray::from_iter(values.iter().copied()).into_array(), + $bit_width, + &mut ctx, + )?; + let threshold = unpack_chunk_threshold::<$T>(); + + for selected in [threshold, threshold + 1] { + let indices = (0..selected) + .map(|index| index * 1_024 / selected) + .collect::>(); + let actual = filter_with_indices::<$T>(&packed, &indices); + let expected = indices + .iter() + .map(|&index| values[index]) + .collect::>(); + assert_eq!(actual.as_slice(), expected); + } + }}; + } + + check_type!(u8, 7); + check_type!(u16, 15); + check_type!(u32, 31); + check_type!(u64, 63); + Ok(()) + } + + #[test] + fn sparse_extraction_supports_zero_width() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let packed = BitPackedData::encode( + &PrimitiveArray::from_iter([0u32; 2_048]).into_array(), + 0, + &mut ctx, + )?; + let actual = filter_with_indices::(&packed, &[0, 17, 1_023, 1_024, 2_047]); + assert_eq!(actual.as_slice(), &[0; 5]); + Ok(()) + } + #[test] fn take_indices() { let mut ctx = SESSION.create_execution_ctx(); diff --git a/encodings/fastlanes/src/bitpacking/compute/mod.rs b/encodings/fastlanes/src/bitpacking/compute/mod.rs index 38f86f781bb..fba1d8fd403 100644 --- a/encodings/fastlanes/src/bitpacking/compute/mod.rs +++ b/encodings/fastlanes/src/bitpacking/compute/mod.rs @@ -1,6 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::mem::size_of; + +use fastlanes::BitPacking; +use vortex_array::dtype::NativePType; +use vortex_buffer::BufferMut; + mod between; mod cast; mod compare; @@ -11,6 +17,35 @@ mod slice; mod stream_predicate; mod take; +const fn unpack_chunk_threshold() -> usize { + // FastLanes and Vortex benchmarks set conservative crossovers for each physical type. + match size_of::() { + 1 => 16, + 2 => 32, + 4 => 64, + 8 => 160, + _ => unreachable!(), + } +} + +fn unpack_indices_into( + output: &mut BufferMut, + bit_width: usize, + packed: &[T], + indices: &[usize], +) { + let output_len = output.len(); + let destination = &mut output.spare_capacity_mut()[..indices.len()]; + + // SAFETY: `bit_width` comes from validated data and fits `T`. + // `packed` contains one complete block, and each index is block-relative. + // The destination length equals the index length, and `output` reserves enough space. + unsafe { + T::unchecked_unpack_indices(bit_width, packed, indices, destination); + output.set_len(output_len + indices.len()); + } +} + // TODO(connor): This is duplicated in `encodings/fastlanes/src/bitpacking/kernels/mod.rs`. fn chunked_indices( mut indices: impl Iterator, diff --git a/encodings/fastlanes/src/bitpacking/compute/take.rs b/encodings/fastlanes/src/bitpacking/compute/take.rs index 86e97623cf6..3c9bf79ab1c 100644 --- a/encodings/fastlanes/src/bitpacking/compute/take.rs +++ b/encodings/fastlanes/src/bitpacking/compute/take.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::mem; use std::mem::MaybeUninit; use fastlanes::BitPacking; @@ -23,15 +22,11 @@ use vortex_error::VortexExpect as _; use vortex_error::VortexResult; use super::chunked_indices; +use super::unpack_chunk_threshold; +use super::unpack_indices_into; use crate::BitPacked; use crate::BitPackedArrayExt; -use crate::bitpack_decompress; - -// TODO(connor): This is duplicated in `encodings/fastlanes/src/bitpacking/kernels/mod.rs`. -/// assuming the buffer is already allocated (which will happen at most once) then unpacking -/// all 1024 elements takes ~8.8x as long as unpacking a single element on an M2 Macbook Air. -/// see -pub(super) const UNPACK_CHUNK_THRESHOLD: usize = 8; +const FULL_ARRAY_DECODE_RATIO: usize = 8; impl TakeExecute for BitPacked { fn take( @@ -40,7 +35,7 @@ impl TakeExecute for BitPacked { ctx: &mut ExecutionCtx, ) -> VortexResult> { // If the indices are large enough, it's faster to flatten and take the primitive array. - if indices.len() * UNPACK_CHUNK_THRESHOLD > array.len() { + if indices.len() * FULL_ARRAY_DECODE_RATIO > array.len() { let prim = array.array().clone().execute::(ctx)?; return prim.into_array().take(indices.clone()).map(Some); } @@ -98,41 +93,22 @@ fn take_primitive( chunked_indices(indices_iter, offset, |chunk_idx, indices_within_chunk| { let packed = &packed[chunk_idx * chunk_len..][..chunk_len]; - let mut have_unpacked = false; - let (offset_chunks, remainder) = indices_within_chunk.as_chunks::(); - - // this loop only runs if we have at least UNPACK_CHUNK_THRESHOLD offsets - for offset_chunk in offset_chunks { - if !have_unpacked { - unsafe { - let dst: &mut [MaybeUninit] = &mut unpacked; - let dst: &mut [T] = mem::transmute(dst); - BitPacking::unchecked_unpack(bit_width, packed, dst); - } - have_unpacked = true; - } - - for &index in offset_chunk { - output.push(unsafe { unpacked[index].assume_init() }); - } - } - - // if we have a remainder (i.e., < UNPACK_CHUNK_THRESHOLD leftover offsets), we need to handle it - if !remainder.is_empty() { - if have_unpacked { - // we already bulk unpacked this chunk, so we can just push the remaining elements - for &index in remainder { - output.push(unsafe { unpacked[index].assume_init() }); - } - } else { - // we had fewer than UNPACK_CHUNK_THRESHOLD offsets in the first place, - // so we need to unpack each one individually - for &index in remainder { - output.push(unsafe { - bitpack_decompress::unpack_single_primitive::(packed, bit_width, index) - }); - } + if indices_within_chunk.len() > unpack_chunk_threshold::() { + // SAFETY: The validated bit width fits `T`. The source and destination contain one + // complete FastLanes block. The call initializes every destination value. + unsafe { + let dst: &mut [MaybeUninit] = &mut unpacked; + let dst: &mut [T] = std::mem::transmute(dst); + BitPacking::unchecked_unpack(bit_width, packed, dst); } + output.extend_trusted( + indices_within_chunk + .iter() + // SAFETY: The preceding unpack initialized the complete temporary block. + .map(|&index| unsafe { unpacked.get_unchecked(index).assume_init() }), + ); + } else { + unpack_indices_into(&mut output, bit_width, packed, indices_within_chunk); } }); @@ -157,7 +133,7 @@ fn take_primitive( #[cfg(test)] #[expect(clippy::cast_possible_truncation)] -mod test { +mod tests { use std::sync::LazyLock; use rand::RngExt; @@ -172,12 +148,14 @@ mod test { use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; + use vortex_error::VortexResult; use vortex_session::VortexSession; use crate::BitPackedArray; use crate::BitPackedData; use crate::bitpacking::array::BitPackedArrayExt; use crate::bitpacking::compute::take::take_primitive; + use crate::bitpacking::compute::take::unpack_chunk_threshold; static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session(); @@ -185,6 +163,76 @@ mod test { session }); + #[test] + fn sparse_extraction_covers_batch_and_full_chunk_paths() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + + macro_rules! check_type { + ($T:ty, $bit_width:expr) => {{ + let values = (0..2_048) + .map(|index| (index % 127) as $T) + .collect::>(); + let packed = BitPackedData::encode( + &PrimitiveArray::from_iter(values.iter().copied()).into_array(), + $bit_width, + &mut ctx, + )?; + let threshold = unpack_chunk_threshold::<$T>(); + + for selected in [threshold, threshold + 1] { + let indices = (0..selected) + .map(|index| (index * 1_024 / selected) as u32) + .collect::>(); + let actual = take_primitive::<$T, u32>( + packed.as_view(), + &PrimitiveArray::from_iter(indices.iter().copied()), + Validity::NonNullable, + &mut ctx, + )?; + let expected = indices + .iter() + .map(|&index| values[index as usize]) + .collect::>(); + assert_eq!(actual.as_slice::<$T>(), expected); + } + }}; + } + + check_type!(u8, 7); + check_type!(u16, 15); + check_type!(u32, 31); + check_type!(u64, 63); + Ok(()) + } + + #[test] + fn sparse_take_preserves_nullable_order_and_duplicates() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = (0..8_192).map(|index| index as u32).collect::>(); + let packed = BitPackedData::encode( + &PrimitiveArray::from_iter(values.iter().copied()).into_array(), + 13, + &mut ctx, + )?; + let indices = [ + Some(3_073u32), + Some(2), + None, + Some(2), + Some(1_025), + Some(3_073), + Some(0), + ]; + let actual = packed + .take(PrimitiveArray::from_option_iter(indices).into_array())? + .execute::(&mut ctx)?; + let expected = PrimitiveArray::from_option_iter( + indices.map(|index| index.map(|index| values[index as usize])), + ); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + #[test] fn take_indices() { let mut ctx = SESSION.create_execution_ctx();