Struct GroupedHashAggregateStream 

pub(crate) struct GroupedHashAggregateStream {

    schema: SchemaRef,
    input: SendableRecordBatchStream,
    mode: AggregateMode,
    aggregate_arguments: Vec<Vec<Arc<dyn PhysicalExpr>>>,
    filter_expressions: Vec<Option<Arc<dyn PhysicalExpr>>>,
    group_by: PhysicalGroupBy,
    batch_size: usize,
    group_values_soft_limit: Option<usize>,
    exec_state: ExecutionState,
    input_done: bool,
    group_values: Box<dyn GroupValues>,
    current_group_indices: Vec<usize>,
    accumulators: Vec<Box<dyn GroupsAccumulator>>,
    group_ordering: GroupOrdering,
    spill_state: SpillState,
    skip_aggregation_probe: Option<SkipAggregationProbe>,
    reservation: MemoryReservation,
    baseline_metrics: BaselineMetrics,
    group_by_metrics: GroupByMetrics,
    reduction_factor: Option<RatioMetrics>,
}

HashTable based Grouping Aggregator

Design Goals

This structure is designed so that updating the aggregates can be vectorized (done in a tight loop) without allocations. The accumulator state is not managed by this operator (e.g in the hash table) and instead is delegated to the individual accumulators which have type specialized inner loops that perform the aggregation.

Architecture

    Assigns a consecutive group           internally stores aggregate values
    index for each unique set                     for all groups
        of group values

        ┌────────────┐              ┌──────────────┐       ┌──────────────┐
        │ ┌────────┐ │              │┌────────────┐│       │┌────────────┐│
        │ │  "A"   │ │              ││accumulator ││       ││accumulator ││
        │ ├────────┤ │              ││     0      ││       ││     N      ││
        │ │  "Z"   │ │              ││ ┌────────┐ ││       ││ ┌────────┐ ││
        │ └────────┘ │              ││ │ state  │ ││       ││ │ state  │ ││
        │            │              ││ │┌─────┐ │ ││  ...  ││ │┌─────┐ │ ││
        │    ...     │              ││ │├─────┤ │ ││       ││ │├─────┤ │ ││
        │            │              ││ │└─────┘ │ ││       ││ │└─────┘ │ ││
        │            │              ││ │        │ ││       ││ │        │ ││
        │ ┌────────┐ │              ││ │  ...   │ ││       ││ │  ...   │ ││
        │ │  "Q"   │ │              ││ │        │ ││       ││ │        │ ││
        │ └────────┘ │              ││ │┌─────┐ │ ││       ││ │┌─────┐ │ ││
        │            │              ││ │└─────┘ │ ││       ││ │└─────┘ │ ││
        └────────────┘              ││ └────────┘ ││       ││ └────────┘ ││
                                    │└────────────┘│       │└────────────┘│
                                    └──────────────┘       └──────────────┘

        group_values                             accumulators

For example, given a query like COUNT(x), SUM(y) ... GROUP BY z, group_values will store the distinct values of z. There will be one accumulator for COUNT(x), specialized for the data type of x and one accumulator for SUM(y), specialized for the data type of y.

Discussion

group_values does not store any aggregate state inline. It only assigns “group indices”, one for each (distinct) group value. The accumulators manage the in-progress aggregate state for each group, with the group values themselves are stored in group_values at the corresponding group index.

The accumulator state (e.g partial sums) is managed by and stored by a [GroupsAccumulator] accumulator. There is one accumulator per aggregate expression (COUNT, AVG, etc) in the stream. Internally, each GroupsAccumulator manages the state for multiple groups, and is passed group_indexes during update. Note The accumulator state is not managed by this operator (e.g in the hash table).

Partial Aggregate and multi-phase grouping

As described on Accumulator::state, this operator is used in the context “multi-phase” grouping when the mode is AggregateMode::Partial.

An important optimization for multi-phase partial aggregation is to skip partial aggregation when it is not effective enough to warrant the memory or CPU cost, as is often the case for queries many distinct groups (high cardinality group by). Memory is particularly important because each Partial aggregator must store the intermediate state for each group.

If the ratio of the number of groups to the number of input rows exceeds a threshold, and [GroupsAccumulator::supports_convert_to_state] is supported, this operator will stop applying Partial aggregation and directly pass the input rows to the next aggregation phase.

Spilling (to disk)

The sizes of group values and accumulators can become large. Before that causes out of memory, this hash aggregator outputs partial states early for partial aggregation or spills to local disk using Arrow IPC format for final aggregation. For every input [RecordBatch], the memory manager checks whether the new input size meets the memory configuration. If not, outputting or spilling happens. For outputting, the final aggregation takes care of re-grouping. For spilling, later stream-merge sort on reading back the spilled data does re-grouping. Note the rows cannot be grouped once spilled onto disk, the read back data needs to be re-grouped again. In addition, re-grouping may cause out of memory again. Thus, re-grouping has to be a sort based aggregation.

