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_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_nullable.rs b/vortex-array/benches/aggregate_nullable.rs new file mode 100644 index 00000000000..4c58c871f09 --- /dev/null +++ b/vortex-array/benches/aggregate_nullable.rs @@ -0,0 +1,186 @@ +// 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::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; +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() + }); +} + +fn decimal_i64_values() -> Buffer { + let mut rng = StdRng::seed_from_u64(3); + (0..N) + .map(|_| rng.random_range(-100_000..100_000)) + .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() + }); +} + +#[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 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()); + bench_with( + bencher, + || PrimitiveArray::new(values.clone(), validity.clone()).into_array(), + |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/benches/aggregate_sum.rs b/vortex-array/benches/aggregate_sum.rs index fd7a7944ed4..5d7ca97f6d9 100644 --- a/vortex-array/benches/aggregate_sum.rs +++ b/vortex-array/benches/aggregate_sum.rs @@ -267,3 +267,52 @@ 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/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 ce3e938a8fd..168f075e957 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,63 @@ 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(()) + } + + /// 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( @@ -801,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 263e94ffc8e..7b3bf7becf5 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/primitive.rs @@ -1,7 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use itertools::Itertools; +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; @@ -23,67 +27,42 @@ pub(super) fn accumulate_primitive( p: &PrimitiveArray, 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 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, args.options.skip_nans).map(min_max_result) + } + ); + partial.merge(args, local); + Ok(()) } -fn compute_min_max_with_validity( - array: &PrimitiveArray, - ctx: &mut ExecutionCtx, - skip_nans: bool, -) -> VortexResult> +/// Min/max of the valid values of an integer slice. +/// +/// 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, - 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::(); - // 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) - } else { - compute_min_max( - v.slices() - .iter() - .flat_map(|&(start, end)| slice[start..end].iter()), - skip_nans, - ) - } - } - }, - ) + 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. @@ -105,51 +84,171 @@ where Some((min, max)) } -fn min_max_result((min, max): (T, T)) -> MinMaxResult +/// 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, - PValue: From, { - MinMaxResult { - min: Scalar::primitive(min, NonNullable), - max: Scalar::primitive(max, NonNullable), + 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 total order, or `None` if there are +/// none. +/// +/// 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()?] + }; + if Float::is_nan(seed) { + return Some((seed, seed)); } + + 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) + { + 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 + }; + let max = if max.is_zero() { + zero_with_sign(!has_zero_with_sign(values, valid, false)) + } else { + max + }; + Some((min, max)) } -fn compute_min_max<'a, T>( - iter: impl Iterator, - skip_nans: bool, -) -> Option +/// 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>, + null: T, + init: A, + f: impl Fn(A, T) -> A, + merge: impl Fn(A, A) -> A, +) -> A where T: NativePType, - PValue: From, + A: Copy, { - 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) + 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), + } +} + +fn add(a: T, b: T) -> T { + a + b +} + +/// `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 } +} + +/// `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 } +} + +/// 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()) +} + +/// 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) +} + +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 minmax_by_total_order<'a, T>(iter: impl Iterator) -> Option +fn min_max_result((min, max): (T, T)) -> MinMaxResult 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, - }) - } - itertools::MinMaxResult::MinMax(&min, &max) => Some(MinMaxResult { - min: Scalar::primitive(min, NonNullable), - max: Scalar::primitive(max, NonNullable), - }), + 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/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 a63edacf6b3..617b866a407 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs @@ -1,12 +1,18 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ops::Add; +use std::ops::BitOr; + 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; @@ -21,17 +27,23 @@ 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, 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), } } @@ -69,37 +81,71 @@ 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; } @@ -113,15 +159,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; } @@ -129,25 +170,37 @@ 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` 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 +/// however the valid values are spread. 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 sum_unsigned_all(acc, &values[start..end]) { - return Ok(true); - } + if T::PTYPE == PType::U64 { + return Ok(sum_wide(acc, values, Some(valid.bit_buffer()))); } - Ok(false) + 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") } @@ -155,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::(); - for &(start, end) in slices { - if sum_signed_all(acc, &values[start..end]) { - return Ok(true); - } - } - Ok(false) + Ok(sum_signed_valid(acc, p.as_slice::(), valid.bit_buffer())) }, floating: |_T| { vortex_panic!("signed sum state with float input") } ), @@ -169,10 +216,7 @@ fn accumulate_primitive_valid( 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 slices { - sum_float_all(acc, &values[start..end], skip_nans); - } + sum_float(acc, p.as_slice::(), Some(valid.bit_buffer()), skip_nans); Ok(false) } ), @@ -182,10 +226,214 @@ 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_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; + use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; @@ -219,6 +467,44 @@ 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(()) + } + + /// 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(); @@ -448,6 +734,112 @@ 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-array/src/validity.rs b/vortex-array/src/validity.rs index 7a817c8264b..1c34b1c379a 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,31 @@ 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 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_bits.to_bit_buffer() & &rhs_bits.to_bit_buffer(), + Validity::NonNullable, + ) + .into_array(), + )); + } + Validity::Array( + Binary::try_new(lhs, rhs, Operator::And)? + .into_array() + .optimize()?, + ) + } }) } @@ -628,12 +650,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 +665,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 +866,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(()) 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/Cargo.toml b/vortex-compute/Cargo.toml index fc90b7029df..53c28ec33da 100644 --- a/vortex-compute/Cargo.toml +++ b/vortex-compute/Cargo.toml @@ -36,3 +36,11 @@ workspace = true [[bench]] name = "lane_kernels" 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..bd59735d614 --- /dev/null +++ b/vortex-compute/benches/fold_lanes.rs @@ -0,0 +1,334 @@ +// 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/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..a301f857ed4 --- /dev/null +++ b/vortex-compute/src/lane_kernels/fold.rs @@ -0,0 +1,466 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! 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; + +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 + } + + /// 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. +#[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 +}; + +/// 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)] +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; + use crate::lane_kernels::LaneZip; + + #[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_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]; + 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..f2bc5719fc4 100644 --- a/vortex-compute/src/lane_kernels/mod.rs +++ b/vortex-compute/src/lane_kernels/mod.rs @@ -19,12 +19,16 @@ //! the sink itself). //! - [`collect_bits`](mod@collect_bits) — allocating wrappers that bit-pack a predicate into a //! `BitBuffer`. +//! - [`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 //! 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 +37,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;