Native PartialMerge over input sorted on a prefix of its keys fails to spill on DataFusion 55.2 #6880

Describe the bug

#6867 fixes #6857 by running a native PartialMerge as DataFusion Single, which spills instead of emitting a key twice. When the PartialMerge's input is sorted on a prefix of its grouping keys, DataFusion 55.2 doesn't run SingleHashAggregateStream: should_use_single_hash_stream requires InputOrderMode::Linear, so the aggregate falls back to the legacy GroupedHashAggregateStream. Under memory pressure, that stream's spill asks for more than its share of the pool to sort the batch it is about to write, and the task fails:

Failed to reserve memory for sort during spill: Resources exhausted: Additional allocation failed for GroupedHashAggregateStream[0] (col_0) ...
Error: Failed to acquire 745281 bytes where this consumer already holds 8396 bytes and the fair limit is 419430 bytes, 5 registered (0 bytes overcommitted)

Before #6867, this shape ran as DataFusion Partial. Its OrderedPartialAggregateStream emits all its groups when a reservation fails, so it returned the wrong count(DISTINCT) of #6857. A task failure is better than a wrong answer, but the stage should spill.

The shape needs the sort and both aggregates in one native plan, with no exchange between the Partial and the PartialMerge. That happens when the child's hash partitioning on g already satisfies the clustered distribution on (g, d): after a native sort-merge join on g, for example, or in the repro below.

Steps to reproduce

On the #6867 branch (2ce73e5), in a CometTestBase suite (2g off-heap):

spark
  .range(0, 125000, 1, 1)
  .selectExpr("0L AS g", "lpad(cast(id % 62500 AS STRING), 128, 'k') AS k", "id AS v")
  .write
  .parquet(path)

withSQLConf(
  SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> EliminateSorts.ruleName,
  CometConf.COMET_BATCH_SIZE.key -> "1024",
  CometConf.COMET_RESPECT_DATAFUSION_CONFIGS.key -> "true",
  "spark.comet.datafusion.execution.sort_spill_reservation_bytes" -> "65536",
  CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key -> (2.0 / 2048).toString) {
  spark.read
    .parquet(path)
    .repartition(1, col("g"))
    .sortWithinPartitions("g")
    .groupBy("g")
    .agg(expr("max(v)"), count_distinct(col("k")))
    .collect()
}

It fails with the error above. Of the pools from 2 to 48 MiB that I tried (pool fraction n / 2048), only 24 and 48 MiB passed. A larger pool doesn't make it go away, because the sort reservation grows with it: about 0.75 MB at 2 MiB, 2.5 MB at 8 and 12 MiB, 4.9 MB at 14 to 20 MiB and 9.6 MB at 32 MiB, between 1.0 and 1.8 times the stream's fair share.

Expected behavior

The PartialMerge spills and reads its spill files back, as it does when its input is unsorted.

Additional context

• SpillReplayPool doesn't help. The refused request comes from the spill itself, while the stream holds only 8 KB, not from reading the spill files back.
• DataFusion main replaces this fallback with OrderedSingleAggregateStream (datafusion/physical-plan/src/aggregates/ordered_single_stream.rs), so this should be rechecked when Comet moves to DataFusion 56 ([EPIC] DataFusion 56 upgrade: regressions found by tracking DataFusion main #6410).
• If it has to be fixed before then, Comet could run a PartialMerge with sorted input as DataFusion Final. Its OrderedFinalAggregateStream spills, and the MergeAsPartial struct output works under Final too. Final reads the states right after the grouping keys, which holds for a pure PartialMerge, and it spills with its input schema, which is native state when the input is sorted by a native operator.