Skip to content

Commit 37031b1

Browse files
authored
Compute Widen aggregates on the values (#10890)
Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent 2fb259b commit 37031b1

10 files changed

Lines changed: 553 additions & 6 deletions

File tree

‎encodings/fastlanes/tests/widen.rs‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
//! `Widen` operations over a `BitPacked` child.
5+
//!
6+
//! Selection and comparison keep the stored width of the child, aggregates read the child, and
7+
//! canonical execution widens the values to the logical dtype of the `Widen`.
8+
9+
#![cfg(test)]
10+
11+
use std::sync::LazyLock;
12+
13+
use rstest::rstest;
14+
use vortex_array::ArrayRef;
15+
use vortex_array::IntoArray;
16+
use vortex_array::VortexSessionExecute;
17+
use vortex_array::aggregate_fn::NumericalAggregateOpts;
18+
use vortex_array::aggregate_fn::fns::is_constant::is_constant;
19+
use vortex_array::aggregate_fn::fns::is_sorted::is_sorted;
20+
use vortex_array::aggregate_fn::fns::is_sorted::is_strict_sorted;
21+
use vortex_array::aggregate_fn::fns::min_max::min_max;
22+
use vortex_array::aggregate_fn::fns::sum::sum;
23+
use vortex_array::arrays::BoolArray;
24+
use vortex_array::arrays::ConstantArray;
25+
use vortex_array::arrays::PrimitiveArray;
26+
use vortex_array::arrays::Widen;
27+
use vortex_array::arrays::WidenArray;
28+
use vortex_array::arrays::widen::WidenArraySlotsExt;
29+
use vortex_array::assert_arrays_eq;
30+
use vortex_array::builtins::ArrayBuiltins;
31+
use vortex_array::dtype::DType;
32+
use vortex_array::dtype::PType;
33+
use vortex_array::scalar::Scalar;
34+
use vortex_array::scalar_fn::fns::operators::Operator;
35+
use vortex_buffer::buffer;
36+
use vortex_error::VortexResult;
37+
use vortex_fastlanes::BitPacked;
38+
use vortex_mask::Mask;
39+
use vortex_session::VortexSession;
40+
41+
static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
42+
let session = VortexSession::empty();
43+
vortex_fastlanes::initialize(&session);
44+
session
45+
});
46+
47+
/// Returns a `u64` `Widen` over `values` packed at `bit_width` bits.
48+
fn widen_bitpacked(values: PrimitiveArray, bit_width: u8) -> VortexResult<ArrayRef> {
49+
let mut ctx = SESSION.create_execution_ctx();
50+
let packed = BitPacked::encode(&values.into_array(), bit_width, &mut ctx)?.into_array();
51+
let dtype = DType::Primitive(PType::U64, packed.dtype().nullability());
52+
53+
Ok(WidenArray::try_new(packed, dtype)?.into_array())
54+
}
55+
56+
#[test]
57+
fn test_operations() -> VortexResult<()> {
58+
let mut ctx = SESSION.create_execution_ctx();
59+
let array = widen_bitpacked(PrimitiveArray::from_iter([0u8, 1, 3, 7]), 3)?;
60+
61+
let selected = array.filter(Mask::from_iter([true, false, true, true]))?;
62+
assert!(selected.is::<Widen>());
63+
assert_eq!(selected.as_::<Widen>().values().dtype(), &PType::U8.into());
64+
assert_arrays_eq!(selected, buffer![0u64, 3, 7].into_array(), &mut ctx);
65+
66+
assert_arrays_eq!(
67+
array.take(buffer![3u32, 1].into_array())?,
68+
buffer![7u64, 1].into_array(),
69+
&mut ctx
70+
);
71+
assert_arrays_eq!(
72+
array.binary(
73+
ConstantArray::new(Scalar::from(3u64), 4).into_array(),
74+
Operator::Lt
75+
)?,
76+
BoolArray::from_iter([true, true, false, false]),
77+
&mut ctx
78+
);
79+
assert_eq!(
80+
array.execute::<PrimitiveArray>(&mut ctx)?.as_slice::<u64>(),
81+
&[0, 1, 3, 7]
82+
);
83+
84+
Ok(())
85+
}
86+
87+
#[rstest]
88+
#[case::sorted(PrimitiveArray::from_iter([0u8, 1, 3, 7]))]
89+
#[case::constant(PrimitiveArray::from_iter([5u8, 5, 5, 5]))]
90+
#[case::partially_valid(PrimitiveArray::from_option_iter([None, Some(6u8), Some(2), None]))]
91+
fn test_aggregates_match_canonical(#[case] values: PrimitiveArray) -> VortexResult<()> {
92+
let mut ctx = SESSION.create_execution_ctx();
93+
let canonical = widen_bitpacked(values.clone(), 3)?
94+
.execute::<PrimitiveArray>(&mut ctx)?
95+
.into_array();
96+
let options = NumericalAggregateOpts::default();
97+
98+
// Each aggregate gets a fresh `Widen`. A statistic cached by an earlier aggregate, such as the
99+
// min and max that `min_max` caches, would otherwise answer a later one before its kernel runs.
100+
let widen = || widen_bitpacked(values.clone(), 3);
101+
102+
assert_eq!(
103+
min_max(&widen()?, &mut ctx, options)?,
104+
min_max(&canonical, &mut ctx, options)?
105+
);
106+
assert_eq!(
107+
is_constant(&widen()?, &mut ctx)?,
108+
is_constant(&canonical, &mut ctx)?
109+
);
110+
assert_eq!(
111+
is_sorted(&widen()?, &mut ctx)?,
112+
is_sorted(&canonical, &mut ctx)?
113+
);
114+
assert_eq!(
115+
is_strict_sorted(&widen()?, &mut ctx)?,
116+
is_strict_sorted(&canonical, &mut ctx)?
117+
);
118+
assert_eq!(sum(&widen()?, &mut ctx)?, sum(&canonical, &mut ctx)?);
119+
120+
Ok(())
121+
}

