Understanding Trino’s PushPartialAggregationThroughExchange Optimizer

As distributed SQL engines evolved to handle massive datasets, one of the biggest challenges has been efficiently aggregating data across multiple nodes. Trino’s PushPartialAggregationThroughExchange optimizer rule represents a sophisticated solution to this problem, transforming how aggregations are executed in distributed environments.

The Problem: Naive Distributed Aggregation

Imagine you’re running a query like SELECT region, COUNT(*) FROM orders GROUP BY region on a table with millions of rows spread across multiple worker nodes. The naive approach would be:

  1. Scan all raw data on each node
  2. Shuffle all data to coordinator/aggregation nodes based on region
  3. Perform aggregation on the shuffled data

This approach has serious problems - you’re moving potentially terabytes of raw data across the network just to count rows!

The Solution: Push Partial Aggregation Through Exchange

The PushPartialAggregationThroughExchange rule implements a much smarter strategy. Instead of shuffling raw data, it:

  1. Partially aggregates data at the source (close to where data lives)
  2. Shuffles only the intermediate aggregation results
  3. Finally aggregates the partial results

This dramatically reduces network traffic and improves query performance.

Before Optimization:               After Optimization:
┌─────────────────────┐           ┌─────────────────────┐
│ AggregationNode     │           │ AggregationNode     │
│ (SINGLE)            │           │ (FINAL)             │
│ COUNT(*), SUM(price)│           │                     │
└─────────────────────┘           └─────────────────────┘
           │                                  │
┌─────────────────────┐           ┌─────────────────────┐
│ ExchangeNode        │    →      │ ExchangeNode        │
│ (shuffle all data)  │           │ (shuffle partial    │
│                     │           │  results only)      │
└─────────────────────┘           └─────────────────────┘
           │                                  │
┌─────────────────────┐           ┌─────────────────────┐
│ TableScan           │           │ AggregationNode     │
│                     │           │ (PARTIAL)           │
└─────────────────────┘           │ partial_count,      │
                                  │ partial_sum         │
                                  └─────────────────────┘
                                             │
                                  ┌─────────────────────┐
                                  │ TableScan           │
                                  └─────────────────────┘

How the Optimizer Works

When Does the Rule Apply?

The PushPartialAggregationThroughExchange rule is quite selective. It looks for a specific pattern in query plans:

AggregationNode
└── ExchangeNode (data shuffle operation)
    └── [data source]

But it doesn’t apply to every aggregation over an exchange. The rule has several requirements:

  • Functions must be decomposable - The aggregation functions need to support being split into partial and final phases (more on this later)
  • Exchange type matters - Only works with GATHER or REPARTITION exchanges, not REPLICATE
  • Grouping compatibility - For repartitioned exchanges, the partitioning columns must be a subset of the GROUP BY columns
  • Session preference - The prefer_partial_aggregation session property must be enabled

The Two Main Transformations

The optimizer handles two key scenarios:

1. Splitting Single-Step Aggregations

This is the most common case. When the optimizer sees a single-step aggregation that could benefit from being distributed, it splits it:

Before: 
AggregationNode(SINGLE: COUNT(*))
└── ExchangeNode(REPARTITION on region)
    └── TableScan(orders)

After:
AggregationNode(FINAL: combine partial counts)
└── ExchangeNode(REPARTITION on region)  
    └── AggregationNode(PARTIAL: count rows locally)
        └── TableScan(orders)

2. Pushing Existing Partial Aggregations

Sometimes the query plan already has partial aggregations that can be pushed closer to the data:

Before:
AggregationNode(PARTIAL)
└── ExchangeNode
    ├── Source1
    ├── Source2  
    └── Source3

After:
ExchangeNode
├── AggregationNode(PARTIAL) → Source1
├── AggregationNode(PARTIAL) → Source2
└── AggregationNode(PARTIAL) → Source3

The Challenging Case: Distinct Aggregations

One of the most complex aspects of this optimizer is handling DISTINCT aggregations. Traditional aggregation functions like COUNT(*) and SUM(price) are “decomposable” - they can be easily split into partial and final phases. But COUNT(DISTINCT customer) is different - you need to see all the data to properly eliminate duplicates.

