You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "comphead (via GitHub)" <gi...@apache.org> on 2026/04/20 17:21:14 UTC
[PR] feat: support `PartialMerge` [datafusion-comet]
comphead opened a new pull request, #4003:
URL: https://github.com/apache/datafusion-comet/pull/4003
## Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
-->
Closes #2870 .
## Rationale for this change
<!--
Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.
-->
## What changes are included in this PR?
<!--
There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
-->
## How are these changes tested?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example, are they covered by existing tests)?
-->
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: support `PartialMerge` aggregation mode [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #4003:
URL: https://github.com/apache/datafusion-comet/pull/4003#issuecomment-4374074012
Thanks @mbutrovich for the review
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: support `PartialMerge` aggregation mode [datafusion-comet]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on PR #4003:
URL: https://github.com/apache/datafusion-comet/pull/4003#issuecomment-4353922011
Thanks @comphead! The `MergeAsPartialUDF` wrapper is a clever way to get `PartialMerge` semantics without reaching into DataFusion's aggregation logic.
### Streaming shim asymmetry
I checked the Spark source at `v3.5.8`, `v4.0.2`, and `v4.1.1`. `StreamSourceAwareSparkPlan` doesn't exist before 4.x, so needing a different 3.x approach is unavoidable. The shapes of the two checks do differ in a way worth flagging though.
On 3.x the shim does `plan.logicalLink.exists(_.isStreaming)`. That only looks at the logical link of the root physical node. `LogicalPlan.isStreaming` does recurse through logical children, but if the root physical node has `logicalLink = None` (a rule built a fresh node without copying the link) the check silently returns `false`.
On 4.x the shim does `plan.exists { case p: StreamSourceAwareSparkPlan => p.getStream.isDefined }`. `SparkPlan.exists` walks the whole physical tree, and the leaf check does not depend on logical link preservation.
Would `plan.exists(_.logicalLink.exists(_.isStreaming))` on 3.x be a safer analogue, or is there a reason you only want to inspect the root? Totally fine if this is out of scope for the PartialMerge work, just wanted to surface it.
### `sum_int.rs` merge_batch rewrite
The prior `merge_batch` only read row 0 of the state array, which looks like a latent bug whenever a merge batch had more than one row. Switching to `update_batch(states)` fixes that. Do you think it's worth adding a small direct test that exercises `merge_batch` with a multi-row state array independent of the `PartialMerge` plumbing, so the fix is captured even if the `MergeAsPartial` path changes later?
### `collect_set` in the description
The description lists `FIRST` / `LAST` / `collect_set` as falling back, but the serde only filters `First` and `Last`. Is `collect_set` covered entirely by the streaming plan fallback? If so, a line of comment near the `First`/`Last` filter explaining the split would help future readers.
### `distinct` test lost its aggregate count assertion
The previous version used `checkSparkAnswerAndNumOfAggregates` to assert how many native aggregate stages ran. The new version drops that. Since this PR is specifically about enabling more stages to stay native, would it be worth keeping a mode count assertion in one or two of the cases so a future regression that silently falls back one stage still shows up? Same thread as the existing review comment.
### Test name mismatch
`partialMerge - distinct + non-distinct with first() FILTER (Expand pattern)` doesn't use `FIRST` or `FILTER` in the queries it runs. Was that the intended coverage, or did the queries drift from the original intent?
### Leftover configs
In `partialMerge - cnt distinct + sum`, are `spark.comet.cast.allowIncompatible` and `spark.comet.expression.Cast.allowIncompatible` actually needed? Looks like they might be leftover from local debugging.
### AVG(DISTINCT) coverage
The serde relaxation also admits `SUM(DISTINCT x)` and `AVG(DISTINCT x)` for single column. The `distinct` test covers `COUNT` and `SUM` variants but I didn't see `AVG(DISTINCT)` exercised. Could we add one case, ideally mixed with a non distinct aggregate like `AVG(DISTINCT a) + SUM(b)`, so the Expand plan is exercised for a non count distinct?
### `planner.rs` mode literals
The match on `agg.mode` uses `0`, `1`, `2` directly. Would it be cleaner to match on the generated `AggregateMode` enum so the proto and native sides can't desync silently?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: support `PartialMerge` aggregation mode [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead merged PR #4003:
URL: https://github.com/apache/datafusion-comet/pull/4003
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: support `PartialMerge` aggregation mode [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #4003:
URL: https://github.com/apache/datafusion-comet/pull/4003#discussion_r3150573805
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -650,6 +650,87 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
}
}
+ // FIRST/LAST are order-dependent aggregates whose merge result depends on hash table
+ // processing order. In PartialMerge mode, DataFusion's hash table may process rows
+ // in a different order than Spark's, so we fall back to Spark for correctness.
+ test("partialMerge - FIRST/LAST with distinct aggregates falls back") {
+ val numValues = 10000
+ Seq(100).foreach { numGroups =>
+ Seq(128).foreach { batchSize =>
+ withSQLConf(
+ SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",
+ CometConf.COMET_BATCH_SIZE.key -> batchSize.toString) {
+ withParquetTable(
+ (0 until numValues).map(i => (i, Random.nextInt() % numGroups)),
+ "tbl",
+ false) {
+ withView("v") {
+ sql("CREATE TEMP VIEW v AS SELECT _1, _2 FROM tbl ORDER BY _1")
+// checkSparkAnswerAndFallbackReason(
Review Comment:
oh this is a good catch, FIRST/LAST should correctly fallback
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: support `PartialMerge` aggregation mode [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #4003:
URL: https://github.com/apache/datafusion-comet/pull/4003#issuecomment-4306609090
Oh, interesting there are some correctness issues, checking it
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: support `PartialMerge` aggregation mode [datafusion-comet]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #4003:
URL: https://github.com/apache/datafusion-comet/pull/4003#discussion_r3150339487
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -650,6 +650,87 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
}
}
+ // FIRST/LAST are order-dependent aggregates whose merge result depends on hash table
+ // processing order. In PartialMerge mode, DataFusion's hash table may process rows
+ // in a different order than Spark's, so we fall back to Spark for correctness.
+ test("partialMerge - FIRST/LAST with distinct aggregates falls back") {
+ val numValues = 10000
+ Seq(100).foreach { numGroups =>
+ Seq(128).foreach { batchSize =>
+ withSQLConf(
+ SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",
+ CometConf.COMET_BATCH_SIZE.key -> batchSize.toString) {
+ withParquetTable(
+ (0 until numValues).map(i => (i, Random.nextInt() % numGroups)),
+ "tbl",
+ false) {
+ withView("v") {
+ sql("CREATE TEMP VIEW v AS SELECT _1, _2 FROM tbl ORDER BY _1")
+// checkSparkAnswerAndFallbackReason(
Review Comment:
Is this test doing anything?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: support `PartialMerge` aggregation mode [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #4003:
URL: https://github.com/apache/datafusion-comet/pull/4003#issuecomment-4339167389
To regenerate for 4.1 filed https://github.com/apache/datafusion-comet/pull/4134
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: support `PartialMerge` aggregation mode [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #4003:
URL: https://github.com/apache/datafusion-comet/pull/4003#issuecomment-4309689008
moving to draft while resolving correctness issue
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: support `PartialMerge` aggregation mode [datafusion-comet]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on code in PR #4003:
URL: https://github.com/apache/datafusion-comet/pull/4003#discussion_r3150342646
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -1108,54 +1189,37 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
}
}
- // TODO enable once https://github.com/apache/datafusion-comet/issues/1267 is implemented
- ignore("distinct") {
+ test("distinct") {
withSQLConf(CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true") {
Seq("native", "jvm").foreach { cometShuffleMode =>
withSQLConf(CometConf.COMET_SHUFFLE_MODE.key -> cometShuffleMode) {
Seq(true, false).foreach { dictionary =>
withSQLConf("parquet.enable.dictionary" -> dictionary.toString) {
- val cometColumnShuffleEnabled = cometShuffleMode == "jvm"
val table = "test"
withTable(table) {
sql(s"create table $table(col1 int, col2 int, col3 int) using parquet")
sql(
s"insert into $table values(1, 1, 1), (1, 1, 1), (1, 3, 1), (1, 4, 2), (5, 3, 2)")
- var expectedNumOfCometAggregates = 2
-
- checkSparkAnswerAndNumOfAggregates(
- s"SELECT DISTINCT(col2) FROM $table",
- expectedNumOfCometAggregates)
-
- expectedNumOfCometAggregates = 4
+ checkSparkAnswerAndOperator(s"SELECT DISTINCT(col2) FROM $table")
- checkSparkAnswerAndNumOfAggregates(
- s"SELECT COUNT(distinct col2) FROM $table",
- expectedNumOfCometAggregates)
+ checkSparkAnswerAndOperator(s"SELECT COUNT(distinct col2) FROM $table")
- checkSparkAnswerAndNumOfAggregates(
- s"SELECT COUNT(distinct col2), col1 FROM $table group by col1",
- expectedNumOfCometAggregates)
+ checkSparkAnswerAndOperator(
+ s"SELECT COUNT(distinct col2), col1 FROM $table group by col1")
- checkSparkAnswerAndNumOfAggregates(
- s"SELECT SUM(distinct col2) FROM $table",
- expectedNumOfCometAggregates)
+ checkSparkAnswerAndOperator(s"SELECT SUM(distinct col2) FROM $table")
- checkSparkAnswerAndNumOfAggregates(
- s"SELECT SUM(distinct col2), col1 FROM $table group by col1",
- expectedNumOfCometAggregates)
+ checkSparkAnswerAndOperator(
+ s"SELECT SUM(distinct col2), col1 FROM $table group by col1")
- checkSparkAnswerAndNumOfAggregates(
+ checkSparkAnswerAndOperator(
"SELECT COUNT(distinct col2), SUM(distinct col2), col1, COUNT(distinct col2)," +
- s" SUM(distinct col2) FROM $table group by col1",
- expectedNumOfCometAggregates)
+ s" SUM(distinct col2) FROM $table group by col1")
- expectedNumOfCometAggregates = if (cometColumnShuffleEnabled) 2 else 1
- checkSparkAnswerAndNumOfAggregates(
+ checkSparkAnswerAndOperator(
Review Comment:
Can we not use `checkSparkAnswerAndNumOfAggregates` anymore?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: support `PartialMerge` aggregation mode [datafusion-comet]
Posted by "mbutrovich (via GitHub)" <gi...@apache.org>.
mbutrovich commented on PR #4003:
URL: https://github.com/apache/datafusion-comet/pull/4003#issuecomment-4331360054
Can you break out the streaming fallback logic for #4115? It would make it easier to understand logic separation, testing, etc.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: support `PartialMerge` aggregation mode [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #4003:
URL: https://github.com/apache/datafusion-comet/pull/4003#issuecomment-4336605559
> Can you break out the streaming fallback logic for #4115? It would make it easier to understand logic separation, testing, etc.
I can yes, I'm not sure if I can easily reproduce the streaming use case without this PR, let me try
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: support `PartialMerge` aggregation mode [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #4003:
URL: https://github.com/apache/datafusion-comet/pull/4003#issuecomment-4355653046
need to regenerate for 4.2
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: support `PartialMerge` aggregation mode [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on PR #4003:
URL: https://github.com/apache/datafusion-comet/pull/4003#issuecomment-4306124701
This PR also fixes the test from https://github.com/apache/datafusion-comet/issues/1267
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org