Skip to content

Commit 424c31a

Browse files
peterxclipeterxcli
authored
fix: enable FIRST/LAST partial merge (#5041)
1 parent 5627ab8 commit 424c31a

6 files changed

Lines changed: 55 additions & 94 deletions

File tree

‎spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala‎

Lines changed: 1 addition & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ import scala.collection.mutable.ListBuffer
2323

2424
import org.apache.spark.sql.SparkSession
2525
import org.apache.spark.sql.catalyst.expressions.{Divide, DoubleLiteral, EqualNullSafe, EqualTo, Expression, FloatLiteral, GreaterThan, GreaterThanOrEqual, KnownFloatingPointNormalized, LessThan, LessThanOrEqual, NamedExpression, Remainder}
26-
import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, Final, First, Last, Partial, PartialMerge}
26+
import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, Final, Partial, PartialMerge}
2727
import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero
2828
import org.apache.spark.sql.catalyst.rules.Rule
2929
import org.apache.spark.sql.catalyst.trees.TreeNodeTag
@@ -1087,11 +1087,6 @@ case class CometExecRule(session: SparkSession)
10871087
if (consumesBuffers &&
10881088
!QueryPlanSerde.allAggsSupportMixedExecution(agg.aggregateExpressions) &&
10891089
!canAggregateBeConverted(agg, consumerMode)) {
1090-
// This pass deliberately records the consumer diagnostic. Once the Partial is tagged
1091-
// below, the consumer can be skipped because its child is no longer native and may not
1092-
// reach doConvert, which normally records the same reason.
1093-
unsupportedPartialMergeFallbackReason(agg).foreach(reason =>
1094-
withFallbackReason(agg, reason))
10951090
findPartialAggInPlan(agg.child).foreach { partial =>
10961091
// Only tag if the Partial would otherwise have been converted. If the Partial itself
10971092
// cannot be converted (e.g. an incompatible input type or a map-typed grouping key),
@@ -1212,11 +1207,6 @@ case class CometExecRule(session: SparkSession)
12121207
expectedMode == PartialMerge && modes.toSet == Set(Partial, PartialMerge)
12131208
if (!mixedPartialMerge && modes != Seq(expectedMode)) return false
12141209

1215-
// FIRST/LAST cannot merge native partial states in a distinct-aggregate rewrite. Predict
1216-
// this refusal before an earlier exchange materializes incompatible buffers for other
1217-
// functions in the same aggregate (for example percentile). Mirror doConvert's restriction.
1218-
if (unsupportedPartialMergeFallbackReason(agg).nonEmpty) return false
1219-
12201210
// Only Partial binds input attributes; Final and PartialMerge consume intermediate buffers.
12211211
// Mixed distinct-aggregate stages need the same per-expression binding as doConvert.
12221212
if (!aggregateExpressions.forall { e =>
@@ -1237,19 +1227,4 @@ case class CometExecRule(session: SparkSession)
12371227
}
12381228
}
12391229

1240-
private def unsupportedPartialMergeFallbackReason(agg: BaseAggregateExec): Option[String] = {
1241-
val unsupportedMerges = agg.aggregateExpressions.filter { expression =>
1242-
expression.mode == PartialMerge &&
1243-
(expression.aggregateFunction.isInstanceOf[First] ||
1244-
expression.aggregateFunction.isInstanceOf[Last])
1245-
}
1246-
if (unsupportedMerges.nonEmpty) {
1247-
Some(
1248-
"PartialMerge not supported for aggregates: " +
1249-
unsupportedMerges.map(_.aggregateFunction.prettyName).mkString(", "))
1250-
} else {
1251-
None
1252-
}
1253-
}
1254-
12551230
}

‎spark/src/main/scala/org/apache/spark/sql/comet/operators.scala‎

Lines changed: 1 addition & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ import org.apache.spark.internal.Logging
3131
import org.apache.spark.rdd.RDD
3232
import org.apache.spark.sql.catalyst.InternalRow
3333
import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeSet, Expression, ExpressionSet, Generator, NamedExpression, SortOrder}
34-
import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateMode, CollectList, CollectSet, Final, First, Last, Partial, PartialMerge, Percentile}
34+
import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateMode, CollectList, CollectSet, Final, Partial, PartialMerge, Percentile}
3535
import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide}
3636
import org.apache.spark.sql.catalyst.plans._
3737
import org.apache.spark.sql.catalyst.plans.physical._
@@ -1770,25 +1770,6 @@ trait CometBaseAggregate {
17701770
}
17711771
}
17721772

