From f0e1cdeb6b0516d14175e98b06fea84f3ca1de85 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 19:42:20 +0000 Subject: [PATCH 1/7] Fold nullable integer min, max and sum with a masked lane kernel With nulls present, the primitive min/max and sum kernels walked the runs of valid values one at a time, so a nullable column with short runs, such as one with half its rows null, reduced thousands of one- or two-element slices. Add a masked reduction lane kernel, IndexedFoldExt::fold_masked, which folds every lane into an accumulator and reads null lanes as an identity of the fold. Lanes are selected against constant bits of each 32-lane half of the validity word, which needs no per-lane variable shift and so vectorizes on baseline x86-64, and halves with no valid lane are skipped. Integer min/max read null lanes as a valid value, and integer sums narrower than 64 bits read them as zero, so both cost the same at any validity density. 64-bit integer sums keep their per-element overflow check, and floats keep their in-order runs. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01X1HnweRGxeDWgQwXXM6f5s --- .../src/aggregate_fn/fns/min_max/mod.rs | 23 +++ .../src/aggregate_fn/fns/min_max/primitive.rs | 46 ++++-- .../src/aggregate_fn/fns/sum/primitive.rs | 79 +++++++-- vortex-compute/Cargo.toml | 4 + vortex-compute/benches/fold_masked.rs | 94 +++++++++++ vortex-compute/src/lane_kernels/fold.rs | 150 ++++++++++++++++++ vortex-compute/src/lane_kernels/mod.rs | 3 + 7 files changed, 373 insertions(+), 26 deletions(-) create mode 100644 vortex-compute/benches/fold_masked.rs create mode 100644 vortex-compute/src/lane_kernels/fold.rs diff --git a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs index ce3e938a8fd..d08b5a6c0db 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs @@ -684,6 +684,29 @@ mod tests { Ok(()) } + /// The null rows hold the type's extremes, which the masked integer kernel must not read, + /// in a full 64-row word and in the partial word after it. + #[test] + fn integer_min_max_ignores_values_under_nulls() -> VortexResult<()> { + let values: Buffer = (0..100) + .map(|i| match i % 6 { + 0 => i16::MIN, + 3 => i16::MAX, + _ => i, + }) + .collect(); + let array = PrimitiveArray::new(values, Validity::from_iter((0..100).map(|i| i % 3 != 0))); + let result = min_max( + &array.into_array(), + &mut SESSION.create_execution_ctx(), + NumericalAggregateOpts::default(), + )? + .vortex_expect("has valid rows"); + assert_eq!(i16::try_from(&result.min)?, 1); + assert_eq!(i16::try_from(&result.max)?, 98); + Ok(()) + } + #[test] fn test_prim_nan() -> VortexResult<()> { let array = PrimitiveArray::new( diff --git a/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs b/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs index 263e94ffc8e..e69f425010a 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs @@ -2,6 +2,9 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use itertools::Itertools; +use vortex_buffer::BitBuffer; +use vortex_compute::lane_kernels::IndexedFoldExt; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; @@ -59,20 +62,17 @@ where Mask::AllFalse(_) => None, Mask::Values(v) => { let slice = array.as_slice::(); - // Each `[start, end)` run is fully valid, so integers can reuse the vectorized - // packed min/max per run and fold the run results; floats chain the runs through - // the NaN-filtering reduction. if T::PTYPE.is_int() { - v.slices() - .iter() - .filter_map(|&(start, end)| integer_min_max_raw(&slice[start..end])) - .reduce(|(amin, amax), (rmin, rmax)| { - ( - if rmin.is_lt(amin) { rmin } else { amin }, - if rmax.is_gt(amax) { rmax } else { amax }, - ) - }) - .map(min_max_result) + let valid = v.bit_buffer(); + let first = valid + .set_indices() + .next() + .vortex_expect("a mask with values has a valid row"); + Some(min_max_result(integer_min_max_masked( + slice, + valid, + slice[first], + ))) } else { compute_min_max( v.slices() @@ -105,6 +105,26 @@ where Some((min, max)) } +/// Min/max of the valid values of an integer slice, reading null lanes as `valid_value`, a value of +/// some valid lane, which leaves both extrema unchanged. Selecting per lane rather than walking the +/// runs of valid values keeps the packed min/max of [`integer_min_max_raw`] however short the runs. +fn integer_min_max_masked(slice: &[T], valid: &BitBuffer, valid_value: T) -> (T, T) +where + T: NativePType, +{ + slice.fold_masked( + valid, + valid_value, + (valid_value, valid_value), + |(min, max), v| { + ( + if v.is_lt(min) { v } else { min }, + if v.is_gt(max) { v } else { max }, + ) + }, + ) +} + fn min_max_result((min, max): (T, T)) -> MinMaxResult where T: NativePType, diff --git a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs index a63edacf6b3..632b22e8199 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs @@ -1,12 +1,17 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ops::Add; + use num_traits::AsPrimitive; use num_traits::ToPrimitive; +use vortex_buffer::BitBuffer; +use vortex_compute::lane_kernels::IndexedFoldExt; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_panic; -use vortex_mask::AllOr; +use vortex_mask::Mask; +use vortex_mask::MaskValues; use super::SumState; use super::checked_add_i64; @@ -27,11 +32,10 @@ pub(crate) fn accumulate_primitive( ctx: &mut ExecutionCtx, skip_nans: bool, ) -> VortexResult { - let mask = p.as_ref().validity()?.execute_mask(p.as_ref().len(), ctx)?; - match mask.slices() { - AllOr::None => Ok(false), - AllOr::All => accumulate_primitive_all(inner, p, skip_nans), - AllOr::Some(slices) => accumulate_primitive_valid(inner, p, slices, skip_nans), + match p.as_ref().validity()?.execute_mask(p.as_ref().len(), ctx)? { + Mask::AllFalse(_) => Ok(false), + Mask::AllTrue(_) => accumulate_primitive_all(inner, p, skip_nans), + Mask::Values(valid) => accumulate_primitive_valid(inner, p, &valid, skip_nans), } } @@ -129,20 +133,28 @@ where false } -/// Sum the valid elements, described as contiguous `[start, end)` runs of set validity bits. Each -/// run is a slice of fully-valid values, so it reuses the same vectorized reduction as the -/// all-valid path instead of a per-element validity branch. +/// Sum the valid elements. +/// +/// Integers narrower than 64 bits read null lanes as zero, so the sum vectorizes the same way as +/// the all-valid path however the valid values are spread. Wider integers and floats sum each +/// contiguous `[start, end)` run of valid values through the all-valid reduction, since they +/// check every element for overflow or must add in order. fn accumulate_primitive_valid( inner: &mut SumState, p: &PrimitiveArray, - slices: &[(usize, usize)], + valid: &MaskValues, skip_nans: bool, ) -> VortexResult { match inner { SumState::Unsigned(acc) => match_each_native_ptype!(p.ptype(), unsigned: |T| { let values = p.as_slice::(); - for &(start, end) in slices { + if T::PTYPE != PType::U64 { + return Ok(sum_masked(values, valid.bit_buffer(), |chunk_sum| { + checked_add_u64(acc, chunk_sum) + })); + } + for &(start, end) in valid.slices() { if sum_unsigned_all(acc, &values[start..end]) { return Ok(true); } @@ -156,7 +168,12 @@ fn accumulate_primitive_valid( unsigned: |_T| { vortex_panic!("signed sum state with unsigned input") }, signed: |T| { let values = p.as_slice::(); - for &(start, end) in slices { + if T::PTYPE != PType::I64 { + return Ok(sum_masked(values, valid.bit_buffer(), |chunk_sum| { + checked_add_i64(acc, chunk_sum) + })); + } + for &(start, end) in valid.slices() { if sum_signed_all(acc, &values[start..end]) { return Ok(true); } @@ -170,7 +187,7 @@ fn accumulate_primitive_valid( signed: |_T| { vortex_panic!("float sum state with signed input") }, floating: |T| { let values = p.as_slice::(); - for &(start, end) in slices { + for &(start, end) in valid.slices() { sum_float_all(acc, &values[start..end], skip_nans); } Ok(false) @@ -182,10 +199,31 @@ fn accumulate_primitive_valid( } } +/// Sum the valid values of an integer slice narrower than 64 bits in chunks of [`SUM_CHUNK`], +/// reading null lanes as zero, and hand each chunk's sum to `add`, which returns `true` on +/// overflow. A chunk cannot overflow its 64-bit sum, as in [`sum_unsigned_all`]. +fn sum_masked(values: &[T], valid: &BitBuffer, mut add: impl FnMut(S) -> bool) -> bool +where + T: NativePType + AsPrimitive, + S: NativePType + Add, +{ + (0..values.len()).step_by(SUM_CHUNK).any(|start| { + let end = (start + SUM_CHUNK).min(values.len()); + let chunk_sum = values[start..end].fold_masked( + &valid.slice(start..end), + T::zero(), + S::zero(), + |sum, v| sum + v.as_(), + ); + add(chunk_sum) + }) +} + #[cfg(test)] mod tests { use std::sync::LazyLock; + use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; @@ -219,6 +257,21 @@ mod tests { Ok(()) } + /// The null rows hold the type's maximum, which the masked kernel must not add, in a full + /// 64-row word and in the partial word after it. + #[test] + fn sum_ignores_values_under_nulls() -> VortexResult<()> { + let values: Buffer = (0..100u8) + .map(|i| if i % 3 == 0 { u8::MAX } else { i }) + .collect(); + let validity = Validity::from_iter((0..100).map(|i| i % 3 != 0)); + let arr = PrimitiveArray::new(values, validity).into_array(); + let expected: u64 = (0..100u64).filter(|i| i % 3 != 0).sum(); + let result = sum(&arr, &mut SESSION.create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(expected)); + Ok(()) + } + #[test] fn sum_u8() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![10u8, 20, 30], Validity::NonNullable).into_array(); diff --git a/vortex-compute/Cargo.toml b/vortex-compute/Cargo.toml index fc90b7029df..7b90bfb9cd5 100644 --- a/vortex-compute/Cargo.toml +++ b/vortex-compute/Cargo.toml @@ -36,3 +36,7 @@ workspace = true [[bench]] name = "lane_kernels" harness = false + +[[bench]] +name = "fold_masked" +harness = false diff --git a/vortex-compute/benches/fold_masked.rs b/vortex-compute/benches/fold_masked.rs new file mode 100644 index 00000000000..48827fa27af --- /dev/null +++ b/vortex-compute/benches/fold_masked.rs @@ -0,0 +1,94 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Masked integer reductions with [`IndexedFoldExt::fold_masked`], against the same reduction over +//! an all-valid slice and the arrow-rs nullable aggregate kernels over the same data. + +#![expect(clippy::unwrap_used)] + +use arrow_arith::aggregate; +use arrow_array::Int32Array; +use arrow_buffer::NullBuffer; +use divan::Bencher; +use mimalloc::MiMalloc; +use rand::SeedableRng; +use rand::prelude::*; +use rand::rngs::StdRng; +use vortex_buffer::BitBuffer; +use vortex_compute::lane_kernels::IndexedFoldExt; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +fn main() { + divan::main(); +} + +const N: usize = 15_000; +const DENSITIES: &[f64] = &[0.5, 0.9, 0.99]; + +fn values() -> Vec { + let mut rng = StdRng::seed_from_u64(1); + (0..N).map(|_| rng.random_range(-1000..1000)).collect() +} + +fn valid(density: f64) -> BitBuffer { + let mut rng = StdRng::seed_from_u64(2); + BitBuffer::from_iter((0..N).map(|_| rng.random_bool(density))) +} + +fn arrow(density: f64) -> Int32Array { + let valid = valid(density); + Int32Array::new( + values().into(), + Some(NullBuffer::from_iter((0..N).map(|i| valid.value(i)))), + ) +} + +/// The minimum of every value, ignoring validity, as the bound for a masked minimum. +#[divan::bench] +fn min_all_valid(bencher: Bencher) { + let values = values(); + bencher.bench(|| divan::black_box(&values).iter().copied().min()); +} + +#[divan::bench(args = DENSITIES)] +fn min_fold_masked(bencher: Bencher, density: f64) { + let (values, valid) = (values(), valid(density)); + bencher.bench(|| { + divan::black_box(&values[..]).fold_masked(&valid, i32::MAX, i32::MAX, |acc, v| acc.min(v)) + }); +} + +#[divan::bench(args = DENSITIES)] +fn min_max_fold_masked(bencher: Bencher, density: f64) { + let (values, valid) = (values(), valid(density)); + bencher.bench(|| { + divan::black_box(&values[..]).fold_masked( + &valid, + values[0], + (values[0], values[0]), + |(min, max), v| (min.min(v), max.max(v)), + ) + }); +} + +#[divan::bench(args = DENSITIES)] +fn sum_fold_masked(bencher: Bencher, density: f64) { + let (values, valid) = (values(), valid(density)); + bencher.bench(|| { + divan::black_box(&values[..]).fold_masked(&valid, 0, 0i64, |acc, v| acc + i64::from(v)) + }); +} + +#[divan::bench(args = DENSITIES)] +fn min_arrow(bencher: Bencher, density: f64) { + let array = arrow(density); + bencher.bench(|| aggregate::min(divan::black_box(&array))); +} + +#[divan::bench(args = DENSITIES)] +fn sum_arrow(bencher: Bencher, density: f64) { + let array = arrow(density); + bencher.bench(|| aggregate::sum(divan::black_box(&array)).unwrap()); +} diff --git a/vortex-compute/src/lane_kernels/fold.rs b/vortex-compute/src/lane_kernels/fold.rs new file mode 100644 index 00000000000..b040fdee0d9 --- /dev/null +++ b/vortex-compute/src/lane_kernels/fold.rs @@ -0,0 +1,150 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Masked reduction kernels: fold every lane of an [`IndexedSource`] into one accumulator. + +use vortex_buffer::BitBuffer; + +use crate::lane_kernels::source::IndexedSource; + +/// Extension trait providing reduction kernels on any [`IndexedSource`]. +/// +/// Bring the trait into scope (`use vortex_compute::lane_kernels::IndexedFoldExt;`) to call +/// them with method syntax: `values.fold_masked(&mask, i32::MAX, i32::MAX, |acc, v| acc.min(v))`. +pub trait IndexedFoldExt: IndexedSource + Sized { + /// Fold every lane into `init` with `f`, reading null lanes as `null`. + /// + /// `null` must be an identity of `f`, such as zero for a sum or the type's maximum for a + /// minimum, so that null lanes leave the accumulator unchanged. Every lane of a 64-lane chunk + /// is read and selected against its validity bit, with no branch per lane, so a reduction the + /// compiler can reassociate, such as an integer sum, minimum or maximum, vectorizes the same + /// way as an unmasked one. Runs of 32 lanes without a valid lane are skipped. + /// + /// # Panics + /// + /// Panics if `self.len() != mask.len()`. + #[inline] + fn fold_masked(self, mask: &BitBuffer, null: Self::Item, init: A, f: F) -> A + where + Self::Item: Copy, + F: Fn(A, Self::Item) -> A, + { + let values = self; + let len = values.len(); + assert_eq!(len, mask.len(), "values and mask must have the same length"); + + let chunks = mask.chunks(); + let mut acc = init; + for (chunk_idx, word) in chunks.iter().enumerate() { + let base = chunk_idx * 64; + for (half, bits) in halves(word) { + if bits != 0 { + acc = fold_lanes(&values, null, acc, &f, bits, base + half, 32); + } + } + } + let remainder = len % 64; + let word = chunks.remainder_bits(); + let base = len - remainder; + for (half, bits) in halves(word) { + let lanes = remainder.saturating_sub(half).min(32); + if bits != 0 && lanes != 0 { + acc = fold_lanes(&values, null, acc, &f, bits, base + half, lanes); + } + } + acc + } +} + +/// The two 32-lane halves of a 64-lane validity word, each with the offset of its first lane. +#[expect( + clippy::cast_possible_truncation, + reason = "each half keeps exactly its 32 bits" +)] +fn halves(word: u64) -> [(usize, u32); 2] { + [(0, word as u32), (32, (word >> 32) as u32)] +} + +/// `LANE_BITS[i]` selects lane `i` of a 32-lane validity word. Testing a lane against a constant +/// bit, rather than shifting the word by the lane index, needs no per-lane variable shift, which +/// baseline x86-64 lacks, so the selection vectorizes on every target. +const LANE_BITS: [u32; 32] = { + let mut bits = [0; 32]; + let mut lane = 0; + while lane < 32 { + bits[lane] = 1 << lane; + lane += 1; + } + bits +}; + +/// Fold `lanes` lanes from `base`, reading the lanes whose bit in `bits` is unset as `null`. +#[allow(clippy::inline_always)] +#[inline(always)] +fn fold_lanes( + values: &S, + null: S::Item, + mut acc: A, + f: &F, + bits: u32, + base: usize, + lanes: usize, +) -> A +where + S: IndexedSource, + S::Item: Copy, + F: Fn(A, S::Item) -> A, +{ + for (lane, lane_bit) in LANE_BITS.iter().enumerate().take(lanes) { + // SAFETY: the caller guarantees base + lanes <= len. + let value = unsafe { values.get_unchecked(base + lane) }; + acc = f(acc, if bits & lane_bit != 0 { value } else { null }); + } + acc +} + +impl IndexedFoldExt for S {} + +#[cfg(test)] +mod tests { + use vortex_buffer::BitBuffer; + + use super::IndexedFoldExt; + + #[test] + fn fold_masked_reads_null_lanes_as_identity() { + let values: Vec = (0..130).map(|v| v - 65).collect(); + let mask = BitBuffer::from_iter((0..130).map(|i| i % 3 == 0)); + let expected: i64 = values + .iter() + .enumerate() + .filter(|(i, _)| i % 3 == 0) + .map(|(_, &v)| i64::from(v)) + .sum(); + + let sum = values[..].fold_masked(&mask, 0, 0i64, |acc, v| acc + i64::from(v)); + let min = values[..].fold_masked(&mask, i32::MAX, i32::MAX, |acc, v| acc.min(v)); + + assert_eq!(sum, expected); + assert_eq!(min, -65); + } + + #[test] + fn fold_masked_handles_unaligned_masks() { + let values = [5u8, 1, 9, 2, 7]; + let mask = BitBuffer::from_iter([false, true, true, false, true, true, false]).slice(2..7); + + let max = values[..].fold_masked(&mask, u8::MIN, u8::MIN, |acc, v| acc.max(v)); + + // Valid lanes 0, 2 and 3 hold 5, 9 and 2. + assert_eq!(max, 9); + } + + #[test] + fn fold_masked_of_all_null_lanes_is_init() { + let values = [3i64; 70]; + let mask = BitBuffer::new_unset(70); + + assert_eq!(values[..].fold_masked(&mask, 0, 11, |acc, v| acc + v), 11); + } +} diff --git a/vortex-compute/src/lane_kernels/mod.rs b/vortex-compute/src/lane_kernels/mod.rs index 5aae0ee0ef8..e3d512859eb 100644 --- a/vortex-compute/src/lane_kernels/mod.rs +++ b/vortex-compute/src/lane_kernels/mod.rs @@ -19,12 +19,14 @@ //! the sink itself). //! - [`collect_bits`](mod@collect_bits) — allocating wrappers that bit-pack a predicate into a //! `BitBuffer`. +//! - [`fold`] — masked reductions via [`IndexedFoldExt`], reading null lanes as an identity. //! //! The kernels outside [`collect_bits`](mod@collect_bits) never allocate. Both kernel families //! handle a mask with a non-byte-aligned offset and with a logical `len` shorter than the //! underlying byte buffer, via `BitBuffer::chunks`. pub mod collect_bits; +pub mod fold; pub mod map_in_place; pub mod map_into; pub mod sink; @@ -33,6 +35,7 @@ pub mod source; pub use collect_bits::bit_buffer_from_words; pub use collect_bits::collect_bits; pub use collect_bits::collect_zip_bits; +pub use fold::IndexedFoldExt; pub use map_in_place::IndexedSinkExt; pub use map_into::IndexedSourceExt; pub use sink::IndexedSink; From b8475a1ecf123f2fdc9e349595856b2b001fde3b Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 20:14:30 +0000 Subject: [PATCH 2/7] Fold float min/max, NaN counts and float sums with the masked kernel Apply fold_masked to the other aggregate kernels that walked runs of valid values or zipped values with validity bits: - Float min/max fold total-order keys, integers that order exactly as total_cmp does, so the reduction is plain integer min/max and its result is bit-for-bit that of the previous total_cmp reduction, all-valid or not. - Decimal min/max read null lanes as a valid value, like integers. - Float NaN counts read null lanes as zero, which is not NaN. - Float sums read null lanes as zero, adding the values in the same order, which leaves the result unchanged. Decimal sums and the grouped sums keep their paths: decimal sums are bound by their checked adds, and grouped sums over short groups pay more to set up the fold per group than the fold saves. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01X1HnweRGxeDWgQwXXM6f5s --- vortex-array/Cargo.toml | 4 + vortex-array/benches/aggregate_nullable.rs | 136 +++++++++ .../src/aggregate_fn/fns/min_max/decimal.rs | 30 +- .../src/aggregate_fn/fns/min_max/mod.rs | 34 +++ .../src/aggregate_fn/fns/min_max/primitive.rs | 260 ++++++++++-------- .../aggregate_fn/fns/nan_count/primitive.rs | 11 +- .../src/aggregate_fn/fns/sum/primitive.rs | 121 +++++--- 7 files changed, 410 insertions(+), 186 deletions(-) create mode 100644 vortex-array/benches/aggregate_nullable.rs diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 10733c48d8b..76880c218f1 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -103,6 +103,10 @@ vortex-bench-support = { workspace = true } name = "aggregate_max" harness = false +[[bench]] +name = "aggregate_nullable" +harness = false + [[bench]] name = "aggregate_sum" harness = false diff --git a/vortex-array/benches/aggregate_nullable.rs b/vortex-array/benches/aggregate_nullable.rs new file mode 100644 index 00000000000..b07c0bade84 --- /dev/null +++ b/vortex-array/benches/aggregate_nullable.rs @@ -0,0 +1,136 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Aggregates over columns whose nulls are scattered at random, so their valid values form runs +//! of one or two rows rather than the long runs of the clustered benchmarks. +//! +//! Each iteration gets a fresh array, so no statistic cached by a previous iteration is reused. + +use std::sync::LazyLock; + +use divan::Bencher; +use mimalloc::MiMalloc; +use rand::prelude::*; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::fns::min_max::min_max; +use vortex_array::aggregate_fn::fns::nan_count::nan_count; +use vortex_array::aggregate_fn::fns::sum::sum; +use vortex_array::aggregate_fn::session::AggregateFnSession; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DecimalDType; +use vortex_array::optimizer::kernels::KernelSession; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_session::VortexSession; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +// Sized to keep the CodSpeed simulation under 1ms per benchmark. +const N: usize = 15_000; + +static SESSION: LazyLock = LazyLock::new(|| { + VortexSession::empty() + .with::() + .with::() +}); + +/// Half the rows, chosen at random, are valid. +fn scattered_validity() -> Validity { + let mut rng = StdRng::seed_from_u64(1); + Validity::from_iter((0..N).map(|_| rng.random_bool(0.5))) +} + +fn f64_values() -> Buffer { + let mut rng = StdRng::seed_from_u64(2); + (0..N).map(|_| rng.random_range(-1000.0..1000.0)).collect() +} + +fn bench_with( + bencher: Bencher, + input: impl Fn() -> ArrayRef + Sync, + f: impl Fn(&ArrayRef) -> O + Sync, +) { + bencher.with_inputs(&input).bench_refs(|array| f(array)); +} + +fn bench_min_max(bencher: Bencher, input: impl Fn() -> ArrayRef + Sync) { + bench_with(bencher, input, |array| { + min_max( + array, + &mut SESSION.create_execution_ctx(), + NumericalAggregateOpts::default(), + ) + }); +} + +fn bench_sum(bencher: Bencher, input: impl Fn() -> ArrayRef + Sync) { + bench_with(bencher, input, |array| { + sum(array, &mut SESSION.create_execution_ctx()) + }); +} + +#[divan::bench] +fn min_max_f64_all_valid(bencher: Bencher) { + let values = f64_values(); + bench_min_max(bencher, || { + PrimitiveArray::new(values.clone(), Validity::NonNullable).into_array() + }); +} + +#[divan::bench] +fn min_max_f64_nulls_scattered(bencher: Bencher) { + let (values, validity) = (f64_values(), scattered_validity()); + bench_min_max(bencher, || { + PrimitiveArray::new(values.clone(), validity.clone()).into_array() + }); +} + +#[divan::bench] +fn min_max_decimal_i64_nulls_scattered(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(3); + let values: Buffer = (0..N) + .map(|_| rng.random_range(-100_000..100_000)) + .collect(); + let validity = scattered_validity(); + bench_min_max(bencher, || { + DecimalArray::new(values.clone(), DecimalDType::new(18, 2), validity.clone()).into_array() + }); +} + +#[divan::bench] +fn sum_i32_nulls_scattered(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(4); + let values: Buffer = (0..N).map(|_| rng.random_range(-1000..1000)).collect(); + let validity = scattered_validity(); + bench_sum(bencher, || { + PrimitiveArray::new(values.clone(), validity.clone()).into_array() + }); +} + +#[divan::bench] +fn sum_f64_nulls_scattered(bencher: Bencher) { + let (values, validity) = (f64_values(), scattered_validity()); + bench_sum(bencher, || { + PrimitiveArray::new(values.clone(), validity.clone()).into_array() + }); +} + +#[divan::bench] +fn nan_count_f64_nulls_scattered(bencher: Bencher) { + let (values, validity) = (f64_values(), scattered_validity()); + bench_with( + bencher, + || PrimitiveArray::new(values.clone(), validity.clone()).into_array(), + |array| nan_count(array, &mut SESSION.create_execution_ctx()), + ); +} diff --git a/vortex-array/src/aggregate_fn/fns/min_max/decimal.rs b/vortex-array/src/aggregate_fn/fns/min_max/decimal.rs index 2fa1e137ffb..87fab423dc1 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/decimal.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use itertools::Itertools; +use vortex_compute::lane_kernels::IndexedFoldExt; use vortex_error::VortexResult; use vortex_mask::Mask; @@ -79,25 +79,17 @@ pub fn valid_range(values: &[T], validity: &Mask) -> Optio "valid_range requires one validity bit per value" ); + let min_max = |(min, max): (T, T), v: T| (min.min(v), max.max(v)); match validity { - Mask::AllTrue(_) => compute_min_max(values.iter()), + Mask::AllTrue(_) => values + .first() + .map(|&first| values.iter().copied().fold((first, first), min_max)), Mask::AllFalse(_) => None, - Mask::Values(validity) => compute_min_max( - values - .iter() - .zip(validity.bit_buffer().iter()) - .filter_map(|(value, valid)| valid.then_some(value)), - ), - } -} - -fn compute_min_max<'a, T>(iter: impl Iterator) -> Option<(T, T)> -where - T: NativeDecimalType, -{ - match iter.minmax_by(|a, b| a.cmp(b)) { - itertools::MinMaxResult::NoElements => None, - itertools::MinMaxResult::OneElement(&value) => Some((value, value)), - itertools::MinMaxResult::MinMax(&min, &max) => Some((min, max)), + // Null lanes read as a valid value, which leaves both extrema unchanged, so the reduction + // needs no branch per lane. + Mask::Values(validity) => validity.bit_buffer().set_indices().next().map(|first| { + let first = values[first]; + values.fold_masked(validity.bit_buffer(), first, (first, first), min_max) + }), } } diff --git a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs index d08b5a6c0db..e07c94f3ec5 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs @@ -707,6 +707,40 @@ mod tests { Ok(()) } + /// The null rows hold infinities and NaNs, which the masked float kernel must not read, in a + /// full 64-row word and in the partial word after it. Valid NaNs are skipped by default. + #[test] + fn float_min_max_ignores_values_under_nulls() -> VortexResult<()> { + let values: Buffer = (0..100) + .map(|i| match i % 6 { + 0 => f64::NEG_INFINITY, + 3 => f64::INFINITY, + 4 if i == 4 => f64::NAN, + _ => f64::from(i) - 50.5, + }) + .collect(); + let validity = Validity::from_iter((0..100).map(|i| i % 3 != 0)); + let array = PrimitiveArray::new(values, validity).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + + let result = min_max(&array, &mut ctx, NumericalAggregateOpts::default())? + .vortex_expect("has valid rows"); + assert_eq!(f64::try_from(&result.min)?, -49.5); + assert_eq!(f64::try_from(&result.max)?, 47.5); + + // Including NaNs, the valid NaN poisons both extrema. + let mut acc = Accumulator::try_new( + MinMax, + NumericalAggregateOpts::include_nans(), + array.dtype().clone(), + )?; + acc.accumulate(&array, &mut ctx)?; + let result = MinMaxResult::from_scalar(&acc.finish()?)?.vortex_expect("has valid rows"); + assert!(f64::try_from(&result.min)?.is_nan()); + assert!(f64::try_from(&result.max)?.is_nan()); + Ok(()) + } + #[test] fn test_prim_nan() -> VortexResult<()> { let array = PrimitiveArray::new( diff --git a/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs b/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs index e69f425010a..b3ad1e957de 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs @@ -1,10 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use itertools::Itertools; +use half::f16; use vortex_buffer::BitBuffer; use vortex_compute::lane_kernels::IndexedFoldExt; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; @@ -27,104 +26,126 @@ pub(super) fn accumulate_primitive( ctx: &mut ExecutionCtx, ) -> VortexResult<()> { let skip_nans = args.options.skip_nans; - match_each_native_ptype!(p.ptype(), |T| { - let local = compute_min_max_with_validity::(p, ctx, skip_nans)?; - partial.merge(args, local); - Ok(()) + let validity = p.as_ref().validity()?.execute_mask(p.as_ref().len(), ctx)?; + let valid = match &validity { + Mask::AllTrue(_) => None, + Mask::AllFalse(_) => return Ok(()), + Mask::Values(v) => Some(v.bit_buffer()), + }; + let local = match_each_native_ptype!(p.ptype(), + integral: |T| { integer_min_max(p.as_slice::(), valid).map(min_max_result) }, + floating: |T| { float_min_max(p.as_slice::(), valid, skip_nans).map(min_max_result) } + ); + partial.merge(args, local); + Ok(()) +} + +/// Min/max of the valid values of an integer slice, given its validity unless every value is +/// valid. +/// +/// Null lanes read as a valid value, which leaves both extrema unchanged, so the reduction +/// autovectorizes to packed min/max however the valid values are spread. +fn integer_min_max(slice: &[T], valid: Option<&BitBuffer>) -> Option<(T, T)> { + let first = slice[first_valid(slice.len(), valid, |_| true)?]; + let min_max = |(min, max): (T, T), v: T| { + ( + if v.is_lt(min) { v } else { min }, + if v.is_gt(max) { v } else { max }, + ) + }; + Some(match valid { + None => slice.iter().copied().fold((first, first), min_max), + Some(valid) => slice.fold_masked(valid, first, (first, first), min_max), }) } -fn compute_min_max_with_validity( - array: &PrimitiveArray, - ctx: &mut ExecutionCtx, +/// Min/max of the valid values of a float slice under the IEEE 754 total order, skipping NaN +/// values when `skip_nans` is set. +/// +/// Without `skip_nans`, NaN values sort to the ends of the total order, so a NaN extremum lets the +/// partial's merge poison the result. The reduction runs over [`TotalOrderKey`]s, integers that +/// order exactly as the floats do, so it autovectorizes to packed integer min/max; null and +/// skipped lanes read as the key of a value that is neither, which leaves both extrema unchanged. +fn float_min_max( + slice: &[T], + valid: Option<&BitBuffer>, skip_nans: bool, -) -> VortexResult> -where - T: NativePType, - PValue: From, -{ - Ok( - match array - .as_ref() - .validity()? - .execute_mask(array.as_ref().len(), ctx)? - { - Mask::AllTrue(_) => { - let slice = array.as_slice::(); - // Integers have no NaNs, so a plain min/max reduction is correct and, unlike the - // `itertools::minmax_by` + NaN-filter path, autovectorizes to packed min/max. - if T::PTYPE.is_int() { - integer_min_max_raw(slice).map(min_max_result) - } else { - compute_min_max(slice.iter(), skip_nans) - } - } - Mask::AllFalse(_) => None, - Mask::Values(v) => { - let slice = array.as_slice::(); - if T::PTYPE.is_int() { - let valid = v.bit_buffer(); - let first = valid - .set_indices() - .next() - .vortex_expect("a mask with values has a valid row"); - Some(min_max_result(integer_min_max_masked( - slice, - valid, - slice[first], - ))) - } else { - compute_min_max( - v.slices() - .iter() - .flat_map(|&(start, end)| slice[start..end].iter()), - skip_nans, - ) - } - } - }, - ) +) -> Option<(T, T)> { + let first = slice[first_valid(slice.len(), valid, |i| !(skip_nans && slice[i].is_nan()))?]; + let first_key = first.key(); + let min_max = |(min, max): (T::Key, T::Key), v: T| { + let key = if skip_nans && v.is_nan() { + first_key + } else { + v.key() + }; + (min.min(key), max.max(key)) + }; + let (min, max) = match valid { + None => slice.iter().copied().fold((first_key, first_key), min_max), + Some(valid) => slice.fold_masked(valid, first, (first_key, first_key), min_max), + }; + Some((T::from_key(min), T::from_key(max))) } -/// Min/max of an all-valid integer slice as native values. Autovectorizes to packed min/max. -fn integer_min_max_raw(slice: &[T]) -> Option<(T, T)> -where - T: NativePType, -{ - let (&first, rest) = slice.split_first()?; - let mut min = first; - let mut max = first; - for &v in rest { - if v.is_lt(min) { - min = v; - } - if v.is_gt(max) { - max = v; - } +/// The index of the first valid lane for which `include` holds, or `None` if there is none. +fn first_valid( + len: usize, + valid: Option<&BitBuffer>, + include: impl Fn(usize) -> bool, +) -> Option { + match valid { + None => (0..len).find(|&i| include(i)), + Some(valid) => valid.set_indices().find(|&i| include(i)), } - Some((min, max)) } -/// Min/max of the valid values of an integer slice, reading null lanes as `valid_value`, a value of -/// some valid lane, which leaves both extrema unchanged. Selecting per lane rather than walking the -/// runs of valid values keeps the packed min/max of [`integer_min_max_raw`] however short the runs. -fn integer_min_max_masked(slice: &[T], valid: &BitBuffer, valid_value: T) -> (T, T) -where - T: NativePType, -{ - slice.fold_masked( - valid, - valid_value, - (valid_value, valid_value), - |(min, max), v| { - ( - if v.is_lt(min) { v } else { min }, - if v.is_gt(max) { v } else { max }, - ) - }, - ) +/// A float whose IEEE 754 total order, the order of `total_cmp`, is the order of a signed +/// integer key derived from its bits. +/// +/// Flipping every bit but the sign of a negative float's bits reverses the order of the negative +/// floats, so the bits order as `total_cmp` does when read as a signed integer. This is how +/// `total_cmp` itself is implemented. The transform keeps the sign bit, so it is its own inverse. +trait TotalOrderKey: NativePType { + type Key: Copy + Ord; + + fn key(self) -> Self::Key; + + fn from_key(key: Self::Key) -> Self; } +macro_rules! total_order_key { + ($float:ty, $signed:ty, $unsigned:ty) => { + impl TotalOrderKey for $float { + type Key = $signed; + + #[inline] + #[expect( + clippy::cast_possible_wrap, + reason = "the key reinterprets the float's bits" + )] + fn key(self) -> $signed { + let bits = self.to_bits() as $signed; + bits ^ ((((bits >> (<$signed>::BITS - 1)) as $unsigned) >> 1) as $signed) + } + + #[inline] + #[expect( + clippy::cast_sign_loss, + reason = "the key reinterprets the float's bits" + )] + fn from_key(key: $signed) -> Self { + let bits = key ^ ((((key >> (<$signed>::BITS - 1)) as $unsigned) >> 1) as $signed); + <$float>::from_bits(bits as $unsigned) + } + } + }; +} + +total_order_key!(f16, i16, u16); +total_order_key!(f32, i32, u32); +total_order_key!(f64, i64, u64); + fn min_max_result((min, max): (T, T)) -> MinMaxResult where T: NativePType, @@ -136,40 +157,35 @@ where } } -fn compute_min_max<'a, T>( - iter: impl Iterator, - skip_nans: bool, -) -> Option -where - T: NativePType, - PValue: From, -{ - if skip_nans { - minmax_by_total_order(iter.filter(|v| !v.is_nan())) - } else { - // Compute extrema under the total order (where NaNs sort to the ends) and let the - // partial's merge poison the result if either end is NaN. - minmax_by_total_order(iter) - } -} +#[cfg(test)] +mod tests { + #![allow( + clippy::cast_possible_truncation, + reason = "narrows the f64 cases to f32" + )] -fn minmax_by_total_order<'a, T>(iter: impl Iterator) -> Option -where - T: NativePType, - PValue: From, -{ - match iter.minmax_by(|a, b| a.total_compare(**b)) { - itertools::MinMaxResult::NoElements => None, - itertools::MinMaxResult::OneElement(&x) => { - let scalar = Scalar::primitive(x, NonNullable); - Some(MinMaxResult { - min: scalar.clone(), - max: scalar, - }) + use half::f16; + use rstest::rstest; + + use super::TotalOrderKey; + + /// Keys order exactly as `total_cmp` does, and convert back to the same bits. + #[rstest] + #[case(&[f64::NAN, -f64::NAN, f64::INFINITY, f64::NEG_INFINITY, 0.0, -0.0, 1.5, -1.5, f64::MIN_POSITIVE, -f64::MIN_POSITIVE, f64::MAX, f64::MIN])] + fn total_order_keys_order_as_total_cmp(#[case] values: &[f64]) { + for &a in values { + assert_eq!(f64::from_key(a.key()).to_bits(), a.to_bits()); + assert_eq!( + f32::from_key((a as f32).key()).to_bits(), + (a as f32).to_bits() + ); + let h = f16::from_f64(a); + assert_eq!(f16::from_key(h.key()).to_bits(), h.to_bits()); + for &b in values { + assert_eq!(a.key().cmp(&b.key()), a.total_cmp(&b), "{a} vs {b}"); + let (af, bf) = (a as f32, b as f32); + assert_eq!(af.key().cmp(&bf.key()), af.total_cmp(&bf), "{af} vs {bf}"); + } } - itertools::MinMaxResult::MinMax(&min, &max) => Some(MinMaxResult { - min: Scalar::primitive(min, NonNullable), - max: Scalar::primitive(max, NonNullable), - }), } } diff --git a/vortex-array/src/aggregate_fn/fns/nan_count/primitive.rs b/vortex-array/src/aggregate_fn/fns/nan_count/primitive.rs index bf380bc4a3f..40ad3a7dcd3 100644 --- a/vortex-array/src/aggregate_fn/fns/nan_count/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/nan_count/primitive.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_compute::lane_kernels::IndexedFoldExt; use vortex_error::VortexResult; use vortex_mask::Mask; @@ -27,12 +28,10 @@ fn compute_nan_count_with_validity(values: &[T], validity: Mask) match validity { Mask::AllTrue(_) => values.iter().filter(|v| v.is_nan()).count(), Mask::AllFalse(_) => 0, - Mask::Values(v) => values - .iter() - .zip(v.bit_buffer().iter()) - .filter_map(|(v, m)| m.then_some(v)) - .filter(|v| v.is_nan()) - .count(), + // Null lanes read as zero, which is not NaN. + Mask::Values(v) => values.fold_masked(v.bit_buffer(), T::zero(), 0, |count, v| { + count + usize::from(v.is_nan()) + }), } } diff --git a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs index 632b22e8199..fb075d192a9 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs @@ -133,63 +133,30 @@ where false } -/// Sum the valid elements. -/// -/// Integers narrower than 64 bits read null lanes as zero, so the sum vectorizes the same way as -/// the all-valid path however the valid values are spread. Wider integers and floats sum each -/// contiguous `[start, end)` run of valid values through the all-valid reduction, since they -/// check every element for overflow or must add in order. +/// Sum the valid elements of `p`, whose validity is `valid`. fn accumulate_primitive_valid( inner: &mut SumState, p: &PrimitiveArray, valid: &MaskValues, skip_nans: bool, ) -> VortexResult { + let valid = valid.bit_buffer(); match inner { SumState::Unsigned(acc) => match_each_native_ptype!(p.ptype(), - unsigned: |T| { - let values = p.as_slice::(); - if T::PTYPE != PType::U64 { - return Ok(sum_masked(values, valid.bit_buffer(), |chunk_sum| { - checked_add_u64(acc, chunk_sum) - })); - } - for &(start, end) in valid.slices() { - if sum_unsigned_all(acc, &values[start..end]) { - return Ok(true); - } - } - Ok(false) - }, + unsigned: |T| { Ok(sum_unsigned_masked(acc, p.as_slice::(), valid)) }, signed: |_T| { vortex_panic!("unsigned sum state with signed input") }, floating: |_T| { vortex_panic!("unsigned sum state with float input") } ), SumState::Signed(acc) => match_each_native_ptype!(p.ptype(), unsigned: |_T| { vortex_panic!("signed sum state with unsigned input") }, - signed: |T| { - let values = p.as_slice::(); - if T::PTYPE != PType::I64 { - return Ok(sum_masked(values, valid.bit_buffer(), |chunk_sum| { - checked_add_i64(acc, chunk_sum) - })); - } - for &(start, end) in valid.slices() { - if sum_signed_all(acc, &values[start..end]) { - return Ok(true); - } - } - Ok(false) - }, + signed: |T| { Ok(sum_signed_masked(acc, p.as_slice::(), valid)) }, floating: |_T| { vortex_panic!("signed sum state with float input") } ), SumState::Float(acc) => match_each_native_ptype!(p.ptype(), unsigned: |_T| { vortex_panic!("float sum state with unsigned input") }, signed: |_T| { vortex_panic!("float sum state with signed input") }, floating: |T| { - let values = p.as_slice::(); - for &(start, end) in valid.slices() { - sum_float_all(acc, &values[start..end], skip_nans); - } + sum_float_masked(acc, p.as_slice::(), valid, skip_nans); Ok(false) } ), @@ -199,10 +166,63 @@ fn accumulate_primitive_valid( } } +/// Sum the values of `values` that `valid` marks valid into a `u64` accumulator. Returns `true` on +/// overflow. +/// +/// Values narrower than 64 bits read null lanes as zero, so the sum vectorizes the same way as +/// [`sum_unsigned_all`] however the valid values are spread. `u64` values check every element +/// for overflow, so they sum each contiguous run of valid values through [`sum_unsigned_all`]. +fn sum_unsigned_masked(acc: &mut u64, values: &[T], valid: &BitBuffer) -> bool +where + T: NativePType + AsPrimitive, +{ + if T::PTYPE == PType::U64 { + return valid + .set_slices() + .any(|(start, end)| sum_unsigned_all(acc, &values[start..end])); + } + sum_chunks_masked(values, valid, |chunk_sum| checked_add_u64(acc, chunk_sum)) +} + +/// Signed counterpart of [`sum_unsigned_masked`]. +fn sum_signed_masked(acc: &mut i64, values: &[T], valid: &BitBuffer) -> bool +where + T: NativePType + AsPrimitive, +{ + if T::PTYPE == PType::I64 { + return valid + .set_slices() + .any(|(start, end)| sum_signed_all(acc, &values[start..end])); + } + sum_chunks_masked(values, valid, |chunk_sum| checked_add_i64(acc, chunk_sum)) +} + +/// Sum the values of `values` that `valid` marks valid into an `f64` accumulator, skipping NaN +/// values when `skip_nans` is set, as [`sum_float_all`] does. +/// +/// Null and skipped lanes read as zero. The values are added in order, as by [`sum_float_all`], +/// and adding zero leaves the accumulator unchanged: it starts at positive zero, so it is never +/// negative zero, the one value adding zero would change. The result is therefore exactly that of +/// summing the valid values alone, without walking their runs. +fn sum_float_masked( + acc: &mut f64, + values: &[T], + valid: &BitBuffer, + skip_nans: bool, +) { + *acc = values.fold_masked(valid, T::zero(), *acc, |sum, v| { + if skip_nans && v.is_nan() { + sum + } else { + sum + ToPrimitive::to_f64(&v).vortex_expect("float to f64") + } + }); +} + /// Sum the valid values of an integer slice narrower than 64 bits in chunks of [`SUM_CHUNK`], /// reading null lanes as zero, and hand each chunk's sum to `add`, which returns `true` on /// overflow. A chunk cannot overflow its 64-bit sum, as in [`sum_unsigned_all`]. -fn sum_masked(values: &[T], valid: &BitBuffer, mut add: impl FnMut(S) -> bool) -> bool +fn sum_chunks_masked(values: &[T], valid: &BitBuffer, mut add: impl FnMut(S) -> bool) -> bool where T: NativePType + AsPrimitive, S: NativePType + Add, @@ -272,6 +292,29 @@ mod tests { Ok(()) } + /// The null rows hold infinities and NaNs, which the masked float kernel must not add, and + /// the valid NaN is skipped by default. + #[test] + fn float_sum_ignores_values_under_nulls() -> VortexResult<()> { + let values: Buffer = (0..100u16) + .map(|i| match i % 3 { + 0 if i % 2 == 0 => f32::INFINITY, + 0 => f32::NAN, + _ if i == 50 => f32::NAN, + _ => f32::from(i) - 50.0, + }) + .collect(); + let validity = Validity::from_iter((0..100).map(|i| i % 3 != 0)); + let arr = PrimitiveArray::new(values, validity).into_array(); + let expected: f64 = (0..100u16) + .filter(|&i| i % 3 != 0 && i != 50) + .map(|i| f64::from(i) - 50.0) + .sum(); + let result = sum(&arr, &mut SESSION.create_execution_ctx())?; + assert_eq!(result.as_primitive().typed_value::(), Some(expected)); + Ok(()) + } + #[test] fn sum_u8() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![10u8, 20, 30], Validity::NonNullable).into_array(); From a3f0e9dc8bce71babf6c7ecf00d778084f8bf0cd Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 28 Oct 2026 21:41:53 +0000 Subject: [PATCH 3/7] Fold aggregates into lane-parallel accumulators Float sums, float min/max and 64-bit integer sums did not vectorize: the compiler cannot reorder float adds, float compares or checked adds, so each ran one element at a time, and with nulls each run of valid values ran separately. Add IndexedFoldExt::fold_lanes and fold_lanes_masked. They fold lane i into accumulator i % LANES and merge the accumulators at the end, so one row of lanes is one vector step for any fold. The masked form tests each lane against a constant bit of the validity word, with no branch. An opaque barrier before each row stops the vectorizers from packing one lane of several rows with strided loads. Use the folds in the aggregates: - Float sums use lane folds. NaN skipping runs a second pass only when the first sum is NaN. - Float min/max use float compares, one fold per extremum. A cold pass settles a valid NaN and the sign of a zero extremum. - 64-bit sums add with wrapping and OR together a bound on each value's magnitude. Values past the bound get an exact 128-bit sum, so a sum overflows only if its exact total does not fit. - Narrower integer sums use 16 lanes, except signed values on x86-64 without SSE4.1, which has no vector sign extension. On Graviton4, Sapphire Rapids, Zen 4 and Apple M4 the float sums are 4.6x to 8.8x faster, and up to 30x with scattered nulls. Float max is 2.4x to 8.8x faster, and 64-bit sums are 1.8x to 3.4x faster. Signed-off-by: Robert Kruszewski Claude-Session: https://claude.ai/code/session_018SWbEqgj32dnKrPu3naMzN --- vortex-array/benches/aggregate_max.rs | 30 ++ vortex-array/benches/aggregate_sum.rs | 46 ++ .../src/aggregate_fn/fns/min_max/mod.rs | 81 ++++ .../src/aggregate_fn/fns/min_max/primitive.rs | 307 +++++++----- .../src/aggregate_fn/fns/sum/primitive.rs | 455 ++++++++++++++---- .../src/aggregate_fn/fns/sum_v2/tests.rs | 4 +- vortex-compute/Cargo.toml | 4 + vortex-compute/benches/fold_lanes.rs | 330 +++++++++++++ vortex-compute/src/lane_kernels/fold.rs | 307 +++++++++++- vortex-compute/src/lane_kernels/mod.rs | 4 +- 10 files changed, 1354 insertions(+), 214 deletions(-) create mode 100644 vortex-compute/benches/fold_lanes.rs diff --git a/vortex-array/benches/aggregate_max.rs b/vortex-array/benches/aggregate_max.rs index dc3159e6428..93d9127d279 100644 --- a/vortex-array/benches/aggregate_max.rs +++ b/vortex-array/benches/aggregate_max.rs @@ -179,6 +179,36 @@ fn max_i32_nulls_scattered(bencher: Bencher) { .bench_refs(|(a, ctx)| a.statistics().compute_max::(ctx)); } +#[divan::bench] +fn max_f32(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(6); + let data: Vec = (0..N).map(|_| rng.random::()).collect(); + bencher + .with_inputs(|| { + ( + PrimitiveArray::from_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| a.statistics().compute_max::(ctx)); +} + +#[divan::bench] +fn max_f64_nulls_scattered(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(7); + let data: Vec> = (0..N) + .map(|_| rng.random_bool(0.5).then(|| rng.random::())) + .collect(); + bencher + .with_inputs(|| { + ( + PrimitiveArray::from_option_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| a.statistics().compute_max::(ctx)); +} + // String extrema. Every view carries its value's first four bytes inline, so the shape that // matters is how far into a value its order is settled: `distinct_prefixes` and // `low_cardinality` never have to look past the view, while the log-shaped columns (every value diff --git a/vortex-array/benches/aggregate_sum.rs b/vortex-array/benches/aggregate_sum.rs index fd7a7944ed4..b1745a068dc 100644 --- a/vortex-array/benches/aggregate_sum.rs +++ b/vortex-array/benches/aggregate_sum.rs @@ -267,3 +267,49 @@ fn sum_v2_i32_nulls_scattered(bencher: Bencher) { }) .bench_refs(|(a, ctx)| sum_v2(a, ctx)); } + +#[divan::bench] +fn sum_u64(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(8); + let data: Vec = (0..N).map(|_| rng.random_range(0..2000)).collect(); + bencher + .with_inputs(|| { + ( + PrimitiveArray::from_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| a.statistics().compute_as::(Stat::Sum, ctx)); +} + +#[divan::bench] +fn sum_i64_nulls_scattered(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(9); + let data: Vec> = (0..N) + .map(|_| rng.random_bool(0.5).then(|| rng.random_range(-1000..1000))) + .collect(); + bencher + .with_inputs(|| { + ( + PrimitiveArray::from_option_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| a.statistics().compute_as::(Stat::Sum, ctx)); +} + +#[divan::bench] +fn sum_f64_nulls_scattered(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(10); + let data: Vec> = (0..N) + .map(|_| rng.random_bool(0.5).then(|| rng.random_range(-1000.0..1000.0))) + .collect(); + bencher + .with_inputs(|| { + ( + PrimitiveArray::from_option_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| a.statistics().compute_as::(Stat::Sum, ctx)); +} diff --git a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs index e07c94f3ec5..01a4ca74b2e 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs @@ -858,6 +858,87 @@ mod tests { Ok(()) } + /// The null rows hold NaNs and extremes, which the masked float kernel must not read, in a + /// full 64-row word and in the partial word after it, whether NaNs are skipped or not. + #[test] + fn float_min_max_ignores_nans_under_nulls() -> VortexResult<()> { + let values: Buffer = (0..100) + .map(|i| match i % 6 { + 0 => f64::NEG_INFINITY, + 3 => f64::NAN, + _ => f64::from(i), + }) + .collect(); + let array = PrimitiveArray::new(values, Validity::from_iter((0..100).map(|i| i % 3 != 0))); + let mut ctx = SESSION.create_execution_ctx(); + for options in [NumericalAggregateOpts::default(), KEEP_NANS] { + let result = + min_max(&array.clone().into_array(), &mut ctx, options)?.vortex_expect("valid"); + assert_eq!(f64::try_from(&result.min)?, 1.0); + assert_eq!(f64::try_from(&result.max)?, 98.0); + } + Ok(()) + } + + /// The total order puts `-0.0` below `0.0`, although they compare equal as floats. + #[rstest] + #[case::negative_zero_is_the_min(&[0.0, 1.0, -0.0, 0.0], -0.0, 1.0)] + #[case::positive_zero_is_the_max(&[-1.0, -0.0, 0.0, -0.0], -1.0, 0.0)] + #[case::negative_zero_alone(&[-0.0, -0.0], -0.0, -0.0)] + #[case::positive_zero_alone(&[0.0, 0.0], 0.0, 0.0)] + fn float_min_max_orders_signed_zeros( + #[case] values: &[f32], + #[case] min: f32, + #[case] max: f32, + ) -> VortexResult<()> { + let array = PrimitiveArray::from_iter(values.iter().copied()).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let result = min_max(&array, &mut ctx, NumericalAggregateOpts::default())? + .vortex_expect("valid"); + assert_eq!(f32::try_from(&result.min)?.to_bits(), min.to_bits()); + assert_eq!(f32::try_from(&result.max)?.to_bits(), max.to_bits()); + Ok(()) + } + + /// Infinities are values, not NaNs, so they never poison the result. + #[test] + fn float_min_max_of_infinities_is_not_poisoned() -> VortexResult<()> { + let array = PrimitiveArray::new( + buffer![1.0f64, f64::INFINITY, -2.0, f64::NEG_INFINITY], + Validity::from_iter([true, true, true, false]), + ) + .into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let result = min_max(&array, &mut ctx, KEEP_NANS)?.vortex_expect("valid"); + assert_eq!(f64::try_from(&result.min)?, -2.0); + assert_eq!(f64::try_from(&result.max)?, f64::INFINITY); + Ok(()) + } + + /// A valid NaN after the first valid value still poisons, and only NaN values is no result + /// when they are skipped. + #[test] + fn float_min_max_masked_nans() -> VortexResult<()> { + let validity = || Validity::from_iter((0..70).map(|i| i % 2 == 0)); + let late_nan: Buffer = (0..70) + .map(|i| if i == 66 { f64::NAN } else { f64::from(i) }) + .collect(); + let only_nans: Buffer = (0..70) + .map(|i| if i % 2 == 0 { f64::NAN } else { f64::from(i) }) + .collect(); + let mut ctx = SESSION.create_execution_ctx(); + + let late_nan = PrimitiveArray::new(late_nan, validity()).into_array(); + assert_poisoned(min_max(&late_nan, &mut ctx, KEEP_NANS)?)?; + let skipped = min_max(&late_nan, &mut ctx, NumericalAggregateOpts::default())? + .vortex_expect("valid"); + assert_eq!(f64::try_from(&skipped.max)?, 68.0); + + let only_nans = PrimitiveArray::new(only_nans, validity()).into_array(); + assert!(min_max(&only_nans, &mut ctx, NumericalAggregateOpts::default())?.is_none()); + Ok(()) + } + #[test] fn test_prim_nan_not_skipping() -> VortexResult<()> { let array = PrimitiveArray::new( diff --git a/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs b/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs index b3ad1e957de..fe729d46751 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs @@ -1,9 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use half::f16; +use itertools::Either; +use num_traits::Float; use vortex_buffer::BitBuffer; use vortex_compute::lane_kernels::IndexedFoldExt; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; @@ -25,126 +27,214 @@ pub(super) fn accumulate_primitive( p: &PrimitiveArray, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { - let skip_nans = args.options.skip_nans; - let validity = p.as_ref().validity()?.execute_mask(p.as_ref().len(), ctx)?; - let valid = match &validity { - Mask::AllTrue(_) => None, + let mask = p.as_ref().validity()?.execute_mask(p.as_ref().len(), ctx)?; + let valid = match &mask { Mask::AllFalse(_) => return Ok(()), + Mask::AllTrue(_) => None, Mask::Values(v) => Some(v.bit_buffer()), }; + let local = match_each_native_ptype!(p.ptype(), integral: |T| { integer_min_max(p.as_slice::(), valid).map(min_max_result) }, - floating: |T| { float_min_max(p.as_slice::(), valid, skip_nans).map(min_max_result) } + floating: |T| { + float_min_max(p.as_slice::(), valid, args.options.skip_nans).map(min_max_result) + } ); partial.merge(args, local); Ok(()) } -/// Min/max of the valid values of an integer slice, given its validity unless every value is -/// valid. +/// Min/max of the valid values of an integer slice. /// -/// Null lanes read as a valid value, which leaves both extrema unchanged, so the reduction -/// autovectorizes to packed min/max however the valid values are spread. -fn integer_min_max(slice: &[T], valid: Option<&BitBuffer>) -> Option<(T, T)> { - let first = slice[first_valid(slice.len(), valid, |_| true)?]; - let min_max = |(min, max): (T, T), v: T| { - ( - if v.is_lt(min) { v } else { min }, - if v.is_gt(max) { v } else { max }, - ) - }; - Some(match valid { - None => slice.iter().copied().fold((first, first), min_max), - Some(valid) => slice.fold_masked(valid, first, (first, first), min_max), - }) +/// Kept out of line so each integer type gets a function of its own, as the float types do: inlined +/// into the dispatch over every type, the scalar 64-bit loop of baseline x86-64 ran 40% slower. +#[inline(never)] +fn integer_min_max(slice: &[T], valid: Option<&BitBuffer>) -> Option<(T, T)> +where + T: NativePType, +{ + match valid { + None => integer_min_max_raw(slice), + Some(valid) => { + let first = valid + .set_indices() + .next() + .vortex_expect("a mask with values has a valid row"); + Some(integer_min_max_masked(slice, valid, slice[first])) + } + } +} + +/// Min/max of an all-valid integer slice as native values. Autovectorizes to packed min/max. +fn integer_min_max_raw(slice: &[T]) -> Option<(T, T)> +where + T: NativePType, +{ + let (&first, rest) = slice.split_first()?; + let mut min = first; + let mut max = first; + for &v in rest { + if v.is_lt(min) { + min = v; + } + if v.is_gt(max) { + max = v; + } + } + Some((min, max)) +} + +/// Min/max of the valid values of an integer slice, reading null lanes as `valid_value`, a value of +/// some valid lane, which leaves both extrema unchanged. Selecting per lane rather than walking the +/// runs of valid values keeps the packed min/max of [`integer_min_max_raw`] however short the runs. +fn integer_min_max_masked(slice: &[T], valid: &BitBuffer, valid_value: T) -> (T, T) +where + T: NativePType, +{ + slice.fold_masked( + valid, + valid_value, + (valid_value, valid_value), + |(min, max), v| { + ( + if v.is_lt(min) { v } else { min }, + if v.is_gt(max) { v } else { max }, + ) + }, + ) } -/// Min/max of the valid values of a float slice under the IEEE 754 total order, skipping NaN -/// values when `skip_nans` is set. +/// Min/max of the valid values of a float slice under the total order, or `None` if there are +/// none. /// -/// Without `skip_nans`, NaN values sort to the ends of the total order, so a NaN extremum lets the -/// partial's merge poison the result. The reduction runs over [`TotalOrderKey`]s, integers that -/// order exactly as the floats do, so it autovectorizes to packed integer min/max; null and -/// skipped lanes read as the key of a value that is neither, which leaves both extrema unchanged. -fn float_min_max( - slice: &[T], - valid: Option<&BitBuffer>, - skip_nans: bool, -) -> Option<(T, T)> { - let first = slice[first_valid(slice.len(), valid, |i| !(skip_nans && slice[i].is_nan()))?]; - let first_key = first.key(); - let min_max = |(min, max): (T::Key, T::Key), v: T| { - let key = if skip_nans && v.is_nan() { - first_key - } else { - v.key() - }; - (min.min(key), max.max(key)) +/// With `skip_nans`, NaN values are ignored. Otherwise a valid NaN makes both extrema NaN, which +/// poisons the partial state. +/// +/// The extrema are folded with float comparisons, which every target compares in vector registers, +/// unlike the 64-bit integer keys of the total order, which baseline x86-64 cannot. A comparison +/// with a NaN is false, so the folds skip NaN values for free, and null lanes read `seed`, a valid +/// non-NaN value, which leaves both extrema unchanged. Each extremum is a fold of its own: a fold +/// into one float per lane vectorizes on every target, where a fold into a pair of floats can stay +/// scalar. Two cases the comparisons cannot see are settled after the folds, each with another +/// pass that runs only when needed: +/// +/// - Without `skip_nans`, a third fold sums `v - v`, which turns NaN only for a NaN or infinite +/// value. When it does, [`find_nan`] looks for a valid NaN. +/// - The total order puts `-0.0` below `0.0`, but the comparisons see them as equal. When an +/// extremum is zero, [`has_zero_with_sign`] finds which signs are present. +fn float_min_max(values: &[T], valid: Option<&BitBuffer>, skip_nans: bool) -> Option<(T, T)> +where + T: NativePType + Float, +{ + let mut indices = valid_indices(values.len(), valid); + let seed = if skip_nans { + indices.map(|i| values[i]).find(|&v| !Float::is_nan(v))? + } else { + values[indices.next()?] }; - let (min, max) = match valid { - None => slice.iter().copied().fold((first_key, first_key), min_max), - Some(valid) => slice.fold_masked(valid, first, (first_key, first_key), min_max), + if Float::is_nan(seed) { + return Some((seed, seed)); + } + + if !skip_nans { + let non_finite = fold_floats(values, valid, seed, T::zero(), |e, v| e + (v - v), add); + if Float::is_nan(non_finite) + && let Some(nan) = find_nan(values, valid) + { + return Some((nan, nan)); + } + } + + let min = fold_floats(values, valid, seed, seed, lower, lower); + let max = fold_floats(values, valid, seed, seed, higher, higher); + + let min = if min.is_zero() { + zero_with_sign(has_zero_with_sign(values, valid, true)) + } else { + min }; - Some((T::from_key(min), T::from_key(max))) + let max = if max.is_zero() { + zero_with_sign(!has_zero_with_sign(values, valid, false)) + } else { + max + }; + Some((min, max)) } -/// The index of the first valid lane for which `include` holds, or `None` if there is none. -fn first_valid( - len: usize, +/// Fold the valid float values with [`IndexedFoldExt::fold_lanes`], or its masked form reading null +/// lanes as `null`, with as many lanes as fill 64 bytes per accumulator. +#[inline] +fn fold_floats( + values: &[T], valid: Option<&BitBuffer>, - include: impl Fn(usize) -> bool, -) -> Option { - match valid { - None => (0..len).find(|&i| include(i)), - Some(valid) => valid.set_indices().find(|&i| include(i)), + null: T, + init: A, + f: impl Fn(A, T) -> A, + merge: impl Fn(A, A) -> A, +) -> A +where + T: NativePType, + A: Copy, +{ + macro_rules! fold { + ($lanes:literal) => { + match valid { + None => values.fold_lanes::<$lanes, _, _, _>(init, f, merge), + Some(valid) => values.fold_lanes_masked::<$lanes, _, _, _>(valid, null, init, f, merge), + } + }; + } + + match size_of::() { + 2 => fold!(32), + 4 => fold!(16), + _ => fold!(8), } } -/// A float whose IEEE 754 total order, the order of `total_cmp`, is the order of a signed -/// integer key derived from its bits. -/// -/// Flipping every bit but the sign of a negative float's bits reverses the order of the negative -/// floats, so the bits order as `total_cmp` does when read as a signed integer. This is how -/// `total_cmp` itself is implemented. The transform keeps the sign bit, so it is its own inverse. -trait TotalOrderKey: NativePType { - type Key: Copy + Ord; +fn add(a: T, b: T) -> T { + a + b +} - fn key(self) -> Self::Key; +/// `v` if it compares below `min`, else `min`. NaN never compares below. +#[inline] +fn lower(min: T, v: T) -> T { + if v < min { v } else { min } +} - fn from_key(key: Self::Key) -> Self; +/// `v` if it compares above `max`, else `max`. NaN never compares above. +#[inline] +fn higher(max: T, v: T) -> T { + if v > max { v } else { max } } -macro_rules! total_order_key { - ($float:ty, $signed:ty, $unsigned:ty) => { - impl TotalOrderKey for $float { - type Key = $signed; - - #[inline] - #[expect( - clippy::cast_possible_wrap, - reason = "the key reinterprets the float's bits" - )] - fn key(self) -> $signed { - let bits = self.to_bits() as $signed; - bits ^ ((((bits >> (<$signed>::BITS - 1)) as $unsigned) >> 1) as $signed) - } +/// The first valid NaN. +#[cold] +#[inline(never)] +fn find_nan(values: &[T], valid: Option<&BitBuffer>) -> Option { + valid_indices(values.len(), valid) + .map(|i| values[i]) + .find(|v| v.is_nan()) +} - #[inline] - #[expect( - clippy::cast_sign_loss, - reason = "the key reinterprets the float's bits" - )] - fn from_key(key: $signed) -> Self { - let bits = key ^ ((((key >> (<$signed>::BITS - 1)) as $unsigned) >> 1) as $signed); - <$float>::from_bits(bits as $unsigned) - } - } - }; +/// Whether a valid value is a zero of the given sign. +#[cold] +#[inline(never)] +fn has_zero_with_sign(values: &[T], valid: Option<&BitBuffer>, negative: bool) -> bool { + valid_indices(values.len(), valid) + .any(|i| values[i].is_zero() && values[i].is_sign_negative() == negative) } -total_order_key!(f16, i16, u16); -total_order_key!(f32, i32, u32); -total_order_key!(f64, i64, u64); +fn zero_with_sign(negative: bool) -> T { + if negative { -T::zero() } else { T::zero() } +} + +fn valid_indices(len: usize, valid: Option<&BitBuffer>) -> impl Iterator + '_ { + match valid { + None => Either::Left(0..len), + Some(valid) => Either::Right(valid.set_indices()), + } +} fn min_max_result((min, max): (T, T)) -> MinMaxResult where @@ -156,36 +246,3 @@ where max: Scalar::primitive(max, NonNullable), } } - -#[cfg(test)] -mod tests { - #![allow( - clippy::cast_possible_truncation, - reason = "narrows the f64 cases to f32" - )] - - use half::f16; - use rstest::rstest; - - use super::TotalOrderKey; - - /// Keys order exactly as `total_cmp` does, and convert back to the same bits. - #[rstest] - #[case(&[f64::NAN, -f64::NAN, f64::INFINITY, f64::NEG_INFINITY, 0.0, -0.0, 1.5, -1.5, f64::MIN_POSITIVE, -f64::MIN_POSITIVE, f64::MAX, f64::MIN])] - fn total_order_keys_order_as_total_cmp(#[case] values: &[f64]) { - for &a in values { - assert_eq!(f64::from_key(a.key()).to_bits(), a.to_bits()); - assert_eq!( - f32::from_key((a as f32).key()).to_bits(), - (a as f32).to_bits() - ); - let h = f16::from_f64(a); - assert_eq!(f16::from_key(h.key()).to_bits(), h.to_bits()); - for &b in values { - assert_eq!(a.key().cmp(&b.key()), a.total_cmp(&b), "{a} vs {b}"); - let (af, bf) = (a as f32, b as f32); - assert_eq!(af.key().cmp(&bf.key()), af.total_cmp(&bf), "{af} vs {bf}"); - } - } - } -} diff --git a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs index fb075d192a9..fe80d8d324a 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::ops::Add; +use std::ops::BitOr; use num_traits::AsPrimitive; use num_traits::ToPrimitive; @@ -26,6 +27,13 @@ use crate::match_each_native_ptype; /// than 64 bits cannot overflow the 64-bit accumulator: `2^16 * (2^32 - 1) < 2^64`. const SUM_CHUNK: usize = 1 << 16; +/// Partial sums of a sum of integers narrower than 64 bits, each widened to a 64-bit lane. +const NARROW_SUM_LANES: usize = 16; + +/// Partial sums of a float sum. Sixteen `f64` lanes fill eight NEON, four AVX2 or two AVX-512 +/// registers, enough to hide the latency of a vector add on each. +const FLOAT_SUM_LANES: usize = 16; + pub(crate) fn accumulate_primitive( inner: &mut SumState, p: &PrimitiveArray, @@ -73,37 +81,66 @@ fn accumulate_primitive_all( /// are skipped to match the scalar `sum` semantics; otherwise any NaN poisons the accumulator to /// NaN. Floats cannot overflow the accumulator, so this never reports saturation. pub(crate) fn sum_float_all(acc: &mut f64, slice: &[T], skip_nans: bool) { - if skip_nans { - for &v in slice { - if !v.is_nan() { - *acc += ToPrimitive::to_f64(&v).vortex_expect("float to f64"); - } - } + sum_float(acc, slice, None, skip_nans); +} + +/// Sum the values of a float slice, or only its `valid` values, into an `f64` accumulator. +/// +/// The values are added in [`FLOAT_SUM_LANES`] independent partial sums, so the sum vectorizes +/// although float addition does not reassociate. The result can differ from an in-order sum only +/// in rounding. +/// +/// With `skip_nans`, the values are first summed as they are: a sum that is not NaN saw no NaN. +/// Only a NaN sum, which a NaN or infinities of both signs make, is summed again reading NaNs as +/// zero. A select per lane to skip NaNs on top of the select for null lanes leaves the masked +/// fold scalar on NEON. +fn sum_float(acc: &mut f64, values: &[T], valid: Option<&BitBuffer>, skip_nans: bool) { + let sum = fold_float(values, valid, |v| v); + *acc += if skip_nans && sum.is_nan() { + fold_float(values, valid, non_nan) } else { - for &v in slice { - *acc += ToPrimitive::to_f64(&v).vortex_expect("float to f64"); - } + sum + }; +} + +/// Sum `read` of each value, or of each `valid` value, into [`FLOAT_SUM_LANES`] partial sums. +#[inline] +fn fold_float( + values: &[T], + valid: Option<&BitBuffer>, + read: impl Fn(f64) -> f64, +) -> f64 { + let add = |sum: f64, v: T| sum + read(ToPrimitive::to_f64(&v).vortex_expect("float to f64")); + match valid { + None => values.fold_lanes::(0.0, add, |a, b| a + b), + Some(valid) => values.fold_lanes_masked::( + valid, + T::zero(), + 0.0, + add, + |a, b| a + b, + ), } } +/// `v`, or zero when `v` is NaN. +#[inline] +fn non_nan(v: f64) -> f64 { + if v.is_nan() { 0.0 } else { v } +} + /// Sum all values into a `u64` accumulator. For types narrower than 64 bits, values are summed in -/// chunks of [`SUM_CHUNK`] with a single checked add per chunk, which lets the inner loop vectorize -/// to packed widening adds. `u64` input keeps a per-element checked add since a chunk of `u64`s -/// could itself overflow. Returns `true` on overflow. +/// chunks of [`SUM_CHUNK`] by [`narrow_sum`] with a single checked add per chunk. A chunk of `u64`s could itself overflow, so `u64` input is summed by +/// [`sum_wide`]. Returns `true` on overflow. pub(crate) fn sum_unsigned_all(acc: &mut u64, slice: &[T]) -> bool where T: NativePType + AsPrimitive, { if T::PTYPE == PType::U64 { - for &v in slice { - if checked_add_u64(acc, v.as_()) { - return true; - } - } - return false; + return sum_wide(acc, slice, None); } for chunk in slice.chunks(SUM_CHUNK) { - let chunk_sum: u64 = chunk.iter().map(|&v| v.as_()).sum(); + let chunk_sum = narrow_sum::(chunk); if checked_add_u64(acc, chunk_sum) { return true; } @@ -117,15 +154,10 @@ where T: NativePType + AsPrimitive, { if T::PTYPE == PType::I64 { - for &v in slice { - if checked_add_i64(acc, v.as_()) { - return true; - } - } - return false; + return sum_wide(acc, slice, None); } for chunk in slice.chunks(SUM_CHUNK) { - let chunk_sum: i64 = chunk.iter().map(|&v| v.as_()).sum(); + let chunk_sum = narrow_sum::(chunk); if checked_add_i64(acc, chunk_sum) { return true; } @@ -133,30 +165,48 @@ where false } -/// Sum the valid elements of `p`, whose validity is `valid`. +/// Sum the valid elements. +/// +/// Every kernel reads null lanes as zero, so the sum vectorizes the same way as the all-valid path +/// however the valid values are spread. fn accumulate_primitive_valid( inner: &mut SumState, p: &PrimitiveArray, valid: &MaskValues, skip_nans: bool, ) -> VortexResult { - let valid = valid.bit_buffer(); match inner { SumState::Unsigned(acc) => match_each_native_ptype!(p.ptype(), - unsigned: |T| { Ok(sum_unsigned_masked(acc, p.as_slice::(), valid)) }, + unsigned: |T| { + let values = p.as_slice::(); + if T::PTYPE == PType::U64 { + return Ok(sum_wide(acc, values, Some(valid.bit_buffer()))); + } + Ok(sum_masked(values, valid.bit_buffer(), |chunk_sum| { + checked_add_u64(acc, chunk_sum) + })) + }, signed: |_T| { vortex_panic!("unsigned sum state with signed input") }, floating: |_T| { vortex_panic!("unsigned sum state with float input") } ), SumState::Signed(acc) => match_each_native_ptype!(p.ptype(), unsigned: |_T| { vortex_panic!("signed sum state with unsigned input") }, - signed: |T| { Ok(sum_signed_masked(acc, p.as_slice::(), valid)) }, + signed: |T| { + let values = p.as_slice::(); + if T::PTYPE == PType::I64 { + return Ok(sum_wide(acc, values, Some(valid.bit_buffer()))); + } + Ok(sum_masked(values, valid.bit_buffer(), |chunk_sum| { + checked_add_i64(acc, chunk_sum) + })) + }, floating: |_T| { vortex_panic!("signed sum state with float input") } ), SumState::Float(acc) => match_each_native_ptype!(p.ptype(), unsigned: |_T| { vortex_panic!("float sum state with unsigned input") }, signed: |_T| { vortex_panic!("float sum state with signed input") }, floating: |T| { - sum_float_masked(acc, p.as_slice::(), valid, skip_nans); + sum_float(acc, p.as_slice::(), Some(valid.bit_buffer()), skip_nans); Ok(false) } ), @@ -166,79 +216,209 @@ fn accumulate_primitive_valid( } } -/// Sum the values of `values` that `valid` marks valid into a `u64` accumulator. Returns `true` on -/// overflow. -/// -/// Values narrower than 64 bits read null lanes as zero, so the sum vectorizes the same way as -/// [`sum_unsigned_all`] however the valid values are spread. `u64` values check every element -/// for overflow, so they sum each contiguous run of valid values through [`sum_unsigned_all`]. -fn sum_unsigned_masked(acc: &mut u64, values: &[T], valid: &BitBuffer) -> bool -where - T: NativePType + AsPrimitive, -{ - if T::PTYPE == PType::U64 { - return valid - .set_slices() - .any(|(start, end)| sum_unsigned_all(acc, &values[start..end])); - } - sum_chunks_masked(values, valid, |chunk_sum| checked_add_u64(acc, chunk_sum)) -} - -/// Signed counterpart of [`sum_unsigned_masked`]. -fn sum_signed_masked(acc: &mut i64, values: &[T], valid: &BitBuffer) -> bool -where - T: NativePType + AsPrimitive, -{ - if T::PTYPE == PType::I64 { - return valid - .set_slices() - .any(|(start, end)| sum_signed_all(acc, &values[start..end])); - } - sum_chunks_masked(values, valid, |chunk_sum| checked_add_i64(acc, chunk_sum)) -} - -/// Sum the values of `values` that `valid` marks valid into an `f64` accumulator, skipping NaN -/// values when `skip_nans` is set, as [`sum_float_all`] does. -/// -/// Null and skipped lanes read as zero. The values are added in order, as by [`sum_float_all`], -/// and adding zero leaves the accumulator unchanged: it starts at positive zero, so it is never -/// negative zero, the one value adding zero would change. The result is therefore exactly that of -/// summing the valid values alone, without walking their runs. -fn sum_float_masked( - acc: &mut f64, - values: &[T], - valid: &BitBuffer, - skip_nans: bool, -) { - *acc = values.fold_masked(valid, T::zero(), *acc, |sum, v| { - if skip_nans && v.is_nan() { - sum - } else { - sum + ToPrimitive::to_f64(&v).vortex_expect("float to f64") - } - }); -} - /// Sum the valid values of an integer slice narrower than 64 bits in chunks of [`SUM_CHUNK`], /// reading null lanes as zero, and hand each chunk's sum to `add`, which returns `true` on /// overflow. A chunk cannot overflow its 64-bit sum, as in [`sum_unsigned_all`]. -fn sum_chunks_masked(values: &[T], valid: &BitBuffer, mut add: impl FnMut(S) -> bool) -> bool +fn sum_masked(values: &[T], valid: &BitBuffer, mut add: impl FnMut(S) -> bool) -> bool where T: NativePType + AsPrimitive, S: NativePType + Add, { (0..values.len()).step_by(SUM_CHUNK).any(|start| { let end = (start + SUM_CHUNK).min(values.len()); - let chunk_sum = values[start..end].fold_masked( + let chunk_sum = values[start..end].fold_lanes_masked::( &valid.slice(start..end), T::zero(), S::zero(), |sum, v| sum + v.as_(), + |a, b| a + b, ); add(chunk_sum) }) } +/// Sum a chunk of at most [`SUM_CHUNK`] integers narrower than 64 bits, which cannot overflow `S`. +/// +/// The compiler would vectorize an in-order sum on its own, but into a loop of two adds that runs +/// at the mercy of where it lands in the instruction cache: on x86-64 without SSE4.1 the same loop +/// ran 30% slower once it straddled a 64-byte boundary. [`NARROW_SUM_LANES`] partial sums make a +/// loop body large enough not to care, and are as fast or faster on every target, with one +/// exception: x86-64 without SSE4.1 has no vector sign extension, and widening signed values into +/// that many lanes costs more than the in-order loop does. +fn narrow_sum(chunk: &[T]) -> S +where + T: NativePType + AsPrimitive, + S: NativePType + Add, +{ + if NO_SIGN_EXTENSION && T::PTYPE.is_signed_int() { + return chunk.iter().map(|&v| v.as_()).fold(S::zero(), |a, b| a + b); + } + chunk.fold_lanes::(S::zero(), |sum, v| sum + v.as_(), |a, b| a + b) +} + +/// Whether the target lacks a vector sign extension of narrow integers to 64 bits. +const NO_SIGN_EXTENSION: bool = cfg!(all(target_arch = "x86_64", not(target_feature = "sse4.1"))); + +/// Sum 64-bit integers, or only the `valid` ones, into `acc`. Returns `true` on overflow. +/// +/// A checked add per element does not vectorize, so the values are added with wrapping, and +/// evidence of each value's magnitude is ORed together alongside. Both reassociate, so the +/// compiler vectorizes the fold in lane order with the width and interleave it picks, which beats +/// independent per-lane partial sums on every target. If every value is small enough that +/// `values.len()` of them cannot overflow, which holds for most data, the sum is exact. Otherwise the values are summed again in +/// [`sum_exact`]. Either way the result overflows only if the exact sum does not fit, whatever the +/// order of the values. +fn sum_wide(acc: &mut S, values: &[T], valid: Option<&BitBuffer>) -> bool +where + T: NativePType + AsPrimitive, + S: WrappingSum, +{ + let Some(bound) = S::Bound::for_len(values.len()) else { + return false; + }; + + let add = |(sum, evidence): (S, S), v: T| { + let v = v.as_(); + (sum.wrapping_add(v), evidence | bound.evidence(v)) + }; + let zero = (S::zero(), S::zero()); + + let (sum, evidence) = match valid { + None => values.iter().fold(zero, |acc, &v| add(acc, v)), + Some(valid) => values.fold_masked(valid, T::zero(), zero, add), + }; + + if !bound.holds(evidence) { + return sum_exact(acc, values, valid); + } + S::checked_add_into(acc, sum) +} + +/// Sum the values, or only the `valid` ones, into `acc` in 128-bit arithmetic, which no slice of +/// 64-bit integers can overflow. Returns `true` if the result does not fit `S`. +#[cold] +#[inline(never)] +fn sum_exact(acc: &mut S, values: &[T], valid: Option<&BitBuffer>) -> bool +where + T: NativePType, + S: WrappingSum, +{ + let sum: i128 = match valid { + None => values.iter().map(|&v| to_i128(v)).sum(), + Some(valid) => valid.set_indices().map(|i| to_i128(values[i])).sum(), + }; + + match S::from(to_i128(*acc) + sum) { + Some(total) => { + *acc = total; + false + } + None => true, + } +} + +fn to_i128(v: T) -> i128 { + v.to_i128().vortex_expect("64-bit integers fit i128") +} + +/// A 64-bit sum accumulator that adds with wrapping, with a bound on its values under which a +/// wrapping sum cannot overflow. +trait WrappingSum: NativePType + BitOr { + /// The magnitude bound for this type. + type Bound: SumBound; + + /// Add with wrapping. + fn wrapping_add(self, v: Self) -> Self; + + /// Checked add of `v` into `acc`, returning `true` on overflow. + fn checked_add_into(acc: &mut Self, v: Self) -> bool; +} + +/// A bound on the magnitude of the values of a sum, tested with one add and one OR per value. +trait SumBound: Copy { + /// The bound under which `len` values cannot overflow their sum, or `None` if `len` is zero. + fn for_len(len: usize) -> Option; + + /// Evidence of the magnitude of `v`, to OR together over all values. + fn evidence(self, v: S) -> S; + + /// Whether the ORed evidence shows every value is within the bound. + fn holds(self, evidence: S) -> bool; +} + +/// Signed values within `[-2^bits, 2^bits)`, where `len` of them sum to less than `2^63` in +/// magnitude. A value is within the bound iff adding `2^bits` leaves it in `[0, 2^(bits + 1))`. +#[derive(Clone, Copy)] +struct SignedBound { + bits: u32, +} + +impl SumBound for SignedBound { + fn for_len(len: usize) -> Option { + (len != 0).then(|| Self { + bits: 63 - (usize::BITS - len.leading_zeros()), + }) + } + + #[inline] + fn evidence(self, v: i64) -> i64 { + v.wrapping_add(1 << self.bits) + } + + fn holds(self, evidence: i64) -> bool { + evidence >> (self.bits + 1) == 0 + } +} + +/// Unsigned values below `2^bits`, where `len` of them sum to less than `2^64`. +#[derive(Clone, Copy)] +struct UnsignedBound { + bits: u32, +} + +impl SumBound for UnsignedBound { + fn for_len(len: usize) -> Option { + (len != 0).then(|| Self { + bits: 64 - (usize::BITS - len.leading_zeros()), + }) + } + + #[inline] + fn evidence(self, v: u64) -> u64 { + v + } + + fn holds(self, evidence: u64) -> bool { + evidence >> self.bits == 0 + } +} + +impl WrappingSum for i64 { + type Bound = SignedBound; + + #[inline] + fn wrapping_add(self, v: i64) -> i64 { + i64::wrapping_add(self, v) + } + + fn checked_add_into(acc: &mut i64, v: i64) -> bool { + checked_add_i64(acc, v) + } +} + +impl WrappingSum for u64 { + type Bound = UnsignedBound; + + #[inline] + fn wrapping_add(self, v: u64) -> u64 { + u64::wrapping_add(self, v) + } + + fn checked_add_into(acc: &mut u64, v: u64) -> bool { + checked_add_u64(acc, v) + } +} + #[cfg(test)] mod tests { use std::sync::LazyLock; @@ -544,6 +724,109 @@ mod tests { Ok(()) } + /// Values past the magnitude bound take the exact sum, which overflows only if the total does + /// not fit, however the values are ordered. + #[test] + fn sum_i64_near_the_limits_is_exact() -> VortexResult<()> { + let fits = + PrimitiveArray::new(buffer![i64::MAX, 1, i64::MIN, -1, 5], Validity::NonNullable) + .into_array(); + let overflows = + PrimitiveArray::new(buffer![i64::MAX - 1, 1, 1, 0], Validity::NonNullable).into_array(); + let underflows = + PrimitiveArray::new(buffer![i64::MIN, 0, -1], Validity::NonNullable).into_array(); + + let ctx = &mut SESSION.create_execution_ctx(); + assert_eq!(sum(&fits, ctx)?.as_primitive().typed_value::(), Some(4)); + assert!(sum(&overflows, ctx)?.is_null()); + assert!(sum(&underflows, ctx)?.is_null()); + Ok(()) + } + + /// The null rows hold values that would overflow, in a full 64-row word and in the partial + /// word after it, on both the bounded and the exact path. + #[test] + fn sum_64_bit_ignores_values_under_nulls() -> VortexResult<()> { + let validity = || Validity::from_iter((0..100).map(|i| i % 3 != 0)); + let unsigned: Buffer = (0..100u64) + .map(|i| if i % 3 == 0 { u64::MAX } else { i }) + .collect(); + let signed: Buffer = (0..100i64) + .map(|i| match i % 6 { + 0 => i64::MIN, + 3 => i64::MAX, + _ => -i, + }) + .collect(); + let large: Buffer = (0..100i64) + .map(|i| if i % 3 == 0 { i64::MAX } else { i64::MAX / 100 }) + .collect(); + + let ctx = &mut SESSION.create_execution_ctx(); + let unsigned = sum(&PrimitiveArray::new(unsigned, validity()).into_array(), ctx)?; + let signed = sum(&PrimitiveArray::new(signed, validity()).into_array(), ctx)?; + let large = sum(&PrimitiveArray::new(large, validity()).into_array(), ctx)?; + + let valid = (0..100i64).filter(|i| i % 3 != 0); + assert_eq!( + unsigned.as_primitive().typed_value::(), + Some(valid.clone().map(i64::cast_unsigned).sum()) + ); + assert_eq!( + signed.as_primitive().typed_value::(), + Some(-valid.clone().sum::()) + ); + assert_eq!( + large.as_primitive().typed_value::(), + Some(valid.count() as i64 * (i64::MAX / 100)) + ); + Ok(()) + } + + /// NaNs under nulls are never read, and valid NaNs are skipped or poison the sum as asked. + #[test] + fn sum_masked_float_nans() -> VortexResult<()> { + let values: Buffer = (0..100) + .map(|i| match i % 5 { + 0 => f64::NAN, + 4 if i < 50 => f64::NAN, + _ => f64::from(i), + }) + .collect(); + let under_nulls = PrimitiveArray::new( + values.clone(), + Validity::from_iter((0..100).map(|i| i % 5 != 0 && !(i % 5 == 4 && i < 50))), + ) + .into_array(); + let valid_nans = + PrimitiveArray::new(values, Validity::from_iter((0..100).map(|i| i % 5 != 0))) + .into_array(); + + let expected: f64 = (0..100) + .filter(|&i| i % 5 != 0 && !(i % 5 == 4 && i < 50)) + .map(f64::from) + .sum(); + assert_eq!( + sum_with_options(&under_nulls, NumericalAggregateOpts::include_nans())? + .as_primitive() + .typed_value::(), + Some(expected) + ); + assert_eq!( + sum_with_options(&valid_nans, NumericalAggregateOpts::default())? + .as_primitive() + .typed_value::(), + Some(expected) + ); + assert!( + sum_with_options(&valid_nans, NumericalAggregateOpts::include_nans())? + .as_primitive() + .typed_value::() + .is_some_and(f64::is_nan) + ); + Ok(()) + } + #[test] fn sum_checked_overflow_is_saturated() -> VortexResult<()> { let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); diff --git a/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs b/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs index 20eb19abd1c..47203b9e1ed 100644 --- a/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs +++ b/vortex-array/src/aggregate_fn/fns/sum_v2/tests.rs @@ -181,13 +181,15 @@ fn overflow_is_explicit_and_absorbing() -> VortexResult<()> { accumulator.accumulate(&batch, &mut SESSION.create_execution_ctx())?; assert!(accumulator.is_saturated()); + // A batch whose sum overflows is not added, so the partial keeps the sum from before it. The + // flag masks that sum in every result. let partial = accumulator.partial_scalar()?; let fields = partial.as_struct(); assert_eq!( fields .field("sum") .and_then(|sum| sum.as_primitive().typed_value::()), - Some(i64::MAX) + Some(0) ); assert_eq!( fields diff --git a/vortex-compute/Cargo.toml b/vortex-compute/Cargo.toml index 7b90bfb9cd5..53c28ec33da 100644 --- a/vortex-compute/Cargo.toml +++ b/vortex-compute/Cargo.toml @@ -40,3 +40,7 @@ harness = false [[bench]] name = "fold_masked" harness = false + +[[bench]] +name = "fold_lanes" +harness = false diff --git a/vortex-compute/benches/fold_lanes.rs b/vortex-compute/benches/fold_lanes.rs new file mode 100644 index 00000000000..95012ec7122 --- /dev/null +++ b/vortex-compute/benches/fold_lanes.rs @@ -0,0 +1,330 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Reductions that the compiler cannot reassociate on its own, folded with +//! [`IndexedFoldExt::fold_lanes`] and [`IndexedFoldExt::fold_lanes_masked`], against in-order +//! folds and the arrow-rs aggregate kernels over the same data. +//! +//! The folds are the shapes the aggregate functions use: a float sum, float extrema by float +//! comparison with one fold per extremum, and a 64-bit sum with a magnitude bound in place of a +//! checked add per element. The +//! `masked_runs` baselines walk the runs of valid values in order, as the aggregates did before. + +#![expect(clippy::unwrap_used)] + +use arrow_arith::aggregate; +use arrow_array::Float64Array; +use arrow_buffer::NullBuffer; +use divan::Bencher; +use mimalloc::MiMalloc; +use rand::SeedableRng; +use rand::prelude::*; +use rand::rngs::StdRng; +use vortex_buffer::BitBuffer; +use vortex_compute::lane_kernels::IndexedFoldExt; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +fn main() { + divan::main(); +} + +const N: usize = 15_000; +const DENSITIES: &[f64] = &[0.5, 0.99]; + +fn floats() -> Vec { + let mut rng = StdRng::seed_from_u64(1); + (0..N).map(|_| rng.random_range(-1000.0..1000.0)).collect() +} + +fn ints() -> Vec { + let mut rng = StdRng::seed_from_u64(1); + (0..N).map(|_| rng.random_range(-1_000_000..1_000_000)).collect() +} + +fn valid(density: f64) -> BitBuffer { + let mut rng = StdRng::seed_from_u64(2); + BitBuffer::from_iter((0..N).map(|_| rng.random_bool(density))) +} + +/// The `[start, end)` runs of valid lanes. +fn runs(valid: &BitBuffer) -> Vec<(usize, usize)> { + let mut runs = Vec::new(); + let mut start = None; + for i in 0..valid.len() { + match (valid.value(i), start) { + (true, None) => start = Some(i), + (false, Some(s)) => { + runs.push((s, i)); + start = None; + } + _ => {} + } + } + if let Some(s) = start { + runs.push((s, valid.len())); + } + runs +} + +/// `1 << BITS` bounds values that `N` of can sum without overflow. +const BITS: u32 = 63 - (usize::BITS - N.leading_zeros()); + +/// Add `v` to a wrapping sum and OR its offset magnitude into `evidence`, which stays below +/// `2^(BITS + 1)` while every value is within `[-2^BITS, 2^BITS)`. +fn add_bounded((sum, evidence): (i64, i64), v: i64) -> (i64, i64) { + (sum.wrapping_add(v), evidence | v.wrapping_add(1 << BITS)) +} + +fn merge_bounded((a, ea): (i64, i64), (b, eb): (i64, i64)) -> (i64, i64) { + (a.wrapping_add(b), ea | eb) +} + +fn lower(min: f64, v: f64) -> f64 { + if v < min { v } else { min } +} + +fn higher(max: f64, v: f64) -> f64 { + if v > max { v } else { max } +} + +#[divan::bench] +fn sum_f64_in_order(bencher: Bencher) { + let values = floats(); + bencher.bench(|| divan::black_box(&values).iter().sum::()); +} + +#[divan::bench] +fn sum_f64_lanes(bencher: Bencher) { + let values = floats(); + bencher.bench(|| { + divan::black_box(&values[..]).fold_lanes::<16, _, _, _>(0.0, |a, v| a + v, |a, b| a + b) + }); +} + +#[divan::bench] +fn sum_f64_arrow(bencher: Bencher) { + let array = Float64Array::from(floats()); + bencher.bench(|| aggregate::sum(divan::black_box(&array)).unwrap()); +} + +#[divan::bench(args = DENSITIES)] +fn sum_f64_masked_runs(bencher: Bencher, density: f64) { + let (values, runs) = (floats(), runs(&valid(density))); + bencher.bench(|| { + divan::black_box(&runs) + .iter() + .map(|&(start, end)| values[start..end].iter().sum::()) + .sum::() + }); +} + +#[divan::bench(args = DENSITIES)] +fn sum_f64_masked_lanes(bencher: Bencher, density: f64) { + let (values, valid) = (floats(), valid(density)); + bencher.bench(|| { + divan::black_box(&values[..]).fold_lanes_masked::<16, _, _, _>( + &valid, + 0.0, + 0.0, + |a, v| a + v, + |a, b| a + b, + ) + }); +} + +#[divan::bench(args = DENSITIES)] +fn sum_f64_masked_arrow(bencher: Bencher, density: f64) { + let valid = valid(density); + let array = Float64Array::new( + floats().into(), + Some(NullBuffer::from_iter((0..N).map(|i| valid.value(i)))), + ); + bencher.bench(|| aggregate::sum(divan::black_box(&array)).unwrap()); +} + +/// Float extrema under the total order, as the aggregate computed them before. +#[divan::bench] +fn min_max_f64_in_order(bencher: Bencher) { + let values = floats(); + bencher.bench(|| { + divan::black_box(&values) + .iter() + .fold((f64::NAN, -f64::NAN), |(min, max), &v| { + ( + if v.total_cmp(&min).is_lt() { v } else { min }, + if v.total_cmp(&max).is_gt() { v } else { max }, + ) + }) + }); +} + +/// One fold per extremum: a fold into a pair of floats stays scalar on NEON. +#[divan::bench] +fn min_max_f64_lanes(bencher: Bencher) { + let values = floats(); + bencher.bench(|| { + let values = divan::black_box(&values[..]); + let seed = values[0]; + ( + values.fold_lanes::<8, _, _, _>(seed, lower, lower), + values.fold_lanes::<8, _, _, _>(seed, higher, higher), + ) + }); +} + +#[divan::bench(args = DENSITIES)] +fn min_max_f64_masked_lanes(bencher: Bencher, density: f64) { + let (values, valid) = (floats(), valid(density)); + bencher.bench(|| { + let values = divan::black_box(&values[..]); + let seed = values[0]; + ( + values.fold_lanes_masked::<8, _, _, _>(&valid, seed, seed, lower, lower), + values.fold_lanes_masked::<8, _, _, _>(&valid, seed, seed, higher, higher), + ) + }); +} + +#[divan::bench] +fn checked_sum_i64_in_order(bencher: Bencher) { + let values = ints(); + bencher.bench(|| { + divan::black_box(&values) + .iter() + .try_fold(0i64, |a, &v| a.checked_add(v)) + }); +} + +#[divan::bench] +fn checked_sum_i64_lanes(bencher: Bencher) { + let values = ints(); + bencher.bench(|| { + divan::black_box(&values[..]).fold_lanes::<8, _, _, _>((0, 0), add_bounded, merge_bounded) + }); +} + +#[divan::bench(args = DENSITIES)] +fn checked_sum_i64_masked_runs(bencher: Bencher, density: f64) { + let (values, runs) = (ints(), runs(&valid(density))); + bencher.bench(|| { + divan::black_box(&runs).iter().try_fold(0i64, |acc, &(start, end)| { + values[start..end] + .iter() + .try_fold(acc, |a, &v| a.checked_add(v)) + }) + }); +} + +#[divan::bench(args = DENSITIES)] +fn checked_sum_i64_masked_lanes(bencher: Bencher, density: f64) { + let (values, valid) = (ints(), valid(density)); + bencher.bench(|| { + divan::black_box(&values[..]).fold_lanes_masked::<8, _, _, _>( + &valid, + 0, + (0, 0), + add_bounded, + merge_bounded, + ) + }); +} + +/// The bounded sum in lane order, which the compiler may reassociate since integer addition and OR +/// do, so it picks its own vector width and interleave. +#[divan::bench] +fn checked_sum_i64_in_order_bounded(bencher: Bencher) { + let values = ints(); + bencher.bench(|| { + divan::black_box(&values) + .iter() + .fold((0, 0), |acc, &v| add_bounded(acc, v)) + }); +} + +#[divan::bench(args = DENSITIES)] +fn checked_sum_i64_masked_in_order_bounded(bencher: Bencher, density: f64) { + let (values, valid) = (ints(), valid(density)); + bencher.bench(|| divan::black_box(&values[..]).fold_masked(&valid, 0, (0, 0), add_bounded)); +} + +fn narrow() -> Vec { + let mut rng = StdRng::seed_from_u64(3); + (0..N).map(|_| rng.random_range(0..2000)).collect() +} + +/// A sum of narrow integers widened to 64 bits in lane order, the loop the aggregate runs. +#[divan::bench] +fn sum_u32_in_order(bencher: Bencher) { + let values = narrow(); + bencher.bench(|| { + divan::black_box(&values) + .iter() + .map(|&v| u64::from(v)) + .sum::() + }); +} + +#[divan::bench] +fn sum_u32_lanes(bencher: Bencher) { + let values = narrow(); + bencher.bench(|| { + divan::black_box(&values[..]).fold_lanes::<16, _, _, _>( + 0u64, + |a, v| a + u64::from(v), + |a, b| a + b, + ) + }); +} + +#[divan::bench(args = DENSITIES)] +fn sum_u32_masked_in_order(bencher: Bencher, density: f64) { + let (values, valid) = (narrow(), valid(density)); + bencher.bench(|| { + divan::black_box(&values[..]).fold_masked(&valid, 0, 0u64, |a, v| a + u64::from(v)) + }); +} + +#[divan::bench(args = DENSITIES)] +fn sum_u32_masked_lanes(bencher: Bencher, density: f64) { + let (values, valid) = (narrow(), valid(density)); + bencher.bench(|| { + divan::black_box(&values[..]).fold_lanes_masked::<16, _, _, _>( + &valid, + 0, + 0u64, + |a, v| a + u64::from(v), + |a, b| a + b, + ) + }); +} + +fn signed_narrow() -> Vec { + let mut rng = StdRng::seed_from_u64(4); + (0..N).map(|_| rng.random_range(-1000..1000)).collect() +} + +/// A sum of `i32` widened to 64 bits in lane order, the loop the aggregate ran before. +#[divan::bench] +fn sum_i32_in_order(bencher: Bencher) { + let values = signed_narrow(); + bencher.bench(|| { + divan::black_box(&values) + .iter() + .map(|&v| i64::from(v)) + .sum::() + }); +} + +#[divan::bench(consts = [8, 16])] +fn sum_i32_lanes(bencher: Bencher) { + let values = signed_narrow(); + bencher.bench(|| { + divan::black_box(&values[..]).fold_lanes::( + 0i64, + |a, v| a + i64::from(v), + |a, b| a + b, + ) + }); +} diff --git a/vortex-compute/src/lane_kernels/fold.rs b/vortex-compute/src/lane_kernels/fold.rs index b040fdee0d9..f728acb2fcc 100644 --- a/vortex-compute/src/lane_kernels/fold.rs +++ b/vortex-compute/src/lane_kernels/fold.rs @@ -1,7 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Masked reduction kernels: fold every lane of an [`IndexedSource`] into one accumulator. +//! Reduction kernels: fold every lane of an [`IndexedSource`] into one accumulator. +//! +//! [`IndexedFoldExt::fold_masked`] folds in lane order into a single accumulator, which vectorizes +//! only when the compiler may reorder the fold, such as an integer sum or minimum. +//! [`IndexedFoldExt::fold_lanes`] and [`IndexedFoldExt::fold_lanes_masked`] fold into independent +//! per-lane accumulators and merge them at the end, which vectorizes any fold, such as a float sum. use vortex_buffer::BitBuffer; @@ -54,6 +59,218 @@ pub trait IndexedFoldExt: IndexedSource + Sized { } acc } + + /// Fold every lane into `LANES` independent accumulators with `f`, then combine them with + /// `merge`. + /// + /// Lane `i` goes into accumulator `i % LANES`, so each accumulator is a separate chain and one + /// row of `LANES` lanes is a single vector step. The compiler needs no permission to reorder + /// `f`, so a fold it cannot reassociate, such as a float sum, a float minimum or an integer + /// sum that also collects overflow evidence, vectorizes the same way as an integer sum. The + /// result can differ from an in-order fold only in the order `f` sees the lanes, such as the + /// rounding of a float sum. + /// + /// `init` starts every accumulator, so it must be an identity of `merge`. Pick `LANES` so that + /// the accumulators fill a few vector registers: `LANES * size_of::()` of 64 to 128 bytes + /// hides the latency of `f` on NEON, AVX2 and AVX-512 alike. `LANES` must divide 64. + /// + /// Prefer an accumulator of one primitive. A pair of integers vectorizes, but a pair of + /// floats, such as a minimum and a maximum, stays scalar on NEON and spills on SSE2, where two + /// folds of one float each vectorize on every target. + #[inline] + fn fold_lanes(self, init: A, f: F, merge: M) -> A + where + A: Copy, + F: Fn(A, Self::Item) -> A, + M: Fn(A, A) -> A, + { + const { assert!(LANES > 0 && 64 % LANES == 0, "LANES must divide 64") }; + + let values = self; + let len = values.len(); + + // Short of one row, such as a small group of a grouped aggregate, an in-order fold into one + // accumulator skips the merge of the others, which `init` leaves unchanged. + if len < LANES { + let mut acc = init; + for i in 0..len { + // SAFETY: i < len. + acc = f(acc, unsafe { values.get_unchecked(i) }); + } + return acc; + } + + let mut acc = [init; LANES]; + let rows = len / LANES; + for row in 0..rows { + row_barrier(); + fold_row(&values, &mut acc, &f, row * LANES, LANES); + } + fold_row(&values, &mut acc, &f, rows * LANES, len % LANES); + + merge_lanes(acc, merge) + } + + /// Fold the valid lanes into `LANES` independent accumulators with `f`, reading null lanes as + /// `null`, then combine the accumulators with `merge`. + /// + /// This is [`fold_lanes`] with a validity mask. `null` must be an identity of `f`, as in + /// [`fold_masked`]. Each lane is selected against a constant bit of its row of the validity + /// word, with no branch per lane and no per-lane variable shift, so the fold vectorizes on + /// every target at the same cost at any validity density. Words with no valid lane are + /// skipped. + /// + /// [`fold_lanes`]: IndexedFoldExt::fold_lanes + /// [`fold_masked`]: IndexedFoldExt::fold_masked + /// + /// # Panics + /// + /// Panics if `self.len() != mask.len()`. + #[inline] + fn fold_lanes_masked( + self, + mask: &BitBuffer, + null: Self::Item, + init: A, + f: F, + merge: M, + ) -> A + where + Self::Item: Copy, + A: Copy, + F: Fn(A, Self::Item) -> A, + M: Fn(A, A) -> A, + { + const { assert!(LANES > 0 && 64 % LANES == 0, "LANES must divide 64") }; + + let values = self; + let len = values.len(); + assert_eq!(len, mask.len(), "values and mask must have the same length"); + + let chunks = mask.chunks(); + let mut acc = [init; LANES]; + + for (chunk_idx, word) in chunks.iter().enumerate() { + if word == 0 { + continue; + } + + let base = chunk_idx * 64; + for row in 0..64 / LANES { + row_barrier(); + let bits = word >> (row * LANES); + fold_row_masked(&values, null, &mut acc, &f, bits, base + row * LANES, LANES); + } + } + + let remainder = len % 64; + let word = chunks.remainder_bits(); + if word != 0 { + let base = len - remainder; + for row in 0..remainder.div_ceil(LANES) { + let lanes = (remainder - row * LANES).min(LANES); + let bits = word >> (row * LANES); + fold_row_masked(&values, null, &mut acc, &f, bits, base + row * LANES, lanes); + } + } + + merge_lanes(acc, merge) + } +} + +/// Keep the vectorizers from packing one lane of several rows. +/// +/// Each accumulator of [`IndexedFoldExt::fold_lanes`] is a reduction over the rows, and the +/// vectorizers take a reduction they recognize, such as an integer minimum or sum, and vectorize it +/// across rows: one vector holds the same lane of several rows and fills it with a strided load of +/// each lane. The loop vectorizer does so over the loop of rows, and the SLP vectorizer over the +/// rows of a 64-lane chunk, which the masked fold unrolls. The opaque barrier before each row is a +/// point no load moves across, which leaves each row to the SLP vectorizer to pack the row's +/// contiguous lanes into full vectors. It emits no instruction. +#[allow(clippy::inline_always)] +#[inline(always)] +fn row_barrier() { + std::hint::black_box(()); +} + +/// Fold `lanes` lanes from `base`, lane `i` into `acc[i]`. +#[allow(clippy::inline_always)] +#[inline(always)] +fn fold_row( + values: &S, + acc: &mut [A; LANES], + f: &F, + base: usize, + lanes: usize, +) where + S: IndexedSource, + A: Copy, + F: Fn(A, S::Item) -> A, +{ + for (lane, acc) in acc.iter_mut().enumerate().take(lanes) { + // SAFETY: the caller guarantees base + lanes <= len. + let value = unsafe { values.get_unchecked(base + lane) }; + *acc = f(*acc, value); + } +} + +/// Fold `lanes` lanes from `base`, lane `i` into `acc[i]`, reading the lanes whose bit in `bits` +/// is unset as `null`. +/// +/// Lanes of at most 32 bits test a 32-bit copy of the row bits, so the lane mask is as wide as the +/// lane and needs no narrowing. Wider lanes, and rows wider than 32 lanes, test all 64 bits, except +/// where [`NARROW_WIDE_LANES`] says otherwise. +#[allow(clippy::inline_always)] +#[inline(always)] +fn fold_row_masked( + values: &S, + null: S::Item, + acc: &mut [A; LANES], + f: &F, + bits: u64, + base: usize, + lanes: usize, +) where + S: IndexedSource, + S::Item: Copy, + A: Copy, + F: Fn(A, S::Item) -> A, +{ + #[expect( + clippy::cast_possible_truncation, + reason = "rows of at most 32 lanes keep their low 32 bits" + )] + let narrow_bits = bits as u32; + let narrow = LANES <= 32 && (size_of::() <= 4 || NARROW_WIDE_LANES); + + for (lane, acc) in acc.iter_mut().enumerate().take(lanes) { + // SAFETY: the caller guarantees base + lanes <= len. + let value = unsafe { values.get_unchecked(base + lane) }; + let valid = if narrow { + narrow_bits & LANE_BITS[lane % 32] != 0 + } else { + bits & WIDE_LANE_BITS[lane] != 0 + }; + *acc = f(*acc, if valid { value } else { null }); + } +} + +/// Combine the accumulators pairwise, halving their count each step. +#[allow(clippy::inline_always)] +#[inline(always)] +fn merge_lanes(mut acc: [A; LANES], merge: M) -> A +where + A: Copy, + M: Fn(A, A) -> A, +{ + let mut width = LANES; + while width > 1 { + width /= 2; + for lane in 0..width { + acc[lane] = merge(acc[lane], acc[lane + width]); + } + } + acc[0] } /// The two 32-lane halves of a 64-lane validity word, each with the offset of its first lane. @@ -78,6 +295,23 @@ const LANE_BITS: [u32; 32] = { bits }; +/// Whether 64-bit lanes test 32-bit row bits and widen the lane mask. x86-64 below SSE4.1 has no +/// 64-bit vector compare, which makes the widening the cheaper of the two, by up to a third on +/// masked `f64` folds. Every other target compares 64-bit lanes faster. +const NARROW_WIDE_LANES: bool = cfg!(all(target_arch = "x86_64", not(target_feature = "sse4.1"))); + +/// `WIDE_LANE_BITS[i]` selects lane `i` of a 64-lane validity word, as [`LANE_BITS`] does for a +/// 32-lane one. +const WIDE_LANE_BITS: [u64; 64] = { + let mut bits = [0; 64]; + let mut lane = 0; + while lane < 64 { + bits[lane] = 1 << lane; + lane += 1; + } + bits +}; + /// Fold `lanes` lanes from `base`, reading the lanes whose bit in `bits` is unset as `null`. #[allow(clippy::inline_always)] #[inline(always)] @@ -110,6 +344,7 @@ mod tests { use vortex_buffer::BitBuffer; use super::IndexedFoldExt; + use crate::lane_kernels::LaneZip; #[test] fn fold_masked_reads_null_lanes_as_identity() { @@ -140,6 +375,76 @@ mod tests { assert_eq!(max, 9); } + #[test] + fn fold_lanes_matches_an_in_order_fold() { + for len in [0, 1, 7, 8, 63, 64, 65, 130, 1000] { + let values: Vec = (0..len).map(|v| v * 7 - 300).collect(); + let expected: i64 = values.iter().sum(); + + assert_eq!(values[..].fold_lanes::<8, _, _, _>(0, |a, v| a + v, |a, b| a + b), expected); + assert_eq!(values[..].fold_lanes::<64, _, _, _>(0, |a, v| a + v, |a, b| a + b), expected); + } + } + + #[test] + fn fold_lanes_masked_reads_null_lanes_as_identity() { + for len in [0, 1, 31, 33, 64, 65, 130, 1000] { + let values: Vec = (0..len).map(|v| v - 65).collect(); + let mask = BitBuffer::from_iter((0..len).map(|i| i % 3 == 0 || i % 7 == 0)); + let expected: i64 = values + .iter() + .enumerate() + .filter(|(i, _)| mask.value(*i)) + .map(|(_, &v)| i64::from(v)) + .sum(); + + let narrow = + values[..].fold_lanes_masked::<16, _, _, _>(&mask, 0, 0, |a, v| a + v, |a, b| a + b); + let wide = values[..].fold_lanes_masked::<64, _, _, _>( + &mask, + 0, + 0i64, + |a, v| a + i64::from(v), + |a, b| a + b, + ); + + assert_eq!(i64::from(narrow), expected); + assert_eq!(wide, expected); + } + } + + #[test] + fn fold_lanes_masked_handles_unaligned_masks() { + let values = [5.0f64, 1.0, 9.0, 2.0, 7.0]; + let mask = BitBuffer::from_iter([false, true, true, false, true, true, false]).slice(2..7); + + let max = values[..].fold_lanes_masked::<4, _, _, _>( + &mask, + f64::NEG_INFINITY, + f64::NEG_INFINITY, + f64::max, + f64::max, + ); + + // Valid lanes 0, 2 and 3 hold 5, 9 and 2. + assert_eq!(max, 9.0); + } + + #[test] + fn fold_lanes_drives_a_lane_zip() { + let a: Vec = (0..100).collect(); + let b: Vec = (0..100).map(|v| 100 - v).collect(); + let expected: i64 = a.iter().zip(&b).map(|(x, y)| x * y).sum(); + + let dot = LaneZip::new(&a[..], &b[..]).fold_lanes::<8, _, _, _>( + 0, + |acc, (x, y)| acc + x * y, + |x, y| x + y, + ); + + assert_eq!(dot, expected); + } + #[test] fn fold_masked_of_all_null_lanes_is_init() { let values = [3i64; 70]; diff --git a/vortex-compute/src/lane_kernels/mod.rs b/vortex-compute/src/lane_kernels/mod.rs index e3d512859eb..f2bc5719fc4 100644 --- a/vortex-compute/src/lane_kernels/mod.rs +++ b/vortex-compute/src/lane_kernels/mod.rs @@ -19,7 +19,9 @@ //! the sink itself). //! - [`collect_bits`](mod@collect_bits) — allocating wrappers that bit-pack a predicate into a //! `BitBuffer`. -//! - [`fold`] — masked reductions via [`IndexedFoldExt`], reading null lanes as an identity. +//! - [`fold`] — reductions via [`IndexedFoldExt`]: a masked fold the compiler may reassociate, +//! and lane-parallel folds into independent accumulators for those it may not, such as float +//! sums. Masked forms read null lanes as an identity. //! //! The kernels outside [`collect_bits`](mod@collect_bits) never allocate. Both kernel families //! handle a mask with a non-byte-aligned offset and with a logical `len` shorter than the From 2e8bff5f4609fcb0ce9c0705ae5de7414ee7c429 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 29 Oct 2026 11:02:40 +0000 Subject: [PATCH 4/7] Fix lints and make the nullable distinct count benchmark smaller Format the changed crates with the pinned nightly. Let clippy accept v - v in the float min/max kernel, because it finds NaN and infinite values. Run integer_stats_all_nullable on 16K values instead of 64K. Signed-off-by: Robert Kruszewski --- vortex-array/benches/aggregate_sum.rs | 5 +++- .../src/aggregate_fn/fns/min_max/mod.rs | 8 +++--- .../src/aggregate_fn/fns/min_max/primitive.rs | 8 +++++- .../src/aggregate_fn/fns/sum/primitive.rs | 12 +++++++-- vortex-compressor/benches/integer_stats.rs | 27 ++++++++++--------- vortex-compute/benches/fold_lanes.rs | 16 ++++++----- vortex-compute/src/lane_kernels/fold.rs | 19 ++++++++++--- 7 files changed, 65 insertions(+), 30 deletions(-) diff --git a/vortex-array/benches/aggregate_sum.rs b/vortex-array/benches/aggregate_sum.rs index b1745a068dc..5d7ca97f6d9 100644 --- a/vortex-array/benches/aggregate_sum.rs +++ b/vortex-array/benches/aggregate_sum.rs @@ -302,7 +302,10 @@ fn sum_i64_nulls_scattered(bencher: Bencher) { fn sum_f64_nulls_scattered(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(10); let data: Vec> = (0..N) - .map(|_| rng.random_bool(0.5).then(|| rng.random_range(-1000.0..1000.0))) + .map(|_| { + rng.random_bool(0.5) + .then(|| rng.random_range(-1000.0..1000.0)) + }) .collect(); bencher .with_inputs(|| { diff --git a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs index 01a4ca74b2e..168f075e957 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs @@ -893,8 +893,8 @@ mod tests { ) -> VortexResult<()> { let array = PrimitiveArray::from_iter(values.iter().copied()).into_array(); let mut ctx = SESSION.create_execution_ctx(); - let result = min_max(&array, &mut ctx, NumericalAggregateOpts::default())? - .vortex_expect("valid"); + let result = + min_max(&array, &mut ctx, NumericalAggregateOpts::default())?.vortex_expect("valid"); assert_eq!(f32::try_from(&result.min)?.to_bits(), min.to_bits()); assert_eq!(f32::try_from(&result.max)?.to_bits(), max.to_bits()); Ok(()) @@ -930,8 +930,8 @@ mod tests { let late_nan = PrimitiveArray::new(late_nan, validity()).into_array(); assert_poisoned(min_max(&late_nan, &mut ctx, KEEP_NANS)?)?; - let skipped = min_max(&late_nan, &mut ctx, NumericalAggregateOpts::default())? - .vortex_expect("valid"); + let skipped = + min_max(&late_nan, &mut ctx, NumericalAggregateOpts::default())?.vortex_expect("valid"); assert_eq!(f64::try_from(&skipped.max)?, 68.0); let only_nans = PrimitiveArray::new(only_nans, validity()).into_array(); diff --git a/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs b/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs index fe729d46751..7b3bf7becf5 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs @@ -137,6 +137,10 @@ where } if !skip_nans { + #[expect( + clippy::eq_op, + reason = "v - v is NaN only for a NaN or infinite value" + )] let non_finite = fold_floats(values, valid, seed, T::zero(), |e, v| e + (v - v), add); if Float::is_nan(non_finite) && let Some(nan) = find_nan(values, valid) @@ -180,7 +184,9 @@ where ($lanes:literal) => { match valid { None => values.fold_lanes::<$lanes, _, _, _>(init, f, merge), - Some(valid) => values.fold_lanes_masked::<$lanes, _, _, _>(valid, null, init, f, merge), + Some(valid) => { + values.fold_lanes_masked::<$lanes, _, _, _>(valid, null, init, f, merge) + } } }; } diff --git a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs index fe80d8d324a..e00d3987776 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs @@ -94,7 +94,12 @@ pub(crate) fn sum_float_all(acc: &mut f64, slice: &[T], skip_nan /// Only a NaN sum, which a NaN or infinities of both signs make, is summed again reading NaNs as /// zero. A select per lane to skip NaNs on top of the select for null lanes leaves the masked /// fold scalar on NEON. -fn sum_float(acc: &mut f64, values: &[T], valid: Option<&BitBuffer>, skip_nans: bool) { +fn sum_float( + acc: &mut f64, + values: &[T], + valid: Option<&BitBuffer>, + skip_nans: bool, +) { let sum = fold_float(values, valid, |v| v); *acc += if skip_nans && sum.is_nan() { fold_float(values, valid, non_nan) @@ -737,7 +742,10 @@ mod tests { PrimitiveArray::new(buffer![i64::MIN, 0, -1], Validity::NonNullable).into_array(); let ctx = &mut SESSION.create_execution_ctx(); - assert_eq!(sum(&fits, ctx)?.as_primitive().typed_value::(), Some(4)); + assert_eq!( + sum(&fits, ctx)?.as_primitive().typed_value::(), + Some(4) + ); assert!(sum(&overflows, ctx)?.is_null()); assert!(sum(&underflows, ctx)?.is_null()); Ok(()) diff --git a/vortex-compressor/benches/integer_stats.rs b/vortex-compressor/benches/integer_stats.rs index 068f6f410c0..2a418c38888 100644 --- a/vortex-compressor/benches/integer_stats.rs +++ b/vortex-compressor/benches/integer_stats.rs @@ -30,34 +30,37 @@ static SESSION: LazyLock = LazyLock::new(|| { const LEN: usize = 64 * 1024; +/// Counting distinct values over a nullable array is the slowest case, so it runs on less data. +const SMALL_LEN: usize = 16 * 1024; + fn main() { LazyLock::force(&SESSION); divan::main(); } /// Runs of up to 8 values, drawn from 1024 distinct values. -fn values() -> Buffer { +fn values(len: usize) -> Buffer { let mut rng = StdRng::seed_from_u64(0); - let mut values = Vec::with_capacity(LEN); - while values.len() < LEN { + let mut values = Vec::with_capacity(len); + while values.len() < len { let value: u32 = rng.random_range(0..1024); let run = rng.random_range(1..=8); - values.extend(std::iter::repeat_n(value, run.min(LEN - values.len()))); + values.extend(std::iter::repeat_n(value, run.min(len - values.len()))); } Buffer::from(values) } /// One value in ten null, at random, so almost every chunk of 64 values has a null. -fn sparse_nulls() -> Validity { +fn sparse_nulls(len: usize) -> Validity { let mut rng = StdRng::seed_from_u64(1); - Validity::from(BitBuffer::from_iter((0..LEN).map(|_| rng.random_bool(0.9)))) + Validity::from(BitBuffer::from_iter((0..len).map(|_| rng.random_bool(0.9)))) } /// Computes the statistics of a fresh array each iteration, so that the bounds cached on the array /// by a previous iteration are not reused. With `count_distinct_values`, every statistic is /// computed, as when a dictionary scheme is enabled. -fn bench_stats(bencher: Bencher, validity: Validity, count_distinct_values: bool) { - let values = values(); +fn bench_stats(bencher: Bencher, len: usize, validity: Validity, count_distinct_values: bool) { + let values = values(len); bencher .with_inputs(|| { ( @@ -78,20 +81,20 @@ fn bench_stats(bencher: Bencher, validity: Validity, count_distinct_values: bool #[divan::bench] fn integer_stats_non_null(bencher: Bencher) { - bench_stats(bencher, Validity::NonNullable, false); + bench_stats(bencher, LEN, Validity::NonNullable, false); } #[divan::bench] fn integer_stats_nullable(bencher: Bencher) { - bench_stats(bencher, sparse_nulls(), false); + bench_stats(bencher, LEN, sparse_nulls(LEN), false); } #[divan::bench] fn integer_stats_all_non_null(bencher: Bencher) { - bench_stats(bencher, Validity::NonNullable, true); + bench_stats(bencher, LEN, Validity::NonNullable, true); } #[divan::bench] fn integer_stats_all_nullable(bencher: Bencher) { - bench_stats(bencher, sparse_nulls(), true); + bench_stats(bencher, SMALL_LEN, sparse_nulls(SMALL_LEN), true); } diff --git a/vortex-compute/benches/fold_lanes.rs b/vortex-compute/benches/fold_lanes.rs index 95012ec7122..bd59735d614 100644 --- a/vortex-compute/benches/fold_lanes.rs +++ b/vortex-compute/benches/fold_lanes.rs @@ -40,7 +40,9 @@ fn floats() -> Vec { fn ints() -> Vec { let mut rng = StdRng::seed_from_u64(1); - (0..N).map(|_| rng.random_range(-1_000_000..1_000_000)).collect() + (0..N) + .map(|_| rng.random_range(-1_000_000..1_000_000)) + .collect() } fn valid(density: f64) -> BitBuffer { @@ -209,11 +211,13 @@ fn checked_sum_i64_lanes(bencher: Bencher) { fn checked_sum_i64_masked_runs(bencher: Bencher, density: f64) { let (values, runs) = (ints(), runs(&valid(density))); bencher.bench(|| { - divan::black_box(&runs).iter().try_fold(0i64, |acc, &(start, end)| { - values[start..end] - .iter() - .try_fold(acc, |a, &v| a.checked_add(v)) - }) + divan::black_box(&runs) + .iter() + .try_fold(0i64, |acc, &(start, end)| { + values[start..end] + .iter() + .try_fold(acc, |a, &v| a.checked_add(v)) + }) }); } diff --git a/vortex-compute/src/lane_kernels/fold.rs b/vortex-compute/src/lane_kernels/fold.rs index f728acb2fcc..a301f857ed4 100644 --- a/vortex-compute/src/lane_kernels/fold.rs +++ b/vortex-compute/src/lane_kernels/fold.rs @@ -381,8 +381,14 @@ mod tests { let values: Vec = (0..len).map(|v| v * 7 - 300).collect(); let expected: i64 = values.iter().sum(); - assert_eq!(values[..].fold_lanes::<8, _, _, _>(0, |a, v| a + v, |a, b| a + b), expected); - assert_eq!(values[..].fold_lanes::<64, _, _, _>(0, |a, v| a + v, |a, b| a + b), expected); + assert_eq!( + values[..].fold_lanes::<8, _, _, _>(0, |a, v| a + v, |a, b| a + b), + expected + ); + assert_eq!( + values[..].fold_lanes::<64, _, _, _>(0, |a, v| a + v, |a, b| a + b), + expected + ); } } @@ -398,8 +404,13 @@ mod tests { .map(|(_, &v)| i64::from(v)) .sum(); - let narrow = - values[..].fold_lanes_masked::<16, _, _, _>(&mask, 0, 0, |a, v| a + v, |a, b| a + b); + let narrow = values[..].fold_lanes_masked::<16, _, _, _>( + &mask, + 0, + 0, + |a, v| a + v, + |a, b| a + b, + ); let wide = values[..].fold_lanes_masked::<64, _, _, _>( &mask, 0, From d5045ac36941690748668d2b7404bd55ac03fb3b Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 29 Oct 2026 13:36:10 +0000 Subject: [PATCH 5/7] Sum decimals stored in at most 64 bits with the vectorized integer sums Decimal values stored in i8 to i64 use the i64 sum kernels, with and without nulls. Wider values, or an i64 sum that overflows, use the checked loop in the result type. Signed-off-by: Robert Kruszewski --- vortex-array/benches/aggregate_nullable.rs | 34 ++++++-- .../src/aggregate_fn/fns/sum/decimal.rs | 82 +++++++++++++++++++ .../src/aggregate_fn/fns/sum/primitive.rs | 19 +++-- 3 files changed, 123 insertions(+), 12 deletions(-) diff --git a/vortex-array/benches/aggregate_nullable.rs b/vortex-array/benches/aggregate_nullable.rs index b07c0bade84..59bb776b95a 100644 --- a/vortex-array/benches/aggregate_nullable.rs +++ b/vortex-array/benches/aggregate_nullable.rs @@ -95,13 +95,16 @@ fn min_max_f64_nulls_scattered(bencher: Bencher) { }); } -#[divan::bench] -fn min_max_decimal_i64_nulls_scattered(bencher: Bencher) { +fn decimal_i64_values() -> Buffer { let mut rng = StdRng::seed_from_u64(3); - let values: Buffer = (0..N) + (0..N) .map(|_| rng.random_range(-100_000..100_000)) - .collect(); - let validity = scattered_validity(); + .collect() +} + +#[divan::bench] +fn min_max_decimal_i64_nulls_scattered(bencher: Bencher) { + let (values, validity) = (decimal_i64_values(), scattered_validity()); bench_min_max(bencher, || { DecimalArray::new(values.clone(), DecimalDType::new(18, 2), validity.clone()).into_array() }); @@ -125,6 +128,27 @@ fn sum_f64_nulls_scattered(bencher: Bencher) { }); } +#[divan::bench] +fn sum_decimal_i64_all_valid(bencher: Bencher) { + let values = decimal_i64_values(); + bench_sum(bencher, || { + DecimalArray::new( + values.clone(), + DecimalDType::new(18, 2), + Validity::NonNullable, + ) + .into_array() + }); +} + +#[divan::bench] +fn sum_decimal_i64_nulls_scattered(bencher: Bencher) { + let (values, validity) = (decimal_i64_values(), scattered_validity()); + bench_sum(bencher, || { + DecimalArray::new(values.clone(), DecimalDType::new(18, 2), validity.clone()).into_array() + }); +} + #[divan::bench] fn nan_count_f64_nulls_scattered(bencher: Bencher) { let (values, validity) = (f64_values(), scattered_validity()); diff --git a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs index 83bb2e0d65e..22da1fa4182 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs @@ -13,11 +13,16 @@ use vortex_error::vortex_panic; use vortex_mask::Mask; use super::SumState; +use super::checked_add_sum_states; +use super::primitive::sum_signed_valid; +use super::sum_signed_all; use crate::ExecutionCtx; use crate::arrays::DecimalArray; use crate::dtype::DType; use crate::dtype::DecimalDType; +use crate::dtype::DecimalType; use crate::dtype::NativeDecimalType; +use crate::dtype::NativePType; use crate::match_each_decimal_value_type; use crate::scalar::DecimalValue; @@ -38,6 +43,10 @@ pub(crate) fn accumulate_decimal( } }; + if let Some(sum) = sum_narrow_decimal(d, validity) { + return checked_add_sum_states(inner, return_dtype, &SumState::Decimal(sum)); + } + let SumState::Decimal(value) = inner else { vortex_panic!("expected decimal sum state for decimal input"); }; @@ -60,6 +69,33 @@ pub(crate) fn accumulate_decimal( }) } +/// Sum values stored in at most 64 bits with the vectorized `i64` integer sums. +/// +/// Returns `None` for wider values, or if the `i64` sum overflows, which a wider result type can +/// still hold. The caller then sums the values in the result type. +fn sum_narrow_decimal(d: &DecimalArray, validity: Option<&BitBuffer>) -> Option { + match d.values_type() { + DecimalType::I8 => sum_narrow(&d.buffer::(), validity), + DecimalType::I16 => sum_narrow(&d.buffer::(), validity), + DecimalType::I32 => sum_narrow(&d.buffer::(), validity), + DecimalType::I64 => sum_narrow(&d.buffer::(), validity), + DecimalType::I128 | DecimalType::I256 => None, + } +} + +fn sum_narrow(values: &[T], validity: Option<&BitBuffer>) -> Option +where + T: NativePType + AsPrimitive, +{ + let mut sum = 0i64; + let overflow = match validity { + Some(valid) => sum_signed_valid(&mut sum, values, valid), + None => sum_signed_all(&mut sum, values), + }; + + (!overflow).then(|| DecimalValue::from(sum)) +} + fn sum_decimal_value( initial: I, values: Buffer, @@ -113,6 +149,7 @@ where mod tests { use std::sync::LazyLock; + use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -246,6 +283,51 @@ mod tests { Ok(()) } + #[test] + fn sum_decimal_ignores_values_under_nulls() -> VortexResult<()> { + // 100 rows cover a full 64-row validity word and the partial word after it. + let values = (0..100i32).map(|i| if i % 3 == 0 { i32::MAX } else { i }); + let validity = Validity::from_iter((0..100).map(|i| i % 3 != 0)); + let expected: i64 = (0..100).filter(|i| i % 3 != 0).sum(); + let decimal = DecimalArray::new( + values.collect::>(), + DecimalDType::new(9, 0), + validity, + ); + + let result = sum(&decimal.into_array(), &mut SESSION.create_execution_ctx())?; + + let expected = Scalar::try_new( + DType::Decimal(DecimalDType::new(19, 0), Nullable), + Some(ScalarValue::from(DecimalValue::from(expected))), + )?; + + assert_eq!(result, expected); + Ok(()) + } + + #[test] + fn sum_decimal_i64_overflow_fits_result() -> VortexResult<()> { + // Eleven valid values of 10^18 - 1 overflow i64 but fit the i128 storage of precision 28. + let max = 999_999_999_999_999_999i64; + let values = (0..12).map(|i| if i == 5 { i64::MIN } else { max }); + let decimal = DecimalArray::new( + values.collect::>(), + DecimalDType::new(18, 0), + Validity::from_iter((0..12).map(|i| i != 5)), + ); + + let result = sum(&decimal.into_array(), &mut SESSION.create_execution_ctx())?; + + let expected = Scalar::try_new( + DType::Decimal(DecimalDType::new(28, 0), Nullable), + Some(ScalarValue::from(DecimalValue::from(i128::from(max) * 11))), + )?; + + assert_eq!(result, expected); + Ok(()) + } + #[test] fn sum_decimal_preserves_scale() -> VortexResult<()> { let decimal = DecimalArray::new( diff --git a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs index e00d3987776..617b866a407 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs @@ -170,6 +170,17 @@ where false } +/// Sum the `valid` values of a signed integer slice into `acc`. Returns `true` on overflow. +pub(super) fn sum_signed_valid(acc: &mut i64, values: &[T], valid: &BitBuffer) -> bool +where + T: NativePType + AsPrimitive, +{ + if T::PTYPE == PType::I64 { + return sum_wide(acc, values, Some(valid)); + } + sum_masked(values, valid, |chunk_sum| checked_add_i64(acc, chunk_sum)) +} + /// Sum the valid elements. /// /// Every kernel reads null lanes as zero, so the sum vectorizes the same way as the all-valid path @@ -197,13 +208,7 @@ fn accumulate_primitive_valid( SumState::Signed(acc) => match_each_native_ptype!(p.ptype(), unsigned: |_T| { vortex_panic!("signed sum state with unsigned input") }, signed: |T| { - let values = p.as_slice::(); - if T::PTYPE == PType::I64 { - return Ok(sum_wide(acc, values, Some(valid.bit_buffer()))); - } - Ok(sum_masked(values, valid.bit_buffer(), |chunk_sum| { - checked_add_i64(acc, chunk_sum) - })) + Ok(sum_signed_valid(acc, p.as_slice::(), valid.bit_buffer())) }, floating: |_T| { vortex_panic!("signed sum state with float input") } ), From 0d5dd39c2ce410e749ef0f51d331e2b5c2eba12a Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 29 Oct 2026 16:45:22 +0000 Subject: [PATCH 6/7] AND two bool validities into bits at once MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Validity::and of two non-nullable bool validities, such as a nullable array masked by another, built a lazy Binary(And) that every reader of the validity evaluated again. Both sides already hold their bits, so AND them word by word into a bool validity up front. Masking a nullable i32 column with scattered nulls and counting its nulls takes 1.36 µs instead of 1.62 µs. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01X1HnweRGxeDWgQwXXM6f5s --- vortex-array/benches/aggregate_nullable.rs | 26 ++++++++++++ vortex-array/src/validity.rs | 48 +++++++++++++++++++--- 2 files changed, 68 insertions(+), 6 deletions(-) diff --git a/vortex-array/benches/aggregate_nullable.rs b/vortex-array/benches/aggregate_nullable.rs index 59bb776b95a..4c58c871f09 100644 --- a/vortex-array/benches/aggregate_nullable.rs +++ b/vortex-array/benches/aggregate_nullable.rs @@ -17,10 +17,13 @@ use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::NumericalAggregateOpts; use vortex_array::aggregate_fn::fns::min_max::min_max; use vortex_array::aggregate_fn::fns::nan_count::nan_count; +use vortex_array::aggregate_fn::fns::null_count::null_count; use vortex_array::aggregate_fn::fns::sum::sum; use vortex_array::aggregate_fn::session::AggregateFnSession; +use vortex_array::arrays::BoolArray; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DecimalDType; use vortex_array::optimizer::kernels::KernelSession; use vortex_array::validity::Validity; @@ -158,3 +161,26 @@ fn nan_count_f64_nulls_scattered(bencher: Bencher) { |array| nan_count(array, &mut SESSION.create_execution_ctx()), ); } + +/// Mask a nullable column, as a parent's nulls mask a field, and count its nulls, which reads the +/// conjunction of the two validities. +#[divan::bench] +fn null_count_i32_masked(bencher: Bencher) { + let values = i32_values(); + let validity = scattered_validity(); + let mut rng = StdRng::seed_from_u64(6); + let mask = BoolArray::from_iter((0..N).map(|_| rng.random_bool(0.5))).into_array(); + bench_with( + bencher, + || PrimitiveArray::new(values.clone(), validity.clone()).into_array(), + |array| { + let masked = array.clone().mask(mask.clone())?; + null_count(&masked, &mut SESSION.create_execution_ctx()) + }, + ); +} + +fn i32_values() -> Buffer { + let mut rng = StdRng::seed_from_u64(7); + (0..N).map(|_| rng.random_range(-1000..1000)).collect() +} diff --git a/vortex-array/src/validity.rs b/vortex-array/src/validity.rs index 7a817c8264b..1685f118b5d 100644 --- a/vortex-array/src/validity.rs +++ b/vortex-array/src/validity.rs @@ -24,9 +24,11 @@ use crate::ArrayRef; use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; +use crate::arrays::Bool; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; +use crate::arrays::bool::BoolArrayExt; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; @@ -309,11 +311,27 @@ impl Validity { | (Validity::AllValid, Validity::NonNullable) | (Validity::AllValid, Validity::AllValid) => Validity::AllValid, // Here we actually have to do some work - (Validity::Array(lhs), Validity::Array(rhs)) => Validity::Array( - Binary::try_new(lhs, rhs, Operator::And)? - .into_array() - .optimize()?, - ), + (Validity::Array(lhs), Validity::Array(rhs)) => { + // Both sides already hold their bits, e.g. a validity masked by a demand: a word-wise + // AND costs less than evaluating a lazy conjunction wherever it is read. + if let (Some(lhs), Some(rhs)) = (lhs.as_opt::(), rhs.as_opt::()) + && !lhs.dtype().is_nullable() + && !rhs.dtype().is_nullable() + { + return Ok(Validity::Array( + BoolArray::new( + &lhs.to_bit_buffer() & &rhs.to_bit_buffer(), + Validity::NonNullable, + ) + .into_array(), + )); + } + Validity::Array( + Binary::try_new(lhs, rhs, Operator::And)? + .into_array() + .optimize()?, + ) + } }) } @@ -628,12 +646,14 @@ mod tests { use rstest::rstest; use vortex_buffer::Buffer; use vortex_buffer::buffer; + use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; + use crate::arrays::Bool; use crate::arrays::PrimitiveArray; use crate::dtype::Nullability; use crate::validity::BoolArray; @@ -641,6 +661,22 @@ mod tests { static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + /// Two validities that already hold their bits are conjoined at once, into bits. + #[test] + fn and_of_two_bool_validities_is_their_word_wise_and() -> VortexResult<()> { + let lhs = Validity::Array(BoolArray::from_iter([true, true, false, true]).into_array()); + let rhs = Validity::Array(BoolArray::from_iter([true, false, false, true]).into_array()); + + let and = lhs.and(rhs)?; + + assert!(matches!(&and, Validity::Array(array) if array.is::())); + assert_eq!( + and.execute_mask(4, &mut SESSION.create_execution_ctx())?, + Mask::from_iter([true, false, false, true]) + ); + Ok(()) + } + #[rstest] #[case(Validity::AllValid, 5, &[2, 4], Validity::AllValid, Validity::AllValid)] #[case( @@ -826,7 +862,7 @@ mod tests { #[case] lhs: Validity, #[case] rhs: Validity, #[case] expected: bool, - ) -> vortex_error::VortexResult<()> { + ) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); assert_eq!(lhs.mask_eq(&rhs, 3, &mut ctx)?, expected); Ok(()) From 1df80b07227cdbce5594f8c430c93a48bcbb0615 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 29 Oct 2026 18:05:32 +0000 Subject: [PATCH 7/7] Keep device validities lazy when ANDing them Validity::and read both validities' bits on the host to AND them, which panics for bits held on a device. Take the word-wise AND only when both validities are host-resident, and keep a lazy conjunction otherwise. Signed-off-by: Claude Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01X1HnweRGxeDWgQwXXM6f5s --- vortex-array/src/validity.rs | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/vortex-array/src/validity.rs b/vortex-array/src/validity.rs index 1685f118b5d..1c34b1c379a 100644 --- a/vortex-array/src/validity.rs +++ b/vortex-array/src/validity.rs @@ -312,15 +312,19 @@ impl Validity { | (Validity::AllValid, Validity::AllValid) => Validity::AllValid, // Here we actually have to do some work (Validity::Array(lhs), Validity::Array(rhs)) => { - // Both sides already hold their bits, e.g. a validity masked by a demand: a word-wise - // AND costs less than evaluating a lazy conjunction wherever it is read. - if let (Some(lhs), Some(rhs)) = (lhs.as_opt::(), rhs.as_opt::()) - && !lhs.dtype().is_nullable() - && !rhs.dtype().is_nullable() + // Both sides already hold their bits in host memory, e.g. a validity masked by a + // demand: a word-wise AND costs less than evaluating a lazy conjunction wherever it + // is read. Bits on a device stay lazy, so they are not copied to the host here. + if let (Some(lhs_bits), Some(rhs_bits)) = + (lhs.as_opt::(), rhs.as_opt::()) + && !lhs_bits.dtype().is_nullable() + && !rhs_bits.dtype().is_nullable() + && lhs.is_host() + && rhs.is_host() { return Ok(Validity::Array( BoolArray::new( - &lhs.to_bit_buffer() & &rhs.to_bit_buffer(), + &lhs_bits.to_bit_buffer() & &rhs_bits.to_bit_buffer(), Validity::NonNullable, ) .into_array(),