Skip to content

Commit 5a33f11

Browse files
korowakorowa
committed
rfc: optional skipping partial aggregation
1 parent f80dde0 commit 5a33f11

9 files changed

Lines changed: 667 additions & 4 deletions

File tree

‎datafusion/common/src/config.rs‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -321,6 +321,15 @@ config_namespace! {
321321

322322
/// Should DataFusion keep the columns used for partition_by in the output RecordBatches
323323
pub keep_partition_by_columns: bool, default = false
324+
325+
/// Aggregation ratio (number of distinct groups / number of input rows)
326+
/// threshold for skipping partial aggregation. If the value is greater
327+
/// then partial aggregation will skip aggregation for further input
328+
pub skip_partial_aggregation_probe_ratio_threshold: f64, default = 0.8
329+
330+
/// Number of input rows partial aggregation partition should process, before
331+
/// aggregation ratio check and trying to switch to skipping aggregation mode
332+
pub skip_partial_aggregation_probe_rows_threshold: usize, default = 100_000
324333
}
325334
}
326335

‎datafusion/expr/src/groups_accumulator.rs‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
//! Vectorized [`GroupsAccumulator`]
1919
2020
use arrow_array::{ArrayRef, BooleanArray};
21-
use datafusion_common::Result;
21+
use datafusion_common::{not_impl_err, Result};
2222

2323
/// Describes how many rows should be emitted during grouping.
2424
#[derive(Debug, Clone, Copy)]
@@ -158,6 +158,24 @@ pub trait GroupsAccumulator: Send {
158158
total_num_groups: usize,
159159
) -> Result<()>;
160160

161+
/// Converts input batch to intermediate aggregate state,
162+
/// without grouping (each input row considered as a separate
163+
/// group).
164+
fn convert_to_state(
165+
&self,
166+
_values: &[ArrayRef],
167+
_opt_filter: Option<&BooleanArray>,
168+
) -> Result<Vec<ArrayRef>> {
169+
not_impl_err!("Input batch conversion to state not implemented")
170+
}
171+
172+
/// Returns `true` is groups accumulator supports input batch
173+
/// to intermediate aggregate state conversion (`convert_to_state`
174+
/// method is implemented).
175+
fn convert_to_state_supported(&self) -> bool {
176+
false
177+
}
178+
161179
/// Amount of memory used to store the state of this accumulator,
162180
/// in bytes. This function is called once per batch, so it should
163181
/// be `O(n)` to compute, not `O(num_groups)`

‎datafusion/functions-aggregate/src/count.rs‎

Lines changed: 44 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ use arrow::{
3333
};
3434

3535
use arrow::{
36-
array::{Array, BooleanArray, Int64Array, PrimitiveArray},
36+
array::{Array, BooleanArray, Int64Array, Int64Builder, PrimitiveArray},
3737
buffer::BooleanBuffer,
3838
};
3939
use datafusion_common::{
@@ -432,6 +432,49 @@ impl GroupsAccumulator for CountGroupsAccumulator {
432432
Ok(vec![Arc::new(counts) as ArrayRef])
433433
}
434434

435+
fn convert_to_state(
436+
&self,
437+
values: &[ArrayRef],
438+
opt_filter: Option<&BooleanArray>,
439+
) -> Result<Vec<ArrayRef>> {
440+
let values = &values[0];
441+
442+
let state_array = match (values.logical_nulls(), opt_filter) {
443+
(Some(nulls), None) => {
444+
let mut builder = Int64Builder::with_capacity(values.len());
445+
nulls
446+
.into_iter()
447+
.for_each(|is_valid| builder.append_value(is_valid as i64));
448+
builder.finish()
449+
}
450+
(Some(nulls), Some(filter)) => {
451+
let mut builder = Int64Builder::with_capacity(values.len());
452+
nulls.into_iter().zip(filter.iter()).for_each(
453+
|(is_valid, filter_value)| {
454+
builder.append_value(
455+
(is_valid && filter_value.is_some_and(|val| val)) as i64,
456+
)
457+
},
458+
);
459+
builder.finish()
460+
}
461+
(None, Some(filter)) => {
462+
let mut builder = Int64Builder::with_capacity(values.len());
463+
filter.into_iter().for_each(|filter_value| {
464+
builder.append_value(filter_value.is_some_and(|val| val) as i64)
465+
});
466+
builder.finish()
467+
}
468+
(None, None) => Int64Array::from_value(1, values.len()),
469+
};
470+
471+
Ok(vec![Arc::new(state_array)])
472+
}
473+
474+
fn convert_to_state_supported(&self) -> bool {
475+
true
476+
}
477+
435478
fn size(&self) -> usize {
436479
self.counts.capacity() * std::mem::size_of::<usize>()
437480
}

‎datafusion/physical-expr-common/src/aggregate/groups_accumulator/prim_op.rs‎

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818
use std::sync::Arc;
1919

20-
use arrow::array::{ArrayRef, AsArray, BooleanArray, PrimitiveArray};
20+
use arrow::array::{ArrayRef, AsArray, BooleanArray, PrimitiveArray, PrimitiveBuilder};
2121
use arrow::datatypes::ArrowPrimitiveType;
2222
use arrow::datatypes::DataType;
2323
use datafusion_common::Result;
@@ -134,6 +134,46 @@ where
134134
self.update_batch(values, group_indices, opt_filter, total_num_groups)
135135
}
136136

137+
fn convert_to_state(
138+
&self,
139+
values: &[ArrayRef],
140+
opt_filter: Option<&BooleanArray>,
141+
) -> Result<Vec<ArrayRef>> {
142+
let values = values[0].as_primitive::<T>();
143+
let mut state = PrimitiveBuilder::<T>::with_capacity(values.len())
144+
.with_data_type(self.data_type.clone());
145+
146+
match opt_filter {
147+
Some(filter) => {
148+
values
149+
.iter()
150+
.zip(filter.iter())
151+
.for_each(|(val, filter_val)| match (val, filter_val) {
152+
(Some(val), Some(true)) => {
153+
let mut state_val = self.starting_value;
154+
(self.prim_fn)(&mut state_val, val);
155+
state.append_value(state_val);
156+
}
157+
(_, _) => state.append_null(),
158+
})
159+
}
160+
None => values.iter().for_each(|val| match val {
161+
Some(val) => {
162+
let mut state_val = self.starting_value;
163+
(self.prim_fn)(&mut state_val, val);
164+
state.append_value(state_val);
165+
}
166+
None => state.append_null(),
167+
}),
168+
};
169+
170+
Ok(vec![Arc::new(state.finish())])
171+
}
172+
173+
fn convert_to_state_supported(&self) -> bool {
174+
true
175+
}
176+
137177
fn size(&self) -> usize {
138178
self.values.capacity() * std::mem::size_of::<T::Native>() + self.null_state.size()
139179
}

‎datafusion/physical-plan/src/aggregates/mod.rs‎

Lines changed: 181 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2442,4 +2442,185 @@ mod tests {
24422442

24432443
Ok(())
24442444
}
2445+
2446+
#[tokio::test]
2447+
async fn test_skip_aggregation_after_first_batch() -> Result<()> {
2448+
let schema = Arc::new(Schema::new(vec![
2449+
Field::new("key", DataType::Int32, true),
2450+
Field::new("val", DataType::Int32, true),
2451+
]));
2452+
2453+
let group_by =
2454+
PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]);
2455+
2456+
let aggr_expr: Vec<Arc<dyn AggregateExpr>> = vec![create_aggregate_expr(
2457+
&count_udaf(),
2458+
&[col("val", &schema)?],
2459+
&[datafusion_expr::col("val")],
2460+
&[],
2461+
&[],
2462+
&schema,
2463+
"COUNT(val)",
2464+
false,
2465+
false,
2466+
false,
2467+
)?];
2468+
2469+
let input_data = vec![
2470+
RecordBatch::try_new(
2471+
Arc::clone(&schema),
2472+
vec![
2473+
Arc::new(Int32Array::from(vec![1, 2, 3])),
2474+
Arc::new(Int32Array::from(vec![0, 0, 0])),
2475+
],
2476+
)
2477+
.unwrap(),
2478+
RecordBatch::try_new(
2479+
Arc::clone(&schema),
2480+
vec![
2481+
Arc::new(Int32Array::from(vec![2, 3, 4])),
2482+
Arc::new(Int32Array::from(vec![0, 0, 0])),
2483+
],
2484+
)
2485+
.unwrap(),
2486+
];
2487+
2488+
let input = Arc::new(MemoryExec::try_new(
2489+
&[input_data],
2490+
Arc::clone(&schema),
2491+
None,
2492+
)?);
2493+
let aggregate_exec = Arc::new(AggregateExec::try_new(
2494+
AggregateMode::Partial,
2495+
group_by,
2496+
aggr_expr,
2497+
vec![None],
2498+
Arc::clone(&input) as Arc<dyn ExecutionPlan>,
2499+
schema,
2500+
)?);
2501+
2502+
let mut session_config = SessionConfig::default();
2503+
session_config = session_config.set(
2504+
"datafusion.execution.skip_partial_aggregation_probe_rows_threshold",
2505+
ScalarValue::Int64(Some(2)),
2506+
);
2507+
session_config = session_config.set(
2508+
"datafusion.execution.skip_partial_aggregation_probe_ratio_threshold",
2509+
ScalarValue::Float64(Some(0.1)),
2510+
);
2511+
2512+
let ctx = TaskContext::default().with_session_config(session_config);
2513+
let output = collect(aggregate_exec.execute(0, Arc::new(ctx))?).await?;
2514+
2515+
let expected = [
2516+
"+-----+-------------------+",
2517+
"| key | COUNT(val)[count] |",
2518+
"+-----+-------------------+",
2519+
"| 1 | 1 |",
2520+
"| 2 | 1 |",
2521+
"| 3 | 1 |",
2522+
"| 2 | 1 |",
2523+
"| 3 | 1 |",
2524+
"| 4 | 1 |",
2525+
"+-----+-------------------+",
2526+
];
2527+
assert_batches_eq!(expected, &output);
2528+
2529+
Ok(())
2530+
}
2531+
2532+
#[tokio::test]
2533+
async fn test_skip_aggregation_after_threshold() -> Result<()> {
2534+
let schema = Arc::new(Schema::new(vec![
2535+
Field::new("key", DataType::Int32, true),
2536+
Field::new("val", DataType::Int32, true),
2537+
]));
2538+
2539+
let group_by =
2540+
PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]);
2541+
2542+
let aggr_expr: Vec<Arc<dyn AggregateExpr>> = vec![create_aggregate_expr(
2543+
&count_udaf(),
2544+
&[col("val", &schema)?],
2545+
&[datafusion_expr::col("val")],
2546+
&[],
2547+
&[],
2548+
&schema,
2549+
"COUNT(val)",
2550+
false,
2551+
false,
2552+
false,
2553+
)?];
2554+
2555+
let input_data = vec![
2556+
RecordBatch::try_new(
2557+
Arc::clone(&schema),
2558+
vec![
2559+
Arc::new(Int32Array::from(vec![1, 2, 3])),
2560+
Arc::new(Int32Array::from(vec![0, 0, 0])),
2561+
],
2562+
)
2563+
.unwrap(),
2564+
RecordBatch::try_new(
2565+
Arc::clone(&schema),
2566+
vec![
2567+
Arc::new(Int32Array::from(vec![2, 3, 4])),
2568+
Arc::new(Int32Array::from(vec![0, 0, 0])),
2569+
],
2570+
)
2571+
.unwrap(),
2572+
RecordBatch::try_new(
2573+
Arc::clone(&schema),
2574+
vec![
2575+
Arc::new(Int32Array::from(vec![2, 3, 4])),
2576+
Arc::new(Int32Array::from(vec![0, 0, 0])),
2577+
],
2578+
)
2579+
.unwrap(),
2580+
];
2581+
2582+
let input = Arc::new(MemoryExec::try_new(
2583+
&[input_data],
2584+
Arc::clone(&schema),
2585+
None,
2586+
)?);
2587+
let aggregate_exec = Arc::new(AggregateExec::try_new(
2588+
AggregateMode::Partial,
2589+
group_by,
2590+
aggr_expr,
2591+
vec![None],
2592+
Arc::clone(&input) as Arc<dyn ExecutionPlan>,
2593+
schema,
2594+
)?);
2595+
2596+
let mut session_config = SessionConfig::default();
2597+
session_config = session_config.set(
2598+
"datafusion.execution.skip_partial_aggregation_probe_rows_threshold",
2599+
ScalarValue::Int64(Some(5)),
2600+
);
2601+
session_config = session_config.set(
2602+
"datafusion.execution.skip_partial_aggregation_probe_ratio_threshold",
2603+
ScalarValue::Float64(Some(0.1)),
2604+
);
2605+
2606+
let ctx = TaskContext::default().with_session_config(session_config);
2607+
let output = collect(aggregate_exec.execute(0, Arc::new(ctx))?).await?;
2608+
2609+
let expected = [
2610+
"+-----+-------------------+",
2611+
"| key | COUNT(val)[count] |",
2612+
"+-----+-------------------+",
2613+
"| 1 | 1 |",
2614+
"| 2 | 2 |",
2615+
"| 3 | 2 |",
2616+
"| 4 | 1 |",
2617+
"| 2 | 1 |",
2618+
"| 3 | 1 |",
2619+
"| 4 | 1 |",
2620+
"+-----+-------------------+",
2621+
];
2622+
assert_batches_eq!(expected, &output);
2623+
2624+
Ok(())
2625+
}
24452626
}

0 commit comments

Comments
 (0)