Skip to content

Commit 6051e06

Browse files
committed
Support out of order grouped aggregate accumulation
Rebases #8379 onto develop and addresses the outstanding review comments. The grouped accumulator now takes caller-assigned dense `GroupIds` (`{ids: ArrayRef, num_groups}`) directly, instead of requiring values to be pre-sorted into complete list groups. Ids may repeat, arrive out of order, or be absent from a batch. ## Changes - Dispatch grouped kernels on both the values encoding and the group-id encoding, with wildcard fallbacks in each dimension. - Cache materialized and validated group ids across clones, so several aggregates over the same ids validate once. - Let each aggregate own its dense grouped state (`GroupedState`), with the generic one-partial-per-group implementation as the fallback. - Store grouped count state as one `u64` per group and grouped sum state in monomorphic `u64`/`i64`/`f64`/decimal vectors, choosing the decimal width once from the sum dtype. - Port `SumV2` onto the same dense sums, adding the per-group empty flag its SQL semantics need, so both sums share one accumulation kernel. - Merge count and sum partial arrays without `execute_scalar` per row. - Adapt between typed scatter for unclustered ids and per-run reduction for clustered ids. - Use one stable counting-sort gather per batch in the universal fallback instead of one `take` per group. - Keep `GroupedArray::dense_input` as the list-group adapter for `list_sum`, which restores list semantics (null, empty, and all-null lists sum to null) on top of the dense result. ## Rebase notes `develop` reverted SQL null-on-empty semantics for `Sum` (#9324) and re-landed them as a separate `SumV2` (#9590), so this branch follows both contracts: `Sum` starts a group at zero and treats a null partial as overflow, while `SumV2` keeps `{sum, is_overflow, is_empty}` and finalizes an untouched group to null. `SumAggregateOpts` is gone, and both aggregates take `NumericalAggregateOpts`. `SumV2`'s list-shaped `PrimitiveGroupedSumV2EncodingKernel` is replaced by a dense kernel that reuses `Sum`'s accumulation and then clears the empty flag of every group that saw a valid value, so `sum_v2` keeps a fast grouped path rather than falling back to per-group accumulators. The single-partial `flush` shortcut from #9597 no longer applies: dense state produces one partial array per flush, so there is no chunk to collapse, and its three tests went with it. ## Review comments - Group ids are an `ArrayRef` so kernels can dispatch on their encoding (Constant, RLE, ...), and `GroupIds` is the "Group struct" the accumulator API talks in. - The hand-rolled `for_each_valid_idx` helpers are gone: the kernels take the mask's bit buffer, handle the all-true and all-false cases directly, and let `BitBuffer::for_each_set_index` walk the rest instead of materializing the mask's index vector. - The run-detection helpers shared by the count and sum kernels moved to `aggregate_fn::kernels` instead of being duplicated per aggregate. - The accumulator tracks `num_groups` itself, so `flush_partials` and `finish` no longer take it; `ensure_groups` is exposed for callers that need trailing groups no batch mentioned. - `merge_group` names its parameters `into_group` / `from_group`. - The benchmark input builder is now `DenseGroupedInput::clustered`. ## list_sum performance `list_sum` is the one caller that has to build dense ids for input that is already grouped, so it pays for the new API. The adapter avoids the two costs that dominated: it reuses the element array as the values array whenever the groups already cover it in order (null lists keep their own ids, and their meaningless sums are masked out afterwards), and it shares one set of group ranges and validity between the adaptation and the mask. What remains is the id array itself, roughly half a nanosecond per element. Medians from `cargo bench -p vortex-array --bench list_sum`, develop (d0a59fe) vs this branch on the same machine: | benchmark | develop | this branch | |---|---:|---:| | `fsl_sum_large` | 5.1 us | 7.3 us | | `fsl_sum_medium` | 3.4 us | 4.8 us | | `fsl_sum_small` | 2.4 us | 3.2 us | | `list_sum_large` | 11.2 us | 13.7 us | | `list_sum_medium` | 8.9 us | 9.9 us | | `list_sum_small` | 7.6 us | 7.5 us | | `list_sum_nullable_elements_large` | 57.6 us | 18.1 us | | `list_sum_nullable_elements_medium` | 31.2 us | 12.3 us | | `listview_sum_large` | 6.9 us | 9.6 us | | `listview_sum_medium` | 4.9 us | 6.0 us | | `listview_sum_small` | 3.8 us | 3.9 us | Lists whose elements are all valid are 20-40% slower; lists with nullable elements are 2.5-3x faster, since the old path sliced the element mask per group. Closing the remaining gap wants a grouped kernel registered for run-encoded group ids, which the new two-encoding kernel lookup already supports - a follow-up rather than part of this change. ## Validation - `cargo +nightly fmt --all` - `cargo build -p vortex-array --all-targets` - `cargo test -p vortex-array --lib` (3,509 passed, 1 ignored) - `cargo test --doc -p vortex-array` - `cargo bench -p vortex-array --bench list_sum` - `cargo bench -p vortex-array --bench aggregate_grouped -- --test` - `cargo clippy -p vortex-array --all-targets --all-features` - `RUSTDOCFLAGS="-D warnings" cargo doc --no-deps -p vortex-array` A workspace-wide build ran out of disk in this environment; the change is confined to `vortex-array`, and no other crate references the grouped accumulator API. Signed-off-by: Nicholas Gates <nick@nickgates.com> Signed-off-by: Robert Kruszewski <robert@spiraldb.com> Co-authored-by: Claude <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01B8VcUXWmDmNH2JjJY4aDQ8
1 parent d0a59fe commit 6051e06

