andygrove commented on code in PR #7222:
URL: https://github.com/apache/datafusion-comet/pull/7222#discussion_r4387944923
##########
native/core/src/execution/planner/aggregate.rs:
##########
@@ -247,6 +247,35 @@ impl OperatorBuilder for HashAggBuilder {
})
.collect();
+ // DataFusion 55 runs Single over input sorted on some of its keys
with the legacy
+ // GroupedHashAggregateStream, whose spill can ask for more than its
share of the pool
+ // (#6880). Its Final runs an OrderedFinalAggregateStream there
instead, which spills, and
+ // MergeAsPartial evaluates to its state under Final too. Two things
make the task succeed
+ // there: the new stream's spill path, and the fair pool exempting
every consumer named
+ // OrderedFinalAggregateStream from its fair limit
(`limited_to_its_share`), the limit the
+ // legacy stream's spill failed against. So this PartialMerge is
checked only against the
+ // pool and Spark's grant.
+ // Final reads each state from the columns after the keys and spills
its groups in the
+ // input's types, so plan the PartialMerge as Final only when its
input is its keys
+ // followed by native states. A PartialMerge over Spark's binary
collect buffers keeps
+ // Single, but that input arrives through a ScanExec with no ordering,
so it is not
+ // expected to be sorted on a key.
+ // DataFusion 56 replaces the legacy stream with an
OrderedSingleAggregateStream, so
+ // recheck this on that upgrade (#6410).
Review Comment:
Done. The comment now names `ClusteredSingleAggregateStream` and
`ClusteredFinalAggregateStream` from apache/datafusion#24697. It says to drop
this switch and update the consumer name prefix in `spill_replay.rs` on the 56
upgrade, and that constant's doc notes the rename too. I'll add the item to
#6410.
##########
native/core/src/execution/planner/aggregate.rs:
##########
@@ -309,6 +338,37 @@ impl OperatorBuilder for HashAggBuilder {
}
}
+/// Whether a PartialMerge's input is its grouping keys, in order, followed by
the states of
+/// its MergeAsPartial aggregates in their native types, which DataFusion's
Final assumes. A
+/// Spark partial aggregate's collect buffers, for example, arrive as binary.
+fn input_is_keys_then_states(
+ agg: &datafusion_comet_proto::spark_operator::HashAggregate,
+ aggr_expr: &[Arc<AggregateFunctionExpr>],
+ input_schema: &SchemaRef,
+) -> Result<bool, ExecutionError> {
+ let num_keys = agg.grouping_exprs.len();
+ let keys_lead = agg.initial_input_buffer_offset as usize == num_keys
+ && agg.grouping_exprs.iter().enumerate().all(|(idx, expr)| {
+ matches!(&expr.expr_struct, Some(ExprStruct::Bound(bound)) if
bound.index as usize == idx)
+ });
+ if !keys_lead {
+ return Ok(false);
+ }
+ let mut input_types = input_schema
+ .fields()
+ .iter()
+ .skip(num_keys)
+ .map(|f| f.data_type());
+ for expr in aggr_expr {
+ for field in expr.state_fields()? {
+ if input_types.next() != Some(field.data_type()) {
+ return Ok(false);
+ }
+ }
+ }
+ Ok(input_types.next().is_none())
Review Comment:
Added three cases that expect `Single` over input sorted on `g`: the keys
swapped (`key(1), key(0)`), `initial_input_buffer_offset` of 1, and an extra
`int64` column after the states. Your two mutations now fail the test, as do
the offset check and the key-order check mutated on their own. The assert names
the case that caught it.
##########
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:
Done. The test now asserts that a 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:
Your reading of `JoinSelectionHelper` is right. The broadcast came from
`CometTestBase`, which sets `spark.sql.adaptive.autoBroadcastJoinThreshold` to
1g for the whole suite, so AQE's runtime check never fell back to -1. The test
now sets both thresholds to -1 and leaves AQE on. With AQE on it keeps the
sort-merge join, passes the shape check, and still fails with the issue's error
without the planner change.
--
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]
