rekbun commented on PR #59522:
URL: https://github.com/apache/spark/pull/59522#issuecomment-6531367774
Moved from the PR description: local synthetic measurements of size-based
packing against `COALESCE`, measured on 4372e37d.
#### Local synthetic measurements against COALESCE
These use `local[4]`, AQE off, a 4 MiB size target, zero opening cost,
suggested minimum four tasks, two warmup rounds and five measured rounds. The
row source reports estimated sizes of 16 bytes per row; no storage IO is
simulated. COALESCE is placed before partial aggregation. Complete aggregate
results match in every case.
| Layout | Original tasks / median ms | Size-packed tasks / median ms |
Count-matched COALESCE tasks / median ms | Bucket-count COALESCE tasks / median
ms |
| --- | --- | --- | --- | --- |
| 512 small inputs, eight keys | 512 / 1,085 | 4 / 57 | 4 / 60 | 8 / 67 |
| 64 larger inputs, one key | 64 / 228 | 16 / 165 | 16 / 188 | 1 / 610 |
| 32 small + four oversized inputs, one key | 36 / 170 | 5 / 105 | 5 / 287 |
1 / 321 |
The small uniform case performs similarly with a suitable COALESCE. In the
skewed case, count-matched COALESCE places all four large inputs in one task;
size packing leaves each oversized input separate. One-task COALESCE removes
the exchange, while the default and packed cases retain the required exchange.
These local timings illustrate the sizing tradeoff and do not establish
production speedups. The measurement source and reproduction command are below.
The generated CSV records trial ranges and every task's row count.
<details>
<summary>Local measurement source, trial ranges and reproduction</summary>
Save the source below as `/tmp/ScanPackingMeasurements.scala`, then run from
this Spark checkout using JDK 25:
```bash
build/sbt -Phive \
'set LocalProject("sql") / Test / unmanagedSources +=
file("/tmp/ScanPackingMeasurements.scala")' \
'sql/Test/runMain org.apache.spark.sql.connector.ScanPackingMeasurements
/tmp/scan-packing-measurements.csv'
```
The SBT source setting applies only to that invocation. The fixture uses
existing Spark test helpers and stays outside the PR source tree. It generates
rows, with no storage IO. Five measured rounds follow two warmup rounds; case
order is shuffled deterministically. A one-task COALESCE legitimately removes
the exchange, and the reproducer checks that shape as well as complete answers.
| Layout | Variant | Min–max ms | Largest task, rows |
| --- | --- | ---: | ---: |
| small_inputs | default | 1063.5–1143.3 | 128 |
| small_inputs | size_packing | 52.2–63.7 | 16,384 |
| small_inputs | coalesce_matching_count | 53.4–64.3 | 16,384 |
| small_inputs | coalesce_key_count | 60.3–74.9 | 8,192 |
| large_single_key | default | 220.2–230.8 | 65,536 |
| large_single_key | size_packing | 163.1–167.9 | 262,144 |
| large_single_key | coalesce_matching_count | 181.6–192.1 | 262,144 |
| large_single_key | coalesce_key_count | 608.7–623.3 | 4,194,304 |
| skewed_inputs | default | 157.2–179.9 | 524,288 |
| skewed_inputs | size_packing | 95.3–123.8 | 524,288 |
| skewed_inputs | coalesce_matching_count | 283.0–290.9 | 2,099,200 |
| skewed_inputs | coalesce_key_count | 318.5–354.4 | 2,113,536 |
```scala
package org.apache.spark.sql.connector
import java.nio.charset.StandardCharsets
import java.nio.file.{Files, Paths}
import java.util.OptionalLong
import scala.collection.mutable
import scala.util.Random
import org.apache.spark.sql.{DataFrame, Row, SparkSession}
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.connector.catalog.Table
import org.apache.spark.sql.connector.expressions.FieldReference
import org.apache.spark.sql.connector.read._
import
org.apache.spark.sql.connector.read.partitioning.{KeyGroupedPartitioning,
Partitioning}
import org.apache.spark.sql.execution.CoalesceExec
import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
import org.apache.spark.sql.functions._
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.util.CaseInsensitiveStringMap
// Local synthetic measurements, outside the PR source tree. No file IO is
simulated.
class ScanPackingMeasurementSource extends TestingV2Source {
override def getTable(options: CaseInsensitiveStringMap): Table = new
SimpleBatchTable {
override def newScanBuilder(options: CaseInsensitiveStringMap):
ScanBuilder =
new SimpleScanBuilder with SupportsReportPartitioning {
private val keys = options.get("keys").toInt
private val inputs = options.get("rows").split(",").zipWithIndex.map
{ case (rows, i) =>
ScanPackingMeasurementInput(i % keys, rows.toInt)
}
override def planInputPartitions(): Array[InputPartition] =
inputs.map(p => p: InputPartition)
override def outputPartitioning(): Partitioning =
new KeyGroupedPartitioning(Array(FieldReference("i")),
inputs.length)
override def createReaderFactory(): PartitionReaderFactory =
ScanPackingMeasurementReaderFactory
}
}
}
case class ScanPackingMeasurementInput(key: Int, rows: Int)
extends HasPartitionStatistics with HasPartitionKey {
override def partitionKey(): InternalRow = InternalRow(key)
override def sizeInBytes(): OptionalLong = OptionalLong.of(rows.toLong *
16)
override def numRows(): OptionalLong = OptionalLong.of(rows)
override def filesCount(): OptionalLong = OptionalLong.empty()
}
object ScanPackingMeasurementReaderFactory extends PartitionReaderFactory {
override def createReader(input: InputPartition):
PartitionReader[InternalRow] = {
val p = input.asInstanceOf[ScanPackingMeasurementInput]
new PartitionReader[InternalRow] {
private var row = -1
override def next(): Boolean = { row += 1; row < p.rows }
override def get(): InternalRow = InternalRow(p.key, row)
override def close(): Unit = {}
}
}
}
object ScanPackingMeasurements {
case class Scenario(name: String, keys: Int, rows: Seq[Int])
case class Variant(name: String, packing: Boolean, coalesce: Option[Int])
def main(args: Array[String]): Unit = {
val spark = SparkSession.builder().master("local[4]")
.appName("SPARK-60265 local scan packing measurements")
.config("spark.ui.enabled", "false")
.config(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "false")
.config(SQLConf.SHUFFLE_PARTITIONS.key, "4")
.config(SQLConf.V2_SCAN_PARTITION_PACKING_MAX_PARTITION_BYTES.key,
"4MB")
.config(SQLConf.V2_SCAN_PARTITION_PACKING_OPEN_COST_IN_BYTES.key, "0")
.config(SQLConf.V2_SCAN_PARTITION_PACKING_MIN_PARTITION_NUM.key, "4")
.getOrCreate()
spark.sparkContext.setLogLevel("ERROR")
val out = mutable.ArrayBuffer.empty[String]
out += "scenario,variant,raw_inputs,scan_stage_tasks,shuffles,rows," +
"median_ms,min_ms,max_ms,task_rows"
try {
val scenarios = Seq(
Scenario("small_inputs", 8, Seq.fill(512)(128)),
Scenario("large_single_key", 1, Seq.fill(64)(65536)),
Scenario("skewed_inputs", 1, Seq.fill(32)(512) ++
Seq.fill(4)(524288)))
scenarios.foreach { scenario =>
def query(v: Variant): DataFrame = {
spark.conf.set(SQLConf.V2_SCAN_PARTITION_PACKING_ENABLED.key,
v.packing)
val loaded =
spark.read.format(classOf[ScanPackingMeasurementSource].getName)
.option("keys", scenario.keys).option("rows",
scenario.rows.mkString(",")).load()
val input = v.coalesce.map(loaded.coalesce).getOrElse(loaded)
input.groupBy((col("j") % 32).as("g"))
.agg(count(lit(1)).as("n"), sum("i").as("k"), sum("j").as("s"))
}
val packed = Variant("size_packing", true, None)
val initial = query(packed)
val packedScan = initial.queryExecution.executedPlan.collect {
case scan: BatchScanExec => scan
}.head
val packedCount = packedScan.inputRDD.getNumPartitions
val variants = Seq(
Variant("default", false, None), packed,
Variant("coalesce_matching_count", false, Some(packedCount)),
Variant("coalesce_key_count", false, Some(scenario.keys)))
val samples = mutable.Map.empty[String,
Vector[Double]].withDefaultValue(Vector.empty)
var expected = Option.empty[Seq[Row]]
// Two warmup rounds, then five interleaved measured rounds with
deterministic ordering.
for (round <- -2 until 5; variant <- new Random(42 +
round).shuffle(variants)) {
val df = query(variant)
val started = System.nanoTime()
val answer = df.collect().toSeq.sortBy(_.getInt(0))
val elapsed = (System.nanoTime() - started) / 1000000.0
assert(answer.map(_.getLong(1)).sum ==
scenario.rows.map(_.toLong).sum)
expected.foreach(rows => assert(answer == rows))
expected = Some(answer)
val shuffles = df.queryExecution.executedPlan.collect {
case e: ShuffleExchangeExec => e
}.size
assert(shuffles == (if (variant.coalesce.contains(1)) 0 else 1))
if (round >= 0) samples(variant.name) = samples(variant.name) :+
elapsed
}
variants.foreach { variant =>
val df = query(variant)
df.collect()
val plan = df.queryExecution.executedPlan
val scan = plan.collect { case scan: BatchScanExec => scan }.head
val counts = plan.collect { case c: CoalesceExec => c }.headOption
match {
case Some(c) => c.execute().mapPartitions(i =>
Iterator.single(i.size)).collect().toSeq
case None => scan.inputRDD.mapPartitions(i =>
Iterator.single(i.size)).collect().toSeq
}
assert(counts.map(_.toLong).sum == scenario.rows.map(_.toLong).sum)
val times = samples(variant.name).sorted
val shuffles = plan.collect { case e: ShuffleExchangeExec => e
}.size
val line = Seq(scenario.name, variant.name, scenario.rows.size,
counts.size, shuffles,
scenario.rows.map(_.toLong).sum, f"${times(times.size / 2)}%.3f",
f"${times.head}%.3f", f"${times.last}%.3f",
counts.mkString(";")).mkString(",")
out += line
println("MEASUREMENT " + line.split(",").take(9).mkString(","))
}
}
Files.write(Paths.get(args(0)), (out.mkString("\n") + "\n")
.getBytes(StandardCharsets.UTF_8))
} finally {
spark.stop()
}
}
}
```
</details>
--
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]
