peter-toth commented on code in PR #60030: URL: https://github.com/apache/spark/pull/60030#discussion_r4456682981
########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala: ########## @@ -46,6 +46,12 @@ import org.apache.spark.util.ArrayImplicits._ * post-scan `Filter` too. * @param appendsCheckedColumn Whether the readers mark the rows they checked in a boolean column, * the last of `output`, see `StorageFilterPlanning.checkedColumn`. + * @param plannedInputPartitions Keeps the input partitions once the node plans them. The copies + * later rules make of the node share it, so they report the layout Review Comment: Right, the copies still read the live conf. That's SPARK-60260, and #59789 fixes it by deciding the flag in `DataSourceV2Strategy`, so I left it to that PR. The comments now claim only the partitions, and the description points to #59789: [474ef79](https://github.com/apache/spark/pull/60030/commits/474ef79bb89401903919eefd1d8495443a7b6843). I checked your recipe with this PR and #59789 together: it passes in both suites, with AQE on and off. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala: ########## @@ -109,7 +116,7 @@ case class BatchScanExec( keyGroupedPartitioning.map(_.filter(_.references.subsetOf(outputSet))) @transient override lazy val inputPartitions: Seq[InputPartition] = - batch.planInputPartitions().toImmutableArraySeq + plannedInputPartitions.getOrPlan(scan)(batch.planInputPartitions().toImmutableArraySeq) Review Comment: Thanks, done in [c016753](https://github.com/apache/spark/pull/60030/commits/c016753d73373f5bdd2a3a0cc65eb9f97f730179), in its final shape in [474ef79](https://github.com/apache/spark/pull/60030/commits/474ef79bb89401903919eefd1d8495443a7b6843). The `PlannedScan` keeps the `Batch` too, and `batch` comes from it. So a copy reads the planned partitions with a reader factory from the `Batch` that planned them. A pushed runtime filter still replans through a new `Batch`, as before, and the `PlannedScan` doc now says so: [c5480e7](https://github.com/apache/spark/pull/60030/commits/c5480e7f08915a0a5e709c15b0dac581b8d6603e). ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala: ########## @@ -181,3 +188,24 @@ case class BatchScanExec( s"BatchScan ${table.name()}".trim } } + +/** + * The input partitions of a [[BatchScanExec]], planned once, by the first node that asks for them. + * Each newly built node gets its own instance, and the copies later rules make of the node share + * that instance. Only the nodes over the `Scan` that planned the partitions use them. A copy given + * another `Scan` plans its own. + */ +class PlannedInputPartitions { + private var planned: Option[(Scan, Seq[InputPartition])] = None + + def getOrPlan(scan: Scan)(plan: => Seq[InputPartition]): Seq[InputPartition] = synchronized { + planned match { + case Some((plannedScan, partitions)) if plannedScan eq scan => partitions + case Some(_) => plan Review Comment: Thanks. In [474ef79](https://github.com/apache/spark/pull/60030/commits/474ef79bb89401903919eefd1d8495443a7b6843), `BatchScanExec` takes a `PlannedScan` in place of its `Scan`, and the `PlannedScan` holds the `Scan` with the `Batch` and the partitions planned from it. So a node over another `Scan` needs another `PlannedScan`, and no per-`Scan` map is needed. ########## sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala: ########## @@ -389,6 +389,95 @@ trait KeyGroupedPartitioningRuntimeFilterTests extends KeyGroupedPartitioningSui checkAnswer(df, Seq(Row(1L, 42.0f), Row(1L, 44.0f))) } } + + 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 these predicates. + 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).map(i => + s"($i, $i.0, cast('2020-01-01' as timestamp))").mkString(", ")) + + def bound(agg: String): String = + s"(SELECT $agg(item_id) FROM testcat.ns.$purchases WHERE price > 4.0)" + Seq(true, false).foreach { aqe => + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) { + Seq(s"i.id = ${bound("max")}" -> Seq(7), s"i.id >= ${bound("min")}" -> (5 to 7)).foreach { + case (filter, ids) => + val df = sql(s"SELECT i.id, p.price FROM testcat.ns.$items i " + + s"JOIN testcat.ns.$purchases p ON i.id = p.item_id WHERE $filter") + checkAnswer(df, ids.map(i => Row(i.toLong, i.toFloat))) + val plan = df.queryExecution.executedPlan + assert(collectShuffles(plan).isEmpty, + s"aqe=$aqe, $filter: the join should stay storage-partitioned:\n$plan") + } + } + } + } + } + + test("SPARK-60752: a nested DPP filter on a DPP subquery's build side does not change the " + Review Comment: You're right, thanks. I ran both cases on branch-4.1 and branch-4.0 with AQE off: - `pushPartValues=false` fails with `List(7, 3)`. - A fully pruned build side fails with `List(7, 0)`. I listed the affected releases in the description, and added 4.1.4 and 4.0.5 to the JIRA. The nested test now covers the fully pruned build side too: [c016753](https://github.com/apache/spark/pull/60030/commits/c016753d73373f5bdd2a3a0cc65eb9f97f730179). ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala: ########## @@ -55,7 +61,8 @@ case class BatchScanExec( @transient table: Table, keyGroupedPartitioning: Option[Seq[Expression]] = None, storageFilters: Seq[Expression] = Nil, - appendsCheckedColumn: Boolean = false + appendsCheckedColumn: Boolean = false, + @transient plannedInputPartitions: PlannedInputPartitions = new PlannedInputPartitions Review Comment: Thanks, good point about the released lines. The cleanest shape we found is to hold the `Scan` in the `PlannedScan`, so a node over another `Scan` needs another `PlannedScan`. That changes the signature, so we can only do it on master, which moves to a new major version, and on branch-4.x, since 4.4.0 is a feature release and `BatchScanExec` changed at feature releases before (3.4.0, 3.5.0, 4.2.0). Master has it now: [474ef79](https://github.com/apache/spark/pull/60030/commits/474ef79bb89401903919eefd1d8495443a7b6843). The backports to branch-4.3 and branch-4.2 will keep the `scan` field and hand the `PlannedScan` over in `makeCopy` and `clone`, as you suggest, so their signatures stay. branch-4.1 and branch-4.0 need a different fix, since they pad to `spjParams`. I'll open separate PRs for those. ########## sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala: ########## @@ -389,6 +389,95 @@ trait KeyGroupedPartitioningRuntimeFilterTests extends KeyGroupedPartitioningSui checkAnswer(df, Seq(Row(1L, 42.0f), Row(1L, 44.0f))) } } + + 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 these predicates. + 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).map(i => + s"($i, $i.0, cast('2020-01-01' as timestamp))").mkString(", ")) + + def bound(agg: String): String = + s"(SELECT $agg(item_id) FROM testcat.ns.$purchases WHERE price > 4.0)" + Seq(true, false).foreach { aqe => + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) { + Seq(s"i.id = ${bound("max")}" -> Seq(7), s"i.id >= ${bound("min")}" -> (5 to 7)).foreach { + case (filter, ids) => + val df = sql(s"SELECT i.id, p.price FROM testcat.ns.$items i " + + s"JOIN testcat.ns.$purchases p ON i.id = p.item_id WHERE $filter") + checkAnswer(df, ids.map(i => Row(i.toLong, i.toFloat))) + val plan = df.queryExecution.executedPlan + assert(collectShuffles(plan).isEmpty, + s"aqe=$aqe, $filter: the join should stay storage-partitioned:\n$plan") + } + } + } + } + } + + 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. + val dim = "dim" + 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(2)(s"($i, $i.0, cast('2020-01-01' as timestamp))")).mkString(", ")) + 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(true, false).foreach { aqe => + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) { + val df = sql(s"SELECT i.id, p.price FROM testcat.ns.$items i JOIN " + + s"(SELECT /*+ BROADCAST(d) */ p.item_id, p.price FROM testcat.ns.$purchases p " + + s"JOIN testcat.ns.$dim d ON p.item_id = d.k WHERE d.flag = 1) p ON i.id = p.item_id") + checkAnswer(df, (5 to 7).flatMap(i => Seq.fill(2)(Row(i.toLong, i.toFloat)))) + val plan = df.queryExecution.executedPlan + assert(collectShuffles(plan).isEmpty, + s"aqe=$aqe: the join should stay storage-partitioned:\n$plan") + } + } + } + } + + test("SPARK-60752: copies of a BatchScanExec share the input partitions it planned") { + createTable(items, itemsColumns, Array(identity("id"))) + sql(s"INSERT INTO testcat.ns.$items VALUES " + + "(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " + + "(2, 'bb', 10.0, cast('2020-01-01' as timestamp))") + def plannedScan(): BatchScanExec = + collectScans(sql(s"SELECT * FROM testcat.ns.$items").queryExecution.executedPlan).head + val scan = plannedScan() + // A copy, like the ones runtime filter rewrites make, uses the partitions the node planned. + assert(scan.copy(runtimeFilters = Nil).inputPartitions eq scan.inputPartitions) + // A copy over another `Scan` plans its own. + assert(!(scan.copy(scan = plannedScan().scan).inputPartitions eq scan.inputPartitions)) Review Comment: Thanks. A node over another `Scan` now has another `PlannedScan`, so nothing can take the holder over ([474ef79](https://github.com/apache/spark/pull/60030/commits/474ef79bb89401903919eefd1d8495443a7b6843)). The test moved to `SparkPlanSuite`, over a `Scan` whose `toBatch` returns a new `Batch`. It checks that the copies through `makeCopy`, `clone` and `copy` share the `Batch` and the partitions. ########## sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala: ########## @@ -389,6 +389,95 @@ trait KeyGroupedPartitioningRuntimeFilterTests extends KeyGroupedPartitioningSui checkAnswer(df, Seq(Row(1L, 42.0f), Row(1L, 44.0f))) } } + + 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 these predicates. + 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).map(i => + s"($i, $i.0, cast('2020-01-01' as timestamp))").mkString(", ")) + + def bound(agg: String): String = + s"(SELECT $agg(item_id) FROM testcat.ns.$purchases WHERE price > 4.0)" + Seq(true, false).foreach { aqe => + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) { + Seq(s"i.id = ${bound("max")}" -> Seq(7), s"i.id >= ${bound("min")}" -> (5 to 7)).foreach { Review Comment: Thanks, dropped `>=` in [c016753](https://github.com/apache/spark/pull/60030/commits/c016753d73373f5bdd2a3a0cc65eb9f97f730179). ########## sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala: ########## @@ -389,6 +389,95 @@ trait KeyGroupedPartitioningRuntimeFilterTests extends KeyGroupedPartitioningSui checkAnswer(df, Seq(Row(1L, 42.0f), Row(1L, 44.0f))) } } + + 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 these predicates. + 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).map(i => + s"($i, $i.0, cast('2020-01-01' as timestamp))").mkString(", ")) + + def bound(agg: String): String = + s"(SELECT $agg(item_id) FROM testcat.ns.$purchases WHERE price > 4.0)" + Seq(true, false).foreach { aqe => + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) { + Seq(s"i.id = ${bound("max")}" -> Seq(7), s"i.id >= ${bound("min")}" -> (5 to 7)).foreach { + case (filter, ids) => + val df = sql(s"SELECT i.id, p.price FROM testcat.ns.$items i " + + s"JOIN testcat.ns.$purchases p ON i.id = p.item_id WHERE $filter") + checkAnswer(df, ids.map(i => Row(i.toLong, i.toFloat))) + val plan = df.queryExecution.executedPlan + assert(collectShuffles(plan).isEmpty, + s"aqe=$aqe, $filter: the join should stay storage-partitioned:\n$plan") + } + } + } + } + } + + 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. + val dim = "dim" + 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(2)(s"($i, $i.0, cast('2020-01-01' as timestamp))")).mkString(", ")) + 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(true, false).foreach { aqe => + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) { + val df = sql(s"SELECT i.id, p.price FROM testcat.ns.$items i JOIN " + + s"(SELECT /*+ BROADCAST(d) */ p.item_id, p.price FROM testcat.ns.$purchases p " + + s"JOIN testcat.ns.$dim d ON p.item_id = d.k WHERE d.flag = 1) p ON i.id = p.item_id") + checkAnswer(df, (5 to 7).flatMap(i => Seq.fill(2)(Row(i.toLong, i.toFloat)))) + val plan = df.queryExecution.executedPlan + assert(collectShuffles(plan).isEmpty, + s"aqe=$aqe: the join should stay storage-partitioned:\n$plan") + } + } + } + } + + test("SPARK-60752: copies of a BatchScanExec share the input partitions it planned") { Review Comment: Thanks, it moved to `SparkPlanSuite` in [c016753](https://github.com/apache/spark/pull/60030/commits/c016753d73373f5bdd2a3a0cc65eb9f97f730179), so it runs once. It needs no table there. -- 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]