Partial Aggregation [batch_size = 2] (max memory = 3 rows)

 INPUTS        PARTIALLY AGGREGATED (UPDATE BATCH)   OUTPUTS
┌─────────┐    ┌─────────────────┐                  ┌─────────────────┐
│ a │ b   │    │ a │    AVG(b)   │                  │ a │    AVG(b)   │
│---│-----│    │   │[count]│[sum]│                  │   │[count]│[sum]│
│ 3 │ 3.0 │ ─▶ │---│-------│-----│                  │---│-------│-----│
│ 2 │ 2.0 │    │ 2 │ 1     │ 2.0 │ ─▶ early emit ─▶ │ 2 │ 1     │ 2.0 │
└─────────┘    │ 3 │ 2     │ 7.0 │               │  │ 3 │ 2     │ 7.0 │
┌─────────┐ ─▶ │ 4 │ 1     │ 8.0 │               │  └─────────────────┘
│ 3 │ 4.0 │    └─────────────────┘               └▶ ┌─────────────────┐
│ 4 │ 8.0 │    ┌─────────────────┐                  │ 4 │ 1     │ 8.0 │
└─────────┘    │ a │    AVG(b)   │               ┌▶ │ 1 │ 1     │ 1.0 │
┌─────────┐    │---│-------│-----│               │  └─────────────────┘
│ 1 │ 1.0 │ ─▶ │ 1 │ 1     │ 1.0 │ ─▶ early emit ─▶ ┌─────────────────┐
│ 3 │ 2.0 │    │ 3 │ 1     │ 2.0 │                  │ 3 │ 1     │ 2.0 │
└─────────┘    └─────────────────┘                  └─────────────────┘


Final Aggregation [batch_size = 2] (max memory = 3 rows)

PARTIALLY INPUTS       FINAL AGGREGATION (MERGE BATCH)       RE-GROUPED (SORTED)
┌─────────────────┐    [keep using the partial schema]       [Real final aggregation
│ a │    AVG(b)   │    ┌─────────────────┐                    output]
│   │[count]│[sum]│    │ a │    AVG(b)   │                   ┌────────────┐
│---│-------│-----│ ─▶ │   │[count]│[sum]│                   │ a │ AVG(b) │
│ 3 │ 3     │ 3.0 │    │---│-------│-----│ ─▶ spill ─┐       │---│--------│
│ 2 │ 2     │ 1.0 │    │ 2 │ 2     │ 1.0 │           │       │ 1 │    4.0 │
└─────────────────┘    │ 3 │ 4     │ 8.0 │           ▼       │ 2 │    1.0 │
┌─────────────────┐ ─▶ │ 4 │ 1     │ 7.0 │     Streaming  ─▶ └────────────┘
│ 3 │ 1     │ 5.0 │    └─────────────────┘     merge sort ─▶ ┌────────────┐
│ 4 │ 1     │ 7.0 │    ┌─────────────────┐            ▲      │ a │ AVG(b) │
└─────────────────┘    │ a │    AVG(b)   │            │      │---│--------│
┌─────────────────┐    │---│-------│-----│ ─▶ memory ─┘      │ 3 │    2.0 │
│ 1 │ 2     │ 8.0 │ ─▶ │ 1 │ 2     │ 8.0 │                   │ 4 │    7.0 │
│ 2 │ 2     │ 3.0 │    │ 2 │ 2     │ 3.0 │                   └────────────┘
└─────────────────┘    └─────────────────┘

Fields

schema: SchemaRef

input: SendableRecordBatchStream

mode: AggregateMode

aggregate_arguments: Vec<Vec<Arc<dyn PhysicalExpr>>>

Arguments to pass to each accumulator.

The arguments in accumulator[i] is passed aggregate_arguments[i]

The argument to each accumulator is itself a Vec because some aggregates such as CORR can accept more than one argument.

filter_expressions: Vec<Option<Arc<dyn PhysicalExpr>>>

Optional filter expression to evaluate, one for each for accumulator. If present, only those rows for which the filter evaluate to true should be included in the aggregate results.

For example, for an aggregate like SUM(x) FILTER (WHERE x >= 100), the filter expression is x > 100.

group_by: PhysicalGroupBy

GROUP BY expressions

batch_size: usize

max rows in output RecordBatches

group_values_soft_limit: Option<usize>

