Skip to content

Commit 0ecb6ec

Browse files
RachelintRachelint
committed
feat: reuse partial aggregate hashes across repartition
1 parent f34a676 commit 0ecb6ec

20 files changed

Lines changed: 483 additions & 156 deletions

File tree

‎datafusion/physical-expr-common/src/binary_map.rs‎

Lines changed: 31 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -298,6 +298,32 @@ where
298298
) where
299299
MP: FnMut(Option<&[u8]>) -> V,
300300
OP: FnMut(V),
301+
{
302+
let mut batch_hashes = std::mem::take(&mut self.hashes_buffer);
303+
batch_hashes.clear();
304+
batch_hashes.resize(values.len(), 0);
305+
create_hashes([values], &self.random_state, &mut batch_hashes)
306+
// hash is supported for all types and create_hashes only
307+
// returns errors for unsupported types
308+
.unwrap();
309+
self.insert_if_new_with_hashes(
310+
values,
311+
&batch_hashes,
312+
make_payload_fn,
313+
observe_payload_fn,
314+
);
315+
self.hashes_buffer = batch_hashes;
316+
}
317+
318+
pub fn insert_if_new_with_hashes<MP, OP>(
319+
&mut self,
320+
values: &ArrayRef,
321+
hashes: &[u64],
322+
make_payload_fn: MP,
323+
observe_payload_fn: OP,
324+
) where
325+
MP: FnMut(Option<&[u8]>) -> V,
326+
OP: FnMut(V),
301327
{
302328
// Sanity array type
303329
match self.output_type {
@@ -308,6 +334,7 @@ where
308334
));
309335
self.insert_if_new_inner::<MP, OP, GenericBinaryType<O>>(
310336
values,
337+
hashes,
311338
make_payload_fn,
312339
observe_payload_fn,
313340
)
@@ -319,6 +346,7 @@ where
319346
));
320347
self.insert_if_new_inner::<MP, OP, GenericStringType<O>>(
321348
values,
349+
hashes,
322350
make_payload_fn,
323351
observe_payload_fn,
324352
)
@@ -338,29 +366,20 @@ where
338366
fn insert_if_new_inner<MP, OP, B>(
339367
&mut self,
340368
values: &ArrayRef,
369+
hashes: &[u64],
341370
mut make_payload_fn: MP,
342371
mut observe_payload_fn: OP,
343372
) where
344373
MP: FnMut(Option<&[u8]>) -> V,
345374
OP: FnMut(V),
346375
B: ByteArrayType,
347376
{
348-
// step 1: compute hashes
349-
let batch_hashes = &mut self.hashes_buffer;
350-
batch_hashes.clear();
351-
batch_hashes.resize(values.len(), 0);
352-
create_hashes([values], &self.random_state, batch_hashes)
353-
// hash is supported for all types and create_hashes only
354-
// returns errors for unsupported types
355-
.unwrap();
356-
357-
// step 2: insert each value into the set, if not already present
358377
let values = values.as_bytes::<B>();
359378

360379
// Ensure lengths are equivalent
361-
assert_eq!(values.len(), batch_hashes.len());
380+
assert_eq!(values.len(), hashes.len());
362381

363-
for (value, &hash) in values.iter().zip(batch_hashes.iter()) {
382+
for (value, &hash) in values.iter().zip(hashes.iter()) {
364383
// handle null value
365384
let Some(value) = value else {
366385
let payload = if let Some(&(payload, _offset)) = self.null.as_ref() {

‎datafusion/physical-expr-common/src/binary_view_map.rs‎

Lines changed: 31 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -211,13 +211,40 @@ where
211211
) where
212212
MP: FnMut(Option<&[u8]>) -> V,
213213
OP: FnMut(V),
214+
{
215+
let mut batch_hashes = std::mem::take(&mut self.hashes_buffer);
216+
batch_hashes.clear();
217+
batch_hashes.resize(values.len(), 0);
218+
create_hashes([values], &self.random_state, &mut batch_hashes)
219+
// hash is supported for all types and create_hashes only
220+
// returns errors for unsupported types
221+
.unwrap();
222+
self.insert_if_new_with_hashes(
223+
values,
224+
&batch_hashes,
225+
make_payload_fn,
226+
observe_payload_fn,
227+
);
228+
self.hashes_buffer = batch_hashes;
229+
}
230+
231+
pub fn insert_if_new_with_hashes<MP, OP>(
232+
&mut self,
233+
values: &ArrayRef,
234+
hashes: &[u64],
235+
make_payload_fn: MP,
236+
observe_payload_fn: OP,
237+
) where
238+
MP: FnMut(Option<&[u8]>) -> V,
239+
OP: FnMut(V),
214240
{
215241
// Sanity check array type
216242
match self.output_type {
217243
OutputType::BinaryView => {
218244
assert!(matches!(values.data_type(), DataType::BinaryView));
219245
self.insert_if_new_inner::<MP, OP, BinaryViewType>(
220246
values,
247+
hashes,
221248
make_payload_fn,
222249
observe_payload_fn,
223250
)
@@ -226,6 +253,7 @@ where
226253
assert!(matches!(values.data_type(), DataType::Utf8View));
227254
self.insert_if_new_inner::<MP, OP, StringViewType>(
228255
values,
256+
hashes,
229257
make_payload_fn,
230258
observe_payload_fn,
231259
)
@@ -245,34 +273,25 @@ where
245273
fn insert_if_new_inner<MP, OP, B>(
246274
&mut self,
247275
values: &ArrayRef,
276+
hashes: &[u64],
248277
mut make_payload_fn: MP,
249278
mut observe_payload_fn: OP,
250279
) where
251280
MP: FnMut(Option<&[u8]>) -> V,
252281
OP: FnMut(V),
253282
B: ByteViewType,
254283
{
255-
// step 1: compute hashes
256-
let batch_hashes = &mut self.hashes_buffer;
257-
batch_hashes.clear();
258-
batch_hashes.resize(values.len(), 0);
259-
create_hashes([values], &self.random_state, batch_hashes)
260-
// hash is supported for all types and create_hashes only
261-
// returns errors for unsupported types
262-
.unwrap();
263-
264-
// step 2: insert each value into the set, if not already present
265284
let values = values.as_byte_view::<B>();
266285

267286
// Get raw views buffer for direct comparison
268287
let input_views = values.views();
269288

270289
// Ensure lengths are equivalent
271-
assert_eq!(values.len(), self.hashes_buffer.len());
290+
assert_eq!(values.len(), hashes.len());
272291

273292
for i in 0..values.len() {
274293
let view_u128 = input_views[i];
275-
let hash = self.hashes_buffer[i];
294+
let hash = hashes[i];
276295

277296
// handle null value via validity bitmap check
278297
if values.is_null(i) {

‎datafusion/physical-plan/benches/dictionary_group_values.rs‎

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ use arrow::datatypes::{DataType, Field, Int32Type, Schema, SchemaRef};
2727
use criterion::{
2828
BatchSize, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main,
2929
};
30+
use datafusion_common::hash_utils::{RandomState, create_hashes};
3031
use datafusion_expr::EmitTo;
3132
use datafusion_physical_plan::aggregates::group_values::new_group_values;
3233
use datafusion_physical_plan::aggregates::order::GroupOrdering;
@@ -109,10 +110,17 @@ fn bench_intern_emit(c: &mut Criterion) {
109110
new_group_values(schema.clone(), &GroupOrdering::None)
110111
.unwrap(),
111112
Vec::<usize>::with_capacity(size),
113+
Vec::<u64>::new(),
112114
)
113115
},
114-
|(gv, groups)| {
115-
gv.intern(std::slice::from_ref(&array), groups).unwrap();
116+
|(gv, groups, hashes)| {
117+
let seed = RandomState::with_seed(15395726432021054657);
118+
hashes.clear();
119+
hashes.resize(array.len(), 0);
120+
create_hashes(std::slice::from_ref(&array), &seed, hashes)
121+
.unwrap();
122+
gv.intern(std::slice::from_ref(&array), groups, hashes)
123+
.unwrap();
116124
black_box(&*groups);
117125
black_box(gv.emit(EmitTo::All).unwrap());
118126
},
@@ -154,11 +162,18 @@ fn bench_repeated_intern_emit(c: &mut Criterion) {
154162
new_group_values(schema.clone(), &GroupOrdering::None)
155163
.unwrap(),
156164
Vec::<usize>::with_capacity(size),
165+
Vec::<u64>::new(),
157166
)
158167
},
159-
|(gv, groups)| {
168+
|(gv, groups, hashes)| {
160169
for arr in &batches {
161-
gv.intern(std::slice::from_ref(arr), groups).unwrap();
170+
let seed = RandomState::with_seed(15395726432021054657);
171+
hashes.clear();
172+
hashes.resize(arr.len(), 0);
173+
create_hashes(std::slice::from_ref(arr), &seed, hashes)
174+
.unwrap();
175+
gv.intern(std::slice::from_ref(arr), groups, hashes)
176+
.unwrap();
162177
black_box(&*groups);
163178
}
164179
black_box(gv.emit(EmitTo::All).unwrap());

‎datafusion/physical-plan/benches/multi_group_by.rs‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
use arrow::array::{ArrayRef, Int32Array};
2929
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
3030
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
31+
use datafusion_common::hash_utils::{RandomState, create_hashes};
3132
use datafusion_physical_plan::aggregates::group_values::GroupValues;
3233
use datafusion_physical_plan::aggregates::group_values::GroupValuesRows;
3334
use datafusion_physical_plan::aggregates::group_values::multi_group_by::GroupValuesColumn;
@@ -95,9 +96,14 @@ fn bench_intern(
9596
batches: &[Vec<ArrayRef>],
9697
groups: &mut Vec<usize>,
9798
) {
99+
let seed = RandomState::with_seed(15395726432021054657);
100+
let mut hashes = Vec::new();
98101
for batch in batches {
99102
groups.clear();
100-
gv.intern(batch, groups).unwrap();
103+
hashes.clear();
104+
hashes.resize(batch[0].len(), 0);
105+
create_hashes(batch, &seed, &mut hashes).unwrap();
106+
gv.intern(batch, groups, &hashes).unwrap();
101107
}
102108
black_box(&*groups);
103109
}

‎datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs‎

Lines changed: 73 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,10 @@
1818
use std::marker::PhantomData;
1919
use std::sync::Arc;
2020

21-
use arrow::array::{ArrayRef, AsArray, new_null_array};
21+
use arrow::array::{ArrayRef, AsArray, UInt64Array, new_null_array};
2222
use arrow::datatypes::SchemaRef;
2323
use arrow::record_batch::RecordBatch;
24+
use datafusion_common::hash_utils::create_hashes;
2425
use datafusion_common::{Result, internal_err};
2526
use datafusion_execution::memory_pool::proxy::VecAllocExt;
2627
use datafusion_expr::{EmitTo, GroupsAccumulator};
@@ -31,7 +32,8 @@ use crate::aggregates::group_values::{GroupByMetrics, GroupValues, new_group_val
3132
use crate::aggregates::grouped_hash_stream::create_group_accumulator;
3233
use crate::aggregates::order::GroupOrdering;
3334
use crate::aggregates::{
34-
AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by,
35+
AGGREGATION_HASH_SEED, AggregateExec, PhysicalGroupBy, aggregate_expressions,
36+
evaluate_group_by, schema_with_group_hash,
3537
};
3638

3739
/// Marker for raw rows -> partial state aggregation.
@@ -136,6 +138,8 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
136138
group_by: Arc::clone(&agg.group_by),
137139
group_values,
138140
batch_group_indices: Default::default(),
141+
group_hashes: Default::default(),
142+
batch_hashes: Default::default(),
139143
accumulators,
140144
}),
141145
_mode: PhantomData,
@@ -183,6 +187,8 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
183187

184188
acc + state.group_values.size()
185189
+ state.batch_group_indices.allocated_size()
190+
+ state.group_hashes.allocated_size()
191+
+ state.batch_hashes.allocated_size()
186192
}
187193
AggregateHashTableState::OutputtingMaterialized(output) => {
188194
output.memory_size()
@@ -212,6 +218,7 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
212218
};
213219

214220
state.batch_group_indices = Vec::new();
221+
state.batch_hashes = Vec::new();
215222
self.state = AggregateHashTableState::Outputting(state);
216223
}
217224
}
@@ -287,6 +294,12 @@ pub(super) struct AggregateHashTableBuffer {
287294
/// accumulator to update that group's aggregate state.
288295
pub(super) batch_group_indices: Vec<usize>,
289296

297+
/// Hash value for each stored group, indexed by group index.
298+
pub(super) group_hashes: Vec<u64>,
299+
300+
/// Scratch hash vector for the current input batch.
301+
pub(super) batch_hashes: Vec<u64>,
302+
290303
/// One item per aggregate expression.
291304
///
292305
/// Example: `COUNT(x), SUM(y)` creates two items. Each item owns the input
@@ -343,6 +356,64 @@ impl MaterializedAggregateOutput {
343356
}
344357
}
345358

359+
pub(super) fn try_new_internal_batch(
360+
output_schema: SchemaRef,
361+
mut columns: Vec<ArrayRef>,
362+
has_group_hash: bool,
363+
) -> Result<RecordBatch> {
364+
if has_group_hash {
365+
Ok(RecordBatch::try_new(
366+
schema_with_group_hash(&output_schema),
367+
columns,
368+
)?)
369+
} else {
370+
Ok(RecordBatch::try_new(
371+
output_schema,
372+
std::mem::take(&mut columns),
373+
)?)
374+
}
375+
}
376+
377+
impl AggregateHashTableBuffer {
378+
pub(super) fn create_batch_hashes(
379+
&mut self,
380+
group_values: &[ArrayRef],
381+
) -> Result<&[u64]> {
382+
self.batch_hashes.clear();
383+
let num_rows = group_values.first().map(|array| array.len()).unwrap_or(0);
384+
self.batch_hashes.resize(num_rows, 0);
385+
create_hashes(group_values, &AGGREGATION_HASH_SEED, &mut self.batch_hashes)?;
386+
Ok(&self.batch_hashes)
387+
}
388+
389+
pub(super) fn reuse_batch_hashes(&mut self, input_hashes: &UInt64Array) -> &[u64] {
390+
self.batch_hashes.clear();
391+
self.batch_hashes.extend_from_slice(input_hashes.values());
392+
&self.batch_hashes
393+
}
394+
395+
pub(super) fn record_new_group_hashes(
396+
&mut self,
397+
starting_num_groups: usize,
398+
total_num_groups: usize,
399+
) {
400+
self.group_hashes.resize(total_num_groups, 0);
401+
for (row, group_index) in self.batch_group_indices.iter().enumerate() {
402+
if *group_index >= starting_num_groups {
403+
self.group_hashes[*group_index] = self.batch_hashes[row];
404+
}
405+
}
406+
}
407+
408+
pub(super) fn emit_hashes(&mut self, emit_to: EmitTo) -> ArrayRef {
409+
let hashes = match emit_to {
410+
EmitTo::All => std::mem::take(&mut self.group_hashes),
411+
EmitTo::First(n) => self.group_hashes.drain(..n).collect(),
412+
};
413+
Arc::new(UInt64Array::from(hashes))
414+
}
415+
}
416+
346417
impl HashAggregateAccumulator {
347418
pub(super) fn new(
348419
aggregate_expr: Arc<AggregateFunctionExpr>,

0 commit comments

Comments
 (0)