Implementation:ArroyoSystems Arroyo Aggregate Extension

Overview

AggregateExtension is a custom DataFusion logical plan extension that represents windowed and non-windowed aggregate operations in Arroyo's streaming SQL engine. It encapsulates the aggregate logic, windowing behavior, key fields, and final projection calculation.

Description

The AggregateExtension struct contains:

• window_behavior: A WindowBehavior enum that indicates whether the window is derived from an operator (tumbling, sliding, session) or is already present in the data (from a prior windowing stage).
• aggregate: The underlying DataFusion LogicalPlan (typically a LogicalPlan::Aggregate) containing the group-by and aggregate expressions.
• key_fields: Indices of the key fields used for partitioning and state management.
• final_calculation: A derived projection that computes the final output schema after the aggregate.

The extension provides window-specific configuration methods:

• tumbling_window_config(): Creates a TumblingWindowAggregateOperator node with binning function and aggregate plan serialization.
• sliding_window_config(): Creates a SlidingWindowAggregateOperator node.
• session_window_config(): Creates a SessionWindowAggregateOperator node.
• instant_window_config(): Creates an instant (non-windowed tumbling) aggregate.

It implements both UserDefinedLogicalNodeCore (for DataFusion integration) and ArroyoExtension (for dataflow graph construction).

Usage

Created by the AggregateRewriter during logical plan rewriting when an Aggregate node is encountered. The extension is later processed by the plan graph builder to create the appropriate streaming aggregate operator.

Code Reference

Source Location

crates/arroyo-planner/src/extension/aggregate.rs

Signature

pub(crate) const AGGREGATE_EXTENSION_NAME: &str = "AggregateExtension";

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub(crate) struct AggregateExtension {
    pub(crate) window_behavior: WindowBehavior,
    pub(crate) aggregate: LogicalPlan,
    pub(crate) schema: DFSchemaRef,
    pub(crate) key_fields: Vec<usize>,
    pub(crate) final_calculation: LogicalPlan,
}

impl AggregateExtension {
    pub fn new(window_behavior: WindowBehavior, aggregate: LogicalPlan, key_fields: Vec<usize>) -> Self
    pub fn tumbling_window_config(&self, planner: &Planner, index: usize, input_schema: DFSchemaRef, width: Duration) -> Result<LogicalNode>
    pub fn sliding_window_config(&self, planner: &Planner, index: usize, input_schema: DFSchemaRef, width: Duration, slide: Duration) -> Result<LogicalNode>
    pub fn session_window_config(&self, planner: &Planner, index: usize, input_schema: DFSchemaRef, gap: Duration) -> Result<LogicalNode>
}

Import

use crate::extension::aggregate::{AggregateExtension, AGGREGATE_EXTENSION_NAME};

I/O Contract

Inputs

Name	Type	Description
window_behavior	WindowBehavior	Specifies whether the window comes from an operator or from input data
aggregate	LogicalPlan	The DataFusion aggregate plan with group-by and aggregate expressions
key_fields	Vec<usize>	Indices of key fields for partitioning

Outputs

Name	Type	Description
LogicalNode	struct	A graph node representing the configured aggregate operator
LogicalEdge	struct	Edge connections to upstream nodes

Usage Examples

let extension = AggregateExtension::new(
    WindowBehavior::FromOperator {
        window: WindowType::Tumbling { width: Duration::from_secs(60) },
        window_field: window_field.clone(),
        window_index: 0,
        is_nested: false,
    },
    aggregate_plan,
    vec![0, 1],  // key field indices
);

Related Pages

• ArroyoSystems_Arroyo_Planner_Extensions - The ArroyoExtension trait that AggregateExtension implements
• ArroyoSystems_Arroyo_Plan_Rewriter - The rewriter that creates AggregateExtension nodes
• ArroyoSystems_Arroyo_Datastream_Types - WindowType enum used for window configuration