[andygrove]

Which issue does this PR close?

Closes #6880.

Rationale for this change

#6867 runs a native PartialMerge as DataFusion Single so that it spills rather than emit a key twice. When the PartialMerge's input is sorted on a prefix of its grouping keys, DataFusion 55.2 doesn't run SingleHashAggregateStream for Single. It falls back to the legacy GroupedHashAggregateStream, whose spill asks for more than its share of the pool to sort the batch it writes, and the task fails with Failed to reserve memory for sort during spill. This is a regression since 1.1.0, which ran this shape as Partial and returned a wrong count(DISTINCT) instead (#6857).

What changes are included in this PR?

• The aggregate planner runs a pure PartialMerge as DataFusion Final when DataFusion sees its input sorted on some of its keys (InputOrderMode is not Linear). Final runs OrderedFinalAggregateStream there, which spills and reads its spill files back. MergeAsPartial evaluates to its state under Final as it does under Single, so the projection that unpacks the states doesn't change.
• Final reads each state from the columns right after the grouping keys, and it writes its spill files with the input's column types. So the planner switches only when the grouping keys are the first input columns, initial_input_buffer_offset equals the number of keys, and the input's state columns have the native state types. Otherwise it keeps Single. For example, a Spark partial's collect_list buffers arrive as binary and keep Single. That input reaches native code through a ScanExec with no ordering, so it isn't expected to be sorted on a key.
• The task now succeeds for two reasons: the ordered final stream's spill path, and the fair pool's exemption. limited_to_its_share exempts every consumer named OrderedFinalAggregateStream from the fair limit, the limit the legacy stream's spill failed against, so this PartialMerge is checked only against the pool and Spark's grant. The limited_to_its_share doc and memory_management.md now say the exemption covers this case.
• SpillReplayPool already covers OrderedFinalAggregateStream, so reading the spill files back past the stream's share works as it does for the A native final aggregate that has spilled can fail the task during its replay #6254 final aggregates. Doc comments in merge_as_partial.rs and spill_replay.rs now mention Final.