1773-
// FIRST/LAST are order-dependent: in PartialMerge mode, DataFusion's hash
1774-
// table may process rows in a different order than Spark's. CollectSet is
1775-
// handled separately (floating-point compat in CometCollectSet; streaming
1776-
// in ShimCometStreaming.isStreamingPlan).
1777-
// https://github.com/apache/datafusion-comet/issues/4131
1778-
if (hasPartialMerge) {
1779-
val unsupportedAggs = aggregateExpressions.filter { a =>
1780-
a.mode == PartialMerge && (a.aggregateFunction.isInstanceOf[First] ||
1781-
a.aggregateFunction.isInstanceOf[Last])
1782-
}
1783-
if (unsupportedAggs.nonEmpty) {
1784-
withFallbackReason(
1785-
aggregate,
1786-
"PartialMerge not supported for aggregates: " +
1787-
unsupportedAggs.map(_.aggregateFunction.prettyName).mkString(", "))
1788-
return None
1789-
}
1790-
}
1791-
17921773
// Per-expression binding: Partial expressions bind to child output,
17931774
// PartialMerge/Final expressions do not (native planner handles their input).
17941775
val output = child.output

‎spark/src/test/resources/sql-tests/expressions/aggregate/partial_merge.sql‎

Lines changed: 13 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -329,52 +329,48 @@ SELECT grp, count(DISTINCT IF(i % 2 = 0, NULL, i)), sum(i)
329329
FROM pm_basic GROUP BY grp ORDER BY grp
330330

331331
-- ############################################################
332-
-- FALLBACK cases
332+
-- FIRST/LAST cases
333333
-- ############################################################
334334

335-
-- FIRST/LAST aggregates in PartialMerge mode are order-dependent and
336-
-- DataFusion's hash table may process rows in a different order than Spark's.
337-
-- See https://github.com/apache/datafusion-comet/issues/4131
338-
339335
-- ============================================================
340-
-- fallback: first + distinct count triggers PartialMerge on first
336+
-- first + distinct count triggers PartialMerge on first
341337
-- ============================================================
342338

343-
query expect_fallback(PartialMerge not supported for aggregates: first)
339+
query
344340
SELECT first(i), count(DISTINCT i) FROM pm_basic
345341

346-
query expect_fallback(PartialMerge not supported for aggregates: first)
342+
query
347343
SELECT grp, first(i), count(DISTINCT i) FROM pm_basic GROUP BY grp ORDER BY grp
348344

349345
-- ============================================================
350-
-- fallback: last + distinct count triggers PartialMerge on last
346+
-- last + distinct count triggers PartialMerge on last
351347
-- ============================================================
352348

353-
query expect_fallback(PartialMerge not supported for aggregates: last)
349+
query
354350
SELECT last(i), count(DISTINCT i) FROM pm_basic
355351

356-
query expect_fallback(PartialMerge not supported for aggregates: last)
352+
query
357353
SELECT grp, last(i), count(DISTINCT i) FROM pm_basic GROUP BY grp ORDER BY grp
358354

359355
-- ============================================================
360-
-- fallback: first and last together with distinct
356+
-- first and last together with distinct
361357
-- ============================================================
362358

363-
query expect_fallback(PartialMerge not supported for aggregates: first, last)
359+
query
364360
SELECT first(i), last(i), count(DISTINCT i) FROM pm_basic
365361

366-
query expect_fallback(PartialMerge not supported for aggregates: first, last)
362+
query
367363
SELECT grp, first(i), last(i), count(DISTINCT i), sum(i)
368364
FROM pm_basic GROUP BY grp ORDER BY grp
369365

370366
-- ============================================================
371-
-- fallback: first/last IGNORE NULLS with distinct
367+
-- first/last IGNORE NULLS with distinct
372368
-- ============================================================
373369

374-
query expect_fallback(PartialMerge not supported for aggregates: first)
370+
query
375371
SELECT grp, first(i) IGNORE NULLS, count(DISTINCT i)
376372
FROM pm_nulls GROUP BY grp ORDER BY grp
377373

378-
query expect_fallback(PartialMerge not supported for aggregates: last)
374+
query
379375
SELECT grp, last(i) IGNORE NULLS, count(DISTINCT i)
380376
FROM pm_nulls GROUP BY grp ORDER BY grp

