alamb commented on code in PR #23324:
URL: https://github.com/apache/datafusion/pull/23324#discussion_r3536183430
##########
datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs:
##########
@@ -171,6 +171,106 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
})
}
+ /// Aggregates one input batch after selecting the mode-specific
accumulator
+ /// operation.
+ ///
+ /// Each aggregation mode chooses a different `aggregate_fn` according to
its
+ /// semantics. For example, partial aggregation takes raw inputs, and
update them
+ /// into stored partial states, so [`GroupsAccumulator::update_batch`] is
used.
+ pub(super) fn aggregate_batch_inner<F>(
Review Comment:
can you please document somewhere what these arguments are? Specifically the
&[usize] and `usize` arguments?
Something like
```rust
/// Function used by [`AggregateHashTable::aggregate_batch_inner`] to update
one
/// accumulator with one evaluated input batch.
///
/// Arguments:
/// * accumulator to update.
/// * accumulator's evaluated arguments and optional filter.
/// * `&[usize]` with one group index per input row, mapping each row to its
interned group
/// * total number of groups currently interned in that buffer, including
newly interned groups
pub(super) type AggregateBatchFn = fn(
&mut AggregateAccumulator,
&EvaluatedAccumulatorArgs,
&[usize],
usize,
) -> Result<()>;
```
So then this would be something like
```rust
pub(super) fn aggregate_batch_inner<F: AggregateBatchFn>(
&mut self,
batch: &RecordBatch,
mut aggregate_fn: F
) -> Result<()>
```
##########
datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs:
##########
@@ -55,91 +49,24 @@ impl AggregateHashTable<PartialReduceMarker> {
pub(in crate::aggregates) fn next_output_batch(
&mut self,
) -> Result<Option<RecordBatch>> {
- let output_schema = Arc::clone(&self.output_schema);
- let batch_size = self.batch_size;
- // Take ownership of the output state. Note
`emit_next_materialized_batch`
- // updates state after it emits a materialized slice.
- match std::mem::replace(&mut self.state,
AggregateHashTableState::Done) {
- AggregateHashTableState::Outputting(state) => {
- if state.group_values.is_empty() {
- return Ok(None);
- }
-
- let output =
- self.materialize_partial_reduce_output(state,
output_schema)?;
- Ok(self.emit_next_materialized_batch(output, batch_size))
- }
- AggregateHashTableState::OutputtingMaterialized(output) => {
- Ok(self.emit_next_materialized_batch(output, batch_size))
- }
- AggregateHashTableState::Done => Ok(None),
- AggregateHashTableState::Building(_) => {
- internal_err!("next_output_batch must be called in the
outputting state")
- }
- }
- }
-
- fn materialize_partial_reduce_output(
- &self,
- mut state: AggregateHashTableBuffer,
- output_schema: SchemaRef,
- ) -> Result<MaterializedAggregateOutput> {
- // `state(EmitTo::All)` consumes accumulator state. Emit all groups
once,
- // then slice the materialized batch on subsequent polls.
- let emit_to_all = EmitTo::All;
- let timer = self.group_by_metrics.emitting_time.timer();
- let mut output = state.group_values.emit(emit_to_all)?;
-
- for acc in state.accumulators.iter_mut() {
- output.extend(acc.state(emit_to_all)?);
- }
- drop(timer);
-
- let batch = RecordBatch::try_new(output_schema, output)?;
- debug_assert!(batch.num_rows() > 0);
- Ok(MaterializedAggregateOutput::new(batch))
- }
-
- fn emit_next_materialized_batch(
- &mut self,
- mut output: MaterializedAggregateOutput,
- batch_size: usize,
- ) -> Option<RecordBatch> {
- let batch = output.next_batch(batch_size);
- if output.is_exhausted() {
- self.state = AggregateHashTableState::Done;
- } else {
- self.state =
AggregateHashTableState::OutputtingMaterialized(output);
- }
- batch
+ self.next_output_batch_inner(|acc, emit_to, output| {
Review Comment:
likewise here I think you can just use the method name directly
```rust
self.next_output_batch_inner(Self::materialize_partial_reduce_output)
```
##########
datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs:
##########
@@ -55,91 +49,24 @@ impl AggregateHashTable<PartialReduceMarker> {
pub(in crate::aggregates) fn next_output_batch(
&mut self,
) -> Result<Option<RecordBatch>> {
- let output_schema = Arc::clone(&self.output_schema);
- let batch_size = self.batch_size;
- // Take ownership of the output state. Note
`emit_next_materialized_batch`
- // updates state after it emits a materialized slice.
- match std::mem::replace(&mut self.state,
AggregateHashTableState::Done) {
- AggregateHashTableState::Outputting(state) => {
- if state.group_values.is_empty() {
- return Ok(None);
- }
-
- let output =
- self.materialize_partial_reduce_output(state,
output_schema)?;
- Ok(self.emit_next_materialized_batch(output, batch_size))
- }
- AggregateHashTableState::OutputtingMaterialized(output) => {
- Ok(self.emit_next_materialized_batch(output, batch_size))
- }
- AggregateHashTableState::Done => Ok(None),
- AggregateHashTableState::Building(_) => {
- internal_err!("next_output_batch must be called in the
outputting state")
- }
- }
- }
-
- fn materialize_partial_reduce_output(
- &self,
- mut state: AggregateHashTableBuffer,
- output_schema: SchemaRef,
- ) -> Result<MaterializedAggregateOutput> {
- // `state(EmitTo::All)` consumes accumulator state. Emit all groups
once,
- // then slice the materialized batch on subsequent polls.
- let emit_to_all = EmitTo::All;
- let timer = self.group_by_metrics.emitting_time.timer();
- let mut output = state.group_values.emit(emit_to_all)?;
-
- for acc in state.accumulators.iter_mut() {
- output.extend(acc.state(emit_to_all)?);
- }
- drop(timer);
-
- let batch = RecordBatch::try_new(output_schema, output)?;
- debug_assert!(batch.num_rows() > 0);
- Ok(MaterializedAggregateOutput::new(batch))
- }
-
- fn emit_next_materialized_batch(
- &mut self,
- mut output: MaterializedAggregateOutput,
- batch_size: usize,
- ) -> Option<RecordBatch> {
- let batch = output.next_batch(batch_size);
- if output.is_exhausted() {
- self.state = AggregateHashTableState::Done;
- } else {
- self.state =
AggregateHashTableState::OutputtingMaterialized(output);
- }
- batch
+ self.next_output_batch_inner(|acc, emit_to, output| {
+ output.extend(acc.state(emit_to)?);
+ Ok(())
+ })
}
+ /// Partial-reduce aggregation consumes partial aggregate states and merges
+ /// them into the table's partial-state accumulators.
pub(in crate::aggregates) fn aggregate_batch(
&mut self,
batch: &RecordBatch,
) -> Result<()> {
- let evaluated_batch = self.evaluate_batch(batch)?;
- let state = self.state.building_mut();
-
- let timer = self.group_by_metrics.aggregation_time.timer();
- for group_values in &evaluated_batch.grouping_set_args {
- state
- .group_values
- .intern(group_values, &mut state.batch_group_indices)?;
- let group_indices = &state.batch_group_indices;
- let total_num_groups = state.group_values.len();
-
- for (acc, values) in state
- .accumulators
- .iter_mut()
- .zip(evaluated_batch.accumulator_args.iter())
- {
- acc.merge_batch(values, group_indices, total_num_groups)?;
- }
- }
- drop(timer);
-
- Ok(())
+ self.aggregate_batch_inner(
+ batch,
+ |acc, values, group_indices, total_num_groups| {
Review Comment:
here too
```rust
self.aggregate_batch_inner(batch, AggregateAccumulator::merge_batch)
```
##########
datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs:
##########
@@ -61,90 +55,24 @@ impl AggregateHashTable<FinalMarker> {
pub(in crate::aggregates) fn next_output_batch(
&mut self,
) -> Result<Option<RecordBatch>> {
- let output_schema = Arc::clone(&self.output_schema);
- let batch_size = self.batch_size;
- // Take ownership of the output state. `emit_next_materialized_batch`
- // restores `self.state` to `OutputtingMaterialized` or `Done`.
- match std::mem::replace(&mut self.state,
AggregateHashTableState::Done) {
- AggregateHashTableState::Outputting(state) => {
- if state.group_values.is_empty() {
- return Ok(None);
- }
-
- let output = self.materialize_final_output(state,
output_schema)?;
- Ok(self.emit_next_materialized_batch(output, batch_size))
- }
- AggregateHashTableState::OutputtingMaterialized(output) => {
- Ok(self.emit_next_materialized_batch(output, batch_size))
- }
- AggregateHashTableState::Done => Ok(None),
- AggregateHashTableState::Building(_) => {
- internal_err!("next_output_batch must be called in the
outputting state")
- }
- }
- }
-
- fn materialize_final_output(
- &self,
- mut state: AggregateHashTableBuffer,
- output_schema: SchemaRef,
- ) -> Result<MaterializedAggregateOutput> {
- // Final aggregate evaluation consumes accumulator state. Evaluate all
- // groups once, then slice the materialized batch on subsequent polls.
- let emit_to = EmitTo::All;
- let timer = self.group_by_metrics.emitting_time.timer();
- let mut output = state.group_values.emit(emit_to)?;
-
- for acc in state.accumulators.iter_mut() {
+ self.next_output_batch_inner(|acc, emit_to, output| {
output.push(acc.evaluate(emit_to)?);
- }
- drop(timer);
-
- let batch = RecordBatch::try_new(output_schema, output)?;
- debug_assert!(batch.num_rows() > 0);
- Ok(MaterializedAggregateOutput::new(batch))
- }
-
- fn emit_next_materialized_batch(
- &mut self,
- mut output: MaterializedAggregateOutput,
- batch_size: usize,
- ) -> Option<RecordBatch> {
- let batch = output.next_batch(batch_size);
- if output.is_exhausted() {
- self.state = AggregateHashTableState::Done;
- } else {
- self.state =
AggregateHashTableState::OutputtingMaterialized(output);
- }
- batch
+ Ok(())
+ })
}
+ /// Final aggregation consumes partial aggregate states and merges them
into
+ /// the table's partial-state accumulators.
pub(in crate::aggregates) fn aggregate_batch(
&mut self,
batch: &RecordBatch,
) -> Result<()> {
- let evaluated_batch = self.evaluate_batch(batch)?;
- let state = self.state.building_mut();
-
- let timer = self.group_by_metrics.aggregation_time.timer();
- for group_values in &evaluated_batch.grouping_set_args {
- state
- .group_values
- .intern(group_values, &mut state.batch_group_indices)?;
- let group_indices = &state.batch_group_indices;
- let total_num_groups = state.group_values.len();
-
- for (acc, values) in state
- .accumulators
- .iter_mut()
- .zip(evaluated_batch.accumulator_args.iter())
- {
- acc.merge_batch(values, group_indices, total_num_groups)?;
- }
- }
- drop(timer);
-
- Ok(())
+ self.aggregate_batch_inner(
+ batch,
+ |acc, values, group_indices, total_num_groups| {
+ acc.merge_batch(values, group_indices, total_num_groups)
+ },
Review Comment:
I think you can use the method name directly here
```rust
pub(in crate::aggregates) fn aggregate_batch(
&mut self,
batch: &RecordBatch,
) -> Result<()> {
self.aggregate_batch_inner(batch, AggregateAccumulator::merge_batch)
}
```
##########
datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs:
##########
@@ -67,60 +65,10 @@ impl AggregateHashTable<PartialMarker> {
pub(in crate::aggregates) fn next_output_batch(
&mut self,
) -> Result<Option<RecordBatch>> {
- let output_schema = Arc::clone(&self.output_schema);
- let batch_size = self.batch_size;
- // Take ownership of the output state. `emit_next_materialized_batch`
- // restores `self.state` to `OutputtingMaterialized` or `Done`.
- match std::mem::replace(&mut self.state,
AggregateHashTableState::Done) {
- AggregateHashTableState::Outputting(state) => {
- if state.group_values.is_empty() {
- return Ok(None);
- }
-
- let output = self.materialize_partial_output(state,
output_schema)?;
- Ok(self.emit_next_materialized_batch(output, batch_size))
- }
- AggregateHashTableState::OutputtingMaterialized(output) => {
- Ok(self.emit_next_materialized_batch(output, batch_size))
- }
- AggregateHashTableState::Done => Ok(None),
- AggregateHashTableState::Building(_) => {
- internal_err!("next_output_batch must be called in the
outputting state")
- }
- }
- }
-
- fn materialize_partial_output(
- &self,
- mut state: AggregateHashTableBuffer,
- output_schema: SchemaRef,
- ) -> Result<MaterializedAggregateOutput> {
- let emit_to = EmitTo::All;
- let timer = self.group_by_metrics.emitting_time.timer();
- let mut output = state.group_values.emit(emit_to)?;
-
- for acc in state.accumulators.iter_mut() {
+ self.next_output_batch_inner(|acc, emit_to, output| {
Review Comment:
```rust
self.next_output_batch_inner(Self::materialize_partial_output)
```
##########
datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs:
##########
@@ -171,6 +171,106 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
})
}
+ /// Aggregates one input batch after selecting the mode-specific
accumulator
+ /// operation.
+ ///
+ /// Each aggregation mode chooses a different `aggregate_fn` according to
its
+ /// semantics. For example, partial aggregation takes raw inputs, and
update them
+ /// into stored partial states, so [`GroupsAccumulator::update_batch`] is
used.
+ pub(super) fn aggregate_batch_inner<F>(
+ &mut self,
+ batch: &RecordBatch,
+ mut aggregate_fn: F,
+ ) -> Result<()>
+ where
+ F: FnMut(
+ &mut HashAggregateAccumulator,
+ &EvaluatedAccumulatorArgs,
+ &[usize],
+ usize,
+ ) -> Result<()>,
+ {
+ let evaluated_batch = self.evaluate_batch(batch)?;
+ let state = self.state.building_mut();
+
+ let _timer = self.group_by_metrics.aggregation_time.timer();
+ for group_values in &evaluated_batch.grouping_set_args {
+ state
+ .group_values
+ .intern(group_values, &mut state.batch_group_indices)?;
+ let group_indices = &state.batch_group_indices;
+ let total_num_groups = state.group_values.len();
+
+ for (acc, values) in state
+ .accumulators
+ .iter_mut()
+ .zip(evaluated_batch.accumulator_args.iter())
+ {
+ aggregate_fn(acc, values, group_indices, total_num_groups)?;
+ }
+ }
+
+ Ok(())
+ }
+
+ /// Materializes the full output once, then returns it downstream
incrementally
+ /// by slicing it into `batch_size` chunks.
+ ///
+ /// Each aggregation mode chooses a different `materialize_accumulator_fn`
+ /// according to its semantics. For example, partial aggregation emits
+ /// partial states to feed the final stage, so it us
[`GroupsAccumulator::state`].
+ ///
+ /// This is a temporary solution until blocked state management is
implemented:
+ /// Issue: <https://github.com/apache/datafusion/issues/7065>
+ pub(super) fn next_output_batch_inner<F>(
Review Comment:
Similarly, here I think a new typedef would help
```rust
/// Function used by [`AggregateHashTable::next_output_batch_inner`] to
/// materialize all currently buffered output for one aggregation mode.
///
/// Arguments
/// * hash table that is producing output
/// *The outputting buffer moved out of the table state
/// * output schema for the materialized [`RecordBatch`].
pub(super) type NextOutputBatchFn<AggrMode> = fn(
&AggregateHashTable<AggrMode>,
AggregateHashTableBuffer,
SchemaRef,
) -> Result<MaterializedAggregateOutput>;
```
##########
datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs:
##########
@@ -161,31 +109,18 @@ impl AggregateHashTable<PartialMarker> {
})
}
+ /// Partial aggregation consumes raw input rows and updates the table's
+ /// partial-state accumulators.
pub(in crate::aggregates) fn aggregate_batch(
&mut self,
batch: &RecordBatch,
) -> Result<()> {
- let evaluated_batch = self.evaluate_batch(batch)?;
- let state = self.state.building_mut();
-
- let _timer = self.group_by_metrics.aggregation_time.timer();
- for group_values in &evaluated_batch.grouping_set_args {
- state
- .group_values
- .intern(group_values, &mut state.batch_group_indices)?;
- let group_indices = &state.batch_group_indices;
- let total_num_groups = state.group_values.len();
-
- for (acc, values) in state
- .accumulators
- .iter_mut()
- .zip(evaluated_batch.accumulator_args.iter())
- {
- acc.update_batch(values, group_indices, total_num_groups)?;
- }
- }
-
- Ok(())
+ self.aggregate_batch_inner(
Review Comment:
```rust
self.aggregate_batch_inner(batch,
HashAggregateAccumulator::update_batch)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
