Skip to content

Commit ae99ee8

Browse files
connortsui20connortsui20
committed
feat: specialize primitive sums for constant and run-end arrays
Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent e3b8eb2 commit ae99ee8

12 files changed

Lines changed: 1320 additions & 4 deletions

File tree

‎encodings/runend/Cargo.toml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,3 +58,7 @@ harness = false
5858
[[bench]]
5959
name = "run_end_filter"
6060
harness = false
61+
62+
[[bench]]
63+
name = "run_end_sum"
64+
harness = false

‎encodings/runend/benches/run_end_sum.rs‎

Lines changed: 181 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,181 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
#![expect(clippy::unwrap_used)]
5+
6+
use std::sync::LazyLock;
7+
8+
use divan::Bencher;
9+
use vortex_array::ArrayRef;
10+
use vortex_array::ArrayVTable;
11+
use vortex_array::ExecutionCtx;
12+
use vortex_array::IntoArray;
13+
use vortex_array::VortexSessionExecute;
14+
use vortex_array::aggregate_fn::AggregateFnRef;
15+
use vortex_array::aggregate_fn::AggregateFnVTable;
16+
use vortex_array::aggregate_fn::DynGroupedAccumulator;
17+
use vortex_array::aggregate_fn::GroupedAccumulator;
18+
use vortex_array::aggregate_fn::GroupedArray;
19+
use vortex_array::aggregate_fn::NumericalAggregateOpts;
20+
use vortex_array::aggregate_fn::fns::sum_v2::SumV2;
21+
use vortex_array::aggregate_fn::fns::sum_v2::sum_v2;
22+
use vortex_array::aggregate_fn::kernels::DynAggregateKernel;
23+
use vortex_array::aggregate_fn::kernels::DynGroupedAggregateKernel;
24+
use vortex_array::aggregate_fn::session::AggregateFnSessionExt;
25+
use vortex_array::arrays::Constant;
26+
use vortex_array::arrays::ConstantArray;
27+
use vortex_array::arrays::FixedSizeListArray;
28+
use vortex_array::arrays::PrimitiveArray;
29+
use vortex_array::scalar::Scalar;
30+
use vortex_array::validity::Validity;
31+
use vortex_error::VortexResult;
32+
use vortex_runend::RunEnd;
33+
use vortex_session::VortexSession;
34+
35+
const LEN: usize = 16_384;
36+
37+
static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
38+
let session = vortex_array::array_session();
39+
vortex_runend::initialize(&session);
40+
session
41+
});
42+
43+
static FALLBACK_SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
44+
let session = vortex_array::array_session();
45+
vortex_runend::initialize(&session);
46+
session
47+
.aggregate_fns()
48+
.register_aggregate_kernel(RunEnd.id(), Some(SumV2.id()), &Decline);
49+
for encoding in [RunEnd.id(), Constant.id()] {
50+
session
51+
.aggregate_fns()
52+
.register_grouped_encoding_kernel(encoding, SumV2.id(), &Decline);
53+
}
54+
session
55+
});
56+
57+
/// Keep the pre-specialization dispatch paths available for benchmark comparisons.
58+
#[derive(Debug)]
59+
struct Decline;
60+
61+
impl DynAggregateKernel for Decline {
62+
fn aggregate(
63+
&self,
64+
_aggregate_fn: &AggregateFnRef,
65+
_batch: &ArrayRef,
66+
_ctx: &mut ExecutionCtx,
67+
) -> VortexResult<Option<Scalar>> {
68+
Ok(None)
69+
}
70+
}
71+
72+
impl DynGroupedAggregateKernel for Decline {
73+
fn grouped_aggregate(
74+
&self,
75+
_aggregate_fn: &AggregateFnRef,
76+
_groups: &GroupedArray,
77+
_ctx: &mut ExecutionCtx,
78+
) -> VortexResult<Option<ArrayRef>> {
79+
Ok(None)
80+
}
81+
}
82+
83+
fn main() {
84+
LazyLock::force(&SESSION);
85+
LazyLock::force(&FALLBACK_SESSION);
86+
divan::main();
87+
}
88+
89+
fn runend(run_length: usize) -> ArrayRef {
90+
let ends =
91+
PrimitiveArray::from_iter((run_length..=LEN).step_by(run_length).map(|end| end as u64));
92+
let values = PrimitiveArray::from_option_iter(
93+
(0..ends.len())
94+
.map(|index| (index % 5 != 0).then_some(i32::try_from(index % 100).unwrap())),
95+
);
96+
RunEnd::try_new(
97+
ends.into_array(),
98+
values.into_array(),
99+
&mut SESSION.create_execution_ctx(),
100+
)
101+
.unwrap()
102+
.into_array()
103+
}
104+
105+
fn bench_sum(bencher: Bencher, run_length: usize, session: &VortexSession) {
106+
let array = runend(run_length);
107+
bencher
108+
.with_inputs(|| session.create_execution_ctx())
109+
.bench_refs(|ctx| sum_v2(&array, ctx).unwrap());
110+
}
111+
112+
fn bench_grouped(bencher: Bencher, elements: ArrayRef, group_size: u32, session: &VortexSession) {
113+
let dtype = elements.dtype().clone();
114+
let groups = FixedSizeListArray::try_new(
115+
elements,
116+
group_size,
117+
Validity::NonNullable,
118+
LEN / group_size as usize,
119+
)
120+
.unwrap()
121+
.into_array();
122+
bencher
123+
.with_inputs(|| {
124+
(
125+
GroupedAccumulator::try_new(
126+
SumV2,
127+
NumericalAggregateOpts::default(),
128+
dtype.clone(),
129+
)
130+
.unwrap(),
131+
session.create_execution_ctx(),
132+
)
133+
})
134+
.bench_refs(|(acc, ctx)| {
135+
acc.accumulate_list(&groups, ctx).unwrap();
136+
acc.finish()
137+
.unwrap()
138+
.execute::<PrimitiveArray>(ctx)
139+
.unwrap()
140+
});
141+
}
142+
143+
#[divan::bench(args = [4, 64, 1024])]
144+
fn sum_runend(bencher: Bencher, run_length: usize) {
145+
bench_sum(bencher, run_length, &SESSION);
146+
}
147+
148+
#[divan::bench(args = [4, 64, 1024])]
149+
fn sum_runend_fallback(bencher: Bencher, run_length: usize) {
150+
bench_sum(bencher, run_length, &FALLBACK_SESSION);
151+
}
152+
153+
#[divan::bench(args = [4, 64, 1024], consts = [1, 2, 8, 128])]
154+
fn grouped_runend<const GROUP_SIZE: u32>(bencher: Bencher, run_length: usize) {
155+
bench_grouped(bencher, runend(run_length), GROUP_SIZE, &SESSION);
156+
}
157+
158+
#[divan::bench(args = [4, 64, 1024], consts = [1, 2, 8, 128])]
159+
fn grouped_runend_fallback<const GROUP_SIZE: u32>(bencher: Bencher, run_length: usize) {
160+
bench_grouped(bencher, runend(run_length), GROUP_SIZE, &FALLBACK_SESSION);
161+
}
162+
163+
#[divan::bench(consts = [1, 2, 8, 128])]
164+
fn grouped_constant<const GROUP_SIZE: u32>(bencher: Bencher) {
165+
bench_grouped(
166+
bencher,
167+
ConstantArray::new(3i32, LEN).into_array(),
168+
GROUP_SIZE,
169+
&SESSION,
170+
);
171+
}
172+
173+
#[divan::bench(consts = [1, 2, 8, 128])]
174+
fn grouped_constant_fallback<const GROUP_SIZE: u32>(bencher: Bencher) {
175+
bench_grouped(
176+
bencher,
177+
ConstantArray::new(3i32, LEN).into_array(),
178+
GROUP_SIZE,
179+
&FALLBACK_SESSION,
180+
);
181+
}

‎encodings/runend/src/compute/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ pub(crate) mod filter;
88
pub(crate) mod is_constant;
99
pub(crate) mod is_sorted;
1010
pub(crate) mod min_max;
11+
pub(crate) mod sum;
1112
pub(crate) mod take;
1213
pub(crate) mod take_from;
1314

‎encodings/runend/src/compute/sum/mod.rs‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
//! Primitive sums over run-end encoded arrays.
5+
//!
6+
//! The kernels decode the run ends and values once, then sum each value weighted by its run
7+
//! length. Grouped sums intersect runs with each group's range. Decimal inputs use the fallback.
8+
//! Float multiplication can round differently from repeated addition, as with constant sums.
9+
10+
mod primitive;
11+
12+
use vortex_array::ArrayRef;
13+
use vortex_array::ExecutionCtx;
14+
use vortex_array::IntoArray;
15+
use vortex_array::aggregate_fn::AggregateFnRef;
16+
use vortex_array::aggregate_fn::GroupedArray;
17+
use vortex_array::aggregate_fn::fns::sum::Sum;
18+
use vortex_array::aggregate_fn::fns::sum_v2::SumV2;
19+
use vortex_array::aggregate_fn::kernels::DynAggregateKernel;
20+
use vortex_array::aggregate_fn::kernels::DynGroupedAggregateKernel;
21+
use vortex_array::scalar::Scalar;
22+
use vortex_error::VortexResult;
23+
24+
use self::primitive::RunEndSums;
25+
use crate::RunEnd;
26+
27+
/// Whole-array and grouped primitive sum kernels for [`RunEnd`].
28+
#[derive(Debug)]
29+
pub(crate) struct RunEndSumKernel;
30+
31+
impl DynAggregateKernel for RunEndSumKernel {
32+
fn aggregate(
33+
&self,
34+
aggregate_fn: &AggregateFnRef,
35+
batch: &ArrayRef,
36+
ctx: &mut ExecutionCtx,
37+
) -> VortexResult<Option<Scalar>> {
38+
let Some(options) = aggregate_fn
39+
.as_opt::<Sum>()
40+
.or_else(|| aggregate_fn.as_opt::<SumV2>())
41+
else {
42+
return Ok(None);
43+
};
44+
let Some(array) = batch.as_opt::<RunEnd>() else {
45+
return Ok(None);
46+
};
47+
if !batch.dtype().is_primitive() {
48+
return Ok(None);
49+
}
50+
51+
let sums = RunEndSums::new(array, ctx, options.skip_nans)?;
52+
let (sum, is_empty) = sums.sum(0..batch.len());
53+
Ok(Some(partial_scalar(aggregate_fn, sum, is_empty)?))
54+
}
55+
}
56+
57+
impl DynGroupedAggregateKernel for RunEndSumKernel {
58+
fn grouped_aggregate(
59+
&self,
60+
aggregate_fn: &AggregateFnRef,
61+
groups: &GroupedArray,
62+
ctx: &mut ExecutionCtx,
63+
) -> VortexResult<Option<ArrayRef>> {
64+
let Some(options) = aggregate_fn
65+
.as_opt::<Sum>()
66+
.or_else(|| aggregate_fn.as_opt::<SumV2>())
67+
else {
68+
return Ok(None);
69+
};
70+
let Some(elements) = groups.elements().as_opt::<RunEnd>() else {
71+
return Ok(None);
72+
};
73+
if !groups.elements().dtype().is_primitive() {
74+
return Ok(None);
75+
}
76+
77+
let ranges = groups.group_ranges(ctx)?;
78+
let validity = groups.group_validity(ctx)?;
79+
let sums = RunEndSums::new(elements, ctx, options.skip_nans)?;
80+
let (results, empty_groups) = sums.grouped_sum(&ranges, &validity);
81+
82+
let results = results.into_array();
83+
if aggregate_fn.is::<SumV2>() {
84+
Ok(Some(SumV2::partials_from_sums(
85+
results,
86+
empty_groups,
87+
validity,
88+
ctx,
89+
)?))
90+
} else {
91+
Ok(Some(results))
92+
}
93+
}
94+
}
95+
96+
fn partial_scalar(
97+
aggregate_fn: &AggregateFnRef,
98+
sum: Scalar,
99+
is_empty: bool,
100+
) -> VortexResult<Scalar> {
101+
if aggregate_fn.is::<SumV2>() {
102+
SumV2::partial_from_sum(sum, is_empty)
103+
} else {
104+
Ok(sum)
105+
}
106+
}
107+
108+
#[cfg(test)]
109+
mod tests;

0 commit comments

Comments
 (0)