peter-toth commented on code in PR #59950:
URL: https://github.com/apache/spark/pull/59950#discussion_r4423516843
##########
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:
**Finding 2.** Without the shuffle, the subquery reads the build side with
one task per partition key (or per bucket), not one per split. I measured it
with this test's `items` table, two keys and eight inserts each, so 16 splits:
- Master reads the scan with 16 tasks, then shuffles.
- This PR reads it with 2 tasks, each going through 8 splits.
The shuffle this saves is small, because the partial aggregate already
dedups each split. The pruned scan waits for this subquery, so a build side
with many splits per key now delays the main query.
It reaches bucketed V1 tables too. With a parquet build table bucketed 4
ways on the key, master's subquery turns off the bucketed scan and shuffles.
Here it keeps the bucketed scan, in 4 tasks, with no shuffle.
A plain `SELECT DISTINCT id` on the same table gets the same 2 tasks on
master. So this matches how aggregates already use a keyed scan. The only
switch for it is `spark.sql.sources.v2.bucketing.enabled` or
`spark.sql.sources.bucketing.enabled`. Those also turn off storage-partitioned
and bucketed joins. Users can see the cost, so please state it in the
description.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -226,6 +228,67 @@ 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 reads only the build side's key column, which is a
partition key of
Review Comment:
**Finding 5.** The subquery also reads `price`, for the `i.price > 40.0`
filter in the build plan. Its scan is `BatchScan testcat.ns.items[id, price]`.
It only groups on, and outputs, the key.
```suggestion
// The subquery groups only on the build side's key column, which
is a partition key of
```
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -226,6 +228,67 @@ 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 reads only 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
+ assert(collect(subPlan) { case s: ShuffleExchangeLike => s
}.isEmpty,
+ s"the DPP subquery should not shuffle:\n$subPlan")
+ assert(collect(subPlan) { case g: GroupPartitionsExec => g
}.nonEmpty,
+ s"the DPP subquery should group the keyed
partitions:\n$subPlan")
+ case other =>
+ fail(s"expected a SubqueryExec, got $other")
+ }
+
+ // The outer query keeps its storage-partitioned join.
Review Comment:
**Finding 4.** `checkAnswer` passes whether or not the subquery's keys reach
the scan, since the join filters the rows anyway. The fixture and the query are
the ones in `data source partitioning + dynamic partition filtering` above.
That test runs with AQE off only, and it checks the pruning with
`assertFilteredPartitions(scans, Seq(5, 5), Seq(0, 2))`. Asserting the same
here also covers the AQE arm, and with it `PlanAdaptiveDynamicPruningFilters`.
It passes in both arms of both runtime-filter suites:
```suggestion
// DPP keeps only the item_id = 1 splits of `purchases`.
assertFilteredPartitions(collectScans(plan), Seq(5, 5), Seq(0, 2))
// The outer query keeps its storage-partitioned join.
```
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -226,6 +228,67 @@ trait KeyGroupedPartitioningRuntimeFilterTests extends
KeyGroupedPartitioningSui
}
}
+ test("SPARK-60741: dynamic partition filtering keeps the build side's keyed
partitioning") {
Review Comment:
**Finding 3.** This test covers a keyed V2 scan. The change also reaches a
build side hash-partitioned on the key. That path has no test. Its common case
is a bucketed V1 table. Something like this in `DynamicPartitionPruningV1Suite`
passes on this PR in both AQE modes. On master it fails with a shuffle in the
subquery:
```scala
test("SPARK-60741: the DPP subquery keeps a bucketed build side's
partitioning") {
withSQLConf(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",
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
withTable("bucketed_dim", "partitioned_fact") {
spark.range(100).selectExpr("id", "id % 10 AS k")
.write.bucketBy(4,
"k").format(tableFormat).saveAsTable("bucketed_dim")
spark.range(100).selectExpr("id AS v", "id % 10 AS p")
.write.partitionBy("p").format(tableFormat).saveAsTable("partitioned_fact")
val df = sql("SELECT sum(f.v) FROM partitioned_fact f JOIN
bucketed_dim d " +
"ON f.p = d.k WHERE d.id < 5")
checkAnswer(df, Row(2350))
// The subquery groups on the bucket column, so the bucketed scan is
kept, and no shuffle
// is needed.
val subqueries =
collectDynamicPruningExpressions(df.queryExecution.executedPlan)
.collect { case i: InSubqueryExec => i.plan }
assert(subqueries.nonEmpty)
subqueries.foreach {
case s: SubqueryExec =>
assert(collect(s.child) { case e: ShuffleExchangeLike => e
}.isEmpty, s.child)
assert(collect(s.child) { case f: FileSourceScanExec =>
f.bucketedScan } === Seq(true),
s.child)
case other => fail(s"expected a SubqueryExec, got $other")
}
}
}
}
```
It needs `ShuffleExchangeLike` in the suite's `exchange` import.
--
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]
