Skip to content
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

Jump to bottom

Chunk based iteration in accumulate_indices #13451

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
+20 −67
Open

Chunk based iteration in accumulate_indices #13451

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
90e8930
filter chunk
jayzhan211 Nov 16, 2024
d0c8e05
fmt
jayzhan211 Nov 16, 2024
f3697b8
acc
jayzhan211 Nov 16, 2024
6d9f26c
BitIndexIterator
jayzhan211 Nov 17, 2024
3721e74
cleanup
jayzhan211 Nov 17, 2024
a889299
count group
jayzhan211 Nov 17, 2024
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
73 changes: 15 additions & 58 deletions datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -395,69 +395,26 @@ pub fn accumulate_indices<F>(
}
}
(None, Some(filter)) => {
assert_eq!(filter.len(), group_indices.len());
// The performance with a filter could be improved by
// iterating over the filter in chunks, rather than a single
// iterator. TODO file a ticket
let iter = group_indices.iter().zip(filter.iter());
for (&group_index, filter_value) in iter {
if let Some(true) = filter_value {
index_fn(group_index)
}
}
debug_assert_eq!(filter.len(), group_indices.len());
filter
.values()
.set_indices()
.for_each(|index| index_fn(group_indices[index]));
}
(Some(valids), None) => {
assert_eq!(valids.len(), group_indices.len());
// This is based on (ahem, COPY/PASTA) arrow::compute::aggregate::sum
// iterate over in chunks of 64 bits for more efficient null checking
let group_indices_chunks = group_indices.chunks_exact(64);
let bit_chunks = valids.inner().bit_chunks();

let group_indices_remainder = group_indices_chunks.remainder();

group_indices_chunks.zip(bit_chunks.iter()).for_each(
|(group_index_chunk, mask)| {
// index_mask has value 1 << i in the loop
let mut index_mask = 1;
group_index_chunk.iter().for_each(|&group_index| {
// valid bit was set, real vale
let is_valid = (mask & index_mask) != 0;
if is_valid {
index_fn(group_index);
}
index_mask <<= 1;
})
},
);

// handle any remaining bits (after the initial 64)
let remainder_bits = bit_chunks.remainder_bits();
group_indices_remainder
.iter()
.enumerate()
.for_each(|(i, &group_index)| {
let is_valid = remainder_bits & (1 << i) != 0;
if is_valid {
index_fn(group_index)
}
});
debug_assert_eq!(valids.len(), group_indices.len());
valids
.valid_indices()
.for_each(|index| index_fn(group_indices[index]));
}

(Some(valids), Some(filter)) => {
assert_eq!(filter.len(), group_indices.len());
assert_eq!(valids.len(), group_indices.len());
// The performance with a filter could likely be improved by
// iterating over the filter in chunks, rather than using
// iterators. TODO file a ticket
filter
.iter()
.zip(group_indices.iter())
.zip(valids.iter())
.for_each(|((filter_value, &group_index), is_valid)| {
if let (Some(true), true) = (filter_value, is_valid) {
index_fn(group_index)
}
})
debug_assert_eq!(filter.len(), group_indices.len());
debug_assert_eq!(valids.len(), group_indices.len());
let valid_and_filter_indices = valids.inner() & filter.values();
valid_and_filter_indices
.set_indices()
.for_each(|index| index_fn(group_indices[index]));
}
}
}
Expand Down
14 changes: 5 additions & 9 deletions datafusion/functions-aggregate/src/count.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -481,15 +481,11 @@ impl GroupsAccumulator for CountGroupsAccumulator {
// Adds the counts with the partial counts
self.counts.resize(total_num_groups, 0);
match opt_filter {
Some(filter) => filter
.iter()
.zip(group_indices.iter())
.zip(partial_counts.iter())
.for_each(|((filter_value, &group_index), partial_count)| {
if let Some(true) = filter_value {
self.counts[group_index] += partial_count;
}
}),
Some(filter) => {
filter.values().set_indices().for_each(|index| {
self.counts[group_indices[index]] += partial_counts[index]
});
}
None => group_indices.iter().zip(partial_counts.iter()).for_each(
|(&group_index, partial_count)| {
self.counts[group_index] += partial_count;
Expand Down