Skip to content

Make AggregateExec state modeling and public updates safe #25563

Description

@2010YOUY01

Motivation

AggregateExec, like several existing physical operators, has implicit coupling between fields. Some field combinations are invalid, but the type system does not prevent them from being constructed.

A simplified example:

pub struct AggregateExec {
    // These modes are mutually exclusive,
    // but the struct can represent both as enabled.
    enable_mode1: bool,
    enable_mode2: bool,
}

impl AggregateExec {
    // Potentially unsafe because changing mode1 may require
    // coordinated changes to other fields.
    pub fn update_mode1(&mut self, enabled: bool) {
        self.enable_mode1 = enabled;
    }
}

This leads to two main problems:

  1. Unsafe public APIs: mutating one field may silently violate invariants involving other fields.

    • Even internal optimizer usage is tricky to get right, and there have been many existing bugs.
  2. Hard to understand: understanding one field requires understanding all of its implicit associations.

Approach 1: Builder Pattern

One way to make construction safer is to validate field combinations in build():

AggregateExecBuilder::new()
    .with_mode1(...)
    .with_mode2(...)
    .build()?;

This centralizes validation in one place. However, it does not address the underlying modeling problem:

  • Unsafe APIs still exist, but are moved to the builder.
  • The struct remains hard to understand because the coupling between fields is still implicit.

This issue describes this approach with specific implementation plan, I think it demonstrates the above-mentioned problems:

Approach 2 (Proposed): Safe Inner States and Public Update APIs

API Design: Atomic and Safe

// Before
pub fn with_limit_options(mut self, limit_options: Option<LimitOptions>) -> Self { ... }
// After: if the optimization is not applicable, just no-op.
pub fn try_optimize_distinct_soft_limit(
    mut self,
    limit: usize,
) -> Result<Transformed<Self>> { ... }

Rather than exposing individual fields for mutation, provide APIs that perform one complete, valid state transition. This should also make the API more intuitive to use.

Inner-struct Modeling: Enum-based

Model mutually exclusive execution modes explicitly using enums, so invalid states are not representable.

For example:

Existing Implementation

/// Flat fields layout reused by multiple execution paths, it causes
/// - Multiplexed fields, that their semantics depends on other fields
/// - Invalid field combination become possible
/// (in short, unsafe and hard to read)
struct AggregateExec {
    mode: AggregateMode,
    input: ExecutionPlan,

    // The execution variant is implicit in the combination of these fields.
    group_by: PhysicalGroupBy,
    aggr_expr: Vec<AggregateExpr>,
    filter_expr: Vec<Option<Expr>>,

    // Represents a distinct soft limit or a TopK bound,
    // depending on the other fields.
    limit_options: Option<LimitOptions>,

    // Common original input schema, output schema, properties, metrics...
}

struct LimitOptions {
    limit: usize,
    // Optional ordering direction used by TopK.
    descending: Option<bool>,
}

Proposed Idea

// This enum represents the underlying model. The existing implementation instead
// uses a flat list of fields, introducing hidden coupling and heavily multiplexed
// fields that make the struct difficult to understand.
enum AggregateKind {
    General {
        group_by: PhysicalGroupBy,
        aggregates: Vec<AggregateExpr>,
        filters: Vec<Option<Expr>>,
    },
    Distinct {
        group_by: PhysicalGroupBy,
        soft_limit: Option<usize>,
    },
    TopK(TopKSpec),
}

struct AggregateExec {
    mode: AggregateMode,
    input: ExecutionPlan,
    kind: AggregateKind,
    // Common original input schema, output schema, properties, metrics...
}

The two problems are distinct but come from the same root cause: AggregateExec has invariants that are not encoded explicitly. The typed inner model makes those invariants structural, so invalid internal states are harder or impossible to represent; the safe public APIs preserve those invariants when optimizers or other callers update the plan. In short, the inner model defines what states are valid, while the public API defines the valid transitions between them.

Together, these changes make invalid states harder to construct and the relationships between fields explicit.

Implementation Plan

The final goal is the enum-based aggregate shape above, together with atomic public update APIs on AggregateExec that optimizers can use safely.

To keep the changes easier to review, this can be split into smaller steps:

  1. Add an enum for the distinct limit optimization, together with a safe API for applying it.
  2. Do the same for TopK.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions