AQE may materialize a non-supported Final-mode HashAggregate

Describe the bug

In cases where we support a HashAggregate's aggregate functions, we will convert the partial stage HashAggregate, execute it in DF, then use native shuffle to forward the results, however, the next stage will not materialize due to AQE waiting for the shuffle and its stats.
Then, if the agg() function contained unsupported expressions wrapping the aggregates themselves, we cannot convert the Final mode.
This causes the Spark HashAggregate to crash as it attempts to access fields that don't exist, and even if they do it expects its own aggregate buffer representation, which we don't have(The error being thrown is the "Not supported on CometListVector"), I'm not sure why the InputAdapter doesn't always work here, I think this can only be supported in cases where no aggregate buffer exists,
This comment implies the possible issue:

datafusion-comet/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala

Lines 491 to 494 in f099e6e

    
           // When Comet shuffle is disabled, we don't want to transform the HashAggregate 
        
           // to CometHashAggregate. Otherwise, we probably get partial Comet aggregation 
        
           // and final Spark aggregation. 
        
           isCometShuffleEnabled(conf) =>

I believe I've seen a few tests that are ignored because of this.
I don't think this is a valid situation, We should not crash based on previously Comet-ran operators if they were successful.

Steps to reproduce

Have a group_by/aggregate that either needs a shuffle or aggregate buffer
Like
.agg(
collect_set(col("my_column"))
)

but wrap that collect_set with an unsupported expression or cast or something
(I am not sure, but I believe I saw something about a simliar behaviour that can be created using Decimal avg, as the Partial aggregate is supported but generates a Sum and Count, but generating results from the intermediate data does is not supported natively, that may no longer be relevant though)

For example:
.agg(
concat(flatten(collect_set(col("my_column"))))
)
can create this behaviour with AQE on.
not sure if this is datatype related.

Expected behavior

Comet should either not convert a Partial HashAggregate whose Final stage cannot be converted.
Or if it already did, should elegantly execute the Final aggregations and let Spark finish the work without breaking the plan.

Additional context

I believe I have a non-invasive solution where if the result expressions are not supported, we convert them into a separate ProjectExec, which will be the parent of the CometHashAggregateExec, which will not have result expressions (Like the Partial stage doesn't), and will have the grouping+aggregate attributes as its output.
We then have a conversion to rows and run a ProjectExec with the unsupported expression, ensuring that even if the rest of the stage cannot be run using Comet, we don't break an already running workflow.

Will open a PR shortly

@viirya

cc @viirya
I forgot the context why we did this in #991

datafusion-comet/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala

Lines 491 to 494 in f099e6e

    
           // When Comet shuffle is disabled, we don't want to transform the HashAggregate 
        
           // to CometHashAggregate. Otherwise, we probably get partial Comet aggregation 
        
           // and final Spark aggregation. 
        
           isCometShuffleEnabled(conf) =>

Because this?

// When Comet shuffle is disabled, we don't want to transform the HashAggregate
// to CometHashAggregate. Otherwise, we probably get partial Comet aggregation
// and final Spark aggregation.

Comet partial aggregation operator might return different data than Spark partial aggregation operator.

@viirya

Thanks @viirya

@EmilyMatt Did you mean that we should be able to run Comet aggregation even Comet shuffle is disabled by

I believe I've seen a few tests that are ignored because of this.
I don't think this is a valid situation, We should not crash based on previously Comet-ran operators if they were successful.

The issue can be dissected as this:

a. There is no reason to have a Partial aggregate and not a Final one, regardless of shuffle, if we support the aggregate expressions, why should we run the Final stage in Spark if it is supported? we should get the performance benefit of running the final aggregate in native as well, and if something can run in Spark and the conversion has to be made, it should be only the result expressions.
This is the less critical aspect.

b. Currently, all the supported aggregate expressions are "naive", as in, they result in some sort of Spark primitive, and they are performed using a primitive aggregate, even ones such as Avg are represented as "Sum" and "Count", but what happens when Spark uses an aggregate buffer that it passes between stages? such as in the case of CollectSet and CollectList.
Despite the expected result probably being an ArrayType(StringType, false), for example, Spark will pass the serialized intermediate representation between stages, which can be observed as the DataType is BinaryType(the attribute is usually called "buf#" or something).
When we reach something like this, using columnar shuffle, the shuffle will throw an error,
because we will output an unexpected intermediate type(ArrayType of non-nullable string), while it expects a BinaryType.
Let's say we fixed this by iterating over the outputs in CometHashAggregateExec, and correcting the DataType to what Comet outputs in the previous stage, now the shuffle will work correctly, as it will call the correct method on the ColumnarBatch and write it to the shuffle file.
However, the Final stage HashAggregate, will still attempt to read the aggregate buffer it expects(binary), not an intermediate value of the result type.
This will cause a crash then.

This cannot be expected before stage materialization as the plan doesn't exist yet and we don't know if the result expressions are not supported, so I think allowing a Spark Final stage with a Comet partial stage is slightly shortsighted, unless in the future we write a conversion(maybe using codegen) between the Comet output and a Spark aggregate binary representation.
I think a better solution is the separation of the result expressions in cases where they are not supported.

Working on this

@EmilyMatt

As per @EmilyMatt ' s comment on the related PR ,this seems to be no longer a bug but a performance boost. @andygrove we might have to remove bug label and remark it as perf

This seems to be a bug rather than a performance issue. If we have a Comet partial aggregate and a Spark final aggregate, and the serialized aggregate buffers are incompatible, the query will fail at runtime. I have seen this in the past, but I do not have a repro for the issue.

It is also possible that the bug has already been fixed, but we really need a test to confirm this, and then we can close this issue.

@andygrove

As @andygrove said, there is definitely a potential bug here, besides that there's a performance boost, I feel like both labels should prob apply here^

@EmilyMatt

@EmilyMatt, now that I have worked on this and become more familiar with the issue, could you confirm that I understand it correctly?

If we have a Comet final and Spark partial agg, or a Spark final and Comet partial agg, we can experience runtime errors if the intermediate buffer representation isn't compatible. In those cases, we should either a) fall back to Spark for both, or b) add additional projections or take other measures to make it still work. Am I understanding this correctly?

@EmilyMatt, now that I have worked on this and become more familiar with the issue, could you confirm that I understand it correctly?

If we have a Comet final and Spark partial agg, or a Spark final and Comet partial agg, we can experience runtime errors if the intermediate buffer representation isn't compatible. In those cases, we should either a) fall back to Spark for both, or b) add additional projections or take other measures to make it still work. Am I understanding this correctly?

Indeed^
The better solution imo is to just separate the result expressions into a Spark ProjectExec, and have the conversion after the aggreagte(assuming that our aggregate is better than Spark ;) )
The conversion will also have less rows because of the group-by, which seems more optimal all-around.
I can give this a go again if you'd like

fix: prevent incompatible-buffer aggregate from splitting across Come…

fix: allow safe mixed Spark/Comet partial/final aggregate execution

Navigation Menu

Repository navigation

AQE may materialize a non-supported Final-mode HashAggregate #1389

Describe the bug

Steps to reproduce

Expected behavior

Additional context

2 remaining items

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	// When Comet shuffle is disabled, we don't want to transform the HashAggregate
	// to CometHashAggregate. Otherwise, we probably get partial Comet aggregation
	// and final Spark aggregation.
	isCometShuffleEnabled(conf) =>

AQE may materialize a non-supported Final-mode HashAggregate #1389

Description

Describe the bug

Steps to reproduce

Expected behavior

Additional context

Activity

kazuyukitanimura commented on Feb 15, 2025

viirya commented on Feb 15, 2025

kazuyukitanimura commented on Feb 15, 2025

EmilyMatt commented on Feb 16, 2025

coderfender commented on Jun 24, 2025

coderfender commented on Jun 24, 2025

andygrove commented on Jun 26, 2025

EmilyMatt commented on Jun 26, 2025

2 remaining items

andygrove commented on Dec 12, 2025

EmilyMatt commented on Dec 12, 2025

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions