dongjoon-hyun commented on code in PR #60030:
URL: https://github.com/apache/spark/pull/60030#discussion_r4462999307
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -389,6 +389,80 @@ trait KeyGroupedPartitioningRuntimeFilterTests extends
KeyGroupedPartitioningSui
checkAnswer(df, Seq(Row(1L, 42.0f), Row(1L, 44.0f)))
}
}
+
+ /**
+ * Runs `f` over `items` and `purchases`, both partitioned by the ids 1 to
7, with `rowsPerKey`
+ * rows in each `purchases` key, under the confs of a DPP subquery that
reuses no broadcast.
+ */
+ private def withDppSubqueryTables(rowsPerKey: Int)(f: => Unit): Unit = {
+ 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") {
+ createTable(items, itemsColumns, Array(identity("id")))
+ sql(s"INSERT INTO testcat.ns.$items VALUES " + (1 to 7).map(i =>
+ s"($i, 'x', 10.0, cast('2020-01-01' as timestamp))").mkString(", "))
+ createTable(purchases, purchasesColumns, Array(identity("item_id")))
+ sql(s"INSERT INTO testcat.ns.$purchases VALUES " + (1 to 7).flatMap(i =>
+ Seq.fill(rowsPerKey)(s"($i, $i.0, cast('2020-01-01' as
timestamp))")).mkString(", "))
+ f
+ }
+ }
+
+ /** Checks the rows of `query`, and that its join plans no shuffle, with AQE
on and off. */
+ private def checkStoragePartitionedJoin(query: String, expected: Seq[Row]):
Unit = {
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ val df = sql(query)
+ checkAnswer(df, expected)
+ val plan = df.queryExecution.executedPlan
+ assert(collectShuffles(plan).isEmpty,
+ s"aqe=$aqe: the join should stay storage-partitioned:\n$plan")
+ }
+ }
+ }
+
+ test("SPARK-60752: a DPP subquery's build side does not change the layout of
the scan it " +
+ "shares with the join") {
+ // A DPP subquery that does not reuse a broadcast plans its build side
again, over the same
+ // `Scan` as the join's `purchases` scan. A scalar subquery filter on
`item_id` is a runtime
+ // filter of both, and the subquery runs it first, which prunes the shared
`Scan`. The join's
+ // scan is a copy made after planning, when the scalar subquery is
replaced by a reused one. Had
+ // it planned its partitions again, it would report fewer of them than
`items`, and the join
+ // would fail with "Can't zip RDDs". Of the two suites that run this, only
the Catalyst one
+ // prunes on this predicate.
+ withDppSubqueryTables(rowsPerKey = 1) {
+ checkStoragePartitionedJoin(
+ s"SELECT i.id, p.price FROM testcat.ns.$items i JOIN
testcat.ns.$purchases p " +
+ s"ON i.id = p.item_id WHERE i.id = " +
+ s"(SELECT max(item_id) FROM testcat.ns.$purchases WHERE price >
4.0)",
+ Seq(Row(7L, 7.0f)))
+ }
+ }
+
+ test("SPARK-60752: a nested DPP filter on a DPP subquery's build side does
not change the " +
+ "layout of the scan it shares with the join") {
+ // As above, with a DPP filter from a broadcast dimension as the build
side's own runtime
+ // filter, which every fixture prunes on. Two rows per key give
`purchases` two splits per key,
+ // so a GroupPartitionsExec groups that side and checks the layout of the
copied scan. With
+ // `flag = 2` the build side's filter prunes every split.
+ val dim = "dim"
+ withDppSubqueryTables(rowsPerKey = 2) {
+ createTable(dim, Array(Column.create("k", LongType),
Column.create("flag", IntegerType)),
+ Array.empty)
+ sql(s"INSERT INTO testcat.ns.$dim VALUES " +
+ (1 to 7).map(i => s"($i, ${if (i >= 5) 1 else 0})").mkString(", "))
+ Seq(1 -> (5 to 7), 2 -> Nil).foreach { case (flag, ids) =>
Review Comment:
**14. [Low] With AQE on, the fully pruned case never runs the join**
The `flag = 2` iteration, which I suggested in item 4, checks nothing with
AQE on. No `dim` row has `flag = 2`, so the `BROADCAST(d)` stage materializes
with 0 rows before the result stage is created (`AdaptiveSparkPlanExec.scala`
L859-L860). `reOptimize` then runs `AQEPropagateEmptyRelation`, which reads
that row count (`AQEPropagateEmptyRelation.scala` L65-L66) and empties
`purchases JOIN dim` and then the outer join, so the final plan is an
`EmptyRelationExec`:
- no `items` or `purchases` scan runs, and neither does the DPP subquery or
a `GroupPartitionsExec`;
- `collectShuffles` only looks under a `SortMergeJoinExec` (L99-L107), so
the assertion at L420 holds trivially;
- `checkAnswer(df, Nil)` holds with or without the fix.
With AQE off the join stays, since the static `PropagateEmptyRelation`
treats only an empty `LocalRelation` as empty (`PropagateEmptyRelation.scala`
L51-L54), so the case is covered there. A `dim` row whose key matches no
`purchases` key, such as `(8, 2)`, would keep the broadcast non-empty with AQE
on too, and the expected rows stay `Nil`:
```scala
sql(s"INSERT INTO testcat.ns.$dim VALUES " +
((1 to 7).map(i => s"($i, ${if (i >= 5) 1 else 0})") :+ "(8,
2)").mkString(", "))
```
Running this case with AQE off only, and saying so in the comment, would
also do.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -389,6 +389,80 @@ trait KeyGroupedPartitioningRuntimeFilterTests extends
KeyGroupedPartitioningSui
checkAnswer(df, Seq(Row(1L, 42.0f), Row(1L, 44.0f)))
}
}
+
+ /**
+ * Runs `f` over `items` and `purchases`, both partitioned by the ids 1 to
7, with `rowsPerKey`
+ * rows in each `purchases` key, under the confs of a DPP subquery that
reuses no broadcast.
+ */
+ private def withDppSubqueryTables(rowsPerKey: Int)(f: => Unit): Unit = {
+ 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") {
+ createTable(items, itemsColumns, Array(identity("id")))
+ sql(s"INSERT INTO testcat.ns.$items VALUES " + (1 to 7).map(i =>
+ s"($i, 'x', 10.0, cast('2020-01-01' as timestamp))").mkString(", "))
+ createTable(purchases, purchasesColumns, Array(identity("item_id")))
+ sql(s"INSERT INTO testcat.ns.$purchases VALUES " + (1 to 7).flatMap(i =>
+ Seq.fill(rowsPerKey)(s"($i, $i.0, cast('2020-01-01' as
timestamp))")).mkString(", "))
+ f
+ }
+ }
+
+ /** Checks the rows of `query`, and that its join plans no shuffle, with AQE
on and off. */
+ private def checkStoragePartitionedJoin(query: String, expected: Seq[Row]):
Unit = {
+ Seq(true, false).foreach { aqe =>
Review Comment:
**15. [Low] The AQE-on half of the new tests passes without the fix**
This extends the last paragraph of item 7, which I only guessed for the
first test. Reading the AQE path, no AQE-on run of either test can fail on the
base commit:
- With AQE on, the copies of the join's `purchases` scan are made in
`optimizeQueryStage` by `PlanAdaptiveDynamicPruningFilters` and
`ReuseAdaptiveSubquery` (`AdaptiveSparkPlanExec.scala` L193-L195).
- `newResultQueryStage` then applies `postStageCreationRules`
(`AdaptiveSparkPlanExec.scala` L998-L1001), whose
`ApplyColumnarRulesAndInsertTransitions` (`AdaptiveSparkPlanExec.scala`
L210-L213) asks each copy `supportsColumnar`. The in-memory scans keep the
default `PARTITION_DEFINED` mode, so that ask plans the partitions of the copy
(`DataSourceV2ScanExecBase.scala` L195-L200).
- The DPP build side calls `filter()` on the shared `Scan` only when the
result stage runs its subqueries, which comes later. So even a copy that plans
its own partitions, as on the base commit, gets the full layout.
So the AQE-on runs, half of the executions of these tests, cannot catch a
regression of this PR, and neither can the AQE-off run of the first test in
`KeyGroupedPartitioningRuntimeFilterSuite`, as the comment at L433-L434 says.
The four older tests of this trait run with AQE off only (L162, L267, L313,
L355). Could `checkStoragePartitionedJoin` run with AQE off only, or say that
its AQE-on run only checks the rows over the AQE copies? The description says
the failure needs AQE off too.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala:
##########
@@ -181,3 +184,21 @@ case class BatchScanExec(
s"BatchScan ${table.name()}".trim
}
}
+
+/**
+ * A `Scan`, with the `Batch` and the input partitions a [[BatchScanExec]]
plans from it once, when
+ * a node first asks for them. The copies later rules make of the node share
it, so they read the
Review Comment:
**16. [Low] Nothing tells a rule that rebuilds the node to keep its
`PlannedScan`**
The description tells code that builds or matches the node to use `new
PlannedScan(scan)` and `case BatchScanExec(_, PlannedScan(scan), ...)`, and
this doc only says that the copies later rules make share the holder. Put
together, they suggest rebuilding a node like this:
```scala
case BatchScanExec(o, PlannedScan(s), rf, ord, t, k, sf, c) =>
BatchScanExec(o, new PlannedScan(s), f(rf), ord, t, k, sf, c)
```
That drops the planned partitions. Say such a rule runs after
`EnsureRequirements` and nothing asks the new holder before execution, e.g. an
injected `ColumnarRule.postColumnarTransitions` with AQE off. Then the first
ask comes at execution, after the DPP build side has pruned the shared `Scan`,
and with the tables and confs of the nested test that is the
`GroupPartitionsExec` INTERNAL_ERROR again, or `List(7, 1)` without a
`GroupPartitionsExec`. `b.copy(runtimeFilters = f(rf))`, or passing
`b.plannedScan`, keeps the sharing. A new holder is right only when the rebuilt
node should plan again, as in the SPARK-60618 test of item 13.
Nothing in this repo rebuilds the node this way, and the base commit planned
again on every rebuild too, so this is only about the guidance. One sentence
here and in the description would cover it, for example "a rule that rebuilds a
node over the same `Scan` passes its `plannedScan` or uses `copy`; `new
PlannedScan` is for a newly planned node or another `Scan`", with `ps @
PlannedScan(scan)` in the pattern example.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala:
##########
@@ -108,8 +112,7 @@ case class BatchScanExec(
@transient lazy val prunedKeyGroupedPartitioning: Option[Seq[Expression]] =
keyGroupedPartitioning.map(_.filter(_.references.subsetOf(outputSet)))
- @transient override lazy val inputPartitions: Seq[InputPartition] =
- batch.planInputPartitions().toImmutableArraySeq
+ override def inputPartitions: Seq[InputPartition] =
plannedScan.inputPartitions
Review Comment:
**13. [Medium] After merging master, CI fails the SPARK-60618 test that
rebuilds a scan with new partitions**
The latest run of this head merges it into master `2c03f1681fb` and fails
one test of `KeyGroupedPartitioningSuite`, `SPARK-60618: scan rebuilt with
different partition keys after alignment`, with `keyedpartitioning(id#1299588,
...) equaled keyedpartitioning(id#1299588, ...), and List(2, 3) did not equal
List(1, 2)`
([job](https://github.com/peter-toth/spark/actions/runs/41663003909/job/126096325517)).
That test came with SPARK-60618 (1a524e44b5a), after the base of this PR. It
changes the data of the in-memory scan and then builds
`leftScan.copy(runtimeFilters = Seq(Literal.TrueLiteral))`, expecting the copy
to plan the new splits ([master
L2136-L2143](https://github.com/apache/spark/blob/a475424a8c56adf84f23e8a6051a385dbb65c50b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala#L2136-L2143)).
With this line the copy reads the partitions the original node planned, so it
keeps the old layout and the `!=` assertion there fails. There is no textual
conflict, so GitHub still shows the PR as mergeable.
The test models "a controlled post-planning rewrite" that changes the split
set of the connector, so after a rebase it can give the copy a new holder, as
item 16 describes:
```scala
val rebuilt = leftScan.copy(plannedScan = new
PlannedScan(leftScan.scan),
runtimeFilters = Seq(Literal.TrueLiteral))
```
The comment above it would need the same change, and the suite would need to
import `PlannedScan`. I traced this from the code and the CI log and did not
run it.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/SparkPlanSuite.scala:
##########
@@ -321,6 +322,30 @@ class SparkPlanSuite extends SharedSparkSession {
assert(canonicalizedTypes == Seq(IntegerType, StringType),
s"Expected [IntegerType, StringType] but got $canonicalizedTypes")
}
+
+ test("SPARK-60752: copies of a BatchScanExec share its Batch and the
partitions it planned") {
+ // A `Scan` that returns a new `Batch` from each `toBatch`, which a
connector may do.
+ class NewBatchScan extends Scan {
+ override def readSchema(): StructType = new StructType()
+ override def toBatch: Batch = new Batch {
+ override def planInputPartitions(): Array[InputPartition] = Array(new
InputPartition {})
+ override def createReaderFactory(): PartitionReaderFactory =
+ throw new UnsupportedOperationException()
+ }
+ }
+ def sharesPlanning(copy: SparkPlan, scan: BatchScanExec): Boolean = copy
match {
Review Comment:
**19. [Nit] `sharesPlanning` takes the same `scan` in every call**
All three calls pass the `scan` of L342, which this parameter also shadows.
It looks like a leftover of c016753, where one call compared another node.
Moving the helper below L342 and closing over `scan` would drop the parameter
and fit the `makeCopy` assertion on one line:
```scala
val scan = BatchScanExec(Seq.empty, new PlannedScan(new NewBatchScan),
Seq.empty, table = null)
def sharesPlanning(copy: SparkPlan): Boolean = copy match {
case c: BatchScanExec =>
(c.batch eq scan.batch) && (c.inputPartitions eq
scan.inputPartitions)
case _ => false
}
// Copies through `makeCopy`, `clone` or `copy` share what the node
planned.
assert(sharesPlanning(scan.makeCopy(scan.productIterator.map(_.asInstanceOf[AnyRef]).toArray)))
assert(sharesPlanning(scan.clone()))
assert(sharesPlanning(scan.copy(runtimeFilters = Nil)))
```
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala:
##########
@@ -181,3 +184,21 @@ case class BatchScanExec(
s"BatchScan ${table.name()}".trim
}
}
+
+/**
+ * A `Scan`, with the `Batch` and the input partitions a [[BatchScanExec]]
plans from it once, when
+ * a node first asks for them. The copies later rules make of the node share
it, so they read the
+ * partitions the first node planned, with a reader factory from the `Batch`
that planned them.
+ * The partitions have to be the planned ones, because the `Scan` can be
shared with another node
+ * whose runtime filter prunes it first, e.g. the build side of a DPP subquery
(SPARK-60752). A
+ * pushed runtime filter replans them from a new `Batch`, but keeps the
planned `Batch`'s reader
+ * factory.
+ */
+final class PlannedScan(val scan: Scan) {
+ lazy val batch: Batch = scan.toBatch
+ lazy val inputPartitions: Seq[InputPartition] =
batch.planInputPartitions().toImmutableArraySeq
+}
+
+object PlannedScan {
Review Comment:
**18. [Nit] The positional test patterns could read the scan through
`BatchScanExec.scan`**
In this repo the extractor is used only by the 11 positional patterns of the
test suites, and they still depend on the 8 fields of `BatchScanExec`. The ORC
one in `OrcV2SchemaPruningSuite` has been rewritten for a change of those
fields 8 times since 2019, most recently by SPARK-60465 a week ago, while its
Parquet twin reads `scan.scan` (`ParquetSchemaPruningSuite.scala` L76) and has
not changed since. Since `scan` stays, `case b: BatchScanExec if
b.scan.isInstanceOf[OrcScan] => ...`, or `getBatchScanExec` in `AvroV2Suite`,
would find the same scans without the field count or the six import changes.
The branch-4.3 and branch-4.2 backports keep the `scan` field too, so code that
builds against several releases cannot use `PlannedScan(...)` anyway. Keeping
the extractor for other code is fine; this is only about the patterns here.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala:
##########
@@ -181,3 +184,21 @@ case class BatchScanExec(
s"BatchScan ${table.name()}".trim
}
}
+
+/**
+ * A `Scan`, with the `Batch` and the input partitions a [[BatchScanExec]]
plans from it once, when
+ * a node first asks for them. The copies later rules make of the node share
it, so they read the
+ * partitions the first node planned, with a reader factory from the `Batch`
that planned them.
+ * The partitions have to be the planned ones, because the `Scan` can be
shared with another node
+ * whose runtime filter prunes it first, e.g. the build side of a DPP subquery
(SPARK-60752). A
+ * pushed runtime filter replans them from a new `Batch`, but keeps the
planned `Batch`'s reader
+ * factory.
+ */
+final class PlannedScan(val scan: Scan) {
+ lazy val batch: Batch = scan.toBatch
+ lazy val inputPartitions: Seq[InputPartition] =
batch.planInputPartitions().toImmutableArraySeq
Review Comment:
**17. [Nit] `readerFactory` is still built once per node**
The holder now shares the `Batch` and the partitions, but `readerFactory`
(L128) is still a lazy val of each node, so each node that is asked calls
`createReaderFactory()` again on the same `Batch`. For a scan in the default
`PARTITION_DEFINED` mode, `DataSourceV2Strategy` asks the planner's node
through `supportsColumnar` (`DataSourceV2Strategy.scala` L242),
`RemoveRedundantProjects` or `ApplyColumnarRulesAndInsertTransitions` asks the
executed clone again, and a copy made by `ReuseExchangeAndSubquery` builds a
third one in `inputRDD`. That is cheap for most sources, but not for all:
- the state data source reader with `joinSide` left or right reads the
operator metadata and the state schema files on each call
(`StateScanBuilder.scala` L134-L155);
- the V2 file scans broadcast the Hadoop conf on each call
(`ParquetScan.scala` L173-L174), once a format is taken off
`spark.sql.sources.useV1SourceList`.
Since the factory comes from the `Batch` alone, a `lazy val readerFactory =
batch.createReaderFactory()` here, with `override def readerFactory =
plannedScan.readerFactory` in the node, would build it once per planned node.
The storage filter and the checked column reach the readers through
`DataSourceRDD`, not the factory, so sharing it changes nothing else, and it
matches "with a reader factory from the `Batch` that planned them" above. L128
is not changed by this PR, so this is only a suggestion.
--
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]