17 files changed

Lines changed: 2735 additions & 1121 deletions

File tree

‎vortex-array/benches/aggregate_grouped.rs‎

Lines changed: 93 additions & 126 deletions
Original file line numberDiff line numberDiff line change
@@ -10,21 +10,20 @@ use divan::Bencher;
1010
use rand::RngExt;
1111
use rand::SeedableRng;
1212
use rand::rngs::StdRng;
13+
use rand::seq::SliceRandom;
1314
use vortex_array::ArrayRef;
1415
use vortex_array::Canonical;
1516
use vortex_array::IntoArray;
1617
use vortex_array::VortexSessionExecute;
1718
use vortex_array::aggregate_fn::AggregateFnVTable;
1819
use vortex_array::aggregate_fn::DynGroupedAccumulator;
20+
use vortex_array::aggregate_fn::GroupIds;
1921
use vortex_array::aggregate_fn::GroupedAccumulator;
20-
use vortex_array::aggregate_fn::NumericalAggregateOpts;
2122
use vortex_array::aggregate_fn::fns::count::Count;
2223
use vortex_array::aggregate_fn::fns::sum::Sum;
2324
use vortex_array::aggregate_fn::fns::sum_v2::SumV2;
24-
use vortex_array::arrays::ListViewArray;
2525
use vortex_array::arrays::PrimitiveArray;
2626
use vortex_array::arrays::VarBinViewArray;
27-
use vortex_array::dtype::DType;
2827
use vortex_array::validity::Validity;
2928
use vortex_buffer::Buffer;
3029
use vortex_session::VortexSession;
@@ -40,6 +39,22 @@ const GROUP_COUNT: usize = 128;
4039
const GROUP_SIZE_SEED: u64 = 42;
4140
const MIN_VALUES_PER_GROUP: usize = 1;
4241
const MAX_VALUES_PER_GROUP: usize = 15;
42+
const CARDINALITY_ELEMENT_COUNT: usize = 1 << 16;
43+
44+
#[derive(Clone, Copy, Debug)]
45+
enum IdOrder {
46+
Clustered,
47+
Shuffled,
48+
}
49+
50+
const CARDINALITY_ARGS: &[(usize, IdOrder)] = &[
51+
(128, IdOrder::Clustered),
52+
(128, IdOrder::Shuffled),
53+
(1 << 12, IdOrder::Clustered),
54+
(1 << 12, IdOrder::Shuffled),
55+
(1 << 16, IdOrder::Clustered),
56+
(1 << 16, IdOrder::Shuffled),
57+
];
4358