When Comet moves to DataFusion 56 (#6410), which replaces this fallback with OrderedSingleAggregateStream, the switch can probably be dropped.

How are these changes tested?

• Two new CometAggregateSuite tests write 125,000 rows over four values of g and check one row per g (max(v) and count_distinct(k)), so a key emitted twice across a change of g shows up in the answer. Both check that the PartialMerge spilled.
  • repartition(1, g), sortWithinPartitions(g) with EliminateSorts excluded, in 2 MiB and 8 MiB pools.
  • A sort-merge join on g with the default rules and AQE off, in a 16 MiB pool, which is the shape users get. The sorts under the join need more than a 2 or 8 MiB pool's share.
    Without the planner change, every case fails with the issue's error. With a 32 MiB pool the ordered final aggregate fits in memory and doesn't spill, so the tests don't use it.
• A new planner unit test checks that a PartialMerge over a native sort on its first key, or on both keys, plans as Final with PartiallySorted or Sorted input and emits the merged max and count states. It also checks that a collect_list PartialMerge over binary Spark buffers stays Single.
• CometAggregateSuite (164 tests) and the native core library tests pass locally.

[viirya]

Thanks for this. Planning a pure PartialMerge as Final when its input is sorted on a key prefix is a neat way to land on OrderedFinalAggregateStream in DataFusion 55.2. Its spill doesn't reserve extra memory to sort what it writes, and MergeAsPartial evaluates to its state there just as it does under Single. I checked the guard against how Final reads its input. merge_expressions takes the states from column group_by.expr.len() on, as_final() reads the keys from columns 0..n, and the ordered final table spills with the input schema. So both checks in input_is_keys_then_states are needed, and they look right to me. The {Partial, PartialMerge} mix is sent with proto mode Partial, so it can't reach this path either.

A few things I'd like to sort out before merge.

The PartialMerge is no longer held to its fair share. limited_to_its_share in fair_pool.rs exempts every consumer whose name starts with OrderedFinalAggregateStream, partially sorted ones included. Once this PartialMerge plans as Final, it is checked only against the pool and Spark's grant. The issue's error was the fair limit refusing the legacy stream's spill, so this exemption is part of why the task now succeeds, along with the new stream's spill path. I think the PR description and the comment in aggregate.rs should say so. The exemption paragraph in docs/source/contributor-guide/memory_management.md and the limited_to_its_share doc both describe it as a workaround for the fully sorted CometSortAggregateExec case. Could they mention that a PartialMerge over partially sorted input now gets it too? Then the next person tuning the fair pool knows which consumers it covers.

The tests don't exercise the path the fix depends on. With g constant, the ordered stream never sees the prefix advance. So nothing exercises groups emitted early before a spill, or the rule that every group waits for the replay once a run has spilled. A key emitted twice across a prefix boundary is the #6857 failure, so I'd like the answer check to be able to catch it. The repro also needs EliminateSorts excluded, so it isn't a shape users get with the default rules. Details are inline.

[viirya]

input_sorted_on_a_key is computed for every aggregate, before the proto_mode == ProtoAggregateMode::PartialMerge check. So every Partial and Final aggregate now runs find_longest_permutation and gains a new ? error path, only for the result to be thrown away. Could it be evaluated only inside the PartialMerge condition, for example as the second operand of the &&?

[andygrove]

Good catch. The sort check now runs last in the && chain, so only a PartialMerge whose input is its keys followed by native states pays for it. Partial and Final aggregates skip it.

[viirya]

Planning this as Final also gives the consumer the OrderedFinalAggregateStream name, which limited_to_its_share in fair_pool.rs exempts by prefix whether the input is fully or partially sorted. So this PartialMerge is no longer checked against its fair limit, the limit the issue's error came from. Could the comment here, and the PR description, say that this is part of why the task now succeeds? Could the exemption paragraph in memory_management.md and the limited_to_its_share doc mention this case too? Both describe the exemption as being for a CometSortAggregateExec over fully sorted input.

[andygrove]

Done. The comment here now says the task succeeds because of the new stream's spill path and because limited_to_its_share exempts the stream from the fair limit. The limited_to_its_share doc, the exemption paragraph in memory_management.md and the PR description say it too.

[viirya]

Is there a real plan where a PartialMerge over Spark's binary collect buffers sees input sorted on only some of its keys? A Spark partial's output reaches native code through a ScanExec that reports no ordering, so it would run as Linear on SingleHashAggregateStream. A sort aggregate sorts on every key, which is Sorted, and the legacy stream doesn't spill there anyway. If there isn't one, could the comment say this case isn't expected in practice, rather than that it can still fail this way? Otherwise a reader may go looking for a failure that can't happen.

[andygrove]

I couldn't construct one. A Spark partial's output arrives through a ScanExec that reports no ordering, so it runs as Linear on the hash stream. The comment now says that input isn't expected to be sorted on a key, and the description no longer says it can still fail.

[viirya]

With 0L AS g, the input is sorted on a key that never changes. So the ordered final stream never emits a finished group before it spills, and the test never exercises the rule that every group waits for the replay once a run has spilled. Could the data use a few g values, for example id % 4 AS g, and check one row per g? A key emitted twice across a prefix boundary would then show up in the answer, and that's the failure #6857 was about.

[andygrove]

Changed. The data now has four values of g (id % 4), and the test checks one row per g with max(v) and count_distinct(k). Both pool sizes still spill and fail with the issue's error when I turn the planner change off.

[viirya]

This shape needs EliminateSorts excluded, because the optimizer otherwise drops the sortWithinPartitions under the aggregate. I'd expect users to hit it after a sort-merge join on g. The join's HashPartitioning(g) satisfies the PartialMerge's clustered distribution on (g, k), so the join, the partial and the PartialMerge run in one native plan, and the PartialMerge sees its input sorted on g. Could the test use that shape with the default rules, or add it as a second case?

[andygrove]

Added as a second test: a sort-merge join on g with the default rules, which plans the join, the partial and the PartialMerge in one native plan. It needs AQE off, or it becomes a broadcast join, and a 16 MiB pool, because the sorts under the join can't get what they need from a 2 or 8 MiB pool's share. It fails with the issue's error without the planner change.

[comphead]

Two points on the head after the first round. One is about the new guard in aggregate.rs, the other about the new planner test.

[comphead]

This collects every aggregate's state fields into a Vec<Vec<_>>, concatenates them, and only then compares. Walking the input types once returns on the first mismatch and builds no vectors (not compiled):

let mut input_types = input_schema.fields().iter().skip(num_keys).map(|f| f.data_type());
for expr in aggr_expr {
    for field in expr.state_fields()? {
        if input_types.next() != Some(field.data_type()) {
            return Ok(false);
        }
    }
}
Ok(input_types.next().is_none())

[andygrove]

Done. The guard walks the input types once against each aggregate's state fields and returns on the first mismatch, then checks that no input columns are left over. It builds no vectors.

[comphead]

The same ScanExec to MemorySourceConfig swap is already at lines 4529, 4632 and 4829, and the int64, v, max and count builders above repeat the ones in partial_merge_runs_as_single_and_emits_merged_states. A helper that takes the planned plan and the input batches and returns the collected output, plus shared builders for the two PartialMerge tests, would leave this test with only what it checks: the sort keys, the mode and input order it expects, and the rows.

[andygrove]

Done. A collect_over_batches helper takes the planned plan and the input batches and returns the collected output. The keyless, list-key, PartialMerge and ordered Final tests use it now. create_int64_datatype, create_unbound_int64, create_max_agg and create_count_agg replace the builders the two PartialMerge tests each repeated.

[viirya]

Thanks for adding this case. The assert only checks that a native sort-merge join exists somewhere in the plan. If an exchange ever lands between the join and the PartialMerge, or the join's ordering stops reaching native code, the PartialMerge would run as Single over unsorted input. It would still spill, so the test would keep passing without covering #6880. Could it check the shape, as the first test does with partial.child.isInstanceOf[CometSortExec]? For example, assert that the spilled PartialMerge's child is the partial aggregate and that the partial's subtree reaches the CometSortMergeJoinExec without a CometShuffleExchangeExec in between.

[viirya]

The helper already sets AUTO_BROADCASTJOIN_THRESHOLD to -1. At runtime, AQE falls back to that value unless spark.sql.adaptive.autoBroadcastJoinThreshold is set (JoinSelectionHelper, plan.stats.isRuntime). So I wouldn't expect AQE to turn this join into a broadcast join. Does the test pass with AQE on? AQE is on by default, so that would make it closer to the shape users get. If AQE does need to be off, could the comment give the reason you saw?

[mbutrovich]

First pass, pretty minor stuff, thanks @andygrove!

[mbutrovich]

As I read DataFusion main, there is no OrderedSingleAggregateStream there anymore. apache/datafusion#24697 renamed the ordered streams to ClusteredSingleAggregateStream and ClusteredFinalAggregateStream, and those are also their memory consumer names. This fix depends on the OrderedFinalAggregateStream[ prefix in spill_replay.rs and the fair pool exemption that reuses it (fair_pool.rs). On 56, neither would match the stream this PartialMerge runs on. #6404 pins a DataFusion revision from before the rename, so its CI hasn't seen this yet.

Could the comment name the stream as it is on main? And since this PR closes #6880, could #6410 get an item for the 56 upgrade that covers both removing this Final switch and updating the consumer name prefix? Its current #6880 entry only says to recheck the failure.

[mbutrovich]

What do you think about adding cases to partial_merge_over_sorted_input_runs_as_final for the inputs this guard rejects? I mutated the guard on the head commit and ran cargo test -p datafusion-comet --lib execution::. Replacing if !keys_lead with a check that never returns, or replacing the last line with Ok(true), leaves every test passing. Only the type comparison is checked today, by the collect_list case. Each rejected input is one where Final would read the wrong columns or spill with a schema that doesn't match its groups, so I think each needs a case that expects Single. For example, the grouping keys in swapped order (key(1), key(0)), initial_input_buffer_offset of 1, and one extra int64 input column after the states.