peter-toth commented on code in PR #60030:
URL: https://github.com/apache/spark/pull/60030#discussion_r4465546078
##########
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:
Thanks, I rebased onto master. The SPARK-60618 test now gives the rebuilt
node a new `PlannedScan`, and the comment above it says why:
[b1cfddb](https://github.com/apache/spark/pull/60030/commits/b1cfddb02bd6a749c355115df833e374ab22e123).
##########
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:
Thanks, the check runs with AQE off only now (item 15), where the fully
pruned case keeps the join:
[b1cfddb](https://github.com/apache/spark/pull/60030/commits/b1cfddb02bd6a749c355115df833e374ab22e123).
##########
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:
Thanks, agreed. `checkStoragePartitionedJoin` runs with AQE off only now,
and its doc says why:
[b1cfddb](https://github.com/apache/spark/pull/60030/commits/b1cfddb02bd6a749c355115df833e374ab22e123).
That matches what I measured: the AQE-on runs pass without the fix 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:
Thanks, done in
[b1cfddb](https://github.com/apache/spark/pull/60030/commits/b1cfddb02bd6a749c355115df833e374ab22e123).
The `PlannedScan` doc and the description say that a rule that rebuilds a node
over the same `Scan` passes on its `plannedScan` or uses `copy`. A new
`PlannedScan` is for a node that is planned again, or one over another `Scan`.
##########
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:
Thanks, done in
[b1cfddb](https://github.com/apache/spark/pull/60030/commits/b1cfddb02bd6a749c355115df833e374ab22e123).
The `PlannedScan` builds the reader factory once, and `readerFactory` reads it
from there.
##########
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:
Thanks. I tried it in
[b1cfddb](https://github.com/apache/spark/pull/60030/commits/b1cfddb02bd6a749c355115df833e374ab22e123),
but I'd rather keep the extractor, and
[c09f3b9](https://github.com/apache/spark/pull/60030/commits/c09f3b9ce21b2cf6d3daadcc71da65e883e5fe46)
restores it with the compact patterns. It gives code that targets this line a
compact form, and `BatchScanExec.scan` stays for code that builds against
several releases.
--
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]
