dongjoon-hyun commented on code in PR #59950:
URL: https://github.com/apache/spark/pull/59950#discussion_r4449347529
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PlanDynamicPruningFilters.scala:
##########
@@ -125,7 +125,12 @@ case class PlanDynamicPruningFilters(sparkSession:
SparkSession) extends Rule[Sp
case None =>
val aliases = broadcastKeyIndices.map(idx =>
Alias(buildKeys(idx), buildKeys(idx).toString)())
- val aggregate = Aggregate(aliases, aliases, buildPlan)
+ // Alias the keys in a Project instead of making them the grouping
expressions.
+ // Grouping on plain attributes lets the aggregate keep the build
side's partitioning
+ // (e.g. KeyedPartitioning) instead of requiring a shuffle. See
+ // AliasAwareOutputExpression.
+ val groupingAttrs = aliases.map(_.toAttribute)
+ val aggregate = Aggregate(groupingAttrs, groupingAttrs,
Project(aliases, buildPlan))
Review Comment:
**1. [Low] The fallback subquery can now fail the planned-partitioning check
of `GroupPartitionsExec`**
Before this change, this subquery always had an `Exchange` above the partial
aggregate. Now a build side that reports an ungrouped `KeyedPartitioning`, i.e.
several splits per key, gets a `GroupPartitionsExec` here instead. That node
checks at execution that its child still reports the partitioning it was
planned over
([L353-L357](https://github.com/apache/spark/blob/e82e6c3d24734c00e80806a0c37feda2402cf782/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala#L353-L357)).
Two things together can break that:
- The build-side scan has a subquery in its runtime filters, for example a
nested DPP from a broadcast dimension joined to the build table, or a scalar
subquery on its partition key. The rule that plans that subquery after
`EnsureRequirements` rebuilds the `BatchScanExec`:
`PlanAdaptiveDynamicPruningFilters` or `ReuseAdaptiveSubquery` with AQE, and
`ReuseExchangeAndSubquery` without it.
- The rebuilt node plans its input partitions again, from a `Scan` that it
shares with the main query's copy of the build scan:
`QueryExecution.prepareExecutedPlan` clones the logical plan, and the clone
keeps the `Scan` instance. If the main query's scan has already pushed its
runtime filter into that `Scan` ([`PushDownUtils.scala`
L205](https://github.com/apache/spark/blob/e82e6c3d24734c00e80806a0c37feda2402cf782/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala#L205)),
the rebuilt node sees fewer partitions, and the query fails with
"GroupPartitionsExec's child no longer reports the partitioning it was planned
over".
This depends on timing, and I did not reproduce it. #60030 handles this kind
of rebuild by sharing the planned input partitions between the copies of a
`BatchScanExec`, so this PR adds a new place where the check can fail rather
than a new mechanism. Could this land after #60030, or could the new test add a
runtime filter on the keyed build side, so that the combination is covered? The
same applies to `PlanAdaptiveDynamicPruningFilters.scala` L128.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -226,6 +228,70 @@ trait KeyGroupedPartitioningRuntimeFilterTests extends
KeyGroupedPartitioningSui
}
}
+ test("SPARK-60741: dynamic partition filtering keeps the build side's keyed
partitioning") {
+ withSQLConf(
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true",
+ SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "false",
+ SQLConf.DYNAMIC_PARTITION_PRUNING_FALLBACK_FILTER_RATIO.key -> "10") {
+ val items_partitions = Array(identity("id"))
+ createTable(items, itemsColumns, items_partitions)
+ sql(s"INSERT INTO testcat.ns.$items VALUES " +
+ s"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " +
+ s"(1, 'aa', 41.0, cast('2020-01-15' as timestamp)), " +
+ s"(2, 'bb', 10.0, cast('2020-01-01' as timestamp)), " +
+ s"(2, 'bb', 10.5, cast('2020-01-01' as timestamp)), " +
+ s"(3, 'cc', 15.5, cast('2020-02-01' as timestamp))")
+
+ val purchases_partitions = Array(identity("item_id"))
+ createTable(purchases, purchasesColumns, purchases_partitions)
+ sql(s"INSERT INTO testcat.ns.$purchases VALUES " +
+ s"(1, 42.0, cast('2020-01-01' as timestamp)), " +
+ s"(1, 44.0, cast('2020-01-15' as timestamp)), " +
+ s"(1, 45.0, cast('2020-01-15' as timestamp)), " +
+ s"(2, 11.0, cast('2020-01-01' as timestamp)), " +
+ s"(3, 19.5, cast('2020-02-01' as timestamp))")
+
+ // The false arm exercises PlanDynamicPruningFilters and the true arm
+ // PlanAdaptiveDynamicPruningFilters.
+ Seq(false, true).foreach { aqeEnabled =>
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key ->
aqeEnabled.toString) {
+ val df = sql(s"SELECT sum(p.price) FROM testcat.ns.$items i, " +
+ s"testcat.ns.$purchases p WHERE i.id = p.item_id AND i.price >
40.0")
+ checkAnswer(df, Seq(Row(131)))
+
+ // The subquery groups only on the build side's key column, which is
a partition key of
+ // `items`, so its aggregate must keep the scan's keyed
partitioning. These keys repeat
+ // across splits, so keeping the partitioning means a
`GroupPartitionsExec` coalesces
+ // them, and no shuffle is needed.
+ val plan = df.queryExecution.executedPlan
+ val subqueries = collectScans(plan)
+ .flatMap(_.runtimeFilters)
+ .flatMap(_.collect { case d: DynamicPruningExpression => d })
+ .flatMap(_.collect { case i: InSubqueryExec => i.plan })
+ assert(subqueries.nonEmpty, s"expected a DPP subquery, got:\n$plan")
+ subqueries.foreach {
+ case subquery: SubqueryExec =>
+ val subPlan = subquery.child
Review Comment:
**3. [Nit] The AQE arm does not check that the subquery is adaptive**
The comment at L255-L256 says that the true arm exercises
`PlanAdaptiveDynamicPruningFilters`, but the assertions accept any
`SubqueryExec`. If AQE stopped applying to this query,
`InsertAdaptiveSparkPlan` would only log a warning and return the plan without
AQE, both arms would go through `PlanDynamicPruningFilters`, and the test would
still pass. `PlanDynamicPruningFilters` never builds an adaptive subquery, so
one assertion pins the arm, with `AdaptiveSparkPlanExec` added to the
`adaptive` import at L55:
```suggestion
val subPlan = subquery.child
assert(subPlan.isInstanceOf[AdaptiveSparkPlanExec] ==
aqeEnabled,
s"the DPP subquery should be adaptive only with
AQE:\n$subPlan")
```
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PlanDynamicPruningFilters.scala:
##########
@@ -125,7 +125,12 @@ case class PlanDynamicPruningFilters(sparkSession:
SparkSession) extends Rule[Sp
case None =>
val aliases = broadcastKeyIndices.map(idx =>
Alias(buildKeys(idx), buildKeys(idx).toString)())
- val aggregate = Aggregate(aliases, aliases, buildPlan)
+ // Alias the keys in a Project instead of making them the grouping
expressions.
+ // Grouping on plain attributes lets the aggregate keep the build
side's partitioning
+ // (e.g. KeyedPartitioning) instead of requiring a shuffle. See
+ // AliasAwareOutputExpression.
Review Comment:
**4. [Nit] The comment does not say why an aliased grouping key loses the
partitioning**
`AliasAwareOutputExpression` names the catalyst trait that only builds the
alias map. The partitioning is carried through the alias by
`PartitioningPreservingUnaryExecNode` in sql/core, which `ProjectExec` mixes
in. The comment also leaves out why the old shape lost it: a partial
aggregate's `resultExpressions` hold only the grouping attributes
(`AggUtils.planAggregateWithoutDistinct`), so an `Alias` among the grouping
expressions never reaches its `outputExpressions`. Without that, someone could
move the alias back into the grouping, and only the two new tests would notice
the shuffle. Maybe something like this, here and at
`PlanAdaptiveDynamicPruningFilters.scala` L123-L126:
```suggestion
// A partial aggregate outputs only its grouping attributes, so
an Alias among the
// grouping expressions hides the build side's partitioning
(e.g. KeyedPartitioning)
// from PartitioningPreservingUnaryExecNode, and
EnsureRequirements adds a shuffle.
// Alias the keys in a Project and group on its attributes
instead, the shape that
// PullOutGroupingExpressions gives complex grouping expressions.
```
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/PlanAdaptiveDynamicPruningFilters.scala:
##########
@@ -120,7 +120,12 @@ case class PlanAdaptiveDynamicPruningFilters(
case None =>
val aliases = indices.map(idx =>
Alias(buildKeys(idx), buildKeys(idx).toString)())
- val aggregate = Aggregate(aliases, aliases, buildPlan)
+ // Alias the keys in a Project instead of making them the grouping
expressions.
+ // Grouping on plain attributes lets the aggregate keep the build
side's partitioning
+ // (e.g. KeyedPartitioning) instead of requiring a shuffle. See
+ // AliasAwareOutputExpression.
+ val groupingAttrs = aliases.map(_.toAttribute)
+ val aggregate = Aggregate(groupingAttrs, groupingAttrs,
Project(aliases, buildPlan))
Review Comment:
**5. [Nit] The new block and its comment are duplicated in both rules**
These eight lines, the comment included, repeat
`PlanDynamicPruningFilters.scala` L126-L133 word for word. The block has
drifted before: SPARK-46946 had to change both copies, and SPARK-52974 brought
the non-AQE copy on par with this one. If a later change reaches only one copy,
the two paths plan different subqueries again, and only one of them keeps the
build side's partitioning. Could one helper build this logical plan for both
rules, for example next to
`DynamicPruningSubquery.usableBroadcastValueProjection`, which both paths
already use?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/PlanAdaptiveDynamicPruningFilters.scala:
##########
@@ -120,7 +120,12 @@ case class PlanAdaptiveDynamicPruningFilters(
case None =>
val aliases = indices.map(idx =>
Alias(buildKeys(idx), buildKeys(idx).toString)())
- val aggregate = Aggregate(aliases, aliases, buildPlan)
+ // Alias the keys in a Project instead of making them the grouping
expressions.
+ // Grouping on plain attributes lets the aggregate keep the build
side's partitioning
+ // (e.g. KeyedPartitioning) instead of requiring a shuffle. See
+ // AliasAwareOutputExpression.
+ val groupingAttrs = aliases.map(_.toAttribute)
+ val aggregate = Aggregate(groupingAttrs, groupingAttrs,
Project(aliases, buildPlan))
Review Comment:
**2. [Low] AQE can no longer split a skewed build-side join inside the
subquery**
When the build side is itself a shuffled join on the DPP key, for example
`fact f JOIN (SELECT d1.k FROM d1 JOIN d2 ON d1.k = d2.k WHERE d2.x < 5) t ON
f.p = t.k`, the `Project` now carries the join's hash partitioning to the
aggregate, and the subquery no longer has an `Exchange` of its own:
```
Before: HashAggregate(Final) <- Exchange <- HashAggregate(Partial) <-
SortMergeJoin
After: HashAggregate(Complete) <- Project <- SortMergeJoin
```
`OptimizeSkewedJoin` reads a split skewed partition with
`UnknownPartitioning`, which no longer satisfies the aggregate's
`ClusteredDistribution`. So `ValidateRequirements` fails, and the rule returns
the original plan
([L259-L275](https://github.com/apache/spark/blob/e82e6c3d24734c00e80806a0c37feda2402cf782/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/OptimizeSkewedJoin.scala#L259-L275)).
A local shuffle read, after AQE turns the join into a broadcast hash join, is
rejected the same way
([L239-L254](https://github.com/apache/spark/blob/e82e6c3d24734c00e80806a0c37feda2402cf782/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala#L239-L254)).
The pruned fact scan then waits on an unsplit skewed task that master avoided.
`spark.sql.adaptive.forceOptimizeSkewedJoin` does not help much:
`CombineAdjacentAggregation` has already merged the two aggregates, so the
shuffle that it adds moves the join output before any deduplication.
`SELECT DISTINCT d1.k FROM d1 JOIN d2 ON d1.k = d2.k` makes the same
trade-off on master, so this is not a bug. Could the code comment or the PR
description mention it? Keeping the old shape when `buildPlan` contains a
shuffled join would also avoid 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: [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]
