rekbun commented on code in PR #59522: URL: https://github.com/apache/spark/pull/59522#discussion_r4498501134
########## 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: Done in dc626a7c49a: `leafNodeDefaultParallelism`, resolved at execution so scan equality doesn't depend on the cluster. ########## 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: Done in dc626a7c49a: merged into `GroupedPartitionCoalescer`, with the tie-break, `@transient` groups and no host for single-partition groups. SPJ locality over a scan is unchanged, which a test with per-input hosts checks. -- 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]
