peter-toth commented on code in PR #59522: URL: https://github.com/apache/spark/pull/59522#discussion_r4470165520
########## sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/PackV2ScanPartitions.scala: ########## @@ -0,0 +1,98 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.exchange + +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.expressions.SortOrder +import org.apache.spark.sql.catalyst.plans.physical.{Distribution, UnspecifiedDistribution} +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.execution._ +import org.apache.spark.sql.execution.aggregate.BaseAggregateExec +import org.apache.spark.sql.execution.command.DataWritingCommandExec +import org.apache.spark.sql.execution.datasources.WriteFilesExec +import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, V2ScanPartitionPacking, V2TableWriteExec} +import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec} +import org.apache.spark.sql.internal.SQLConf + +/** + * Pack V2 scan inputs after distribution and ordering requirements have been planned. + * A shuffle or broadcast exchange consumes neither input property. Otherwise, propagate + * consumers' requirements conservatively through known operators, including a broadcast + * join's streamed side. Keep unknown operators and explicit coalesces unchanged. + * + * Packing happens inside BatchScanExec, before partial aggregation and UnsafeRow conversion. + * No keyed or ordered output is claimed, and scans used by GroupPartitionsExec stay unchanged. + * + * Like `DisableUnnecessaryBucketedScan`, this is a separate rule that runs after + * `EnsureRequirements` has planned every exchange. Under AQE, `requiredDistribution` is the + * distribution the final plan must satisfy, the same one `EnsureRequirements` is given. + */ +case class PackV2ScanPartitions( + requiredDistribution: Option[Distribution] = None) extends Rule[SparkPlan] { + + override def apply(plan: SparkPlan): SparkPlan = { + if (!conf.v2ScanPartitionPackingEnabled) { + return plan + } + lazy val options = V2ScanPartitionPacking( + conf.v2ScanPartitionPackingMaxPartitionBytes, + conf.v2ScanPartitionPackingOpenCostInBytes, + conf.v2ScanPartitionPackingMinPartitionNum.getOrElse( + conf.getConf(SQLConf.LEAF_NODE_DEFAULT_PARALLELISM).getOrElse( + SparkSession.active.sparkContext.defaultParallelism)), + conf.v2ScanPartitionPackingMaxPartitionNum) + + def rewrite(node: SparkPlan, needsPartitioning: Boolean, needsOrdering: Boolean): SparkPlan = { + node match { + case scan: BatchScanExec if !needsPartitioning && !needsOrdering && + scan.partitionPacking.isEmpty && + options.pack(scan.inputPartitions).size < scan.inputPartitions.size => Review Comment: **Finding 9.** This guard calls `planInputPartitions()` on every candidate scan while planning, and `withPartitionPacking` then creates its reader factory. On master, an unkeyed scan whose `columnarSupportMode()` is `SUPPORTED` or `UNSUPPORTED` plans its inputs only at execution (SPARK-44505). With a runtime filter it plans only the filtered scan, since `replanWithRuntimeFilters` takes the original inputs by name (`PushDownUtils.scala:295`). I measured it with an unkeyed source of 8 sized inputs that supports runtime filtering, joined to a small table so that DPP keeps 2 of them. These are its `planInputPartitions()` calls: | | while planning | in total | |---|---|---| | packing off, AQE off or on | 0 | 1 (2 inputs) | | packing on, AQE off | 1 (8 inputs) | 2 | | packing on, AQE on | 1 (8 inputs) | 3 (8, 8, 2) | So with packing on, the scan plans the whole table that DPP is there to skip, and AQE plans it once more when it re-plans. For a connector like Iceberg that is file planning over the whole table. The rows are the same in every run. Nothing above an eligible scan uses its partitioning or ordering, so the rule could mark every eligible scan and leave the grouping to `inputRDD`. It already packs the filtered inputs there, and keeps them as they are when nothing merges. The snapshot then has no planning-time decision to keep consistent, so a packed copy can plan its inputs like any other copy. I tried this by dropping the guard and the snapshot's inputs, keeping only the options in `V2ScanPartitionPackingInfo`: - the query above plans 0 times while planning and once in total, AQE on and off; - 274 of the 279 tests in the packing suites and the three packing re-runs pass; - the 5 others check what the guard and the snapshot do. 3 expect no packing on a scan that can't merge anything, or the same input instance across copies. 2 read a claim on a scan nothing consumes ("non-clustered distribution: single partition", "SPARK-59995: a scan's output ordering drops sort orders that hold a partition transform"), like the tests this PR already adjusts with `copy(partitionPacking = None)`. The cost is an explain line, and withdrawn claims nothing reads on an eligible scan that doesn't merge. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/PackV2ScanPartitions.scala: ########## @@ -0,0 +1,98 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.exchange + +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.expressions.SortOrder +import org.apache.spark.sql.catalyst.plans.physical.{Distribution, UnspecifiedDistribution} +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.execution._ +import org.apache.spark.sql.execution.aggregate.BaseAggregateExec +import org.apache.spark.sql.execution.command.DataWritingCommandExec +import org.apache.spark.sql.execution.datasources.WriteFilesExec +import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, V2ScanPartitionPacking, V2TableWriteExec} +import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec} +import org.apache.spark.sql.internal.SQLConf + +/** + * Pack V2 scan inputs after distribution and ordering requirements have been planned. + * A shuffle or broadcast exchange consumes neither input property. Otherwise, propagate + * consumers' requirements conservatively through known operators, including a broadcast + * join's streamed side. Keep unknown operators and explicit coalesces unchanged. + * + * Packing happens inside BatchScanExec, before partial aggregation and UnsafeRow conversion. + * No keyed or ordered output is claimed, and scans used by GroupPartitionsExec stay unchanged. + * + * Like `DisableUnnecessaryBucketedScan`, this is a separate rule that runs after + * `EnsureRequirements` has planned every exchange. Under AQE, `requiredDistribution` is the + * distribution the final plan must satisfy, the same one `EnsureRequirements` is given. + */ +case class PackV2ScanPartitions( + requiredDistribution: Option[Distribution] = None) extends Rule[SparkPlan] { + + override def apply(plan: SparkPlan): SparkPlan = { + if (!conf.v2ScanPartitionPackingEnabled) { + return plan + } + lazy val options = V2ScanPartitionPacking( + conf.v2ScanPartitionPackingMaxPartitionBytes, + conf.v2ScanPartitionPackingOpenCostInBytes, + conf.v2ScanPartitionPackingMinPartitionNum.getOrElse( + conf.getConf(SQLConf.LEAF_NODE_DEFAULT_PARALLELISM).getOrElse( + SparkSession.active.sparkContext.defaultParallelism)), + conf.v2ScanPartitionPackingMaxPartitionNum) + + def rewrite(node: SparkPlan, needsPartitioning: Boolean, needsOrdering: Boolean): SparkPlan = { + node match { + case scan: BatchScanExec if !needsPartitioning && !needsOrdering && + scan.partitionPacking.isEmpty && Review Comment: **Finding 11.** `FilePartition` doesn't implement `HasPartitionStatistics`, so a V2 file scan always takes the count fallback in `pack` (`V2ScanPartitionPacking.scala:47`). A session that sets `maxPartitionNum` for an unsized connector then also re-chunks its Parquet and ORC scans by count. With `spark.sql.sources.useV1SourceList` empty, a Parquet scan of 8 files planned as 8 `FilePartition`s runs as 2 tasks of 40,000 rows with `maxPartitionNum=2`, ignoring `spark.sql.files.maxPartitionBytes`. `spark.sql.files.*` already sizes those scans, as it does V1 file scans, which this option leaves alone. Leaving `FileScan` out keeps them that way: ```suggestion scan.partitionPacking.isEmpty && !scan.scan.isInstanceOf[FileScan] && ``` with `FileScan` in the `datasources.v2` import. The same scan then runs 8 tasks. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala: ########## @@ -425,16 +426,36 @@ class QueryExecution( !hasShuffle(executedPlan)) { this } else { - val session = SparkSession.getOrCloneSessionWithConfigsOff( - sparkSession, Seq(SQLConf.LOCAL_PIPELINED_SHUFFLE_ENABLED)) - new QueryExecution(session, commandExecuted, mode = mode, - shuffleCleanupModeOpt = shuffleCleanupModeOpt, - refreshPhaseEnabled = refreshPhaseEnabled, analyzerOpt = Some(analyzer)) + withConfigsOff(Seq(SQLConf.LOCAL_PIPELINED_SHUFFLE_ENABLED)) } } private[sql] def withRegularShuffle: QueryExecution = lazyRegularShuffle.get + /** + * A checkpoint reports its plan's output partitioning to the queries that read it later. Its + * root has no consumer when it is planned, so scan packing would withdraw a keyed partitioning + * those queries could use. Plan it with packing off, as well as with regular shuffles. + */ + private val lazyForCheckpoint = LazyTry { + if (!sparkSession.sessionState.conf.v2ScanPartitionPackingEnabled) { Review Comment: **Finding 10.** With packing on, this builds a second `QueryExecution` for every checkpoint, also when no scan in the plan gets packed, e.g. for `spark.range(10).toDF()`. `withRegularShuffle` re-plans only when it has to, and this path has no such check. The second execution also runs its own refresh phase (`lazyTableVersionsRefreshed`, `:296`). So when a Dataset has already run, its checkpoint reloads versioned tables at their latest version and can hold other rows than the ones the Dataset returned. That part is from reading the code, since the in-memory test catalog doesn't pin versions. Re-planning only when a scan was packed keeps every other checkpoint on this execution, as on master. It plans this execution first, which a checkpoint on master does anyway: ```suggestion def hasPackedScan(plan: SparkPlan): Boolean = plan match { case s: BatchScanExec => s.partitionPacking.isDefined case a: AdaptiveSparkPlanExec => hasPackedScan(a.executedPlan) case q: QueryStageExec => hasPackedScan(q.plan) case other => other.children.exists(hasPackedScan) } if (!sparkSession.sessionState.conf.v2ScanPartitionPackingEnabled || !hasPackedScan(executedPlan)) { ``` It also needs `BatchScanExec` in the `datasources.v2` import. With it, `PackV2ScanPartitionsPlanSuite` and the 9 `DatasetSuite` checkpoint tests pass, including "checkpointed keyed scans keep their layout for later joins", which still re-plans. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala: ########## @@ -2644,6 +2644,61 @@ object SQLConf { .booleanConf .createWithDefault(true) + val V2_SCAN_PARTITION_PACKING_ENABLED = Review Comment: **Finding 12.** These five entries sit between `V2_BUCKETING_ENABLED` and `V2_BUCKETING_PUSH_PART_VALUES_ENABLED`, in the middle of the twelve `spark.sql.sources.v2.bucketing.*` entries, although their doc says the feature is independent of V2 bucketing. The repo's `AGENTS.md` asks that a new member go with the code it belongs with rather than split a section. After `V2_BUCKETING_PRESERVE_ORDERING_ON_COALESCE_ENABLED` (`:2833`), or next to `FILES_MAX_PARTITION_NUM` (`:3418`), their file-source counterpart, would keep both groups whole. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/PackV2ScanPartitions.scala: ########## @@ -0,0 +1,98 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.exchange + +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.expressions.SortOrder +import org.apache.spark.sql.catalyst.plans.physical.{Distribution, UnspecifiedDistribution} +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.execution._ +import org.apache.spark.sql.execution.aggregate.BaseAggregateExec +import org.apache.spark.sql.execution.command.DataWritingCommandExec +import org.apache.spark.sql.execution.datasources.WriteFilesExec +import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, V2ScanPartitionPacking, V2TableWriteExec} +import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, BroadcastNestedLoopJoinExec} +import org.apache.spark.sql.internal.SQLConf + +/** + * Pack V2 scan inputs after distribution and ordering requirements have been planned. + * A shuffle or broadcast exchange consumes neither input property. Otherwise, propagate + * consumers' requirements conservatively through known operators, including a broadcast + * join's streamed side. Keep unknown operators and explicit coalesces unchanged. + * + * Packing happens inside BatchScanExec, before partial aggregation and UnsafeRow conversion. + * No keyed or ordered output is claimed, and scans used by GroupPartitionsExec stay unchanged. + * + * Like `DisableUnnecessaryBucketedScan`, this is a separate rule that runs after + * `EnsureRequirements` has planned every exchange. Under AQE, `requiredDistribution` is the + * distribution the final plan must satisfy, the same one `EnsureRequirements` is given. + */ +case class PackV2ScanPartitions( + requiredDistribution: Option[Distribution] = None) extends Rule[SparkPlan] { + + override def apply(plan: SparkPlan): SparkPlan = { + if (!conf.v2ScanPartitionPackingEnabled) { + return plan + } + lazy val options = V2ScanPartitionPacking( + conf.v2ScanPartitionPackingMaxPartitionBytes, + conf.v2ScanPartitionPackingOpenCostInBytes, + conf.v2ScanPartitionPackingMinPartitionNum.getOrElse( + conf.getConf(SQLConf.LEAF_NODE_DEFAULT_PARALLELISM).getOrElse( + SparkSession.active.sparkContext.defaultParallelism)), Review Comment: **Finding 13.** This repeats `SparkSession.leafNodeDefaultParallelism` (`classic/SparkSession.scala:913`), which `FilePartition.maxSplitBytes` uses for the same default: ```suggestion conf.v2ScanPartitionPackingMinPartitionNum.getOrElse( SparkSession.active.leafNodeDefaultParallelism), ``` with `org.apache.spark.sql.classic.SparkSession` imported in place of `org.apache.spark.sql.SparkSession`, and the `SQLConf` import dropped. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanPartitionPacking.scala: ########## @@ -0,0 +1,118 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.datasources.v2 + +import scala.collection.mutable + +import org.apache.spark.rdd.{PartitionCoalescer, PartitionGroup, RDD} +import org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning +import org.apache.spark.sql.connector.read.{Batch, HasPartitionStatistics, InputPartition, PartitionReaderFactory, Scan} + +/** Immutable sizing policy captured when a scan is prepared, before runtime filtering. */ +case class V2ScanPartitionPacking( + maxPartitionBytes: Long, + openCostInBytes: Long, + minPartitionNum: Int, + maxPartitionNum: Option[Int]) { + + /** + * Group adjacent inputs by estimated size, keeping an oversized input on its own. Missing, + * negative or mixed statistics use the explicit count cap, or leave inputs unchanged. + * The cap is only a fallback: it never overrides the size limit for inputs reporting bytes. + */ + def pack(partitions: Seq[InputPartition]): Seq[Seq[Int]] = { + if (partitions.isEmpty) return Seq.empty + val costs = partitions.map { + case p: HasPartitionStatistics => + val bytes = p.sizeInBytes() + Option.when(bytes.isPresent && bytes.getAsLong >= 0)( + add(bytes.getAsLong, openCostInBytes)) + case _ => None + } + if (costs.exists(_.isEmpty)) { + val count = maxPartitionNum.getOrElse(partitions.size).min(partitions.size) + // Balanced consecutive ranges, without integer overflow or empty groups. + (0 until count).map { index => + val start = (index.toLong * partitions.size / count).toInt + val end = ((index.toLong + 1) * partitions.size / count).toInt + start until end + } + } else { + val sizes = costs.map(_.get) + val total = sizes.foldLeft(0L)(add) + val bytesPerTask = total / minPartitionNum + (if (total % minPartitionNum == 0) 0 else 1) + val target = math.min(maxPartitionBytes, math.max(openCostInBytes, bytesPerTask)) + val groups = mutable.ArrayBuffer.empty[Seq[Int]] + val current = mutable.ArrayBuffer.empty[Int] + var currentBytes = 0L + def finish(): Unit = { + if (current.nonEmpty) { + groups += current.toVector + current.clear() + currentBytes = 0L + } + } + sizes.zipWithIndex.foreach { case (size, index) => + if (current.nonEmpty && (size > target || currentBytes > target - size)) finish() + current += index + currentBytes = add(currentBytes, size) + if (size >= target && size > 0) finish() + } + finish() + groups.toSeq + } + } + + private def add(left: Long, right: Long): Long = + if (left > Long.MaxValue - right) Long.MaxValue else left + right +} + +/** + * Driver-side snapshot of a scan that has already planned its inputs. Adding task sizing reuses + * this metadata and factory. Only the sizing policy contributes to scan equality and display. + */ +class V2ScanPartitionPackingInfo( + val options: V2ScanPartitionPacking, + @transient val scan: Scan, + @transient val batch: Batch, + @transient val inputPartitions: Seq[InputPartition], + @transient val readerFactory: PartitionReaderFactory, + @transient val keyedPartitioning: Option[KeyedPartitioning]) extends Serializable { + override def toString: String = options.toString +} + +/** Use the size planner's groups rather than the RDD coalescer's count-based grouping. */ +private[v2] class V2ScanPartitionCoalescer(@transient groups: Seq[Seq[Int]]) Review Comment: **Finding 14.** `GroupedPartitionCoalescer` (`GroupPartitionsExec.scala:869`), in the same package, already coalesces an RDD into given `Seq[Seq[Int]]` groups, with the most common preferred host of each group. This class differs only in how it breaks ties across groups, which `PackV2ScanPartitionsReadSuite` pins (`worker-a`, `worker-b`, `worker-a`, `worker-b`). Adding that tie-break to `GroupedPartitionCoalescer` and using it here would leave one class, and the SPJ groups would get the same balancing. -- 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]
