From 7609b8f8a262a1a33a2448b82ff4e5cf890f9c83 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 4 Jun 2026 17:20:30 +0000 Subject: [PATCH] perf(mask): SIMD popcount for valid_counts_for_indices `Mask::valid_counts_for_indices` walked the validity bit buffer one bit at a time via `BitIterator::next`, accumulating a running valid count. The pco/zstd slice decoders call it with `[slice_start, slice_stop]`, so the cost was a bit-by-bit scan of the entire prefix up to `slice_stop`. Replace the per-bit walk with an incremental SIMD popcount over each gap between consecutive indices. Total scanned bits are unchanged, but the work now goes through the vectorized `count_ones` path. Add `BitBuffer::count_range(start, end)`, which counts set bits in a range directly over the backing buffer without cloning a sliced `BitBuffer`, so the many-indices case has no per-gap allocation overhead. Benchmark (vortex-mask/benches/valid_counts.rs, median): slice_bounds 16384 8.7us -> 51ns (~170x) slice_bounds 262144 138us -> 310ns (~445x) slice_bounds 1M 554us -> 1.86us (~300x) many_indices 16384 10.1us -> 1.68us (~6x) many_indices 1M 613us -> 3.33us (~185x) Signed-off-by: Joe Isaacs Signed-off-by: Robert Kruszewski --- AGENTS.md | 44 ++++++ Cargo.lock | 1 + encodings/runend/Cargo.toml | 4 + encodings/runend/benches/run_end_filter.rs | 113 +++++++++++++++ encodings/runend/src/compute/filter.rs | 24 ++-- encodings/runend/src/lib.rs | 1 + vortex-array/Cargo.toml | 4 + vortex-array/benches/validity_is_valid.rs | 71 ++++++++++ vortex-buffer/src/bit/buf.rs | 110 +++++++++++++++ vortex-duckdb/Cargo.toml | 5 + vortex-duckdb/benches/bool_export.rs | 58 ++++++++ vortex-duckdb/src/exporter/bool.rs | 17 ++- vortex-mask/Cargo.toml | 8 ++ vortex-mask/benches/mask_iteration.rs | 153 +++++++++++++++++++++ vortex-mask/benches/valid_counts.rs | 51 +++++++ vortex-mask/src/lib.rs | 24 ++-- 16 files changed, 662 insertions(+), 26 deletions(-) create mode 100644 encodings/runend/benches/run_end_filter.rs create mode 100644 vortex-array/benches/validity_is_valid.rs create mode 100644 vortex-duckdb/benches/bool_export.rs create mode 100644 vortex-mask/benches/mask_iteration.rs create mode 100644 vortex-mask/benches/valid_counts.rs diff --git a/AGENTS.md b/AGENTS.md index e5c3d0cc13b..30ddeeb9515 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -151,6 +151,47 @@ Notes: self-explanatory code. - Keep public APIs small and consistent with neighboring crates. +## Performance: avoid hidden-cost accessors in hot loops + +The most common performance trap in this codebase is calling a *per-element accessor that +hides non-trivial work* inside an `O(n)` loop, instead of doing the work once over the whole +chunk. The call site looks like a cheap getter, but each call re-pays a cost that is constant +(or amortizable) across the loop, making the loop `O(n · k)`. + +Watch for these accessors used inside `for i in 0..n { ... }`: + +| Per-element accessor (in a loop) | Hidden cost | Bulk replacement | +| --- | --- | --- | +| `Validity::is_valid(i)` / `is_null(i)` | for array-backed validity, **allocates an `ExecutionCtx` and runs a scalar lookup per call** | `validity.execute_mask(len, ctx)?` once, then `Mask::value(i)` | +| `array.scalar_at(i)` / `array.execute_scalar(i, ctx)` | per-element execution through the compute stack | canonicalize once (`execute::` / `as_slice`) then index | +| `BitBuffer::value(i)` / `Mask::value(i)` accumulated into a count | recomputes the byte address each call; defeats popcount | `true_count()`, `BitBuffer::count_range(start, end)`, `set_indices()` | +| `BitIterator::next()` to accumulate a running rank/prefix count | bit-at-a-time | `count_range` over each gap (SIMD popcount) | +| re-deriving a value inside the loop (e.g. `self.validity()?` each iteration) | re-runs the derivation `n` times | hoist it above the loop | + +Decide per site — bulk is not always the answer: + +- **Sequential / contiguous access** over an accessor that hides amortizable work → hoist and + go bulk (materialize once, then index or iterate the chunk). +- **Gather over arbitrary indices** → you cannot amortize a per-element *decode*, but you can + still materialize the backing buffer once (e.g. `execute_mask`) and then do cheap `O(1)` + random reads, avoiding per-call context/allocation. +- **The accessor is already genuinely `O(1)`** (e.g. reading an already-materialized `Mask`/ + slice, or a native bitmap) → leave it; bulk would not help. + +Even after materializing into a `Mask`/`BitBuffer`, **do not loop with `value(i)` per +element** to act on each set bit — the per-element branch dominates. Iterate a `u64` word +at a time with all-set / all-unset fast paths: use [`BitBuffer::for_each_set_index`]. It +beats `for i in 0..len { if buf.value(i) {..} }` by 2-45x (more at low density) and beats +collecting `set_indices()` by ~2x at mid/high density, while self-adapting from sparse to +dense — see `vortex-mask/benches/mask_iteration.rs`. Reach for the cached `indices()` / +`slices()` representations when you need them more than once; otherwise `for_each_set_index` +needs no materialization. + +When you touch such a loop, back the change with a benchmark (see +`vortex-array/benches/validity_is_valid.rs` for the `is_valid` case, +`vortex-mask/benches/valid_counts.rs` for the popcount case, and +`vortex-mask/benches/mask_iteration.rs` for the per-element-vs-word-vs-sparse comparison). + ## Tests - Strongly consider `rstest` cases when parameterizing repetitive test logic. @@ -179,6 +220,9 @@ Check new and modified lines against this list before finishing: - Updating expected test output to match buggy behavior without independently verifying the intended semantics. - Silently reducing the scope of an approved plan when implementation is harder than expected. +- Calling a hidden-cost per-element accessor (`Validity::is_valid`, `scalar_at`, `BitBuffer:: + value` accumulation) inside a hot loop instead of materializing once — see + "Performance: avoid hidden-cost accessors in hot loops". ## Summaries diff --git a/Cargo.lock b/Cargo.lock index 173843bbaad..1dccf84bb34 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10201,6 +10201,7 @@ dependencies = [ "bitvec", "cbindgen", "cc", + "codspeed-divan-compat", "custom-labels", "futures", "geo-types", diff --git a/encodings/runend/Cargo.toml b/encodings/runend/Cargo.toml index 268aa6ac3ca..c7b38d2494f 100644 --- a/encodings/runend/Cargo.toml +++ b/encodings/runend/Cargo.toml @@ -56,3 +56,7 @@ harness = false [[bench]] name = "run_end_take" harness = false + +[[bench]] +name = "run_end_filter" +harness = false diff --git a/encodings/runend/benches/run_end_filter.rs b/encodings/runend/benches/run_end_filter.rs new file mode 100644 index 00000000000..395b41063b1 --- /dev/null +++ b/encodings/runend/benches/run_end_filter.rs @@ -0,0 +1,113 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks for the run-end filter inner loop (`filter_run_end_primitive`). +//! +//! 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. + +#![expect(clippy::cast_possible_truncation)] +#![expect(clippy::cast_precision_loss)] +#![expect(clippy::cast_sign_loss)] +#![expect(clippy::expect_used)] + +use std::fmt; + +use divan::Bencher; +use rand::SeedableRng; +use rand::rngs::StdRng; +use rand::seq::SliceRandom; +use vortex_buffer::BitBuffer; +use vortex_runend::_benchmarking::filter_run_end_primitive; + +fn main() { + divan::main(); +} + +#[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 { + write!( + f, + "len={}_run={}_density={:.1}", + self.length, self.run_length, self.density + ) + } +} + +const FILTER_ARGS: &[FilterBenchArgs] = &[ + FilterBenchArgs { + length: 4_096, + run_length: 16, + density: 0.1, + }, + FilterBenchArgs { + length: 4_096, + run_length: 16, + density: 0.5, + }, + FilterBenchArgs { + length: 4_096, + run_length: 16, + density: 0.9, + }, + FilterBenchArgs { + length: 16_384, + run_length: 16, + density: 0.1, + }, + FilterBenchArgs { + length: 16_384, + run_length: 16, + density: 0.5, + }, + FilterBenchArgs { + length: 16_384, + run_length: 16, + density: 0.9, + }, +]; + +/// 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) + .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 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); + BitBuffer::from(bits) +} + +#[divan::bench(args = FILTER_ARGS)] +fn filter_run_end(bencher: Bencher, args: FilterBenchArgs) { + let run_ends = build_run_ends(args.length, args.run_length); + let mask = build_mask(args.length, args.density); + let length = args.length as u64; + bencher + .with_inputs(|| (run_ends.clone(), mask.clone())) + .bench_refs(|(run_ends, mask)| { + filter_run_end_primitive::(run_ends, 0, length, mask).expect("filter") + }); +} diff --git a/encodings/runend/src/compute/filter.rs b/encodings/runend/src/compute/filter.rs index ca99fdf1912..6d8d4aab6ea 100644 --- a/encodings/runend/src/compute/filter.rs +++ b/encodings/runend/src/compute/filter.rs @@ -5,6 +5,7 @@ use std::cmp::min; use std::ops::AddAssign; use num_traits::AsPrimitive; +use num_traits::NumCast; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; @@ -74,7 +75,7 @@ impl FilterKernel for RunEnd { } // Code adapted from apache arrow-rs https://github.com/apache/arrow-rs/blob/b1f5c250ebb6c1252b4e7c51d15b8e77f4c361fa/arrow-select/src/filter.rs#L425 -fn filter_run_end_primitive + AsPrimitive>( +pub fn filter_run_end_primitive + AsPrimitive>( run_ends: &[R], offset: u64, length: u64, @@ -87,16 +88,21 @@ fn filter_run_end_primitive + AsPrimitiv let mut count = R::zero(); let new_mask: Mask = BitBuffer::collect_bool(run_ends.len(), |i| { - let mut keep = false; let end = min(run_ends[i].as_() - offset, length); - // Safety: predicate must be the same length as the array the ends have been taken from - for pred in (start..end).map(|i| unsafe { - mask.value_unchecked(i.try_into().vortex_expect("index must fit in usize")) - }) { - count += >::from(pred); - keep |= pred - } + // SIMD popcount of the predicate bits in this run. The range matches the + // bit-by-bit `value_unchecked` read it replaces, so `end <= mask.len()`. + let start_usize = start + .try_into() + .vortex_expect("run start index must fit in usize"); + let end_usize = end + .try_into() + .vortex_expect("run end index must fit in usize"); + let run_trues = mask.count_range(start_usize, end_usize); + count += ::from(run_trues) + .vortex_expect("run popcount must fit in run-end native type"); + let keep = run_trues > 0; + // this is to avoid branching new_run_ends[j] = count; j += keep as usize; diff --git a/encodings/runend/src/lib.rs b/encodings/runend/src/lib.rs index 8169f5d8dfc..a0b26b54dab 100644 --- a/encodings/runend/src/lib.rs +++ b/encodings/runend/src/lib.rs @@ -21,6 +21,7 @@ mod rules; #[doc(hidden)] pub mod _benchmarking { + pub use compute::filter::filter_run_end_primitive; pub use compute::take::take_indices_unchecked; use super::*; diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 98de7caf95d..14a959fa45c 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -200,6 +200,10 @@ required-features = ["_test-harness"] name = "dict_mask" harness = false +[[bench]] +name = "validity_is_valid" +harness = false + [[bench]] name = "dict_unreferenced_mask" harness = false diff --git a/vortex-array/benches/validity_is_valid.rs b/vortex-array/benches/validity_is_valid.rs new file mode 100644 index 00000000000..dc22fcb5478 --- /dev/null +++ b/vortex-array/benches/validity_is_valid.rs @@ -0,0 +1,71 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks the cost of checking array-backed [`Validity`] per element. +//! +//! For the `Validity::Array` variant, `Validity::execute_is_valid` runs a scalar lookup through +//! the compute stack on *every* call, and the (now deprecated) `is_valid` additionally spins up a +//! fresh `ExecutionCtx` per call. Either way, calling it in a `for i in 0..n` loop is +//! pathologically slow. The fix used by callers is to materialize the validity into a `Mask` once +//! (`execute_mask`) and then do cheap O(1) bit reads via `Mask::value`. This bench contrasts the two. +//! +//! Sizes are kept small because the per-element variant is intentionally the slow one we are +//! measuring. + +#![allow(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::validity::Validity; +use vortex_buffer::BitBuffer; +use vortex_session::VortexSession; + +fn main() { + divan::main(); +} + +const SIZES: &[usize] = &[256, 1024]; + +/// Build an array-backed validity (~10% nulls so it is `Validity::Array`). +fn array_validity(len: usize) -> Validity { + Validity::from(BitBuffer::from_iter( + (0..len).map(|i| !i.is_multiple_of(10)), + )) +} + +static SESSION: LazyLock = LazyLock::new(array_session); + +/// Per-element validity check over array-backed validity (the antipattern). This mirrors the +/// deprecated `Validity::is_valid(i)`: a fresh `ExecutionCtx` plus a scalar lookup on every call. +#[divan::bench(args = SIZES)] +fn is_valid_per_element(bencher: Bencher, len: usize) { + let validity = array_validity(len); + bencher + .with_inputs(|| (&validity, SESSION.create_execution_ctx())) + .bench_refs(|(validity, ctx)| { + let mut count = 0usize; + for i in 0..len { + count += validity.execute_is_valid(i, ctx).unwrap() as usize; + } + count + }); +} + +/// Materialize the validity into a `Mask` once, then read bits (the fix). +#[divan::bench(args = SIZES)] +fn execute_mask_then_value(bencher: Bencher, len: usize) { + let validity = array_validity(len); + bencher + .with_inputs(|| (&validity, SESSION.create_execution_ctx())) + .bench_refs(|(validity, ctx)| { + let mask = validity.execute_mask(len, ctx).unwrap(); + let mut count = 0usize; + for i in 0..len { + count += mask.value(i) as usize; + } + count + }); +} diff --git a/vortex-buffer/src/bit/buf.rs b/vortex-buffer/src/bit/buf.rs index a8bd9e1929c..d229329fcd2 100644 --- a/vortex-buffer/src/bit/buf.rs +++ b/vortex-buffer/src/bit/buf.rs @@ -379,6 +379,20 @@ impl BitBuffer { count_ones(self.buffer.as_slice(), self.offset, self.len) } + /// Get the number of set bits in the bit range `[start, end)`. + /// + /// Unlike `self.slice(start..end).true_count()`, this counts directly over the + /// existing backing buffer without allocating or cloning a new [`BitBuffer`], + /// making it cheap to call repeatedly over many small ranges. + /// + /// Panics if `start > end` or `end > len`. + #[inline] + pub fn count_range(&self, start: usize, end: usize) -> usize { + assert!(start <= end, "start {start} exceeds end {end}"); + assert!(end <= self.len, "end {end} exceeds len {}", self.len); + count_ones(self.buffer.as_slice(), self.offset + start, end - start) + } + /// Returns the position of the `nth` set bit (0-indexed). /// /// This is the "select" operation on a bitmap: given a rank `nth`, find @@ -410,6 +424,33 @@ impl BitBuffer { BitSliceIterator::new(self.buffer.as_slice(), self.offset, self.len) } + /// Invoke `f(index)` for every set bit, in ascending order, processing a `u64` + /// word at a time. + /// + /// This is the fast way to "do something for each set bit": it skips all-zero + /// words, fast-paths all-one words, and walks the remaining bits with + /// `trailing_zeros`. Prefer it over `for i in 0..len { if buf.value(i) { f(i) } }` + /// (which pays a branch per element) and over collecting [`Self::set_indices`] + /// (whose per-`next` iterator state does not inline as well). + #[inline] + pub fn for_each_set_index(&self, mut f: F) { + let mut base = 0usize; + for word in self.chunks().iter_padded() { + if word == u64::MAX { + for k in 0..64 { + f(base + k); + } + } else { + let mut w = word; + while w != 0 { + f(base + w.trailing_zeros() as usize); + w &= w - 1; + } + } + base += 64; + } + } + /// Created a new BitBuffer with offset reset to 0 pub fn sliced(&self) -> Self { if self.offset.is_multiple_of(8) { @@ -857,6 +898,75 @@ mod tests { } } + #[rstest] + #[case(0, 0)] + #[case(0, 64)] + #[case(5, 70)] + #[case(64, 130)] + #[case(0, 200)] + fn test_count_range(#[case] start: usize, #[case] end: usize) { + let len = 200; + let buf = BitBuffer::collect_bool(len, |i| i % 3 == 0); + let expected = (start..end).filter(|i| i % 3 == 0).count(); + assert_eq!(buf.count_range(start, end), expected); + // Must agree with slicing then counting. + assert_eq!( + buf.count_range(start, end), + buf.slice(start..end).true_count() + ); + } + + #[rstest] + #[case(3)] + #[case(7)] + fn test_count_range_with_offset(#[case] offset: usize) { + let len = 150; + let buf = BitBuffer::collect_bool(offset + len, |i| i % 2 == 0); + let view = BitBuffer::new_with_offset(buf.inner().clone(), len, offset); + for (start, end) in [(0, len), (10, 100), (1, 2), (63, 129)] { + let expected = (offset + start..offset + end) + .filter(|i| i % 2 == 0) + .count(); + assert_eq!(view.count_range(start, end), expected, "[{start}, {end})"); + } + } + + #[rstest] + #[case(0)] + #[case(1)] + #[case(63)] + #[case(64)] + #[case(65)] + #[case(200)] + #[case(1000)] + fn test_for_each_set_index_matches_set_indices(#[case] len: usize) { + let buf = BitBuffer::collect_bool(len, |i| i % 5 == 0 || i % 7 == 0); + let expected: Vec = buf.set_indices().collect(); + let mut got = Vec::new(); + buf.for_each_set_index(|i| got.push(i)); + assert_eq!(got, expected); + } + + #[rstest] + #[case(3, 200)] + #[case(7, 130)] + fn test_for_each_set_index_with_offset(#[case] offset: usize, #[case] len: usize) { + let base = BitBuffer::collect_bool(offset + len, |i| i % 3 == 0); + let view = BitBuffer::new_with_offset(base.inner().clone(), len, offset); + let expected: Vec = view.set_indices().collect(); + let mut got = Vec::new(); + view.for_each_set_index(|i| got.push(i)); + assert_eq!(got, expected); + } + + #[test] + fn test_for_each_set_index_all_set() { + let buf = BitBuffer::new_set(130); + let mut got = Vec::new(); + buf.for_each_set_index(|i| got.push(i)); + assert_eq!(got, (0..130).collect::>()); + } + #[test] fn test_map_cmp_conditional() { // map_cmp with conditional logic based on index and bit value diff --git a/vortex-duckdb/Cargo.toml b/vortex-duckdb/Cargo.toml index b6896910a89..01a0e114fde 100644 --- a/vortex-duckdb/Cargo.toml +++ b/vortex-duckdb/Cargo.toml @@ -45,6 +45,7 @@ vortex-utils = { workspace = true, features = ["dashmap"] } [dev-dependencies] anyhow = { workspace = true } +divan = { workspace = true } geo-types = { workspace = true } jiff = { workspace = true } rstest = { workspace = true } @@ -54,6 +55,10 @@ vortex-runend = { workspace = true } vortex-sequence = { workspace = true } wkb = { workspace = true } +[[bench]] +name = "bool_export" +harness = false + [lints] workspace = true diff --git a/vortex-duckdb/benches/bool_export.rs b/vortex-duckdb/benches/bool_export.rs new file mode 100644 index 00000000000..f8094b554ae --- /dev/null +++ b/vortex-duckdb/benches/bool_export.rs @@ -0,0 +1,58 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Microbenchmark for the bit -> byte-bool unpack performed by `BoolExporter::export`. +//! +//! DuckDB stores booleans as one byte per value (`&mut [bool]`), while Vortex stores them +//! bit-packed in a [`BitBuffer`]. On every boolean column export we unpack a slice of the +//! bit buffer into the destination byte slice. +//! +//! This bench isolates that pure unpacking logic so it can run without a live DuckDB +//! vector: a reused `Vec` stands in for the DuckDB byte-bool destination. It compares +//! the previous implementation (`.iter().collect::>()` followed by +//! `copy_from_slice`) against the new allocation-free direct zip-write. + +use divan::Bencher; +use vortex::buffer::BitBuffer; + +fn main() { + divan::main(); +} + +const SIZES: &[usize] = &[1024, 16_384]; + +/// Bit densities (true ratio), expressed as `(numerator, denominator)`. +const DENSITIES: &[(usize, usize)] = &[(1, 2), (1, 10), (9, 10)]; + +fn make_buffer(len: usize, density: (usize, usize)) -> BitBuffer { + let (num, den) = density; + // `+ 8` so we can slice off a non-byte-aligned offset, matching real exports. + BitBuffer::from_iter((0..len + 8).map(|i| (i % den) < num)) +} + +/// Previous implementation: allocate a throwaway `Vec` then copy into the destination. +#[divan::bench(args = SIZES, consts = [0usize, 1usize, 2usize])] +fn old_collect_copy(bencher: Bencher, len: usize) { + let buffer = make_buffer(len, DENSITIES[DENSITY_IDX]); + bencher + .with_inputs(|| vec![false; len]) + .bench_refs(|dst: &mut Vec| { + // Offset by 1 to exercise a non-byte-aligned slice like the export hot path. + dst.copy_from_slice(&buffer.slice(1..(1 + len)).iter().collect::>()); + divan::black_box(&dst); + }); +} + +/// New implementation: zip the sliced bit iterator directly into the destination slice. +#[divan::bench(args = SIZES, consts = [0usize, 1usize, 2usize])] +fn new_zip_write(bencher: Bencher, len: usize) { + let buffer = make_buffer(len, DENSITIES[DENSITY_IDX]); + bencher + .with_inputs(|| vec![false; len]) + .bench_refs(|dst: &mut Vec| { + for (slot, bit) in dst.iter_mut().zip(buffer.slice(1..(1 + len)).iter()) { + *slot = bit; + } + divan::black_box(&dst); + }); +} diff --git a/vortex-duckdb/src/exporter/bool.rs b/vortex-duckdb/src/exporter/bool.rs index 4d0a2b29883..fc1dd3e10c3 100644 --- a/vortex-duckdb/src/exporter/bool.rs +++ b/vortex-duckdb/src/exporter/bool.rs @@ -46,13 +46,16 @@ impl ColumnExporter for BoolExporter { ) -> VortexResult<()> { // DuckDB uses byte bools, not bit bools. // maybe we can convert into these from a compressed array sometimes?. - unsafe { vector.as_slice_mut(len) }.copy_from_slice( - &self - .bit_buffer - .slice(offset..(offset + len)) - .iter() - .collect::>(), - ); + // Unpack the bits directly into the destination byte slice, avoiding the + // throwaway `Vec` allocation and the extra copy pass. The sliced + // iterator already accounts for the buffer's bit offset. + let dst = unsafe { vector.as_slice_mut(len) }; + for (slot, bit) in dst + .iter_mut() + .zip(self.bit_buffer.slice(offset..(offset + len)).iter()) + { + *slot = bit; + } Ok(()) } diff --git a/vortex-mask/Cargo.toml b/vortex-mask/Cargo.toml index 7aa5bb2ef85..cf681101d43 100644 --- a/vortex-mask/Cargo.toml +++ b/vortex-mask/Cargo.toml @@ -37,5 +37,13 @@ harness = false name = "rank" harness = false +[[bench]] +name = "valid_counts" +harness = false + +[[bench]] +name = "mask_iteration" +harness = false + [lints] workspace = true diff --git a/vortex-mask/benches/mask_iteration.rs b/vortex-mask/benches/mask_iteration.rs new file mode 100644 index 00000000000..84ca37991ef --- /dev/null +++ b/vortex-mask/benches/mask_iteration.rs @@ -0,0 +1,153 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Strategy comparison for "process the valid elements of a mask". +//! +//! The kernel is `sum of values[i] for every set bit i` — a stand-in for any +//! validity-gated loop. We compare: +//! +//! * `per_element_value` — `for i in 0..n { if buf.value(i) {..} }` (what callers +//! typically write after materializing validity into a mask) +//! * `bit_iterator` — arrow's `BitIterator` (tracks a word internally) +//! * `word_trailing_zeros`— iterate `u64` words, fast-path all-set/all-unset, and +//! walk set bits with `trailing_zeros` / `w &= w - 1` +//! * `set_slices` — iterate contiguous true runs (`set_slices`) +//! * `set_indices` — iterate set positions (`set_indices`) +//! +//! Run across densities to expose the dense-vs-sparse crossover. + +#![allow(clippy::unwrap_used, clippy::cast_possible_truncation)] + +use divan::Bencher; +use vortex_buffer::BitBuffer; + +fn main() { + divan::main(); +} + +const ARGS: &[(usize, f64)] = &[ + (16_384, 0.01), + (16_384, 0.10), + (16_384, 0.50), + (16_384, 0.90), + (16_384, 0.99), +]; + +fn make(len: usize, density: f64) -> (BitBuffer, Vec) { + let threshold = (density * 1000.0) as usize; + let buf = BitBuffer::from_iter((0..len).map(|i| (i * 7 + 13) % 1000 < threshold)); + let values = (0..len as u64).collect(); + (buf, values) +} + +#[divan::bench(args = ARGS)] +fn per_element_value(bencher: Bencher, (len, density): (usize, f64)) { + let (buf, values) = make(len, density); + bencher + .with_inputs(|| (&buf, &values)) + .bench_refs(|(buf, values)| { + let mut acc = 0u64; + for i in 0..len { + if buf.value(i) { + acc = acc.wrapping_add(values[i]); + } + } + acc + }); +} + +#[divan::bench(args = ARGS)] +fn bit_iterator(bencher: Bencher, (len, density): (usize, f64)) { + let (buf, values) = make(len, density); + bencher + .with_inputs(|| (&buf, &values)) + .bench_refs(|(buf, values)| { + let mut acc = 0u64; + for (i, set) in buf.iter().enumerate() { + if set { + acc = acc.wrapping_add(values[i]); + } + } + acc + }); +} + +#[divan::bench(args = ARGS)] +fn word_trailing_zeros(bencher: Bencher, (len, density): (usize, f64)) { + let (buf, values) = make(len, density); + bencher + .with_inputs(|| (&buf, &values)) + .bench_refs(|(buf, values)| { + let mut acc = 0u64; + let chunks = buf.chunks(); + let mut base = 0usize; + for word in chunks.iter() { + if word == u64::MAX { + for k in 0..64 { + acc = acc.wrapping_add(values[base + k]); + } + } else { + let mut w = word; + while w != 0 { + let b = w.trailing_zeros() as usize; + acc = acc.wrapping_add(values[base + b]); + w &= w - 1; + } + } + base += 64; + } + let mut w = chunks.remainder_bits(); + let rem = buf.len() - base; + while w != 0 { + let b = w.trailing_zeros() as usize; + if b >= rem { + break; + } + acc = acc.wrapping_add(values[base + b]); + w &= w - 1; + } + acc + }); +} + +#[divan::bench(args = ARGS)] +fn for_each_set_index(bencher: Bencher, (len, density): (usize, f64)) { + let (buf, values) = make(len, density); + bencher + .with_inputs(|| (&buf, &values)) + .bench_refs(|(buf, values)| { + let mut acc = 0u64; + buf.for_each_set_index(|i| acc = acc.wrapping_add(values[i])); + acc + }); +} + +#[divan::bench(args = ARGS)] +fn set_slices(bencher: Bencher, (len, density): (usize, f64)) { + let (buf, values) = make(len, density); + bencher + .with_inputs(|| (&buf, &values)) + .bench_refs(|(buf, values)| { + let mut acc = 0u64; + for (start, end) in buf.set_slices() { + for i in start..end { + acc = acc.wrapping_add(values[i]); + } + } + acc + }); +} + +#[divan::bench(args = ARGS)] +fn set_indices(bencher: Bencher, (len, density): (usize, f64)) { + let (buf, values) = make(len, density); + bencher + .with_inputs(|| (&buf, &values)) + .bench_refs(|(buf, values)| { + let mut acc = 0u64; + for i in buf.set_indices() { + acc = acc.wrapping_add(values[i]); + } + acc + }); +} diff --git a/vortex-mask/benches/valid_counts.rs b/vortex-mask/benches/valid_counts.rs new file mode 100644 index 00000000000..425a021fa6f --- /dev/null +++ b/vortex-mask/benches/valid_counts.rs @@ -0,0 +1,51 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks for `Mask::valid_counts_for_indices`. +//! +//! This mirrors the hot path in the pco/zstd slice decoders, which call +//! `valid_counts_for_indices(&[slice_start, slice_stop])` to translate a row +//! range into a value range. The cost is dominated by counting set bits in the +//! prefix `[0, slice_stop)`, so a SIMD popcount over the bit buffer should beat +//! a bit-by-bit walk handily for large `slice_stop`. + +#![allow(clippy::unwrap_used, clippy::cast_possible_truncation)] + +use divan::Bencher; +use vortex_buffer::BitBuffer; +use vortex_mask::Mask; + +fn main() { + divan::main(); +} + +const BENCH_SIZES: &[(usize, f64)] = &[(4_096, 0.1), (4_096, 0.9), (16_384, 0.1), (16_384, 0.9)]; + +fn create_mask(len: usize, density: f64) -> Mask { + Mask::from_buffer(BitBuffer::from_iter((0..len).map(|i| { + let threshold = (density * 1000.0) as usize; + (i * 7 + 13) % 1000 < threshold + }))) +} + +/// The pco/zstd slice pattern: two indices bracketing a sub-range. The prefix +/// up to `slice_stop` must be counted, so `slice_stop` near the end is worst case. +#[divan::bench(args = BENCH_SIZES)] +fn slice_bounds(bencher: Bencher, (len, density): (usize, f64)) { + let mask = create_mask(len, density); + let indices = [len / 4, len - len / 8]; + bencher + .with_inputs(|| (&mask, indices)) + .bench_refs(|(mask, indices)| mask.valid_counts_for_indices(indices)); +} + +/// Many monotonically increasing indices spread across the whole mask. +#[divan::bench(args = BENCH_SIZES)] +fn many_indices(bencher: Bencher, (len, density): (usize, f64)) { + let mask = create_mask(len, density); + let stride = (len / 256).max(1); + let indices: Vec = (0..len).step_by(stride).collect(); + bencher + .with_inputs(|| (&mask, indices.as_slice())) + .bench_refs(|(mask, indices)| mask.valid_counts_for_indices(indices)); +} diff --git a/vortex-mask/src/lib.rs b/vortex-mask/src/lib.rs index 401077f0b34..1f41c367947 100644 --- a/vortex-mask/src/lib.rs +++ b/vortex-mask/src/lib.rs @@ -628,23 +628,24 @@ impl Mask { /// Given monotonically increasing `indices` in [0, n_rows], returns the /// count of valid elements up to each index. /// - /// This is O(n_rows). + /// This is O(n_rows), but the per-gap counts are computed with a SIMD + /// popcount over the underlying bit buffer rather than walking bit-by-bit. pub fn valid_counts_for_indices(&self, indices: &[usize]) -> Vec { match self { Self::AllTrue(_) => indices.to_vec(), Self::AllFalse(_) => vec![0; indices.len()], Self::Values(values) => { - let mut bool_iter = values.bit_buffer().iter(); + let buffer = values.bit_buffer(); let mut valid_counts = Vec::with_capacity(indices.len()); let mut valid_count = 0; - let mut idx = 0; + let mut prev = 0; for &next_idx in indices { - while idx < next_idx { - idx += 1; - valid_count += bool_iter - .next() - .unwrap_or_else(|| vortex_panic!("Row indices exceed array length")) - as usize; + assert!(next_idx <= buffer.len(), "Row indices exceed array length"); + // `indices` is monotonically increasing, so each gap is counted once; + // the total work across all gaps scans the prefix `[0, last_idx)` once. + if next_idx > prev { + valid_count += buffer.count_range(prev, next_idx); + prev = next_idx; } valid_counts.push(valid_count); } @@ -779,7 +780,10 @@ impl MaskValues { } let mut indices = Vec::with_capacity(self.true_count); - indices.extend(self.buffer.set_indices()); + // Word-at-a-time set-bit walk; faster than collecting `set_indices()`, + // whose per-`next` iterator state inlines less well (see + // `vortex-mask/benches/mask_iteration.rs`). + self.buffer.for_each_set_index(|i| indices.push(i)); debug_assert!(indices.is_sorted()); assert_eq!(indices.len(), self.true_count); indices