Skip to content

Replacing an AggregateExec's aggregate expressions keeps state derived from the old ones #25394

Description

@adriangb

Describe the bug

AggregateExec can have its aggregate expressions replaced, by OptimizeAggregateOrder through AggregateExec::to_builder().with_aggr_exprs(..) (or, before #25376, through the now deprecated with_new_aggr_exprs). Both keep the state the node derives from those very expressions, so the replacement can leave the node describing the expressions it no longer has.

Ordering requirement. required_input_ordering is derived from the aggregate expressions by get_finer_aggregate_exprs_requirement. A replacement keeps the old one, so swapping an order-insensitive aggregate for one with an ORDER BY produces a node that advertises no ordering requirement and can run against unsorted input:

// array_agg(b) requires nothing of its input
let exec = AggregateExec::builder(AggregateMode::Single, input)
    .with_group_by(group_by)
    .with_aggr_exprs(vec![array_agg_b(&schema, /* ordered */ false)?])
    .build()?;
assert!(exec.required_input_ordering().iter().all(Option::is_none));

// array_agg(b ORDER BY b) does, and produces the same output field, so the
// inherited output schema still describes it
let rewritten = exec
    .to_builder()
    .with_aggr_exprs(vec![array_agg_b(&schema, /* ordered */ true)?])
    .build()?;
// FAILS: the requirement of the expressions the node no longer has
assert!(rewritten.required_input_ordering().iter().any(Option::is_some));

SanityCheckPlan runs after OptimizeAggregateOrder, so nothing downstream catches it.

Output schema. A rewrite deliberately keeps the output schema, so that it cannot rename output fields. Nothing checks that the new expressions still produce that schema, so a replacement that changes a field's type or nullability produces a node whose schema lies about its output.

Dynamic filter. AggrDynFilter records which aggregate expressions are MIN/MAX, at which index, and over which column; a MIN pushes down col < bound and a MAX pushes col > bound. A rewrite that turns a MIN into a MAX keeps the old state and pushes down the wrong predicate, pruning rows the aggregate needs. The schema check above cannot see this, because both produce the same output field.

The no-op check. AggregateExecBuilder::with_aggr_exprs currently applies the replacement unconditionally, which is correct. Any future short-circuit must compare by allocation, not with ==: AggregateFunctionExpr's PartialEq compares the name, the return field, the function and the arguments, and ignores DISTINCT, IGNORE NULLS and ORDER BY, so count(b) and count(DISTINCT b) with the same alias compare equal.

To Reproduce

The ordering case is the snippet above, as a unit test in datafusion/physical-plan/src/aggregates/builder.rs. All of these need a hand-built plan: the rewrites DataFusion's own rules perform today preserve the type, the ordering requirement and the MIN/MAX kind, so none of this is reachable from SQL.

Expected behavior

AggregateExecBuilder::build should derive everything it computes from the aggregate expressions again when they are replaced, carrying over only the output schema, and check that the new expressions still describe that schema. The dynamic filter should be kept only while it still describes the same aggregates, and rebuilt otherwise — rebuilding loses the link to whichever child accepted the old filter during pushdown, which costs the optimization but cannot give wrong results, since an abandoned filter is never narrowed.

Additional context

Pre-existing; with_new_aggr_exprs behaves the same way on main. Carved out of #25376, which is a pure refactor and keeps this behaviour unchanged. Related: #25393.

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

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions