Skip to content

Support out of order grouped aggregate accumulation #8379

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
gatesn wants to merge 12 commits into
base: develop
Choose a base branch
from
Draft

Support out of order grouped aggregate accumulation #8379

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
6bdcd32
Support dense grouped aggregate accumulation
gatesn Jun 11, 2026
9bd157f
Clarify dense grouped aggregate ids
gatesn Jun 11, 2026
50701b2
Fix grouped aggregate CI checks
gatesn Jun 11, 2026
adae76e
DCO Remediation Commit for Nicholas Gates <nick@nickgates.com>
gatesn Jun 11, 2026
e999a5f
Merge origin/develop into grouped aggregate
gatesn Jun 12, 2026
c61b7bb
Restore dense grouped aggregate test coverage
gatesn Jun 12, 2026
2ef64b2
Optimize dense grouped aggregates
gatesn Jun 12, 2026
4ba3988
group ids as array ref, multi encoding kernel lookup (#8550)
onursatici Jun 29, 2026
9022490
Merge develop and complete grouped aggregates
gatesn Aug 7, 2026
3560666
Optimize grouped aggregate hot paths
gatesn Aug 7, 2026
f52c15d
Optimize grouped aggregate state storage
gatesn Aug 8, 2026
b27ec6a
Merge remote-tracking branch 'origin/develop' into ngates/grouped-agg…
gatesn Aug 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
170 changes: 82 additions & 88 deletions vortex-array/benches/aggregate_grouped.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,19 +10,19 @@ use divan::Bencher;
use rand::RngExt;
use rand::SeedableRng;
use rand::rngs::StdRng;
use rand::seq::SliceRandom;
use vortex_array::ArrayRef;
use vortex_array::Canonical;
use vortex_array::IntoArray;
use vortex_array::VortexSessionExecute;
use vortex_array::aggregate_fn::AggregateFnVTable;
use vortex_array::aggregate_fn::DynGroupedAccumulator;
use vortex_array::aggregate_fn::GroupIds;
use vortex_array::aggregate_fn::GroupedAccumulator;
use vortex_array::aggregate_fn::fns::count::Count;
use vortex_array::aggregate_fn::fns::sum::Sum;
use vortex_array::arrays::ListViewArray;
use vortex_array::arrays::PrimitiveArray;
use vortex_array::arrays::VarBinViewArray;
use vortex_array::dtype::DType;
use vortex_array::validity::Validity;
use vortex_buffer::Buffer;
use vortex_session::VortexSession;
Expand All @@ -38,6 +38,22 @@ const GROUP_COUNT: usize = 128;
const GROUP_SIZE_SEED: u64 = 42;
const MIN_VALUES_PER_GROUP: usize = 1;
const MAX_VALUES_PER_GROUP: usize = 15;
const CARDINALITY_ELEMENT_COUNT: usize = 1 << 16;

#[derive(Clone, Copy, Debug)]
enum IdOrder {
Clustered,
Shuffled,
}

const CARDINALITY_ARGS: &[(usize, IdOrder)] = &[
(128, IdOrder::Clustered),
(128, IdOrder::Shuffled),
(1 << 12, IdOrder::Clustered),
(1 << 12, IdOrder::Shuffled),
(1 << 16, IdOrder::Clustered),
(1 << 16, IdOrder::Shuffled),
];

fn random_group_sizes() -> Vec<usize> {
let mut rng = StdRng::seed_from_u64(GROUP_SIZE_SEED);
Expand All @@ -50,44 +66,40 @@ fn total_element_count(group_sizes: &[usize]) -> usize {
group_sizes.iter().sum()
}

fn contiguous_list_view(elements: ArrayRef, group_sizes: &[usize]) -> ArrayRef {
let mut offset = 0usize;
let offsets: Buffer<u32> = group_sizes
.iter()
.map(|&size| {
let current_offset = offset;
offset += size;
current_offset as u32
})
.collect();
let sizes: Buffer<u32> = group_sizes.iter().map(|&size| size as u32).collect();
struct DenseGroupedInput {
values: ArrayRef,
group_ids: GroupIds,
}

assert_eq!(elements.len(), total_element_count(group_sizes));
fn dense_grouped_input(values: ArrayRef, group_sizes: &[usize]) -> DenseGroupedInput {

@AdamGS AdamGS Jun 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit - this is basically the constructor for DenseGroupedInput

assert_eq!(values.len(), total_element_count(group_sizes));

ListViewArray::try_new(
elements,
offsets.into_array(),
sizes.into_array(),
Validity::NonNullable,
let group_ids = GroupIds::from_iter(
group_sizes
.iter()
.enumerate()
.flat_map(|(group_id, &size)| std::iter::repeat_n(group_id as u32, size)),
group_sizes.len(),
)
.unwrap()
.into_array()
.unwrap();

DenseGroupedInput { values, group_ids }
}

fn i32_nullable_all_valid_input() -> ArrayRef {
fn i32_nullable_all_valid_input() -> DenseGroupedInput {
let group_sizes = random_group_sizes();
let element_count = total_element_count(&group_sizes);
let values: Buffer<i32> = (0..element_count)
.map(|i| (i % 1024) as i32 - 512)
.collect();
let validity = Validity::from_iter(std::iter::repeat_n(true, element_count));
contiguous_list_view(
dense_grouped_input(
PrimitiveArray::new(values, validity).into_array(),
&group_sizes,
)
}

fn i32_clustered_nulls_input() -> ArrayRef {
fn i32_clustered_nulls_input() -> DenseGroupedInput {
let group_sizes = random_group_sizes();
let element_count = total_element_count(&group_sizes);
let values = (0..element_count).map(|i| {
Expand All @@ -97,26 +109,26 @@ fn i32_clustered_nulls_input() -> ArrayRef {
Some((i % 1024) as i32 - 512)
}
});
contiguous_list_view(
dense_grouped_input(
PrimitiveArray::from_option_iter(values).into_array(),
&group_sizes,
)
}

fn f64_all_valid_input() -> ArrayRef {
fn f64_all_valid_input() -> DenseGroupedInput {
let group_sizes = random_group_sizes();
let element_count = total_element_count(&group_sizes);
let mut rng = StdRng::seed_from_u64(GROUP_SIZE_SEED);
let values: Buffer<f64> = (0..element_count)
.map(|_| rng.random_range(-1000.0..1000.0))
.collect();
contiguous_list_view(
dense_grouped_input(
PrimitiveArray::new(values, Validity::NonNullable).into_array(),
&group_sizes,
)
}

fn f64_clustered_nulls_input() -> ArrayRef {
fn f64_clustered_nulls_input() -> DenseGroupedInput {
let group_sizes = random_group_sizes();
let element_count = total_element_count(&group_sizes);
let mut rng = StdRng::seed_from_u64(GROUP_SIZE_SEED);
Expand All @@ -127,42 +139,61 @@ fn f64_clustered_nulls_input() -> ArrayRef {
Some(rng.random_range(-1000.0f64..1000.0))
}
});
contiguous_list_view(
dense_grouped_input(
PrimitiveArray::from_option_iter(values).into_array(),
&group_sizes,
)
}

fn varbinview_input() -> ArrayRef {
fn varbinview_input() -> DenseGroupedInput {
let group_sizes = random_group_sizes();
let element_count = total_element_count(&group_sizes);
let values: Vec<String> = (0..element_count)
.map(|i| format!("value-{i:06}"))
.collect();
contiguous_list_view(
dense_grouped_input(
VarBinViewArray::from_iter_str(values.iter().map(String::as_str)).into_array(),
&group_sizes,
)
}

fn list_element_dtype(list_view: &ArrayRef) -> DType {
match list_view.dtype() {
DType::List(element_dtype, _) => element_dtype.as_ref().clone(),
dtype => unreachable!("expected List dtype, got {dtype}"),
fn i32_cardinality_input(group_count: usize, order: IdOrder) -> DenseGroupedInput {
let mut rng = StdRng::seed_from_u64(GROUP_SIZE_SEED);
let values: Buffer<i32> = (0..CARDINALITY_ELEMENT_COUNT)
.map(|_| rng.random_range(-512..512))
.collect();
let mut group_ids: Vec<u32> = (0..CARDINALITY_ELEMENT_COUNT)
.map(|idx| (idx * group_count / CARDINALITY_ELEMENT_COUNT) as u32)
.collect();
if matches!(order, IdOrder::Shuffled) {
group_ids.shuffle(&mut rng);
}

DenseGroupedInput {
values: PrimitiveArray::new(values, Validity::NonNullable).into_array(),
group_ids: GroupIds::from_buffer(Buffer::from(group_ids), group_count).unwrap(),
}
}

fn grouped_accumulator<V>(list_view: &ArrayRef, vtable: V) -> ArrayRef
fn grouped_accumulator<V>(input: &DenseGroupedInput, vtable: V) -> ArrayRef
where
V: AggregateFnVTable + Clone,
V::Options: Default,
{
let mut acc =
GroupedAccumulator::try_new(vtable, V::Options::default(), list_element_dtype(list_view))
GroupedAccumulator::try_new(vtable, V::Options::default(), input.values.dtype().clone())
.unwrap();
acc.accumulate_list(list_view, &mut SESSION.create_execution_ctx())
let num_groups = input.group_ids.num_groups();
let mut ctx = SESSION.create_execution_ctx();
acc.accumulate(&input.values, &input.group_ids, &mut ctx)
.unwrap();
divan::black_box(acc.finish().unwrap())
let result = acc
.finish(num_groups)
.unwrap()
.execute::<Canonical>(&mut ctx)
.unwrap()
.into_array();
divan::black_box(result)
}

#[divan::bench]
Expand Down Expand Up @@ -197,70 +228,33 @@ fn sum_f64_clustered_nulls(bencher: Bencher) {
.bench_refs(|input| grouped_accumulator(input, Sum));
}

/// Like [`grouped_accumulator`], but executes the lazy finalize result to canonical so the
/// bench measures the full cost of producing usable sums.
fn grouped_accumulator_canonical<V>(list_view: &ArrayRef, vtable: V) -> ArrayRef
where
V: AggregateFnVTable + Clone,
V::Options: Default,
{
let mut acc =
GroupedAccumulator::try_new(vtable, V::Options::default(), list_element_dtype(list_view))
.unwrap();
let mut ctx = SESSION.create_execution_ctx();
acc.accumulate_list(list_view, &mut ctx).unwrap();
let result = acc
.finish()
.unwrap()
.execute::<Canonical>(&mut ctx)
.unwrap()
.into_array();
divan::black_box(result)
}

#[divan::bench]
fn canonical_sum_i32_nullable_all_valid(bencher: Bencher) {
let input = i32_nullable_all_valid_input();
bencher
.with_inputs(|| &input)
.bench_refs(|input| grouped_accumulator_canonical(input, Sum));
}

#[divan::bench]
fn canonical_sum_i32_clustered_nulls(bencher: Bencher) {
fn count_i32_clustered_nulls(bencher: Bencher) {
let input = i32_clustered_nulls_input();
bencher
.with_inputs(|| &input)
.bench_refs(|input| grouped_accumulator_canonical(input, Sum));
}

#[divan::bench]
fn canonical_sum_f64_all_valid(bencher: Bencher) {
let input = f64_all_valid_input();
bencher
.with_inputs(|| &input)
.bench_refs(|input| grouped_accumulator_canonical(input, Sum));
.bench_refs(|input| grouped_accumulator(input, Count));
}

#[divan::bench]
fn canonical_sum_f64_clustered_nulls(bencher: Bencher) {
let input = f64_clustered_nulls_input();
fn count_varbinview(bencher: Bencher) {
let input = varbinview_input();
bencher
.with_inputs(|| &input)
.bench_refs(|input| grouped_accumulator_canonical(input, Sum));
.bench_refs(|input| grouped_accumulator(input, Count));
}

#[divan::bench]
fn count_i32_clustered_nulls(bencher: Bencher) {
let input = i32_clustered_nulls_input();
#[divan::bench(args = CARDINALITY_ARGS)]
fn sum_i32_cardinality(bencher: Bencher, (group_count, order): (usize, IdOrder)) {
let input = i32_cardinality_input(group_count, order);
bencher
.with_inputs(|| &input)
.bench_refs(|input| grouped_accumulator(input, Count));
.bench_refs(|input| grouped_accumulator(input, Sum));
}

#[divan::bench]
fn count_varbinview(bencher: Bencher) {
let input = varbinview_input();
#[divan::bench(args = CARDINALITY_ARGS)]
fn count_i32_cardinality(bencher: Bencher, (group_count, order): (usize, IdOrder)) {
let input = i32_cardinality_input(group_count, order);
bencher
.with_inputs(|| &input)
.bench_refs(|input| grouped_accumulator(input, Count));
Expand Down
2 changes: 1 addition & 1 deletion vortex-array/src/aggregate_fn/accumulator.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -181,7 +181,7 @@ impl<V: AggregateFnVTable> DynAccumulator for Accumulator<V> {
}

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