peter-toth commented on code in PR #59789:
URL: https://github.com/apache/spark/pull/59789#discussion_r4402755073
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala:
##########
@@ -225,6 +225,12 @@ class DataSourceV2Strategy(session: SparkSession) extends
Strategy with Predicat
// rewritten it to TrueLiteral by the time we get here.
val runtimeFilters = dynamicFilters ++ scalarSubqueryFilters
+ // Decide storage-partitioned join here, once, so a rebuild of the
planned node (for instance
+ // the copy PlanAdaptiveDynamicPruningFilters makes at stage creation)
does not read the
+ // config again and drop the keyed layout the join was planned with
(SPARK-60260).
+ val keyGroupedPartitioning = relation.keyGroupedPartitioning
+ .filter(_ => session.sessionState.conf.v2BucketingEnabled)
Review Comment:
**Finding 10.** This makes a comment in `DSv2PlanMergingSuite.scala:199-201`
stale. It says only `DataSourceV2ScanExecBase.outputPartitioning` reads
`V2_BUCKETING_ENABLED`, while the read is now here. For example:
```scala
// V2_BUCKETING_ENABLED is what makes the preserved report reach the
physical plan, which the
// last assertion checks. DataSourceV2Strategy hands the scan its
key-grouped partitioning
// only when it is on. AQE is off so that `executedPlan` is the plan
the assertions can walk.
```
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -389,6 +389,104 @@ trait KeyGroupedPartitioningRuntimeFilterTests extends
KeyGroupedPartitioningSui
checkAnswer(df, Seq(Row(1L, 42.0f), Row(1L, 44.0f)))
}
}
+
+ test("SPARK-60260: a storage-partitioned join keeps the planned keyed layout
when v2 bucketing " +
+ "is disabled after planning") {
+ // reportedKeyedPartitioning used to read
spark.sql.sources.v2.bucketing.enabled at execution,
+ // so turning it off between planning and execution dropped the keyed
layout on the scan that
+ // PlanAdaptiveDynamicPruningFilters rebuilds at stage creation, and the
sort-merge join then
+ // lined partitions up wrong. The decision is frozen at planning now.
Covered across three
+ // shapes, each planned with bucketing on and then executed with it off,
AQE on.
+ val letters = Seq("a", "b", "c", "d", "e", "f", "g")
+
+ // Plan under bucketing on (so the keyed layout and DPP are decided here),
assert the shape,
+ // then execute with bucketing off and check both the rows and that the
rebuilt keyed scan kept
+ // its KeyedPartitioning. checkToRDD = false runs only the committed plan,
not a new one.
+ def check(filter: String, expected: Seq[Row], groupPartitions: Boolean,
shuffle: Boolean)
+ : Unit = {
+ val df = sql(s"SELECT i.id, i.name, p.price FROM testcat.ns.$items i " +
+ s"JOIN testcat.ns.$purchases p ON i.id = p.item_id WHERE $filter")
+ val plan = df.queryExecution.executedPlan
+ assert(collect(plan) { case j: SortMergeJoinExec => j }.nonEmpty,
+ s"expected a sort-merge join:\n$plan")
+ assert(collectScans(plan).exists(_.runtimeFilters.nonEmpty),
+ s"expected dynamic partition pruning on the keyed scan:\n$plan")
+ assert(collectGroupPartitions(plan).nonEmpty == groupPartitions,
+ s"expected GroupPartitionsExec present == $groupPartitions:\n$plan")
+ assert(collectShuffles(plan).nonEmpty == shuffle,
+ s"expected a shuffle present == $shuffle:\n$plan")
+
+ withSQLConf(SQLConf.V2_BUCKETING_ENABLED.key -> "false") {
+ // Without checkToRDD, so only the committed plan runs, not a newly
planned one.
+ QueryTest.checkAnswer(df, expected, checkToRDD = false)
+ // Re-collected, because under AQE executedPlan only reaches the
rebuilt scans once the
+ // query has run. Without this fix the rebuilt keyed scan reports
UnknownPartitioning here.
+ val executed = df.queryExecution.executedPlan
+ assert(
+
collectScans(executed).exists(_.outputPartitioning.isInstanceOf[KeyedPartitioning]),
+ s"expected the rebuilt scan to keep its keyed layout:\n$executed")
Review Comment:
**Finding 4.** `exists` passes on master in shapes 1 and 2. Both tables are
partitioned there, and the `purchases` scan, which DPP does not rebuild, keeps
the `KeyedPartitioning` it computed at planning. Only shape 3 makes it fail,
since `purchases` is not partitioned there. A `forall` over every scan, as I
first suggested, fails on that unpartitioned `purchases`, which is presumably
why this became `exists`. The rebuilt scan is the one with the runtime filter:
```suggestion
// The scan with the DPP filter is the one
PlanAdaptiveDynamicPruningFilters rebuilds.
val rebuilt =
collectScans(executed).filter(_.runtimeFilters.nonEmpty)
assert(rebuilt.nonEmpty &&
rebuilt.forall(_.outputPartitioning.isInstanceOf[KeyedPartitioning]),
s"expected the rebuilt scan to keep its keyed layout:\n$executed")
```
With it, the test passes here in both suites. On master, with the row check
taken out, it fails in all six cases, while `exists` fails only in shape 3.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBase.scala:
##########
@@ -45,8 +45,14 @@ trait DataSourceV2ScanExecBase
def readerFactory: PartitionReaderFactory
- /** Optional partitioning expressions provided by the V2 data sources,
through
- * `SupportsReportPartitioning` */
+ /**
+ * Optional partitioning expressions provided by the V2 data sources, through
+ * `SupportsReportPartitioning`. `None` means either the source reported no
key, or
+ * storage-partitioned join is off for this scan: the planner gates it on
Review Comment:
**Finding 11.** `V2ScanPartitioningAndOrdering` also leaves the field `None`
when a key cannot be resolved, or when no key survives column pruning:
```suggestion
* `SupportsReportPartitioning`. `None` means the source reported no key,
* `V2ScanPartitioningAndOrdering` dropped it (a key it cannot resolve, or
none left after
* pruning), or storage-partitioned join is off for this scan: the planner
gates it on
```
--
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]
