Avoid dropping aggregate groupings in local plans

[dnhatn]

The local plan optimizer should not change the layout, as it has already been agreed upon. However, CombineProjections can violate this when some grouping elements refer to the same attribute. This occurs when ReplaceFieldWithConstantOrNull replaces missing fields with the same reference for a given data type.

I see several ways to address this issue:

1. Skip CombineProjections in the local plan. This optimization is useful in the global plan, but likely unnecessary in the local plan, especially since we now need to retain the same number of aggregates and groupings.

2. Update RuleUtils#aliasedNulls to generate an eval for each missing field separately, rather than one per data type.

This change adopts option 2 as it is simpler and more robust. However, either option is fine, and I am open to other suggestions.

Closes #128054
Closes #129811

[elasticsearchmachine]

Hi @dnhatn, I've created a changelog YAML for you.

[elasticsearchmachine]

Pinging @elastic/es-analytical-engine (Team:Analytics)

[astefan]

Had a first look at it, but I need more time to understand it.
It seems to be doing what the comment in RuleUtils.aliasedNulls is suggesting, which is good.

When I tested this I really had to use "max_concurrent_shards_per_node":1 (as you have it in your IT test). Why is this bit needed, if you can shed some light?

[bpintea]

Nice find of the cause!

This change adopts option 2 as it is simpler and more robust.

I think skipping CombineProjections in LocalLogicalPlanOptimizer should be (1) even simpler (basically, similar to this line) and likely needed anyways, as indeed, the rule (1) shouldn't normally apply locally, but (2) could also change the layout (as it happened).

So IMO we should at a minimum apply 1. Then, the NULL aliasing should make a difference with many locally missing fields, so even it has some fragility, I'd opt to keep it (at least until we substitute with something else).

[alex-spies]

Hey, nice find @dnhatn !

Update RuleUtils#aliasedNulls to generate an eval for each missing field separately, rather than one per data type.

I believe this is a good quick fix.

I think skipping CombineProjections in LocalLogicalPlanOptimizer should be (1) even simpler (basically, similar to this line) and likely needed anyways, as indeed, the rule (1) shouldn't normally apply locally, but (2) could also change the layout (as it happened)

I wouldn't disable this locally altogether. There are very valid reasons to combine projections on the local node; for one, ProjectAwayColumns introduces a projection that's only present on local nodes; also, ReplaceFieldWithConstantOrNull introduces a projection only locally; combining those makes sense and simplifies our plans.

I think the crux is that on a local node Aggregate can actually mean "this will be the initial aggregate that corresponds to an intermediate/final aggregate downstream and thus mustn't have its output changed".

Thus, I think a more robust approach would be to mark such Aggregates and avoid the deduplication of the groupings here for any aggregates marked in such a way.

I know this requires a new transport version, but we could add a fixedOutput field to Aggregate and check it in CombineProjections.

[bpintea]

I wouldn't disable this locally altogether. There are very valid reasons to combine projections on the local node; for one, ProjectAwayColumns introduces a projection that's only present on local nodes; also, ReplaceFieldWithConstantOrNull introduces a projection only locally; combining those makes sense and simplifies our plans.

I think ProjectAwayColumns should matter later, as a physical rule, but the projections that ReplaceFieldWithConstantOrNull produces should be combined, you're right!

Thus, I think a more robust approach would be to mark such Aggregates and avoid the deduplication of the groupings here for any aggregates marked in such a way.

CombineProjections would need to become aware of a special Aggregate flavour, then, plus have a rule to mark these instances (coming/going from/in a fragment). An alternative might be to just make CombineProjections location-aware (coordinator vs. local) and based on this discriminator just use a list, rather then a set, when building the output of the agg in #combineUpperGroupingsAndLowerProjections? That shouldn't require transport changes.

[alex-spies]

I think ProjectAwayColumns should matter later, as a physical rule,

It adds a logical projection into the fragment here. Currently, this top projection gets normally combined with anything that rules like ReplaceFieldWithConstantOrNull add.

An alternative might be to just make CombineProjections location-aware (coordinator vs. local) and based on this discriminator just use a list, rather then a set, when building the output of the agg in #combineUpperGroupingsAndLowerProjections? That shouldn't require transport changes.

Oh, right - I thought this wasn't an option because there may be other aggs in the local plan, but of course there can be at most one aggregate 🤦 . So, a local flavor of CombineProjections would do the trick!

[bpintea]

Laterally, I was also wondering if we couldn't do a verification of the outputs before and after local replanning.

[astefan]

LGTM. The code looks good to me, maybe a bit more unit tests (not sure exactly what) that show the different "context" (coordinator/data node) of CombineProjections in "action", like LogicalPlan optimizer tests on both coordinator and data node branches.

[dnhatn]

Thanks friends for reviews and feedback!

[dnhatn]

💚 All backports created successfully

Status	Branch	Result
✅	9.0
✅	8.19
✅	8.18
✅	8.17

[alex-spies]

@dnhatn , I realized that 8.17 might also be affected - the optimization of missing fields changed with #126187 (which triggered the bug in CombineProjections to surface on main).

Should we backport one more time, to 8.17?

[dnhatn]

Should we backport one more time, to 8.17?

++ I opened #130128. Thanks Alex!