Optional soft limit on the number of group_values in a batch If the number of group_values in a single batch exceeds this value, the GroupedHashAggregateStream operation immediately switches to output mode and emits all groups.

exec_state: ExecutionState

Tracks if this stream is generating input or output

input_done: bool

Have we seen the end of the input

group_values: Box<dyn GroupValues>

An interning store of group keys

current_group_indices: Vec<usize>

scratch space for the current input [RecordBatch] being processed. Reused across batches here to avoid reallocations

accumulators: Vec<Box<dyn GroupsAccumulator>>

Accumulators, one for each AggregateFunctionExpr in the query

For example, if the query has aggregates, SUM(x), COUNT(y), there will be two accumulators, each one specialized for that particular aggregate and its input types

group_ordering: GroupOrdering

Optional ordering information, that might allow groups to be emitted from the hash table prior to seeing the end of the input

spill_state: SpillState

The spill state object

skip_aggregation_probe: Option<SkipAggregationProbe>

Optional probe for skipping data aggregation, if supported by current stream.

reservation: MemoryReservation

The memory reservation for this grouping

baseline_metrics: BaselineMetrics

Execution metrics

group_by_metrics: GroupByMetrics

Aggregation-specific metrics

reduction_factor: Option<RatioMetrics>

Reduction factor metric, calculated as output_rows/input_rows (only for partial aggregation)

Implementations

impl GroupedHashAggregateStream

pub fn new( agg: &AggregateExec, context: Arc<TaskContext>, partition: usize, ) -> Result<Self>

Create a new GroupedHashAggregateStream

impl GroupedHashAggregateStream

fn group_aggregate_batch(&mut self, batch: RecordBatch) -> Result<()>

Perform group-by aggregation for the given [RecordBatch].

fn update_memory_reservation(&mut self) -> Result<()>

fn emit( &mut self, emit_to: EmitTo, spilling: bool, ) -> Result<Option<RecordBatch>>

Create an output RecordBatch with the group keys and accumulator states/values specified in emit_to

fn spill_previous_if_necessary(&mut self, batch: &RecordBatch) -> Result<()>

Optimistically, Self::group_aggregate_batch allows to exceed the memory target slightly (~ 1 [RecordBatch]) for simplicity. In such cases, spill the data to disk and clear the memory. Currently only GroupOrdering::None is supported for spilling.

fn spill(&mut self) -> Result<()>

Emit all intermediate aggregation states, sort them, and store them on disk. This process helps in reducing memory pressure by allowing the data to be read back with streaming merge.

fn clear_shrink(&mut self, batch: &RecordBatch)

Clear memory and shirk capacities to the size of the batch.

fn clear_all(&mut self)

Clear memory and shirk capacities to zero.

fn emit_early_if_necessary(&mut self) -> Result<()>

Emit if the used memory exceeds the target for partial aggregation. Currently only GroupOrdering::None is supported for early emitting. TODO: support group_ordering for early emitting

fn update_merged_stream(&mut self) -> Result<()>

At this point, all the inputs are read and there are some spills. Emit the remaining rows and create a batch. Conduct a streaming merge sort between the batch and spilled data. Since the stream is fully sorted, set self.group_ordering to Full, then later we can read with [EmitTo::First].

fn hit_soft_group_limit(&self) -> bool

returns true if there is a soft groups limit and the number of distinct groups we have seen is over that limit

fn set_input_done_and_produce_output(&mut self) -> Result<()>

common function for signalling end of processing of the input stream

fn update_skip_aggregation_probe(&mut self, input_rows: usize)

Updates skip aggregation probe state.

Notice: It should only be called in Partial aggregation

fn switch_to_skip_aggregation(&mut self) -> Result<()>

In case the probe indicates that aggregation may be skipped, forces stream to produce currently accumulated output.

Notice: It should only be called in Partial aggregation

fn should_skip_aggregation(&self) -> bool

Returns true if the aggregation probe indicates that aggregation should be skipped.

Notice: It should only be called in Partial aggregation

fn transform_to_states(&self, batch: RecordBatch) -> Result<RecordBatch>

Transforms input batch to intermediate aggregate state, without grouping it

Trait Implementations

impl RecordBatchStream for GroupedHashAggregateStream

fn schema(&self) -> SchemaRef

Returns the schema of this RecordBatchStream.

impl Stream for GroupedHashAggregateStream

type Item = Result<RecordBatch, DataFusionError>

Values yielded by the stream.

fn poll_next( self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll<Option<Self::Item>>

Attempt to pull out the next value of this stream, registering the current task for wakeup if the value is not yet available, and returning None if the stream is exhausted.

fn size_hint(&self) -> (usize, Option<usize>)

Returns the bounds on the remaining length of the stream.