4459
fn random_group_sizes() -> Vec<usize> {
4560
let mut rng = StdRng::seed_from_u64(GROUP_SIZE_SEED);
@@ -52,44 +67,43 @@ fn total_element_count(group_sizes: &[usize]) -> usize {
5267
group_sizes.iter().sum()
5368
}
5469

55-
fn contiguous_list_view(elements: ArrayRef, group_sizes: &[usize]) -> ArrayRef {
56-
let mut offset = 0usize;
57-
let offsets: Buffer<u32> = group_sizes
58-
.iter()
59-
.map(|&size| {
60-
let current_offset = offset;
61-
offset += size;
62-
current_offset as u32
63-
})
64-
.collect();
65-
let sizes: Buffer<u32> = group_sizes.iter().map(|&size| size as u32).collect();
70+
struct DenseGroupedInput {
71+
values: ArrayRef,
72+
group_ids: GroupIds,
73+
}
6674

67-
assert_eq!(elements.len(), total_element_count(group_sizes));
75+
impl DenseGroupedInput {
76+
/// Assign each value the id of the group whose run of `group_sizes` covers it.
77+
fn clustered(values: ArrayRef, group_sizes: &[usize]) -> Self {
78+
assert_eq!(values.len(), total_element_count(group_sizes));
79+
80+
let group_ids = GroupIds::from_iter(
81+
group_sizes
82+
.iter()
83+
.enumerate()
84+
.flat_map(|(group_id, &size)| std::iter::repeat_n(group_id as u32, size)),
85+
group_sizes.len(),
86+
)
87+
.unwrap();
6888

69-
ListViewArray::try_new(
70-
elements,
71-
offsets.into_array(),
72-
sizes.into_array(),
73-
Validity::NonNullable,
74-
)
75-
.unwrap()
76-
.into_array()
89+
Self { values, group_ids }
90+
}
7791
}
7892

79-
fn i32_nullable_all_valid_input() -> ArrayRef {
93+
fn i32_nullable_all_valid_input() -> DenseGroupedInput {
8094
let group_sizes = random_group_sizes();
8195
let element_count = total_element_count(&group_sizes);
8296
let values: Buffer<i32> = (0..element_count)
8397
.map(|i| (i % 1024) as i32 - 512)
8498
.collect();
8599
let validity = Validity::from_iter(std::iter::repeat_n(true, element_count));
86-
contiguous_list_view(
100+
DenseGroupedInput::clustered(
87101
PrimitiveArray::new(values, validity).into_array(),
88102
&group_sizes,
89103
)
90104
}
91105

92-
fn i32_clustered_nulls_input() -> ArrayRef {
106+
fn i32_clustered_nulls_input() -> DenseGroupedInput {
93107
let group_sizes = random_group_sizes();
94108
let element_count = total_element_count(&group_sizes);
95109
let values = (0..element_count).map(|i| {
@@ -99,26 +113,26 @@ fn i32_clustered_nulls_input() -> ArrayRef {
99113
Some((i % 1024) as i32 - 512)
100114
}
101115
});
102-
contiguous_list_view(
116+
DenseGroupedInput::clustered(
103117
PrimitiveArray::from_option_iter(values).into_array(),
104118
&group_sizes,
105119
)
106120
}
107121

108-
fn f64_all_valid_input() -> ArrayRef {
122+
fn f64_all_valid_input() -> DenseGroupedInput {
109123
let group_sizes = random_group_sizes();
110124
let element_count = total_element_count(&group_sizes);
111125
let mut rng = StdRng::seed_from_u64(GROUP_SIZE_SEED);
112126
let values: Buffer<f64> = (0..element_count)
113127
.map(|_| rng.random_range(-1000.0..1000.0))
114128
.collect();
115-
contiguous_list_view(
129+
DenseGroupedInput::clustered(
116130
PrimitiveArray::new(values, Validity::NonNullable).into_array(),
117131
&group_sizes,
118132
)
119133
}
120134

121-
fn f64_clustered_nulls_input() -> ArrayRef {
135+
fn f64_clustered_nulls_input() -> DenseGroupedInput {
122136
let group_sizes = random_group_sizes();
123137
let element_count = total_element_count(&group_sizes);
124138
let mut rng = StdRng::seed_from_u64(GROUP_SIZE_SEED);
@@ -129,44 +143,60 @@ fn f64_clustered_nulls_input() -> ArrayRef {
129143
Some(rng.random_range(-1000.0f64..1000.0))
130144
}
131145
});
132-
contiguous_list_view(
146+
DenseGroupedInput::clustered(
133147
PrimitiveArray::from_option_iter(values).into_array(),
134148
&group_sizes,
135149
)
136150
}
137151

138-
fn varbinview_input() -> ArrayRef {
152+
fn varbinview_input() -> DenseGroupedInput {
139153
let group_sizes = random_group_sizes();
140154
let element_count = total_element_count(&group_sizes);
141155
let values: Vec<String> = (0..element_count)
142156
.map(|i| format!("value-{i:06}"))
143157
.collect();
144-
contiguous_list_view(
158+
DenseGroupedInput::clustered(
145159
VarBinViewArray::from_iter_str(values.iter().map(String::as_str)).into_array(),
146160
&group_sizes,
147161
)
148162
}
149163

150-
fn list_element_dtype(list_view: &ArrayRef) -> DType {
151-
match list_view.dtype() {
152-
DType::List(element_dtype, _) => element_dtype.as_ref().clone(),
153-
dtype => unreachable!("expected List dtype, got {dtype}"),
164+
fn i32_cardinality_input(group_count: usize, order: IdOrder) -> DenseGroupedInput {
165+
let mut rng = StdRng::seed_from_u64(GROUP_SIZE_SEED);
166+
let values: Buffer<i32> = (0..CARDINALITY_ELEMENT_COUNT)
167+
.map(|_| rng.random_range(-512..512))
168+
.collect();
169+
let mut group_ids: Vec<u32> = (0..CARDINALITY_ELEMENT_COUNT)
170+
.map(|idx| (idx * group_count / CARDINALITY_ELEMENT_COUNT) as u32)
171+
.collect();
172+
if matches!(order, IdOrder::Shuffled) {
173+
group_ids.shuffle(&mut rng);
174+
}
175+
176+
DenseGroupedInput {
177+
values: PrimitiveArray::new(values, Validity::NonNullable).into_array(),
178+
group_ids: GroupIds::from_buffer(Buffer::from(group_ids), group_count).unwrap(),
154179
}
155180
}
156181

157-
fn grouped_accumulator<V>(list_view: &ArrayRef, vtable: V) -> ArrayRef
182+
fn grouped_accumulator<V>(input: &DenseGroupedInput, vtable: V) -> ArrayRef
158183
where
159-
V: AggregateFnVTable<Options = NumericalAggregateOpts> + Clone,
184+
V: AggregateFnVTable + Clone,
185+
V::Options: Default,
160186
{
161-
let mut acc = GroupedAccumulator::try_new(
162-
vtable,
163-
NumericalAggregateOpts::default(),
164-
list_element_dtype(list_view),
165-
)
166-
.unwrap();
167-
acc.accumulate_list(list_view, &mut SESSION.create_execution_ctx())
187+
let mut acc =
188+
GroupedAccumulator::try_new(vtable, V::Options::default(), input.values.dtype().clone())
189+
.unwrap();
190+
let mut ctx = SESSION.create_execution_ctx();
191+
acc.accumulate(&input.values, &input.group_ids, &mut ctx)
168192
.unwrap();
169-
divan::black_box(acc.finish().unwrap())
193+
let result = acc
194+
.finish()
195+
.unwrap()
196+
.execute::<Canonical>(&mut ctx)
197+
.unwrap()
198+
.into_array();
199+
divan::black_box(result)
170200
}
171201

172202
#[divan::bench]
@@ -233,104 +263,41 @@ fn sum_v2_f64_clustered_nulls(bencher: Bencher) {
233263
.bench_refs(|input| grouped_accumulator(input, SumV2));
234264
}
235265

236-
/// Execute the lazy finalize result to canonical form so the benchmark includes the full cost of
237-
/// producing usable sums.
238-
fn grouped_accumulator_canonical<V>(list_view: &ArrayRef, vtable: V) -> ArrayRef
239-
where
240-
V: AggregateFnVTable<Options = NumericalAggregateOpts> + Clone,
241-
{
242-
let mut acc = GroupedAccumulator::try_new(
243-
vtable,
244-
NumericalAggregateOpts::default(),
245-
list_element_dtype(list_view),
246-
)
247-
.unwrap();
248-
let mut ctx = SESSION.create_execution_ctx();
249-
acc.accumulate_list(list_view, &mut ctx).unwrap();
250-
let result = acc
251-
.finish()
252-
.unwrap()
253-
.execute::<Canonical>(&mut ctx)
254-
.unwrap()
255-
.into_array();
256-
divan::black_box(result)
257-
}
258-
259-
#[divan::bench]
260-
fn canonical_sum_i32_nullable_all_valid(bencher: Bencher) {
261-
let input = i32_nullable_all_valid_input();
262-
bencher
263-
.with_inputs(|| &input)
264-
.bench_refs(|input| grouped_accumulator_canonical(input, Sum));
265-
}
266-
267266
#[divan::bench]
268-
fn canonical_sum_i32_clustered_nulls(bencher: Bencher) {
269-
let input = i32_clustered_nulls_input();
270-
bencher
271-
.with_inputs(|| &input)
272-
.bench_refs(|input| grouped_accumulator_canonical(input, Sum));
273-
}
274-
275-
#[divan::bench]
276-
fn canonical_sum_f64_all_valid(bencher: Bencher) {
277-
let input = f64_all_valid_input();
278-
bencher
279-
.with_inputs(|| &input)
280-
.bench_refs(|input| grouped_accumulator_canonical(input, Sum));
281-
}
282-
283-
#[divan::bench]
284-
fn canonical_sum_f64_clustered_nulls(bencher: Bencher) {
285-
let input = f64_clustered_nulls_input();
286-
bencher
287-
.with_inputs(|| &input)
288-
.bench_refs(|input| grouped_accumulator_canonical(input, Sum));
289-
}
290-
291-
#[divan::bench]
292-
fn canonical_sum_v2_i32_nullable_all_valid(bencher: Bencher) {
293-
let input = i32_nullable_all_valid_input();
294-
bencher
295-
.with_inputs(|| &input)
296-
.bench_refs(|input| grouped_accumulator_canonical(input, SumV2));
297-
}
298-
299-
#[divan::bench]
300-
fn canonical_sum_v2_i32_clustered_nulls(bencher: Bencher) {
267+
fn count_i32_clustered_nulls(bencher: Bencher) {
301268
let input = i32_clustered_nulls_input();
302269
bencher
303270
.with_inputs(|| &input)
304-
.bench_refs(|input| grouped_accumulator_canonical(input, SumV2));
271+
.bench_refs(|input| grouped_accumulator(input, Count));
305272
}
306273

307274
#[divan::bench]
308-
fn canonical_sum_v2_f64_all_valid(bencher: Bencher) {
309-
let input = f64_all_valid_input();
275+
fn count_varbinview(bencher: Bencher) {
276+
let input = varbinview_input();
310277
bencher
311278
.with_inputs(|| &input)
312-
.bench_refs(|input| grouped_accumulator_canonical(input, SumV2));
279+
.bench_refs(|input| grouped_accumulator(input, Count));
313280
}
314281

315-
#[divan::bench]
316-
fn canonical_sum_v2_f64_clustered_nulls(bencher: Bencher) {
317-
let input = f64_clustered_nulls_input();
282+
#[divan::bench(args = CARDINALITY_ARGS)]
283+
fn sum_i32_cardinality(bencher: Bencher, (group_count, order): (usize, IdOrder)) {
284+
let input = i32_cardinality_input(group_count, order);
318285
bencher
319286
.with_inputs(|| &input)
320-
.bench_refs(|input| grouped_accumulator_canonical(input, SumV2));
287+
.bench_refs(|input| grouped_accumulator(input, Sum));
321288
}
322289

323-
#[divan::bench]
324-
fn count_i32_clustered_nulls(bencher: Bencher) {
325-
let input = i32_clustered_nulls_input();
290+
#[divan::bench(args = CARDINALITY_ARGS)]
291+
fn sum_v2_i32_cardinality(bencher: Bencher, (group_count, order): (usize, IdOrder)) {
292+
let input = i32_cardinality_input(group_count, order);
326293
bencher
327294
.with_inputs(|| &input)
328-
.bench_refs(|input| grouped_accumulator(input, Count));
295+
.bench_refs(|input| grouped_accumulator(input, SumV2));
329296
}
330297

331-
#[divan::bench]
332-
fn count_varbinview(bencher: Bencher) {
333-
let input = varbinview_input();
298+
#[divan::bench(args = CARDINALITY_ARGS)]
299+
fn count_i32_cardinality(bencher: Bencher, (group_count, order): (usize, IdOrder)) {
300+
let input = i32_cardinality_input(group_count, order);
334301
bencher
335302
.with_inputs(|| &input)
336303
.bench_refs(|input| grouped_accumulator(input, Count));

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -172,7 +172,7 @@ impl<V: AggregateFnVTable> DynAccumulator for Accumulator<V> {
172172
}
173173

174174
// 3. Iteratively check the registry against each intermediate encoding, executing one
175-
// step between checks. Mirrors the loop in `GroupedAccumulator::accumulate_list_view`.
175+
// step between checks. Mirrors the loop in `GroupedAccumulator::accumulate`.
176176
// Iteration 0 re-checks the initial encoding — a redundant HashMap miss, the price of
177177
// keeping the loop body uniform. Terminates on `AnyColumnar` (Canonical or Constant)
178178
// since the vtable's `accumulate(&Columnar)` handles both cases directly.

0 commit comments

Comments
 (0)