Skip to content

Commit bf3c8f5

Browse files
committed
perf: avoid grouped sum bitmap appends
Signed-off-by: Matt Katz <mhkatz97@gmail.com>
1 parent 1664756 commit bf3c8f5

1 file changed

Lines changed: 13 additions & 6 deletions

File tree

‎vortex-array/src/aggregate_fn/fns/sum/grouped.rs‎

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -135,12 +135,11 @@ fn collect_sums<T: NativePType, A: NativePType + Default>(
135135
all_valid: bool,
136136
sum_run: impl Fn(&mut A, &[T]) -> bool,
137137
) -> (PrimitiveArray, BitBuffer, BitBuffer) {
138-
let mut is_overflow = BitBufferMut::with_capacity(group_ranges.len());
139-
let mut is_empty = BitBufferMut::with_capacity(group_ranges.len());
138+
let group_count = group_ranges.len();
139+
let mut is_overflow = BitBufferMut::new_unset(group_count);
140+
let mut is_empty = BitBufferMut::new_unset(group_count);
140141
let sums = group_ranges.iter().enumerate().map(|(i, (offset, size))| {
141142
if !group_validity.value(i) {
142-
is_overflow.append(false);
143-
is_empty.append(true);
144143
return A::default();
145144
}
146145
let mut acc = A::default();
@@ -149,8 +148,16 @@ fn collect_sums<T: NativePType, A: NativePType + Default>(
149148
} else {
150149
sum_masked_group(&mut acc, values, offset, size, elem_mask, &sum_run)
151150
};
152-
is_overflow.append(overflow);
153-
is_empty.append(!any_valid);
151+
if overflow {
152+
// SAFETY: `i` comes from enumerating `group_ranges`, and the bitmap has one bit per
153+
// group.
154+
unsafe { is_overflow.set_unchecked(i) };
155+
}
156+
if !any_valid {
157+
// SAFETY: `i` comes from enumerating `group_ranges`, and the bitmap has one bit per
158+
// group.
159+
unsafe { is_empty.set_unchecked(i) };
160+
}
154161
acc
155162
});
156163
let sums = PrimitiveArray::from_iter(sums);

0 commit comments

Comments
 (0)