Skip to content

Commit 963e8af

Browse files
jayzhan211jayzhan211
authored
Chunk based iteration in accumulate_indices (#13451)
* filter chunk Signed-off-by: Jay Zhan <jayzhan211@gmail.com> * fmt Signed-off-by: Jay Zhan <jayzhan211@gmail.com> * acc Signed-off-by: Jay Zhan <jayzhan211@gmail.com> * BitIndexIterator Signed-off-by: Jay Zhan <jayzhan211@gmail.com> * cleanup Signed-off-by: Jay Zhan <jayzhan211@gmail.com> * count group Signed-off-by: Jay Zhan <jayzhan211@gmail.com> * add benches Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * revert to fixed chunk based method instead of iterating set_indices Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * revert count change Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * clippy Signed-off-by: jayzhan211 <jayzhan211@gmail.com> * taplo format Signed-off-by: jayzhan211 <jayzhan211@gmail.com> --------- Signed-off-by: Jay Zhan <jayzhan211@gmail.com> Signed-off-by: jayzhan211 <jayzhan211@gmail.com>
1 parent aef232b commit 963e8af

3 files changed

Lines changed: 190 additions & 22 deletions

File tree

‎datafusion/functions-aggregate-common/Cargo.toml‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,3 +43,10 @@ datafusion-common = { workspace = true }
4343
datafusion-expr-common = { workspace = true }
4444
datafusion-physical-expr-common = { workspace = true }
4545
rand = { workspace = true }
46+
47+
[dev-dependencies]
48+
criterion = "0.5"
49+
50+
[[bench]]
51+
harness = false
52+
name = "accumulate"

‎datafusion/functions-aggregate-common/benches/accumulate.rs‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
extern crate criterion;
19+
20+
use std::sync::Arc;
21+
22+
use arrow::array::{ArrayRef, BooleanArray, Int64Array};
23+
use criterion::{criterion_group, criterion_main, Criterion};
24+
use datafusion_functions_aggregate_common::aggregate::groups_accumulator::accumulate::accumulate_indices;
25+
26+
fn generate_group_indices(len: usize) -> Vec<usize> {
27+
(0..len).collect()
28+
}
29+
30+
fn generate_values(len: usize, has_null: bool) -> ArrayRef {
31+
if has_null {
32+
let values = (0..len)
33+
.map(|i| if i % 7 == 0 { None } else { Some(i as i64) })
34+
.collect::<Vec<_>>();
35+
Arc::new(Int64Array::from(values))
36+
} else {
37+
let values = (0..len).map(|i| Some(i as i64)).collect::<Vec<_>>();
38+
Arc::new(Int64Array::from(values))
39+
}
40+
}
41+
42+
fn generate_filter(len: usize) -> Option<BooleanArray> {
43+
let values = (0..len)
44+
.map(|i| {
45+
if i % 7 == 0 {
46+
None
47+
} else if i % 5 == 0 {
48+
Some(false)
49+
} else {
50+
Some(true)
51+
}
52+
})
53+
.collect::<Vec<_>>();
54+
Some(BooleanArray::from(values))
55+
}
56+
57+
fn criterion_benchmark(c: &mut Criterion) {
58+
let len = 500_000;
59+
let group_indices = generate_group_indices(len);
60+
let rows_count = group_indices.len();
61+
let values = generate_values(len, true);
62+
let opt_filter = generate_filter(len);
63+
let mut counts: Vec<i64> = vec![0; rows_count];
64+
accumulate_indices(
65+
&group_indices,
66+
values.logical_nulls().as_ref(),
67+
opt_filter.as_ref(),
68+
|group_index| {
69+
counts[group_index] += 1;
70+
},
71+
);
72+
73+
c.bench_function("Handle both nulls and filter", |b| {
74+
b.iter(|| {
75+
accumulate_indices(
76+
&group_indices,
77+
values.logical_nulls().as_ref(),
78+
opt_filter.as_ref(),
79+
|group_index| {
80+
counts[group_index] += 1;
81+
},
82+
);
83+
})
84+
});
85+
86+
c.bench_function("Handle nulls only", |b| {
87+
b.iter(|| {
88+
accumulate_indices(
89+
&group_indices,
90+
values.logical_nulls().as_ref(),
91+
None,
92+
|group_index| {
93+
counts[group_index] += 1;
94+
},
95+
);
96+
})
97+
});
98+
99+
let values = generate_values(len, false);
100+
c.bench_function("Handle filter only", |b| {
101+
b.iter(|| {
102+
accumulate_indices(
103+
&group_indices,
104+
values.logical_nulls().as_ref(),
105+
opt_filter.as_ref(),
106+
|group_index| {
107+
counts[group_index] += 1;
108+
},
109+
);
110+
})
111+
});
112+
}
113+
114+
criterion_group!(benches, criterion_benchmark);
115+
criterion_main!(benches);

‎datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs‎

Lines changed: 68 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -395,19 +395,41 @@ pub fn accumulate_indices<F>(
395395
}
396396
}
397397
(None, Some(filter)) => {
398-
assert_eq!(filter.len(), group_indices.len());
399-
// The performance with a filter could be improved by
400-
// iterating over the filter in chunks, rather than a single
401-
// iterator. TODO file a ticket
402-
let iter = group_indices.iter().zip(filter.iter());
403-
for (&group_index, filter_value) in iter {
404-
if let Some(true) = filter_value {
405-
index_fn(group_index)
406-
}
407-
}
398+
debug_assert_eq!(filter.len(), group_indices.len());
399+
let group_indices_chunks = group_indices.chunks_exact(64);
400+
let bit_chunks = filter.values().bit_chunks();
401+
402+
let group_indices_remainder = group_indices_chunks.remainder();
403+
404+
group_indices_chunks.zip(bit_chunks.iter()).for_each(
405+
|(group_index_chunk, mask)| {
406+
// index_mask has value 1 << i in the loop
407+
let mut index_mask = 1;
408+
group_index_chunk.iter().for_each(|&group_index| {
409+
// valid bit was set, real vale
410+
let is_valid = (mask & index_mask) != 0;
411+
if is_valid {
412+
index_fn(group_index);
413+
}
414+
index_mask <<= 1;
415+
})
416+
},
417+
);
418+
419+
// handle any remaining bits (after the initial 64)
420+
let remainder_bits = bit_chunks.remainder_bits();
421+
group_indices_remainder
422+
.iter()
423+
.enumerate()
424+
.for_each(|(i, &group_index)| {
425+
let is_valid = remainder_bits & (1 << i) != 0;
426+
if is_valid {
427+
index_fn(group_index)
428+
}
429+
});
408430
}
409431
(Some(valids), None) => {
410-
assert_eq!(valids.len(), group_indices.len());
432+
debug_assert_eq!(valids.len(), group_indices.len());
411433
// This is based on (ahem, COPY/PASTA) arrow::compute::aggregate::sum
412434
// iterate over in chunks of 64 bits for more efficient null checking
413435
let group_indices_chunks = group_indices.chunks_exact(64);
@@ -444,20 +466,44 @@ pub fn accumulate_indices<F>(
444466
}
445467

446468
(Some(valids), Some(filter)) => {
447-
assert_eq!(filter.len(), group_indices.len());
448-
assert_eq!(valids.len(), group_indices.len());
449-
// The performance with a filter could likely be improved by
450-
// iterating over the filter in chunks, rather than using
451-
// iterators. TODO file a ticket
452-
filter
469+
debug_assert_eq!(filter.len(), group_indices.len());
470+
debug_assert_eq!(valids.len(), group_indices.len());
471+
472+
let group_indices_chunks = group_indices.chunks_exact(64);
473+
let valid_bit_chunks = valids.inner().bit_chunks();
474+
let filter_bit_chunks = filter.values().bit_chunks();
475+
476+
let group_indices_remainder = group_indices_chunks.remainder();
477+
478+
group_indices_chunks
479+
.zip(valid_bit_chunks.iter())
480+
.zip(filter_bit_chunks.iter())
481+
.for_each(|((group_index_chunk, valid_mask), filter_mask)| {
482+
// index_mask has value 1 << i in the loop
483+
let mut index_mask = 1;
484+
group_index_chunk.iter().for_each(|&group_index| {
485+
// valid bit was set, real vale
486+
let is_valid = (valid_mask & filter_mask & index_mask) != 0;
487+
if is_valid {
488+
index_fn(group_index);
489+
}
490+
index_mask <<= 1;
491+
})
492+
});
493+
494+
// handle any remaining bits (after the initial 64)
495+
let remainder_valid_bits = valid_bit_chunks.remainder_bits();
496+
let remainder_filter_bits = filter_bit_chunks.remainder_bits();
497+
group_indices_remainder
453498
.iter()
454-
.zip(group_indices.iter())
455-
.zip(valids.iter())
456-
.for_each(|((filter_value, &group_index), is_valid)| {
457-
if let (Some(true), true) = (filter_value, is_valid) {
499+
.enumerate()
500+
.for_each(|(i, &group_index)| {
501+
let is_valid =
502+
remainder_valid_bits & remainder_filter_bits & (1 << i) != 0;
503+
if is_valid {
458504
index_fn(group_index)
459505
}
460-
})
506+
});
461507
}
462508
}
463509
}

0 commit comments

Comments
 (0)