‎spark/src/test/resources/sql-tests/expressions/aggregate/partial_merge_first_last.sql‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
-- Licensed to the Apache Software Foundation (ASF) under one
2+
-- or more contributor license agreements. See the NOTICE file
3+
-- distributed with this work for additional information
4+
-- regarding copyright ownership. The ASF licenses this file
5+
-- to you under the Apache License, Version 2.0 (the
6+
-- "License"); you may not use this file except in compliance
7+
-- with the License. You may obtain a copy of the License at
8+
--
9+
-- http://www.apache.org/licenses/LICENSE-2.0
10+
--
11+
-- Unless required by applicable law or agreed to in writing,
12+
-- software distributed under the License is distributed on an
13+
-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
-- KIND, either express or implied. See the License for the
15+
-- specific language governing permissions and limitations
16+
-- under the License.
17+
18+
-- Exercise FIRST/LAST partial-state merging across multiple batches.
19+
-- https://github.com/apache/datafusion-comet/issues/4131
20+
-- Config: spark.comet.batchSize=128
21+
-- Config: spark.sql.adaptive.coalescePartitions.enabled=true
22+
-- Config: parquet.enable.dictionary=false
23+
24+
statement
25+
CREATE TABLE pm_first_last(i int, grp int) USING parquet
26+
27+
statement
28+
INSERT INTO pm_first_last
29+
SELECT CAST(id AS int), CAST(id % 100 AS int) FROM range(10000)
30+
31+
-- Hash aggregation does not preserve input order. Use a value constant within
32+
-- each group so FIRST/LAST agree regardless of the engines' processing order.
33+
query
34+
SELECT grp, first(grp), last(grp), count(DISTINCT i)
35+
FROM pm_first_last GROUP BY grp ORDER BY grp

‎spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala‎

Lines changed: 2 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -303,7 +303,7 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
303303
}
304304
}
305305

306-
test("unsupported PartialMerge preserves percentile buffers with local Comet shuffle") {
306+
test("disabled FIRST/LAST preserves percentile buffers with local Comet shuffle") {
307307
for {
308308
adaptive <- Seq(false, true)
309309
percentile <- Seq("percentile", "percentile_approx")
@@ -312,6 +312,7 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
312312
withSQLConf(
313313
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString,
314314
SQLConf.SHUFFLE_PARTITIONS.key -> "4",
315+
s"spark.comet.expression.${mergeFunction.capitalize}.enabled" -> "false",
315316
CometConf.COMET_SHUFFLE_MODE.key -> "native") {
316317
val query = spark
317318
.range(0, 18, 1, 4)
@@ -959,34 +960,6 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
959960
}
960961
}
961962

962-
// FIRST/LAST are order-dependent aggregates whose merge result depends on hash table
963-
// processing order. In PartialMerge mode, DataFusion's hash table may process rows
964-
// in a different order than Spark's, so we fall back to Spark for correctness.
965-
// https://github.com/apache/datafusion-comet/issues/4131
966-
test("partialMerge - FIRST/LAST with distinct aggregates falls back") {
967-
val numValues = 10000
968-
Seq(100).foreach { numGroups =>
969-
Seq(128).foreach { batchSize =>
970-
withSQLConf(
971-
SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",
972-
CometConf.COMET_BATCH_SIZE.key -> batchSize.toString) {
973-
withParquetTable(
974-
(0 until numValues).map(i => (i, Random.nextInt() % numGroups)),
975-
"tbl",
976-
false) {
977-
withView("v") {
978-
sql("CREATE TEMP VIEW v AS SELECT _1, _2 FROM tbl ORDER BY _1")
979-
checkSparkAnswerAndFallbackReason(
980-
"SELECT _2, FIRST(_1), LAST(_1), COUNT(DISTINCT _1)" +
981-
" FROM v GROUP BY _2 ORDER BY _2",
982-
"PartialMerge not supported for aggregates: first, last")
983-
}
984-
}
985-
}
986-
}
987-
}
988-
}
989-
990963
test("partialMerge - cnt distinct + sum") {
991964
withTempDir(dir => {
992965
withSQLConf("spark.comet.enabled" -> "false") {

‎spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -476,13 +476,14 @@ class CometCelebornShufflePlanningSuite extends CometTestBase {
476476
percentile <- Seq("percentile", "percentile_approx")
477477
mergeFunction <- Seq("first", "last")
478478
} {
479-
test(s"unsupported $mergeFunction merge preserves $percentile buffers with AQE=$adaptive") {
479+
test(s"disabled $mergeFunction preserves $percentile buffers with AQE=$adaptive") {
480480
manager.withPlanningSupport(CelebornNativeShufflePlanningSupport()) {
481481
withSQLConf(
482482
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString,
483483
SQLConf.SHUFFLE_PARTITIONS.key -> "4",
484+
s"spark.comet.expression.${mergeFunction.capitalize}.enabled" -> "false",
484485
CometConf.COMET_SHUFFLE_MODE.key -> "native") {
485-
// FIRST/LAST cannot merge natively. Tag the incompatible percentile producer
486+
// Disable FIRST/LAST to exercise fallback. Tag the incompatible percentile producer
486487
// before the first DISTINCT exchange is materialized, not just at the later
487488
// exchange that falls back. Its grouping key makes FIRST/LAST deterministic.
488489
val query = spark

0 commit comments

Comments
 (0)