Skip to content

Commit f6685e9

Browse files
connortsui20connortsui20
committed
refactor: separate run-end sum traversal
Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent 762736d commit f6685e9

5 files changed

Lines changed: 418 additions & 366 deletions

File tree

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

Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
//! Grouped aggregation with traversal selected by the group layout.
5+
//!
6+
//! Fixed-size groups share a forward run cursor. List-view ranges can overlap or arrive out of
7+
//! order, so they locate their runs independently. Both paths weight runs by their intersection
8+
//! with the group, and skip null groups before visiting any runs.
9+
10+
use std::ops::Range;
11+
12+
use itertools::Either;
13+
use vortex_array::ArrayRef;
14+
use vortex_array::ExecutionCtx;
15+
use vortex_array::IntoArray;
16+
use vortex_array::aggregate_fn::AggregateFnRef;
17+
use vortex_array::aggregate_fn::GroupRanges;
18+
use vortex_array::aggregate_fn::GroupedArray;
19+
use vortex_array::aggregate_fn::fns::sum::Sum;
20+
use vortex_array::aggregate_fn::fns::sum_v2::SumV2;
21+
use vortex_array::aggregate_fn::kernels::DynGroupedAggregateKernel;
22+
use vortex_array::arrays::ConstantArray;
23+
use vortex_array::arrays::PrimitiveArray;
24+
use vortex_array::builtins::ArrayBuiltins;
25+
use vortex_array::dtype::IntegerPType;
26+
use vortex_array::dtype::NativePType;
27+
use vortex_array::dtype::Nullability::Nullable;
28+
use vortex_array::match_each_native_ptype;
29+
use vortex_array::match_each_unsigned_integer_ptype;
30+
use vortex_array::validity::Validity;
31+
use vortex_buffer::BitBuffer;
32+
use vortex_buffer::BitBufferMut;
33+
use vortex_error::VortexResult;
34+
use vortex_mask::AllOr;
35+
use vortex_mask::Mask;
36+
37+
use super::RunEndInputs;
38+
use super::RunEndSumKernel;
39+
use super::empty_partial;
40+
use super::runs::add_float_run;
41+
use super::runs::add_signed_run;
42+
use super::runs::add_unsigned_run;
43+
use super::runs::sum_next_range;
44+
use super::runs::sum_range;
45+
use crate::RunEnd;
46+
47+
impl DynGroupedAggregateKernel for RunEndSumKernel {
48+
fn grouped_aggregate(
49+
&self,
50+
aggregate_fn: &AggregateFnRef,
51+
groups: &GroupedArray,
52+
ctx: &mut ExecutionCtx,
53+
) -> VortexResult<Option<ArrayRef>> {
54+
let Some(options) = aggregate_fn
55+
.as_opt::<Sum>()
56+
.or_else(|| aggregate_fn.as_opt::<SumV2>())
57+
else {
58+
return Ok(None);
59+
};
60+
let Some(elements) = groups.elements().as_opt::<RunEnd>() else {
61+
return Ok(None);
62+
};
63+
if !groups.elements().dtype().is_primitive() {
64+
return Ok(None);
65+
}
66+
67+
let validity = groups.group_validity(ctx)?;
68+
let runs = if validity.all_false() {
69+
None
70+
} else {
71+
RunEndInputs::new(elements, ctx)?
72+
};
73+
let Some(runs) = runs else {
74+
let partial = empty_partial(aggregate_fn, groups.elements().dtype())?;
75+
let partials = ConstantArray::new(partial, groups.len()).into_array();
76+
let validity = Validity::from_mask(validity, Nullable).to_array(groups.len());
77+
return Ok(Some(partials.mask(validity)?));
78+
};
79+
80+
let ranges = groups.group_ranges(ctx)?;
81+
let valid_runs = runs.validity.indices();
82+
83+
let (results, empty_groups) = match_each_unsigned_integer_ptype!(runs.ends.ptype(), |E| {
84+
let ends = runs.ends.as_slice::<E>();
85+
match_each_native_ptype!(runs.values.ptype(),
86+
unsigned: |T| {
87+
sum_groups(ends, runs.values.as_slice::<T>(), valid_runs, &ranges,
88+
&validity, runs.offset, add_unsigned_run)
89+
},
90+
signed: |T| {
91+
sum_groups(ends, runs.values.as_slice::<T>(), valid_runs, &ranges,
92+
&validity, runs.offset, add_signed_run)
93+
},
94+
floating: |T| {
95+
sum_groups(ends, runs.values.as_slice::<T>(), valid_runs, &ranges,
96+
&validity, runs.offset,
97+
|sum, value, len| add_float_run(sum, value, len, options.skip_nans))
98+
}
99+
)
100+
});
101+
102+
let results = results.into_array();
103+
if aggregate_fn.is::<SumV2>() {
104+
Ok(Some(SumV2::partials_from_sums(
105+
results,
106+
empty_groups,
107+
validity,
108+
ctx,
109+
)?))
110+
} else {
111+
Ok(Some(results))
112+
}
113+
}
114+
}
115+
116+
fn sum_groups<E: IntegerPType, T: NativePType, A: NativePType>(
117+
ends: &[E],
118+
values: &[T],
119+
validity: AllOr<&[usize]>,
120+
ranges: &GroupRanges,
121+
group_validity: &Mask,
122+
offset: usize,
123+
add_run: impl Fn(A, T, usize) -> Option<A>,
124+
) -> (PrimitiveArray, BitBuffer) {
125+
match ranges {
126+
// Consecutive groups reuse the cursor, including a run split across group boundaries.
127+
GroupRanges::FixedSizeList { .. } => {
128+
let indices = match validity {
129+
AllOr::All => Either::Left(0..ends.len()),
130+
AllOr::None => Either::Left(0..0),
131+
AllOr::Some(indices) => Either::Right(indices.iter().copied()),
132+
};
133+
let mut indices = indices.peekable();
134+
135+
collect_group_sums(ranges, group_validity, offset, |range| {
136+
sum_next_range(ends, values, &mut indices, range, &add_run)
137+
})
138+
}
139+
// These ranges can overlap or go backwards, so each group seeks independently.
140+
GroupRanges::ListView { .. } => {
141+
collect_group_sums(ranges, group_validity, offset, |range| {
142+
sum_range(ends, values, &validity, range, &add_run)
143+
})
144+
}
145+
}
146+
}
147+
148+
fn collect_group_sums<A: NativePType>(
149+
ranges: &GroupRanges,
150+
group_validity: &Mask,
151+
offset: usize,
152+
mut sum_group: impl FnMut(Range<usize>) -> (Option<A>, bool),
153+
) -> (PrimitiveArray, BitBuffer) {
154+
let mut empty_groups = BitBufferMut::new_unset(ranges.len());
155+
let sums =
156+
PrimitiveArray::from_option_iter(ranges.iter().zip(group_validity.iter()).enumerate().map(
157+
|(index, ((start, len), valid))| {
158+
if !valid {
159+
return None;
160+
}
161+
162+
let (sum, is_empty) = sum_group(offset + start..offset + start + len);
163+
empty_groups.set_to(index, is_empty);
164+
sum
165+
},
166+
));
167+
168+
(sums, empty_groups.freeze())
169+
}

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