The PushPartialAggregationThroughExchange rule includes sophisticated logic to handle even these challenging cases.

Simple Case: Only Distinct Aggregations

Let’s start with the simpler case: queries that only have distinct aggregations, like SELECT COUNT(DISTINCT customer) FROM orders GROUP BY region.

The key insight is transforming the COUNT(DISTINCT customer) into a two-step process:

  1. Deduplication: Group by both region and customer to eliminate duplicates
  2. Counting: Count the deduplicated customer-region combinations

This way, we’re only shuffling unique customer-region pairs instead of all the raw order data.

Original Plan:
AggregationNode(COUNT(DISTINCT customer))
└── ExchangeNode
    └── TableScan(orders)

Optimized Plan:
AggregationNode(COUNT(*))  // Step 2: Count deduplicated rows
└── AggregationNode(GROUP BY region, customer)  // Step 1: Deduplication
    └── ExchangeNode
        └── TableScan(orders)

Step-by-step execution:

Step 1: Raw Data
┌──────────┬────────┬────────┐
│ customer │ region │ amount │
├──────────┼────────┼────────┤
│ Alice    │ US     │ 100    │
│ Alice    │ US     │ 200    │  ← Duplicate customer in US
│ Bob      │ EU     │ 150    │
└──────────┴────────┴────────┘

Step 2: Deduplication (GROUP BY region, customer)
┌──────────┬────────┐
│ customer │ region │
├──────────┼────────┤
│ Alice    │ US     │  ← Deduplicated
│ Bob      │ EU     │
└──────────┴────────┘

Step 3: Final Count (GROUP BY region)
┌────────┬───────┐
│ region │ count │
├────────┼───────┤
│ US     │ 1     │
│ EU     │ 1     │
└────────┴───────┘

The Really Complex Case: Mixed Distinct and Non-Distinct Aggregations

But what about queries that mix distinct and non-distinct aggregations? Something like:

SELECT region, COUNT(*), COUNT(DISTINCT customer) FROM orders GROUP BY region

This is where things get really interesting. The optimizer can’t simply apply the deduplication approach because the non-distinct COUNT(*) would be affected by the deduplication - we’d lose the original row count!

The solution involves a sophisticated multi-stage approach using something called a GroupId node. This creates separate “contexts” for processing distinct vs non-distinct aggregations.

Understanding the Multi-Stage Approach

The transformation happens in four distinct stages:

Stage 1: Partial Aggregation with Expanded Grouping

AggregationNode(step=PARTIAL)
├── groupBy: [region, customer]  // Expanded to include distinct columns
└── aggregations: {partial_count: partial_count(*)}

First, we perform partial aggregation but expand the grouping to include the distinct columns. This gives us partial counts for each customer-region combination.

Stage 2: GroupId Node Creates Separate Contexts

GroupIdNode
├── groupingSets: 
│   ├── [region, partial_count]     // Group 0: Non-distinct context
│   └── [region, customer]          // Group 1: Distinct context
└── groupIdSymbol: group

The GroupId node duplicates each row, creating separate contexts. Group 0 contains data for non-distinct aggregations, Group 1 for distinct aggregations.

Stage 3: Final Aggregation for Non-Distinct Functions

AggregationNode(step=FINAL)
├── mask: group = 0  // Only process non-distinct rows
└── aggregations: {count: final_count(partial_count)}

This stage completes the partial→final transformation, but only for the non-distinct context (group = 0).

Stage 4: Outer Aggregation Combining Both Results

AggregationNode(step=SINGLE)
├── aggregations:
│   ├── count: arbitrary(count)           // Extract non-distinct result
│   └── distinct_count: count() MASK      // Count distinct with mask
└── mask: group = 1 for distinct aggregations

The final stage combines both results - using arbitrary() to extract the already-computed non-distinct counts, and masked counting for the distinct aggregation.

Visual Flow: Mixed Aggregations Example

Let’s trace through a concrete example:

