viirya commented on code in PR #7222:
URL: https://github.com/apache/datafusion-comet/pull/7222#discussion_r4381371836
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -301,6 +301,90 @@ class CometAggregateSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
}
}
+ // Writes 125,000 rows over four values of g, with 62,500 distinct k that
each fall in one g,
+ // then runs `f` in each of the pools. The PartialMerge keyed by (g, k) has
to spill several runs
+ // of its 62,500 groups, and its input is sorted on g, a prefix of its keys,
so DataFusion runs
+ // it as an ordered final aggregate (issue #6880). Each g has 15,625
distinct k and the largest v
+ // with id % 4 = g, so a key emitted twice across a change of g shows up in
the answer.
+ private def withPartialMergeOverSortedPrefix(poolsMiB: Seq[Int])(f: String
=> Unit): Unit = {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ spark
+ .range(0, 125000, 1, 1)
+ .selectExpr("id % 4 AS g", "lpad(cast(id % 62500 AS STRING), 128, 'k')
AS k", "id AS v")
+ .write
+ .parquet(path + "/data")
+ spark.range(0, 4).selectExpr("id AS g").write.parquet(path + "/keys")
+ for (poolMiB <- poolsMiB) {
+ withSQLConf(
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ 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 -> (poolMiB /
2048.0).toString) {
+ withClue(s"With a $poolMiB MiB pool: ") {
+ f(path)
+ }
+ }
+ }
+ }
+ }
+
+ private val partialMergeOverSortedPrefixAnswer =
+ (0 until 4).map(g => Row(g.toLong, 124996L + g, 15625L))
+
+ test("PartialMerge over input sorted on a prefix of its keys spills (issue
#6880)") {
+ withPartialMergeOverSortedPrefix(Seq(2, 8)) { path =>
+ // The repartition on g satisfies the PartialMerge keyed by (g, k), so
the sort, the partial
+ // aggregate and the PartialMerge run in one native plan with no
exchange between them. The
+ // optimizer drops the sort under the aggregate unless EliminateSorts is
excluded.
+ withSQLConf(SQLConf.OPTIMIZER_EXCLUDED_RULES.key ->
EliminateSorts.ruleName) {
+ val df = spark.read
+ .parquet(path + "/data")
+ .repartition(1, col("g"))
+ .sortWithinPartitions("g")
+ .groupBy("g")
+ .agg(expr("max(v)"), count_distinct(col("k")))
+ .orderBy("g")
+ checkCometAnswer(df, partialMergeOverSortedPrefixAnswer)
+
+ assert(
+ spilledAggregates(df, PartialMerge).exists(_.child match {
+ case partial: CometHashAggregateExec =>
partial.child.isInstanceOf[CometSortExec]
+ case _ => false
+ }),
+ "Expected a native PartialMerge over a partial aggregate over a
sort:\n" +
+ df.queryExecution.executedPlan)
+ }
+ }
+ }
+
+ test("PartialMerge over a sort-merge join on a prefix of its keys spills
(issue #6880)") {
+ // The sorts under the join need more than a 2 or 8 MiB pool's fair share,
so this uses 16 MiB.
+ withPartialMergeOverSortedPrefix(Seq(16)) { path =>
+ // With the default rules, the join's hash partitioning on g satisfies
the PartialMerge keyed
+ // by (g, k), so the join, the partial aggregate and the PartialMerge
run in one native plan,
+ // and the PartialMerge sees its input sorted on g. AQE would turn the
join into a broadcast
+ // hash join, because the keys are tiny.
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = spark.read
+ .parquet(path + "/data")
+ .join(spark.read.parquet(path + "/keys"), "g")
+ .groupBy("g")
+ .agg(expr("max(v)"), count_distinct(col("k")))
+ .orderBy("g")
+ checkCometAnswer(df, partialMergeOverSortedPrefixAnswer)
+
+ val plan = df.queryExecution.executedPlan
+ assert(
+ collect(plan) { case join: CometSortMergeJoinExec => join }.nonEmpty,
Review Comment:
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.
##########
spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala:
##########
@@ -301,6 +301,90 @@ class CometAggregateSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
}
}
+ // Writes 125,000 rows over four values of g, with 62,500 distinct k that
each fall in one g,
+ // then runs `f` in each of the pools. The PartialMerge keyed by (g, k) has
to spill several runs
+ // of its 62,500 groups, and its input is sorted on g, a prefix of its keys,
so DataFusion runs
+ // it as an ordered final aggregate (issue #6880). Each g has 15,625
distinct k and the largest v
+ // with id % 4 = g, so a key emitted twice across a change of g shows up in
the answer.
+ private def withPartialMergeOverSortedPrefix(poolsMiB: Seq[Int])(f: String
=> Unit): Unit = {
+ withTempPath { dir =>
+ val path = dir.getCanonicalPath
+ spark
+ .range(0, 125000, 1, 1)
+ .selectExpr("id % 4 AS g", "lpad(cast(id % 62500 AS STRING), 128, 'k')
AS k", "id AS v")
+ .write
+ .parquet(path + "/data")
+ spark.range(0, 4).selectExpr("id AS g").write.parquet(path + "/keys")
+ for (poolMiB <- poolsMiB) {
+ withSQLConf(
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ 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 -> (poolMiB /
2048.0).toString) {
+ withClue(s"With a $poolMiB MiB pool: ") {
+ f(path)
+ }
+ }
+ }
+ }
+ }
+
+ private val partialMergeOverSortedPrefixAnswer =
+ (0 until 4).map(g => Row(g.toLong, 124996L + g, 15625L))
+
+ test("PartialMerge over input sorted on a prefix of its keys spills (issue
#6880)") {
+ withPartialMergeOverSortedPrefix(Seq(2, 8)) { path =>
+ // The repartition on g satisfies the PartialMerge keyed by (g, k), so
the sort, the partial
+ // aggregate and the PartialMerge run in one native plan with no
exchange between them. The
+ // optimizer drops the sort under the aggregate unless EliminateSorts is
excluded.
+ withSQLConf(SQLConf.OPTIMIZER_EXCLUDED_RULES.key ->
EliminateSorts.ruleName) {
+ val df = spark.read
+ .parquet(path + "/data")
+ .repartition(1, col("g"))
+ .sortWithinPartitions("g")
+ .groupBy("g")
+ .agg(expr("max(v)"), count_distinct(col("k")))
+ .orderBy("g")
+ checkCometAnswer(df, partialMergeOverSortedPrefixAnswer)
+
+ assert(
+ spilledAggregates(df, PartialMerge).exists(_.child match {
+ case partial: CometHashAggregateExec =>
partial.child.isInstanceOf[CometSortExec]
+ case _ => false
+ }),
+ "Expected a native PartialMerge over a partial aggregate over a
sort:\n" +
+ df.queryExecution.executedPlan)
+ }
+ }
+ }
+
+ test("PartialMerge over a sort-merge join on a prefix of its keys spills
(issue #6880)") {
+ // The sorts under the join need more than a 2 or 8 MiB pool's fair share,
so this uses 16 MiB.
+ withPartialMergeOverSortedPrefix(Seq(16)) { path =>
+ // With the default rules, the join's hash partitioning on g satisfies
the PartialMerge keyed
+ // by (g, k), so the join, the partial aggregate and the PartialMerge
run in one native plan,
+ // and the PartialMerge sees its input sorted on g. AQE would turn the
join into a broadcast
+ // hash join, because the keys are tiny.
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
Review Comment:
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?
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