Lines changed: 41 additions & 84 deletions
Original file line numberDiff line numberDiff line change
@@ -3,113 +3,70 @@
33

44
//! Primitive sums over run-end encoded arrays.
55
//!
6-
//! The kernels check validity before decoding run ends and values, then sum each valid value
7-
//! weighted by its run length. Grouped sums intersect runs with each group's range.
8-
//! Decimal inputs use the fallback.
9-
//! Float multiplication can round differently from repeated addition, as with constant sums.
6+
//! Empty arrays and all-null inputs return before decoding the children. Otherwise, each valid
7+
//! run contributes its value multiplied by the length included in the input.
8+
//!
9+
//! Whole-array sums visit one range, clipped at the array's slice boundaries. Fixed-size groups
10+
//! share a forward cursor. List-view groups can overlap or arrive out of order, so each group
11+
//! locates its first run independently. The shared reduction in [`runs`] clips intersecting runs.
12+
//!
13+
//! Decimal inputs use the fallback. Floating-point multiplication can round differently from
14+
//! repeated addition, as with constant sums.
1015
11-
mod primitive;
16+
mod grouped;
17+
mod runs;
18+
mod whole;
1219

13-
use vortex_array::ArrayRef;
20+
use vortex_array::ArrayView;
1421
use vortex_array::ExecutionCtx;
15-
use vortex_array::IntoArray;
1622
use vortex_array::aggregate_fn::AggregateFnRef;
17-
use vortex_array::aggregate_fn::GroupedArray;
18-
use vortex_array::aggregate_fn::fns::sum::Sum;
1923
use vortex_array::aggregate_fn::fns::sum_v2::SumV2;
20-
use vortex_array::aggregate_fn::kernels::DynAggregateKernel;
21-
use vortex_array::aggregate_fn::kernels::DynGroupedAggregateKernel;
22-
use vortex_array::arrays::ConstantArray;
23-
use vortex_array::builtins::ArrayBuiltins;
24+
use vortex_array::arrays::PrimitiveArray;
2425
use vortex_array::dtype::DType;
25-
use vortex_array::dtype::Nullability::Nullable;
2626
use vortex_array::scalar::Scalar;
27-
use vortex_array::validity::Validity;
2827
use vortex_error::VortexExpect;
2928
use vortex_error::VortexResult;
29+
use vortex_mask::Mask;
3030

31-
use self::primitive::RunEndSums;
3231
use crate::RunEnd;
32+
use crate::RunEndArrayExt;
33+
use crate::RunEndArraySlotsExt;
3334

3435
/// Whole-array and grouped primitive sum kernels for [`RunEnd`].
3536
#[derive(Debug)]
3637
pub(crate) struct RunEndSumKernel;
3738

38-
impl DynAggregateKernel for RunEndSumKernel {
39-
fn aggregate(
40-
&self,
41-
aggregate_fn: &AggregateFnRef,
42-
batch: &ArrayRef,
43-
ctx: &mut ExecutionCtx,
44-
) -> VortexResult<Option<Scalar>> {
45-
let Some(options) = aggregate_fn
46-
.as_opt::<Sum>()
47-
.or_else(|| aggregate_fn.as_opt::<SumV2>())
48-
else {
49-
return Ok(None);
50-
};
51-
let Some(array) = batch.as_opt::<RunEnd>() else {
52-
return Ok(None);
53-
};
54-
if !batch.dtype().is_primitive() {
55-
return Ok(None);
56-
}
57-
58-
let Some(sums) = RunEndSums::new(array, ctx, options.skip_nans)? else {
59-
return Ok(Some(empty_partial(aggregate_fn, batch.dtype())?));
60-
};
61-
let (sum, is_empty) = sums.sum(0..batch.len());
62-
Ok(Some(partial_scalar(aggregate_fn, sum, is_empty)?))
63-
}
39+
struct RunEndInputs {
40+
ends: PrimitiveArray,
41+
values: PrimitiveArray,
42+
validity: Mask,
43+
offset: usize,
6444
}
6545

66-
impl DynGroupedAggregateKernel for RunEndSumKernel {
67-
fn grouped_aggregate(
68-
&self,
69-
aggregate_fn: &AggregateFnRef,
70-
groups: &GroupedArray,
71-
ctx: &mut ExecutionCtx,
72-
) -> VortexResult<Option<ArrayRef>> {
73-
let Some(options) = aggregate_fn
74-
.as_opt::<Sum>()
75-
.or_else(|| aggregate_fn.as_opt::<SumV2>())
76-
else {
77-
return Ok(None);
78-
};
79-
let Some(elements) = groups.elements().as_opt::<RunEnd>() else {
80-
return Ok(None);
81-
};
82-
if !groups.elements().dtype().is_primitive() {
46+
impl RunEndInputs {
47+
/// Skip materializing the children when the array is empty or every run is null.
48+
fn new(array: ArrayView<'_, RunEnd>, ctx: &mut ExecutionCtx) -> VortexResult<Option<Self>> {
49+
if array.is_empty() {
8350
return Ok(None);
8451
}
8552

86-
let validity = groups.group_validity(ctx)?;
87-
let sums = if validity.all_false() {
88-
None
89-
} else {
90-
RunEndSums::new(elements, ctx, options.skip_nans)?
91-
};
92-
let Some(sums) = sums else {
93-
let partial = empty_partial(aggregate_fn, groups.elements().dtype())?;
94-
let partials = ConstantArray::new(partial, groups.len()).into_array();
95-
let validity = Validity::from_mask(validity, Nullable).to_array(groups.len());
96-
return Ok(Some(partials.mask(validity)?));
97-
};
53+
let validity = array
54+
.values()
55+
.validity()?
56+
.execute_mask(array.values().len(), ctx)?;
57+
if validity.all_false() {
58+
return Ok(None);
59+
}
9860

99-
let ranges = groups.group_ranges(ctx)?;
100-
let (results, empty_groups) = sums.grouped_sum(&ranges, &validity);
61+
let ends = array.ends().clone().execute::<PrimitiveArray>(ctx)?;
62+
let values = array.values().clone().execute::<PrimitiveArray>(ctx)?;
10163

102-
let results = results.into_array();
103-
if aggregate_fn.is::<SumV2>() {
104-
Ok(Some(SumV2::partials_from_sums(
105-
results,
106-
empty_groups,
107-
validity,
108-
ctx,
109-
)?))
110-
} else {
111-
Ok(Some(results))
112-
}
64+
Ok(Some(Self {
65+
ends,
66+
values,
67+
validity,
68+
offset: array.offset(),
69+
}))
11370
}
11471
}
11572

0 commit comments

Comments
 (0)