‎vortex-array/src/aggregate_fn/session.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ use crate::arrays::Chunked;
4646
use crate::arrays::Dict;
4747
use crate::arrays::Primitive;
4848
use crate::arrays::RunEnd;
49+
use crate::arrays::Widen;
4950
use crate::arrays::chunked::compute::aggregate::ChunkedArrayAggregate;
5051
use crate::arrays::dict::compute::is_constant::DictIsConstantKernel;
5152
use crate::arrays::dict::compute::is_sorted::DictIsSortedKernel;
@@ -54,6 +55,7 @@ use crate::arrays::runend::compute::is_constant::RunEndIsConstantKernel;
5455
use crate::arrays::runend::compute::is_sorted::RunEndIsSortedKernel;
5556
use crate::arrays::runend::compute::min_max::RunEndMinMaxKernel;
5657
use crate::arrays::runend::compute::sum::RunEndSumKernel;
58+
use crate::arrays::widen::compute::aggregate::WidenArrayAggregate;
5759
use crate::dtype::DType;
5860
use crate::dtype::extension::ExtVTable;
5961
use crate::extension::datetime::Date;
@@ -143,6 +145,12 @@ impl Default for AggregateFnSession {
143145
this.register_aggregate_kernel(RunEnd.id(), Some(IsSorted.id()), &RunEndIsSortedKernel);
144146
this.register_aggregate_kernel(RunEnd.id(), Some(Sum.id()), &RunEndSumKernel);
145147
this.register_aggregate_kernel(RunEnd.id(), Some(SumV2.id()), &RunEndSumKernel);
148+
this.register_aggregate_kernel(Widen.id(), Some(MinMax.id()), &WidenArrayAggregate);
149+
this.register_aggregate_kernel(Widen.id(), Some(Min.id()), &WidenArrayAggregate);
150+
this.register_aggregate_kernel(Widen.id(), Some(Max.id()), &WidenArrayAggregate);
151+
this.register_aggregate_kernel(Widen.id(), Some(IsConstant.id()), &WidenArrayAggregate);
152+
this.register_aggregate_kernel(Widen.id(), Some(IsSorted.id()), &WidenArrayAggregate);
153+
this.register_aggregate_kernel(Widen.id(), Some(Sum.id()), &WidenArrayAggregate);
146154

147155
// Register the built-in grouped aggregate kernels.
148156
this.register_grouped_kernel(Count.id(), &CountGroupedKernel);
Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
use vortex_error::VortexExpect;
5+
use vortex_error::VortexResult;
6+
7+
use crate::ArrayRef;
8+
use crate::ExecutionCtx;
9+
use crate::aggregate_fn::AggregateFnRef;
10+
use crate::aggregate_fn::fns::sum::Sum;
11+
use crate::aggregate_fn::kernels::DynAggregateKernel;
12+
use crate::arrays::Widen;
13+
use crate::arrays::widen::WidenArraySlotsExt;
14+
use crate::dtype::DType;
15+
use crate::expr::stats::Stat;
16+
use crate::scalar::Scalar;
17+
18+
/// Runs an aggregate of a [`Widen`] on its values and widens the partial to the logical dtype.
19+
///
20+
/// The values run the same accumulator code as the canonical array, so edge cases such as an empty
21+
/// batch or a single null match by construction, and the values keep their own cached statistics
22+
/// and kernels. [`Sum`] is served only when [`is_exact_sum`] holds.
23+
#[derive(Debug)]
24+
pub(crate) struct WidenArrayAggregate;
25+
26+
impl DynAggregateKernel for WidenArrayAggregate {
27+
fn aggregate(
28+
&self,
29+
aggregate_fn: &AggregateFnRef,
30+
batch: &ArrayRef,
31+
ctx: &mut ExecutionCtx,
32+
) -> VortexResult<Option<Scalar>> {
33+
let Some(widen) = batch.as_opt::<Widen>() else {
34+
return Ok(None);
35+
};
36+
37+
// NB: Exactness rests on integer decoding. A decimal `Widen` (CD4 in #10733) needs its own
38+
// argument, because `Sum` checks a decimal for overflow against the precision of its input.
39+
if !batch.dtype().is_int() {
40+
return Ok(None);
41+
}
42+
43+
let values = widen.values();
44+
if aggregate_fn.is::<Sum>() && !is_exact_sum(aggregate_fn, batch, values) {
45+
return Ok(None);
46+
}
47+
48+
let mut accumulator = aggregate_fn.accumulator(values.dtype())?;
49+
accumulator.accumulate(values, ctx)?;
50+
51+
let partial_dtype = aggregate_fn.state_dtype(batch.dtype()).vortex_expect(
52+
"an aggregate that accepts the values accepts the logical dtype, which is a wider \
53+
integer of the same signedness",
54+
);
55+
56+
Ok(Some(accumulator.flush()?.cast(&partial_dtype)?))
57+
}
58+
}
59+
60+
/// Returns whether the [`Sum`] of the values of `batch` equals the sum of the canonical batch.
61+
///
62+
/// - The accumulator **must** sum each batch from zero and then add the total to the running sum.
63+
/// It does so for an aggregate with the legacy [`Stat::Sum`] slot. Otherwise it adds each value
64+
/// to the running sum, which a batch total cannot match near the 64-bit bounds: from a running
65+
/// sum of `i64::MAX - 5`, the values `[10, -10]` overflow, but their total of zero does not.
66+
/// - The sum of the batch **must not** be able to overflow. `Sum` adds every integer type into a
67+
/// 64-bit sum, but checks for overflow at different points for different widths: once per block
68+
/// for types narrower than 64 bits, and after every value for 64-bit types. A signed value of
69+
/// `b` bits lies in `[-2^(b-1), 2^(b-1))` and an unsigned one in `[0, 2^b)`, so no sum of at
70+
/// most `2^(64 - b)` of them leaves the range of `i64` or `u64`, and no check fires on either
71+
/// path.
72+
///
73+
/// Only `Sum` with `skip_nans` has the slot. Integers have no NaNs, so `skip_nans` matters here
74+
/// only because it selects the slot.
75+
fn is_exact_sum(aggregate_fn: &AggregateFnRef, batch: &ArrayRef, values: &ArrayRef) -> bool {
76+
let DType::Primitive(values_ptype, _) = values.dtype() else {
77+
return false;
78+
};
79+
let max_len = 1u64 << (64 - values_ptype.bit_width());
80+
81+
Stat::from_aggregate_fn(aggregate_fn) == Some(Stat::Sum) && batch.len() as u64 <= max_len
82+
}

‎vortex-array/src/arrays/widen/compute/mod.rs‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,21 @@
11
// SPDX-License-Identifier: Apache-2.0
22
// SPDX-FileCopyrightText: Copyright the Vortex contributors
33

4-
//! Operations that run on the values of a [`Widen`](crate::arrays::Widen).
4+
//! Operations that run on the values of a [`Widen`](crate::arrays::Widen) instead of on a
5+
//! widened copy.
6+
//!
7+
//! An operation runs on the values only where it gives the canonical result. Decoding a value is a
8+
//! strictly increasing, one-to-one map that keeps validity, so comparisons, sortedness, constancy
9+
//! and extrema carry over from the values, and every value in a result widens exactly to the
10+
//! logical dtype.
11+
//!
12+
//! [`AggregateFnSession::default`] registers the aggregate kernel, not
13+
//! [`initialize`](super::initialize), because the kernel session runs `initialize` on a temporary
14+
//! session and drops its aggregate registry.
15+
//!
16+
//! [`AggregateFnSession::default`]: crate::aggregate_fn::session::AggregateFnSession::default
517
18+
pub(crate) mod aggregate;
619
mod cast;
720
mod compare;
821
mod fill_null;

‎vortex-array/src/arrays/widen/mod.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ pub use array::WidenSlots;
3030
pub use array::WidenSlotsView;
3131
pub use vtable::WidenArray;
3232

33-
mod compute;
33+
pub(crate) mod compute;
3434
mod encode;
3535

3636
mod vtable;

0 commit comments

Comments
 (0)