diff --git a/Cargo.lock b/Cargo.lock index 3bf63e6512b..86c4b79468a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11502,6 +11502,7 @@ dependencies = [ "rand 0.10.2", "rstest", "vortex-array", + "vortex-bench-support", "vortex-buffer", "vortex-error", "vortex-mask", diff --git a/encodings/runend/Cargo.toml b/encodings/runend/Cargo.toml index 5e607b78cc6..b491b970b96 100644 --- a/encodings/runend/Cargo.toml +++ b/encodings/runend/Cargo.toml @@ -35,6 +35,7 @@ mimalloc = { workspace = true } rand = { workspace = true } rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-bench-support = { workspace = true } [features] arbitrary = ["dep:arbitrary", "vortex-array/arbitrary"] diff --git a/encodings/runend/benches/run_end_filter.rs b/encodings/runend/benches/run_end_filter.rs index 395b41063b1..4aa1d967f1a 100644 --- a/encodings/runend/benches/run_end_filter.rs +++ b/encodings/runend/benches/run_end_filter.rs @@ -1,12 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Benchmarks for the run-end filter inner loop (`filter_run_end_primitive`). +//! Benchmarks for the run-end filter inner loop. //! -//! This measures the kernel directly rather than going through the lazy -//! `ArrayRef::filter` (which only builds a `FilterArray` node and does not run -//! the kernel). The hot work is a per-run popcount of the predicate mask, which -//! now uses `BitBuffer::count_range` (SIMD) instead of a bit-by-bit walk. +//! `filter_run_end` retains the historical production benchmark. The two materialization +//! benchmarks compare the range and sequential implementations on each CPU feature runner. #![expect(clippy::cast_possible_truncation)] #![expect(clippy::cast_precision_loss)] @@ -14,32 +12,60 @@ #![expect(clippy::expect_used)] use std::fmt; +use std::sync::LazyLock; use divan::Bencher; +use mimalloc::MiMalloc; use rand::SeedableRng; use rand::rngs::StdRng; use rand::seq::SliceRandom; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::RecursiveCanonical; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::PrimitiveArray; use vortex_buffer::BitBuffer; +use vortex_buffer::Buffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_mask::Mask; use vortex_runend::_benchmarking::filter_run_end_primitive; +use vortex_runend::_benchmarking::filter_run_end_ranges; +use vortex_runend::_benchmarking::filter_run_end_sequential; +use vortex_runend::RunEnd; +use vortex_runend::RunEndArray; +use vortex_runend::RunEndArrayExt; +use vortex_runend::RunEndArraySlotsExt; +use vortex_session::VortexSession; + +// Filtering allocates run ends and value masks inside the timed region. Use the same allocator on +// each benchmark runner so that allocator differences do not obscure the scan comparison. +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; fn main() { divan::main(); } +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_runend::initialize(&session); + session +}); + #[derive(Clone, Copy)] struct FilterBenchArgs { - /// Total logical length of the decoded array. length: usize, - /// Average run length used when building the run-end array. run_length: usize, - /// Fraction of mask bits that are set to `true`. density: f64, } impl fmt::Display for FilterBenchArgs { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { write!( - f, + formatter, "len={}_run={}_density={:.1}", self.length, self.run_length, self.density ) @@ -79,24 +105,17 @@ const FILTER_ARGS: &[FilterBenchArgs] = &[ }, ]; -/// Build the run-end boundaries (cumulative run lengths) for `length` rows. fn build_run_ends(length: usize, run_length: usize) -> Vec { - let n_runs = length.div_ceil(run_length); - (0..n_runs) - .map(|r| (((r + 1) * run_length).min(length)) as u32) + (0..length.div_ceil(run_length)) + .map(|run_index| (((run_index + 1) * run_length).min(length)) as u32) .collect() } -/// Build a predicate mask of `length` bits with approximately `density` set bits, -/// shuffled so the set bits are spread across runs. fn build_mask(length: usize, density: f64) -> BitBuffer { - let n_true = (length as f64 * density).round() as usize; + let selected = (length as f64 * density).round() as usize; let mut bits = vec![false; length]; - for b in bits.iter_mut().take(n_true) { - *b = true; - } - let mut rng = StdRng::seed_from_u64(0x5eed); - bits.shuffle(&mut rng); + bits[..selected].fill(true); + bits.shuffle(&mut StdRng::seed_from_u64(0x5eed)); BitBuffer::from(bits) } @@ -111,3 +130,206 @@ fn filter_run_end(bencher: Bencher, args: FilterBenchArgs) { filter_run_end_primitive::(run_ends, 0, length, mask).expect("filter") }); } + +#[derive(Clone, Copy)] +struct StrategyBenchArgs { + length: usize, + run_length: usize, + density_percent: usize, + shape: ValuesShape, + offset: usize, +} + +impl fmt::Display for StrategyBenchArgs { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + formatter, + "len{}_run{}_density{}_{}_offset{}", + self.length, self.run_length, self.density_percent, self.shape, self.offset + ) + } +} + +#[derive(Clone, Copy)] +enum ValuesShape { + Primitive, + Dictionary, + Irregular, +} + +impl fmt::Display for ValuesShape { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Primitive => formatter.write_str("primitive"), + Self::Dictionary => formatter.write_str("dictionary"), + Self::Irregular => formatter.write_str("irregular"), + } + } +} + +const STRATEGY_ARGS: &[StrategyBenchArgs] = &[ + strategy_case(65_536, 16, 50), + strategy_case(65_536, 32, 50), + strategy_case(65_536, 64, 1), + strategy_case(65_536, 64, 50), + strategy_case(65_536, 64, 95), + strategy_case(65_536, 96, 50), + strategy_case(65_536, 128, 50), + strategy_case(65_536, 256, 50), + strategy_case(65_536, 512, 50), + strategy_case(128, 64, 50), + strategy_case(4_096, 64, 50), + StrategyBenchArgs { + length: 65_536, + run_length: 64, + density_percent: 50, + shape: ValuesShape::Irregular, + offset: 37, + }, + StrategyBenchArgs { + length: 65_536, + run_length: 64, + density_percent: 50, + shape: ValuesShape::Dictionary, + offset: 0, + }, + StrategyBenchArgs { + length: 65_536, + run_length: 128, + density_percent: 50, + shape: ValuesShape::Dictionary, + offset: 0, + }, +]; + +const fn strategy_case( + length: usize, + run_length: usize, + density_percent: usize, +) -> StrategyBenchArgs { + StrategyBenchArgs { + length, + run_length, + density_percent, + shape: ValuesShape::Primitive, + offset: 0, + } +} + +#[vortex_bench_support::cpu_features] +#[divan::bench(args = STRATEGY_ARGS, sample_size = 1)] +fn filter_materialized_range(bencher: Bencher, args: StrategyBenchArgs) { + benchmark_filter(bencher, args, filter_run_end_ranges::); +} + +#[vortex_bench_support::cpu_features] +#[divan::bench(args = STRATEGY_ARGS, sample_size = 1)] +fn filter_materialized_sequential(bencher: Bencher, args: StrategyBenchArgs) { + benchmark_filter(bencher, args, |run_ends, offset, length, mask| { + Ok(filter_run_end_sequential(run_ends, offset, length, mask)) + }); +} + +fn benchmark_filter(bencher: Bencher, args: StrategyBenchArgs, filter_run_ends: F) +where + F: Fn(&[u32], u64, u64, &BitBuffer) -> VortexResult<(PrimitiveArray, Mask)> + + Copy + + Send + + Sync + + 'static, +{ + let array = run_end_array(args); + let mask: Mask = build_mask(args.length, args.density_percent as f64 / 100.0).into(); + bencher + .with_inputs(|| (array.clone(), mask.clone(), SESSION.create_execution_ctx())) + .bench_refs(|(array, mask, execution_ctx)| { + filter_with_strategy(array, mask, filter_run_ends, execution_ctx) + .expect("filter") + .execute::(execution_ctx) + .expect("materialize") + }); +} + +fn run_end_array(args: StrategyBenchArgs) -> RunEndArray { + let source_length = args.length + args.offset; + let mut run_ends = match args.shape { + ValuesShape::Primitive | ValuesShape::Dictionary => { + build_run_ends(source_length, args.run_length) + } + ValuesShape::Irregular => build_irregular_run_ends(source_length, args.run_length), + }; + let first_visible_run = run_ends.partition_point(|&run_end| run_end < args.offset as u32); + run_ends.drain(..first_visible_run); + let values = match args.shape { + ValuesShape::Primitive | ValuesShape::Irregular => { + PrimitiveArray::from_iter(0..run_ends.len() as u64).into_array() + } + ValuesShape::Dictionary => DictArray::try_new( + (0..run_ends.len()) + .map(|run_index| (run_index % 16) as u8) + .collect::>() + .into_array(), + PrimitiveArray::from_iter(0u64..16).into_array(), + ) + .expect("dictionary") + .into_array(), + }; + RunEnd::try_new_offset_length( + PrimitiveArray::from_iter(run_ends).into_array(), + values, + args.offset, + args.length, + &mut SESSION.create_execution_ctx(), + ) + .expect("run-end array") +} + +fn build_irregular_run_ends(length: usize, run_length: usize) -> Vec { + let mut run_ends = Vec::new(); + let mut run_end = 0; + let mut run_index = 0; + while run_end < length { + let next_run_length = if run_index % 2 == 0 { + 1 + } else { + run_length * 2 - 1 + }; + run_end = (run_end + next_run_length).min(length); + run_ends.push(run_end as u32); + run_index += 1; + } + run_ends +} + +fn filter_with_strategy( + array: &RunEndArray, + mask: &Mask, + filter_run_ends: F, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + F: Fn(&[u32], u64, u64, &BitBuffer) -> VortexResult<(PrimitiveArray, Mask)>, +{ + let mask_values = mask + .values() + .vortex_expect("benchmark mask must be non-trivial"); + let primitive_run_ends = array.ends().clone().execute::(ctx)?; + let (filtered_run_ends, values_mask) = filter_run_ends( + primitive_run_ends.as_slice::(), + array.offset() as u64, + array.len() as u64, + mask_values.bit_buffer(), + )?; + let filtered_values = array.values().filter(values_mask)?; + + // SAFETY: Both scan functions return one increasing end for each retained value. + Ok(unsafe { + RunEnd::new_unchecked( + filtered_run_ends.into_array(), + filtered_values, + 0, + mask_values.true_count(), + ) + .into_array() + }) +} diff --git a/encodings/runend/src/compute/filter.rs b/encodings/runend/src/compute/filter.rs index 4dc00ea2aba..9942cb3ef6b 100644 --- a/encodings/runend/src/compute/filter.rs +++ b/encodings/runend/src/compute/filter.rs @@ -12,7 +12,7 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::filter::FilterKernel; -use vortex_array::dtype::NativePType; +use vortex_array::dtype::UnsignedPType; use vortex_array::match_each_unsigned_integer_ptype; use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; @@ -40,6 +40,12 @@ const TAKE_SELECTED_ROWS_PER_RUN_THRESHOLD: f64 = 0.1; /// [#1969]: https://github.com/vortex-data/vortex/pull/1969 const MIN_RUN_FILTER_SELECTED_ROWS: usize = 25; +/// Uses one sequential bitmap cursor when the average run length is at most this value. +/// +/// The general range popcount performs less work for long runs. For short runs, its repeated +/// alignment and range setup costs more than a cursor that keeps the current bitmap word loaded. +const SEQUENTIAL_MASK_SCAN_MAX_AVERAGE_RUN_LENGTH: u64 = 64; + impl FilterKernel for RunEnd { fn filter( array: ArrayView<'_, Self>, @@ -98,7 +104,22 @@ impl FilterKernel for RunEnd { /// selected run value. /// /// Adapted from the [Apache Arrow Rust implementation](https://github.com/apache/arrow-rs/blob/b1f5c250ebb6c1252b4e7c51d15b8e77f4c361fa/arrow-select/src/filter.rs#L425). -pub fn filter_run_end_primitive + AsPrimitive>( +pub fn filter_run_end_primitive + AsPrimitive>( + run_ends: &[R], + offset: u64, + length: u64, + mask: &BitBuffer, +) -> VortexResult<(PrimitiveArray, Mask)> { + if length <= (run_ends.len() as u64).saturating_mul(SEQUENTIAL_MASK_SCAN_MAX_AVERAGE_RUN_LENGTH) + { + return Ok(filter_run_end_sequential(run_ends, offset, length, mask)); + } + + filter_run_end_ranges(run_ends, offset, length, mask) +} + +#[doc(hidden)] +pub fn filter_run_end_ranges + AsPrimitive>( run_ends: &[R], offset: u64, length: u64, @@ -111,7 +132,8 @@ pub fn filter_run_end_primitive + AsPrim let mut filtered_end = R::zero(); let values_mask: Mask = BitBuffer::collect_bool(run_ends.len(), |run_idx| { - let run_end = min(run_ends[run_idx].as_() - offset, length); + let absolute_run_end: u64 = run_ends[run_idx].as_(); + let run_end = min(absolute_run_end - offset, length); // Bulk popcount is SIMD-capable and avoids per-bit reads. The input contract and clamp prove // `run_start_idx <= run_end_idx <= mask.len()`. @@ -144,15 +166,138 @@ pub fn filter_run_end_primitive + AsPrim )) } +/// Recomputes run ends with one sequential cursor over the filter bitmap. +/// +/// The mask must contain `length` bits. The first run end can equal `offset` for a sliced array. +/// Later run ends must increase and cover the filtered range. +/// +/// # Panics +/// +/// Panics if the mask length differs from `length`, or if the run ends violate the documented +/// order. +#[doc(hidden)] +pub fn filter_run_end_sequential + AsPrimitive>( + run_ends: &[R], + offset: u64, + length: u64, + mask: &BitBuffer, +) -> (PrimitiveArray, Mask) { + let mut filtered_run_ends = buffer_mut![R::zero(); run_ends.len()]; + let chunks = mask.chunks(); + let mask_length = usize::try_from(length).vortex_expect("mask length must fit in usize"); + assert_eq!( + mask.len(), + mask_length, + "filter mask length must equal the run-end array length" + ); + let mut mask_cursor = MaskCountCursor::new(chunks.iter_padded(), mask_length); + let mut retained_run_count = 0; + let mut filtered_end = R::zero(); + let mut previous_absolute_run_end = offset; + + let values_mask = BitBuffer::collect_bool(run_ends.len(), |run_index| { + let absolute_run_end = run_ends[run_index].as_(); + if run_index == 0 { + assert!( + absolute_run_end >= offset, + "first run end {absolute_run_end} is before offset {offset}" + ); + } else { + assert!( + absolute_run_end > previous_absolute_run_end, + "run end {absolute_run_end} does not follow {previous_absolute_run_end}" + ); + } + previous_absolute_run_end = absolute_run_end; + let run_end = min(absolute_run_end - offset, length) + .try_into() + .vortex_expect("run end must fit in usize"); + let selected_in_run = mask_cursor.count_to(run_end); + filtered_end += ::from(selected_in_run) + .vortex_expect("run popcount must fit in run-end native type"); + let retain_run = selected_in_run > 0; + filtered_run_ends[retained_run_count] = filtered_end; + retained_run_count += retain_run as usize; + retain_run + }) + .into(); + assert_eq!( + mask_cursor.position, mask_length, + "run ends must cover the filtered range" + ); + + filtered_run_ends.truncate(retained_run_count); + ( + PrimitiveArray::new(filtered_run_ends, Validity::NonNullable), + values_mask, + ) +} + +struct MaskCountCursor { + words: I, + current_word: u64, + position: usize, + length: usize, +} + +impl> MaskCountCursor { + fn new(mut words: I, length: usize) -> Self { + let current_word = if length == 0 { + 0 + } else { + words.next().vortex_expect("mask word must exist") + }; + Self { + words, + current_word, + position: 0, + length, + } + } + + fn count_to(&mut self, end: usize) -> usize { + assert!(end >= self.position); + assert!(end <= self.length); + + let mut selected = 0; + while self.position < end { + let bit_in_word = self.position % 64; + let bits_to_read = (end - self.position).min(64 - bit_in_word); + let bit_mask = if bits_to_read == 64 { + u64::MAX + } else { + (1u64 << bits_to_read) - 1 + }; + selected += ((self.current_word >> bit_in_word) & bit_mask).count_ones() as usize; + self.position += bits_to_read; + if self.position.is_multiple_of(64) && self.position < self.length { + self.current_word = self.words.next().vortex_expect("mask word must exist"); + } + } + selected + } +} + #[cfg(test)] mod tests { + use std::ops::AddAssign; + + use num_traits::AsPrimitive; + use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; + use vortex_array::compute::conformance::filter::test_filter_conformance; + use vortex_array::dtype::UnsignedPType; + use vortex_buffer::BitBuffer; + use vortex_buffer::Buffer; + use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; + use super::filter_run_end_ranges; + use super::filter_run_end_sequential; use crate::RunEnd; use crate::RunEndArray; use crate::tests::SESSION; @@ -165,6 +310,117 @@ mod tests { .unwrap() } + #[rstest] + #[case::one_row_runs(0, 1, 0)] + #[case::short_runs_with_bitmap_offset(1, 4, 0)] + #[case::threshold_runs_with_slice(7, 64, 31)] + #[case::long_runs_with_slice(5, 256, 127)] + fn sequential_scan_matches_range_reference( + #[case] bitmap_offset: usize, + #[case] run_length: usize, + #[case] array_offset: usize, + ) -> VortexResult<()> { + let length = 1_003; + let source_length = array_offset + length; + let run_ends = (0..source_length.div_ceil(run_length)) + .map(|run_index| { + u32::try_from(((run_index + 1) * run_length).min(source_length)) + .vortex_expect("test run end must fit in u32") + }) + .collect::>(); + + assert_scan_matches_reference(&run_ends, array_offset, bitmap_offset, length) + } + + #[test] + fn irregular_sequential_scan_matches_range_reference() -> VortexResult<()> { + assert_scan_matches_reference(&[7u32, 71, 72, 145, 146, 500, 501, 1_008], 5, 3, 1_003) + } + + #[test] + fn sequential_scan_accepts_empty_leading_run() -> VortexResult<()> { + assert_scan_matches_reference(&[64u32, 128, 192], 64, 0, 128) + } + + #[test] + fn sequential_scan_supports_all_run_end_widths() -> VortexResult<()> { + assert_scan_matches_reference(&[7u8, 71, 72, 145, 146, 200], 5, 3, 195)?; + assert_scan_matches_reference(&[7u16, 71, 72, 145, 146, 200], 5, 3, 195)?; + assert_scan_matches_reference(&[7u32, 71, 72, 145, 146, 200], 5, 3, 195)?; + assert_scan_matches_reference(&[7u64, 71, 72, 145, 146, 200], 5, 3, 195) + } + + fn assert_scan_matches_reference( + run_ends: &[R], + array_offset: usize, + bitmap_offset: usize, + length: usize, + ) -> VortexResult<()> + where + R: UnsignedPType + AddAssign + From + AsPrimitive, + { + let backing_mask = BitBuffer::collect_bool(bitmap_offset + length, |index| { + index >= bitmap_offset && !(index - bitmap_offset).is_multiple_of(5) + }); + let mask = backing_mask.slice(bitmap_offset..bitmap_offset + length); + + let (actual_ends, actual_values_mask) = + filter_run_end_sequential(run_ends, array_offset as u64, length as u64, &mask); + let (expected_ends, expected_values_mask) = + filter_run_end_ranges(run_ends, array_offset as u64, length as u64, &mask)?; + + assert_eq!(actual_ends.as_slice::(), expected_ends.as_slice::()); + assert_eq!(actual_values_mask, expected_values_mask); + Ok(()) + } + + #[rstest] + #[case::short_runs(4, false, false)] + #[case::long_nullable_sliced_runs(256, true, true)] + fn filter_conformance( + #[case] run_length: usize, + #[case] nullable: bool, + #[case] sliced: bool, + ) -> VortexResult<()> { + let leading_slice = if sliced { run_length / 2 } else { 0 }; + let length = 1_024; + let source_length = leading_slice + length; + let run_count = source_length.div_ceil(run_length); + let ends = (0..run_count) + .map(|run_index| { + u32::try_from(((run_index + 1) * run_length).min(source_length)) + .vortex_expect("test run end must fit in u32") + }) + .collect::>() + .into_array(); + let values = if nullable { + PrimitiveArray::from_option_iter((0..run_count).map(|run_index| { + (!run_index.is_multiple_of(7)) + .then(|| i32::try_from(run_index).vortex_expect("test value must fit in i32")) + })) + .into_array() + } else { + PrimitiveArray::from_iter((0..run_count).map(|run_index| { + i32::try_from(run_index).vortex_expect("test value must fit in i32") + })) + .into_array() + }; + let array = if sliced { + RunEnd::try_new_offset_length( + ends, + values, + leading_slice, + length, + &mut SESSION.create_execution_ctx(), + )? + } else { + RunEnd::try_new(ends, values, &mut SESSION.create_execution_ctx())? + }; + + test_filter_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); + Ok(()) + } + #[test] fn filter_sliced_run_end() -> VortexResult<()> { let arr = ree_array().slice(2..7)?; diff --git a/encodings/runend/src/lib.rs b/encodings/runend/src/lib.rs index b991609c19c..cd238510bb7 100644 --- a/encodings/runend/src/lib.rs +++ b/encodings/runend/src/lib.rs @@ -23,6 +23,8 @@ mod trace_tests; #[doc(hidden)] pub mod _benchmarking { pub use compute::filter::filter_run_end_primitive; + pub use compute::filter::filter_run_end_ranges; + pub use compute::filter::filter_run_end_sequential; pub use compute::take::take_indices_unchecked; use super::*;