Skip to content

feat: Improved multiple column aggregation performance by using bitmasks rather than Vec<bool> #21886

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

Open
huymq1710 wants to merge 9 commits into
base: main
Choose a base branch
from
Open

feat: Improved multiple column aggregation performance by using bitmasks rather than Vec<bool> #21886

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
31556f6
feat: use bitmasks for multiple column aggregation
huymq1710 Apr 28, 2026
8efed09
Revert "feat: use bitmasks for multiple column aggregation"
huymq1710 Apr 28, 2026
28b5b37
feat: use bitmasks for multiple column aggregation
huymq1710 Apr 28, 2026
4162bb6
feat: use bitmasks for multiple column aggregation
huymq1710 Apr 28, 2026
063a26b
Revert "feat: use bitmasks for multiple column aggregation"
huymq1710 Apr 28, 2026
27dfdb7
feat: use bitmasks for multiple column aggregation
huymq1710 Apr 28, 2026
988aff6
Merge branch 'column-aggregation-bitmasks' of https://github.com/huym…
huymq1710 Apr 28, 2026
f3f0def
chore: add comments
huymq1710 Apr 28, 2026
f878431
chore: run cargo fmt
huymq1710 Apr 28, 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
194 changes: 106 additions & 88 deletions datafusion/physical-plan/src/aggregates/group_values/multi_group_by/boolean.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@ use crate::aggregates::group_values::multi_group_by::{GroupColumn, nulls_equal_t
use crate::aggregates::group_values::null_builder::MaybeNullBufferBuilder;
use arrow::array::{Array as _, ArrayRef, AsArray, BooleanArray, BooleanBufferBuilder};
use datafusion_common::Result;
use itertools::izip;

/// An implementation of [`GroupColumn`] for booleans
///
Expand Down Expand Up @@ -81,32 +80,31 @@ impl<const NULLABLE: bool> GroupColumn for BooleanGroupValueBuilder<NULLABLE> {
lhs_rows: &[usize],
array: &ArrayRef,
rhs_rows: &[usize],
equal_to_results: &mut [bool],
equal_to_results: &mut BooleanBufferBuilder,
) {
let array = array.as_boolean();

let iter = izip!(
lhs_rows.iter(),
rhs_rows.iter(),
equal_to_results.iter_mut(),
);

for (&lhs_row, &rhs_row, equal_to_result) in iter {
// Has found not equal to in previous column, don't need to check
if !*equal_to_result {
for (idx, (&lhs_row, &rhs_row)) in
lhs_rows.iter().zip(rhs_rows.iter()).enumerate()
{
if !equal_to_results.get_bit(idx) {
continue;
}

if NULLABLE {
let exist_null = self.nulls.is_null(lhs_row);
let input_null = array.is_null(rhs_row);
if let Some(result) = nulls_equal_to(exist_null, input_null) {
*equal_to_result = result;
if !result {
equal_to_results.set_bit(idx, false);
}
continue;
}
}

*equal_to_result = self.buffer.get_bit(lhs_row) == array.value(rhs_row);
if self.buffer.get_bit(lhs_row) != array.value(rhs_row) {
equal_to_results.set_bit(idx, false);
}
}
}

Expand Down Expand Up @@ -195,10 +193,20 @@ impl<const NULLABLE: bool> GroupColumn for BooleanGroupValueBuilder<NULLABLE> {

#[cfg(test)]
mod tests {
use arrow::array::NullBufferBuilder;
use arrow::array::{BooleanBufferBuilder, NullBufferBuilder};

use super::*;

fn make_true_buffer(n: usize) -> BooleanBufferBuilder {
let mut buf = BooleanBufferBuilder::new(n);
buf.append_n(n, true);
buf
}

fn to_vec(buf: &BooleanBufferBuilder) -> Vec<bool> {
(0..buf.len()).map(|i| buf.get_bit(i)).collect()
}

#[test]
fn test_nullable_boolean_equal_to() {
let append = |builder: &mut BooleanGroupValueBuilder<true>,
Expand All @@ -209,16 +217,18 @@ mod tests {
}
};

let equal_to = |builder: &BooleanGroupValueBuilder<true>,
lhs_rows: &[usize],
input_array: &ArrayRef,
rhs_rows: &[usize],
equal_to_results: &mut Vec<bool>| {
let iter = lhs_rows.iter().zip(rhs_rows.iter());
for (idx, (&lhs_row, &rhs_row)) in iter.enumerate() {
equal_to_results[idx] = builder.equal_to(lhs_row, input_array, rhs_row);
}
};
let equal_to =
|builder: &BooleanGroupValueBuilder<true>,
lhs_rows: &[usize],
input_array: &ArrayRef,
rhs_rows: &[usize],
equal_to_results: &mut BooleanBufferBuilder| {
let iter = lhs_rows.iter().zip(rhs_rows.iter());
for (idx, (&lhs_row, &rhs_row)) in iter.enumerate() {
equal_to_results
.set_bit(idx, builder.equal_to(lhs_row, input_array, rhs_row));
}
};

test_nullable_boolean_equal_to_internal(append, equal_to);
}
Expand All @@ -233,18 +243,19 @@ mod tests {
.unwrap();
};

let equal_to = |builder: &BooleanGroupValueBuilder<true>,
lhs_rows: &[usize],
input_array: &ArrayRef,
rhs_rows: &[usize],
equal_to_results: &mut Vec<bool>| {
builder.vectorized_equal_to(
lhs_rows,
input_array,
rhs_rows,
equal_to_results,
);
};
let equal_to =
|builder: &BooleanGroupValueBuilder<true>,
lhs_rows: &[usize],
input_array: &ArrayRef,
rhs_rows: &[usize],
equal_to_results: &mut BooleanBufferBuilder| {
builder.vectorized_equal_to(
lhs_rows,
input_array,
rhs_rows,
equal_to_results,
);
};

test_nullable_boolean_equal_to_internal(append, equal_to);
}
Expand All @@ -257,7 +268,7 @@ mod tests {
&[usize],
&ArrayRef,
&[usize],
&mut Vec<bool>,
&mut BooleanBufferBuilder,
),
{
// Will cover such cases:
Expand All @@ -268,7 +279,7 @@ mod tests {
// - exist not null, input not null; values not equal
// - exist not null, input not null; values equal

// Define PrimitiveGroupValueBuilder
// Define BooleanGroupValueBuilder
let mut builder = BooleanGroupValueBuilder::<true>::new();
let builder_array = Arc::new(BooleanArray::from(vec![
None,
Expand All @@ -294,29 +305,30 @@ mod tests {
// explicitly build a null buffer where one of the null values also happens to match
let mut nulls = NullBufferBuilder::new(6);
nulls.append_non_null();
nulls.append_null(); // this sets Some(false) to null above
nulls.append_null();
nulls.append_null();
nulls.append_null();
nulls.append_non_null();
nulls.append_non_null();
let input_array = Arc::new(BooleanArray::new(values, nulls.finish())) as ArrayRef;

// Check
let mut equal_to_results = vec![true; builder.len()];
let mut equal_to_results = make_true_buffer(builder.len());
equal_to(
&builder,
&[0, 1, 2, 3, 4, 5],
&input_array,
&[0, 1, 2, 3, 4, 5],
&mut equal_to_results,
);

assert!(!equal_to_results[0]);
assert!(equal_to_results[1]);
assert!(equal_to_results[2]);
assert!(!equal_to_results[3]);
assert!(!equal_to_results[4]);
assert!(equal_to_results[5]);
let results = to_vec(&equal_to_results);

assert!(!results[0]);
assert!(results[1]);
assert!(results[2]);
assert!(!results[3]);
assert!(!results[4]);
assert!(results[5]);
}

#[test]
Expand All @@ -329,16 +341,18 @@ mod tests {
}
};

let equal_to = |builder: &BooleanGroupValueBuilder<false>,
lhs_rows: &[usize],
input_array: &ArrayRef,
rhs_rows: &[usize],
equal_to_results: &mut Vec<bool>| {
let iter = lhs_rows.iter().zip(rhs_rows.iter());
for (idx, (&lhs_row, &rhs_row)) in iter.enumerate() {
equal_to_results[idx] = builder.equal_to(lhs_row, input_array, rhs_row);
}
};
let equal_to =
|builder: &BooleanGroupValueBuilder<false>,
lhs_rows: &[usize],
input_array: &ArrayRef,
rhs_rows: &[usize],
equal_to_results: &mut BooleanBufferBuilder| {
let iter = lhs_rows.iter().zip(rhs_rows.iter());
for (idx, (&lhs_row, &rhs_row)) in iter.enumerate() {
equal_to_results
.set_bit(idx, builder.equal_to(lhs_row, input_array, rhs_row));
}
};

test_not_nullable_boolean_equal_to_internal(append, equal_to);
}
Expand All @@ -353,18 +367,19 @@ mod tests {
.unwrap();
};

let equal_to = |builder: &BooleanGroupValueBuilder<false>,
lhs_rows: &[usize],
input_array: &ArrayRef,
rhs_rows: &[usize],
equal_to_results: &mut Vec<bool>| {
builder.vectorized_equal_to(
lhs_rows,
input_array,
rhs_rows,
equal_to_results,
);
};
let equal_to =
|builder: &BooleanGroupValueBuilder<false>,
lhs_rows: &[usize],
input_array: &ArrayRef,
rhs_rows: &[usize],
equal_to_results: &mut BooleanBufferBuilder| {
builder.vectorized_equal_to(
lhs_rows,
input_array,
rhs_rows,
equal_to_results,
);
};

test_not_nullable_boolean_equal_to_internal(append, equal_to);
}
Expand All @@ -377,14 +392,14 @@ mod tests {
&[usize],
&ArrayRef,
&[usize],
&mut Vec<bool>,
&mut BooleanBufferBuilder,
),
{
// Will cover such cases:
// - values equal
// - values not equal

// Define PrimitiveGroupValueBuilder
// Define BooleanGroupValueBuilder
let mut builder = BooleanGroupValueBuilder::<false>::new();
let builder_array = Arc::new(BooleanArray::from(vec![
Some(false),
Expand All @@ -403,19 +418,20 @@ mod tests {
])) as ArrayRef;

// Check
let mut equal_to_results = vec![true; builder.len()];
let mut equal_to_results = make_true_buffer(builder.len());
equal_to(
&builder,
&[0, 1, 2, 3],
&input_array,
&[0, 1, 2, 3],
&mut equal_to_results,
);
let results = to_vec(&equal_to_results);

assert!(equal_to_results[0]);
assert!(!equal_to_results[1]);
assert!(!equal_to_results[2]);
assert!(equal_to_results[3]);
assert!(results[0]);
assert!(!results[1]);
assert!(!results[2]);
assert!(results[3]);
}

#[test]
Expand All @@ -432,19 +448,20 @@ mod tests {
.vectorized_append(&all_nulls_input_array, &[0, 1, 2, 3, 4])
.unwrap();

let mut equal_to_results = vec![true; all_nulls_input_array.len()];
let mut equal_to_results = make_true_buffer(all_nulls_input_array.len());
builder.vectorized_equal_to(
&[0, 1, 2, 3, 4],
&all_nulls_input_array,
&[0, 1, 2, 3, 4],
&mut equal_to_results,
);
let results = to_vec(&equal_to_results);

assert!(equal_to_results[0]);
assert!(equal_to_results[1]);
assert!(equal_to_results[2]);
assert!(equal_to_results[3]);
assert!(equal_to_results[4]);
assert!(results[0]);
assert!(results[1]);
assert!(results[2]);
assert!(results[3]);
assert!(results[4]);

// All not nulls input array
let all_not_nulls_input_array = Arc::new(BooleanArray::from(vec![
Expand All @@ -458,18 +475,19 @@ mod tests {
.vectorized_append(&all_not_nulls_input_array, &[0, 1, 2, 3, 4])
.unwrap();

let mut equal_to_results = vec![true; all_not_nulls_input_array.len()];
let mut equal_to_results = make_true_buffer(all_not_nulls_input_array.len());
builder.vectorized_equal_to(
&[5, 6, 7, 8, 9],
&all_not_nulls_input_array,
&[0, 1, 2, 3, 4],
&mut equal_to_results,
);
let results = to_vec(&equal_to_results);

assert!(equal_to_results[0]);
assert!(equal_to_results[1]);
assert!(equal_to_results[2]);
assert!(equal_to_results[3]);
assert!(equal_to_results[4]);
assert!(results[0]);
assert!(results[1]);
assert!(results[2]);
assert!(results[3]);
assert!(results[4]);
}
}
Loading
Loading