Input Data:
┌──────────┬────────┬────────┐
│ customer │ region │ amount │
├──────────┼────────┼────────┤
│ Alice    │ US     │ 100    │
│ Alice    │ US     │ 200    │
│ Bob      │ US     │ 150    │
│ Carol    │ EU     │ 300    │
└──────────┴────────┴────────┘

Stage 1 - Partial Aggregation (group by region, customer):
┌──────────┬────────┬──────────────┐
│ customer │ region │ partial_count│
├──────────┼────────┼──────────────┤
│ Alice    │ US     │ 2            │  ← Alice had 2 orders in US
│ Bob      │ US     │ 1            │  ← Bob had 1 order in US
│ Carol    │ EU     │ 1            │  ← Carol had 1 order in EU
└──────────┴────────┴──────────────┘

Stage 2 - GroupId creates separate contexts:
┌──────────┬────────┬──────────────┬───────┐
│ customer │ region │ partial_count│ group │
├──────────┼────────┼──────────────┼───────┤
│ null     │ US     │ 2            │ 0     │ ← Non-distinct context (Alice's count)
│ null     │ US     │ 1            │ 0     │ ← Non-distinct context (Bob's count)
│ null     │ EU     │ 1            │ 0     │ ← Non-distinct context (Carol's count)
│ Alice    │ US     │ null         │ 1     │ ← Distinct context  
│ Bob      │ US     │ null         │ 1     │ ← Distinct context
│ Carol    │ EU     │ null         │ 1     │ ← Distinct context
└──────────┴────────┴──────────────┴───────┘

Stage 3 - Final Aggregation (FINAL step, mask: group = 0):
┌────────┬───────────┬───────┐
│ region │ count(*)  │ group │
├────────┼───────────┼───────┤
│ US     │ 3         │ 0     │ ← Final count: 2 + 1 = 3 total orders
│ EU     │ 1         │ 0     │ ← Final count: 1 total order
└────────┴───────────┴───────┘

Stage 4 - Outer Aggregation combines both contexts:
┌──────────┬────────┬───────────┬───────┬─────────────────┐
│ customer │ region │ count(*)  │ group │ distinct_mask   │
├──────────┼────────┼───────────┼───────┼─────────────────┤
│ null     │ US     │ 3         │ 0     │ false           │ ← Non-distinct result
│ null     │ EU     │ 1         │ 0     │ false           │ ← Non-distinct result
│ Alice    │ US     │ null      │ 1     │ true            │ ← Distinct context
│ Bob      │ US     │ null      │ 1     │ true            │ ← Distinct context  
│ Carol    │ EU     │ null      │ 1     │ true            │ ← Distinct context
└──────────┴────────┴───────────┴───────┴─────────────────┘

Final Results:
┌────────┬───────────┬──────────────────┐
│ region │ count(*)  │ count(distinct)  │
├────────┼───────────┼──────────────────┤
│ US     │ 3         │ 2                │ ← 3 total orders, 2 unique customers
│ EU     │ 1         │ 1                │ ← 1 total order, 1 unique customer
└────────┴───────────┴──────────────────┘

Why This Approach Works:

  • Non-distinct aggregations get their partial→final optimization through stages 1 and 3
  • Distinct aggregations get proper deduplication through the GroupId separation
  • Both results are combined efficiently in the final stage
  • Network traffic is minimized by only shuffling the aggregated intermediate results

This sophisticated approach allows Trino to optimize even the most complex mixed aggregation queries while maintaining correctness.

The Decision Process

The optimizer doesn’t blindly apply these transformations. It follows a careful decision tree, considering factors like:

  • Whether aggregation functions are decomposable
  • The type of data exchange operation
  • Whether distinct aggregations can be handled
  • The specific mix of aggregation types

This ensures optimizations are only applied when they’ll actually improve performance.

  • Aggregations Overview - Essential background on aggregation concepts, PARTIAL/FINAL steps, and decomposability
  • GroupIdNode - Understanding the row multiplication technique used in mixed aggregation scenarios
  • Aggregations & Exchange - Practical examples showing how these optimizations affect query execution