Skip to content

feat: reuse partial aggregate hashes across repartition #23446

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
Rachelint wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

feat: reuse partial aggregate hashes across repartition #23446

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
43 changes: 31 additions & 12 deletions datafusion/physical-expr-common/src/binary_map.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -298,6 +298,32 @@ where
) where
MP: FnMut(Option<&[u8]>) -> V,
OP: FnMut(V),
{
let mut batch_hashes = std::mem::take(&mut self.hashes_buffer);
batch_hashes.clear();
batch_hashes.resize(values.len(), 0);
create_hashes([values], &self.random_state, &mut batch_hashes)
// hash is supported for all types and create_hashes only
// returns errors for unsupported types
.unwrap();
self.insert_if_new_with_hashes(
values,
&batch_hashes,
make_payload_fn,
observe_payload_fn,
);
self.hashes_buffer = batch_hashes;
}

pub fn insert_if_new_with_hashes<MP, OP>(
&mut self,
values: &ArrayRef,
hashes: &[u64],
make_payload_fn: MP,
observe_payload_fn: OP,
) where
MP: FnMut(Option<&[u8]>) -> V,
OP: FnMut(V),
{
// Sanity array type
match self.output_type {
Expand All @@ -308,6 +334,7 @@ where
));
self.insert_if_new_inner::<MP, OP, GenericBinaryType<O>>(
values,
hashes,
make_payload_fn,
observe_payload_fn,
)
Expand All @@ -319,6 +346,7 @@ where
));
self.insert_if_new_inner::<MP, OP, GenericStringType<O>>(
values,
hashes,
make_payload_fn,
observe_payload_fn,
)
Expand All @@ -338,29 +366,20 @@ where
fn insert_if_new_inner<MP, OP, B>(
&mut self,
values: &ArrayRef,
hashes: &[u64],
mut make_payload_fn: MP,
mut observe_payload_fn: OP,
) where
MP: FnMut(Option<&[u8]>) -> V,
OP: FnMut(V),
B: ByteArrayType,
{
// step 1: compute hashes
let batch_hashes = &mut self.hashes_buffer;
batch_hashes.clear();
batch_hashes.resize(values.len(), 0);
create_hashes([values], &self.random_state, batch_hashes)
// hash is supported for all types and create_hashes only
// returns errors for unsupported types
.unwrap();

// step 2: insert each value into the set, if not already present
let values = values.as_bytes::<B>();

// Ensure lengths are equivalent
assert_eq!(values.len(), batch_hashes.len());
assert_eq!(values.len(), hashes.len());

for (value, &hash) in values.iter().zip(batch_hashes.iter()) {
for (value, &hash) in values.iter().zip(hashes.iter()) {
// handle null value
let Some(value) = value else {
let payload = if let Some(&(payload, _offset)) = self.null.as_ref() {
Expand Down
43 changes: 31 additions & 12 deletions datafusion/physical-expr-common/src/binary_view_map.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -211,13 +211,40 @@ where
) where
MP: FnMut(Option<&[u8]>) -> V,
OP: FnMut(V),
{
let mut batch_hashes = std::mem::take(&mut self.hashes_buffer);
batch_hashes.clear();
batch_hashes.resize(values.len(), 0);
create_hashes([values], &self.random_state, &mut batch_hashes)
// hash is supported for all types and create_hashes only
// returns errors for unsupported types
.unwrap();
self.insert_if_new_with_hashes(
values,
&batch_hashes,
make_payload_fn,
observe_payload_fn,
);
self.hashes_buffer = batch_hashes;
}

pub fn insert_if_new_with_hashes<MP, OP>(
&mut self,
values: &ArrayRef,
hashes: &[u64],
make_payload_fn: MP,
observe_payload_fn: OP,
) where
MP: FnMut(Option<&[u8]>) -> V,
OP: FnMut(V),
{
// Sanity check array type
match self.output_type {
OutputType::BinaryView => {
assert!(matches!(values.data_type(), DataType::BinaryView));
self.insert_if_new_inner::<MP, OP, BinaryViewType>(
values,
hashes,
make_payload_fn,
observe_payload_fn,
)
Expand All @@ -226,6 +253,7 @@ where
assert!(matches!(values.data_type(), DataType::Utf8View));
self.insert_if_new_inner::<MP, OP, StringViewType>(
values,
hashes,
make_payload_fn,
observe_payload_fn,
)
Expand All @@ -245,34 +273,25 @@ where
fn insert_if_new_inner<MP, OP, B>(
&mut self,
values: &ArrayRef,
hashes: &[u64],
mut make_payload_fn: MP,
mut observe_payload_fn: OP,
) where
MP: FnMut(Option<&[u8]>) -> V,
OP: FnMut(V),
B: ByteViewType,
{
// step 1: compute hashes
let batch_hashes = &mut self.hashes_buffer;
batch_hashes.clear();
batch_hashes.resize(values.len(), 0);
create_hashes([values], &self.random_state, batch_hashes)
// hash is supported for all types and create_hashes only
// returns errors for unsupported types
.unwrap();

// step 2: insert each value into the set, if not already present
let values = values.as_byte_view::<B>();

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

// Ensure lengths are equivalent
assert_eq!(values.len(), self.hashes_buffer.len());
assert_eq!(values.len(), hashes.len());

for i in 0..values.len() {
let view_u128 = input_views[i];
let hash = self.hashes_buffer[i];
let hash = hashes[i];

// handle null value via validity bitmap check
if values.is_null(i) {
Expand Down
23 changes: 19 additions & 4 deletions datafusion/physical-plan/benches/dictionary_group_values.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ use arrow::datatypes::{DataType, Field, Int32Type, Schema, SchemaRef};
use criterion::{
BatchSize, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main,
};
use datafusion_common::hash_utils::{RandomState, create_hashes};
use datafusion_expr::EmitTo;
use datafusion_physical_plan::aggregates::group_values::new_group_values;
use datafusion_physical_plan::aggregates::order::GroupOrdering;
Expand Down Expand Up @@ -109,10 +110,17 @@ fn bench_intern_emit(c: &mut Criterion) {
new_group_values(schema.clone(), &GroupOrdering::None)
.unwrap(),
Vec::<usize>::with_capacity(size),
Vec::<u64>::new(),
)
},
|(gv, groups)| {
gv.intern(std::slice::from_ref(&array), groups).unwrap();
|(gv, groups, hashes)| {
let seed = RandomState::with_seed(15395726432021054657);
hashes.clear();
hashes.resize(array.len(), 0);
create_hashes(std::slice::from_ref(&array), &seed, hashes)
.unwrap();
gv.intern(std::slice::from_ref(&array), groups, hashes)
.unwrap();
black_box(&*groups);
black_box(gv.emit(EmitTo::All).unwrap());
},
Expand Down Expand Up @@ -154,11 +162,18 @@ fn bench_repeated_intern_emit(c: &mut Criterion) {
new_group_values(schema.clone(), &GroupOrdering::None)
.unwrap(),
Vec::<usize>::with_capacity(size),
Vec::<u64>::new(),
)
},
|(gv, groups)| {
|(gv, groups, hashes)| {
for arr in &batches {
gv.intern(std::slice::from_ref(arr), groups).unwrap();
let seed = RandomState::with_seed(15395726432021054657);
hashes.clear();
hashes.resize(arr.len(), 0);
create_hashes(std::slice::from_ref(arr), &seed, hashes)
.unwrap();
gv.intern(std::slice::from_ref(arr), groups, hashes)
.unwrap();
black_box(&*groups);
}
black_box(gv.emit(EmitTo::All).unwrap());
Expand Down
8 changes: 7 additions & 1 deletion datafusion/physical-plan/benches/multi_group_by.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
use arrow::array::{ArrayRef, Int32Array};
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
use datafusion_common::hash_utils::{RandomState, create_hashes};
use datafusion_physical_plan::aggregates::group_values::GroupValues;
use datafusion_physical_plan::aggregates::group_values::GroupValuesRows;
use datafusion_physical_plan::aggregates::group_values::multi_group_by::GroupValuesColumn;
Expand Down Expand Up @@ -95,9 +96,14 @@ fn bench_intern(
batches: &[Vec<ArrayRef>],
groups: &mut Vec<usize>,
) {
let seed = RandomState::with_seed(15395726432021054657);
let mut hashes = Vec::new();
for batch in batches {
groups.clear();
gv.intern(batch, groups).unwrap();
hashes.clear();
hashes.resize(batch[0].len(), 0);
create_hashes(batch, &seed, &mut hashes).unwrap();
gv.intern(batch, groups, &hashes).unwrap();
}
black_box(&*groups);
}
Expand Down
75 changes: 73 additions & 2 deletions datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,9 +18,10 @@
use std::marker::PhantomData;
use std::sync::Arc;

use arrow::array::{ArrayRef, AsArray, new_null_array};
use arrow::array::{ArrayRef, AsArray, UInt64Array, new_null_array};
use arrow::datatypes::SchemaRef;
use arrow::record_batch::RecordBatch;
use datafusion_common::hash_utils::create_hashes;
use datafusion_common::{Result, internal_err};
use datafusion_execution::memory_pool::proxy::VecAllocExt;
use datafusion_expr::{EmitTo, GroupsAccumulator};
Expand All @@ -31,7 +32,8 @@ use crate::aggregates::group_values::{GroupByMetrics, GroupValues, new_group_val
use crate::aggregates::grouped_hash_stream::create_group_accumulator;
use crate::aggregates::order::GroupOrdering;
use crate::aggregates::{
AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by,
AGGREGATION_HASH_SEED, AggregateExec, PhysicalGroupBy, aggregate_expressions,
evaluate_group_by, schema_with_group_hash,
};

/// Marker for raw rows -> partial state aggregation.
Expand Down Expand Up @@ -136,6 +138,8 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
group_by: Arc::clone(&agg.group_by),
group_values,
batch_group_indices: Default::default(),
group_hashes: Default::default(),
batch_hashes: Default::default(),
accumulators,
}),
_mode: PhantomData,
Expand Down Expand Up @@ -183,6 +187,8 @@ impl<AggrMode> AggregateHashTable<AggrMode> {

acc + state.group_values.size()
+ state.batch_group_indices.allocated_size()
+ state.group_hashes.allocated_size()
+ state.batch_hashes.allocated_size()
}
AggregateHashTableState::OutputtingMaterialized(output) => {
output.memory_size()
Expand Down Expand Up @@ -212,6 +218,7 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
};

state.batch_group_indices = Vec::new();
state.batch_hashes = Vec::new();
self.state = AggregateHashTableState::Outputting(state);
}
}
Expand Down Expand Up @@ -287,6 +294,12 @@ pub(super) struct AggregateHashTableBuffer {
/// accumulator to update that group's aggregate state.
pub(super) batch_group_indices: Vec<usize>,

/// Hash value for each stored group, indexed by group index.
pub(super) group_hashes: Vec<u64>,

/// Scratch hash vector for the current input batch.
pub(super) batch_hashes: Vec<u64>,

/// One item per aggregate expression.
///
/// Example: `COUNT(x), SUM(y)` creates two items. Each item owns the input
Expand Down Expand Up @@ -343,6 +356,64 @@ impl MaterializedAggregateOutput {
}
}

pub(super) fn try_new_internal_batch(
output_schema: SchemaRef,
mut columns: Vec<ArrayRef>,
has_group_hash: bool,
) -> Result<RecordBatch> {
if has_group_hash {
Ok(RecordBatch::try_new(
schema_with_group_hash(&output_schema),
columns,
)?)
} else {
Ok(RecordBatch::try_new(
output_schema,
std::mem::take(&mut columns),
)?)
}
}

impl AggregateHashTableBuffer {
pub(super) fn create_batch_hashes(
&mut self,
group_values: &[ArrayRef],
) -> Result<&[u64]> {
self.batch_hashes.clear();
let num_rows = group_values.first().map(|array| array.len()).unwrap_or(0);
self.batch_hashes.resize(num_rows, 0);
create_hashes(group_values, &AGGREGATION_HASH_SEED, &mut self.batch_hashes)?;
Ok(&self.batch_hashes)
}

pub(super) fn reuse_batch_hashes(&mut self, input_hashes: &UInt64Array) -> &[u64] {
self.batch_hashes.clear();
self.batch_hashes.extend_from_slice(input_hashes.values());
&self.batch_hashes
}

pub(super) fn record_new_group_hashes(
&mut self,
starting_num_groups: usize,
total_num_groups: usize,
) {
self.group_hashes.resize(total_num_groups, 0);
for (row, group_index) in self.batch_group_indices.iter().enumerate() {
if *group_index >= starting_num_groups {
self.group_hashes[*group_index] = self.batch_hashes[row];
}
}
}

pub(super) fn emit_hashes(&mut self, emit_to: EmitTo) -> ArrayRef {
let hashes = match emit_to {
EmitTo::All => std::mem::take(&mut self.group_hashes),
EmitTo::First(n) => self.group_hashes.drain(..n).collect(),
};
Arc::new(UInt64Array::from(hashes))
}
}

impl HashAggregateAccumulator {
pub(super) fn new(
aggregate_expr: Arc<AggregateFunctionExpr>,
Expand Down
Loading
Loading