Skip to content

AggregateExec physical-plan serde rebuilds the output schema, losing optimizer-preserved field names (and is_reversed) #24202

Description

@adriangb

Describe the bug

AggregateExec deliberately distinguishes two constructors:

  • AggregateExec::try_new derives the output schema with create_schema, whose field names come from aggr_expr[i].field() / state_fields() — i.e. from each aggregate expression's name.
  • AggregateExec::try_new_with_schema (private) takes the schema as given. Its doc comment states why: "a rule may re-write aggregate expressions (e.g. reverse them) during initialization, field names may change inadvertently if one re-creates the schema in such cases."

with_new_aggr_exprs and with_new_children both go through the preserving path. Physical-plan deserialization does not — it calls AggregateExec::try_new, so it rebuilds the schema from the decoded expression names.

That is a round-trip gap whenever the plan's stored schema and its current aggr_expr names have diverged.

How they diverge. OptimizeAggregateOrder (datafusion/physical-optimizer/src/update_aggr_exprs.rs:113) calls with_new_aggr_exprs, keeping the old schema while swapping in rewritten expressions. AggregateFunctionExpr::reverse_expr (datafusion/physical-expr/src/aggregate.rs:938) rewrites the expression's name when the reversed UDAF differs from the original — e.g. last_value(a) ORDER BY [b ASC ...] becomes first_value(a) ORDER BY [b DESC ...]. So after this rule runs, exec.schema() carries the pre-reversal field name while exec.aggr_expr()[i].name() is the post-reversal name.

What serde does with that. The encoder writes aggr_expr_name from the (rewritten) expressions and does not write the schema at all — AggregateExecNode has no schema field. The decoder rebuilds via try_new, so the decoded plan's output field name is the post-reversal name. The original plan's was the pre-reversal one.

Column indices still line up, so this does not usually fail loudly. It surfaces as original.schema() != decoded.schema(), and in any name-based lookup against the aggregate's output.

Scope caveat: reverse_expr only rewrites the name when human_display_alias() is None, i.e. when the aggregate had no explicit SQL alias. SELECT last_value(a ORDER BY b) AS x ... keeps its name and is unaffected; SELECT last_value(a ORDER BY b) ... is affected.

Related gap, same area: PhysicalAggregateExprNode (datafusion/proto-models/proto/datafusion.proto:1089) has no is_reversed field, and the decoder never calls AggregateExprBuilder::with_reversed. A reversed aggregate therefore decodes as non-reversed, so AccumulatorArgs::is_reversed is wrong for any UDAF that branches on it. These look like the same underlying issue — the serialized form does not carry the identity the optimizer established.

To Reproduce

Plan a query that lets OptimizeAggregateOrder reverse an un-aliased order-sensitive aggregate (e.g. last_value(a ORDER BY b) over an input already ordered by b DESC), then round-trip the optimized physical plan through PhysicalPlanNode encode/decode and compare schema() before and after.

I have traced this by inspection rather than executing a failing test; a regression test that drives OptimizeAggregateOrder and round-trips is part of what this issue asks for.

Expected behavior

Round-tripping a physical plan preserves the output schema. A decoded AggregateExec should have the same field names as the plan that was encoded.

Additional context

This is pre-existing, and predates DataFusion 54. Verified at tag 54.0.0:

  • datafusion/proto/src/physical_plan/mod.rs:1352 — decode already called AggregateExec::try_new
  • datafusion/physical-optimizer/src/update_aggr_exprs.rs:113 — already used with_new_aggr_exprs
  • datafusion/physical-plan/src/aggregates/mod.rs:758try_new_with_schema already existed and was already private
  • is_reversed was already absent from the proto

On current main the decode call lives at datafusion/physical-plan/src/aggregates/mod.rs:2447 after the serde hooks moved into physical-plan; the behaviour is unchanged.

Surfaced during review of #24166 (part of #23494), where the exhaustive-destructure refactor made the unserialized schema field explicit. Thanks to @kumarUjjawal for spotting it. It was left out of that PR because fixing it is not a drive-by:

Metadata

Metadata

Assignees

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