Skip to content

Commit cf2ff79

Browse files
andygroveandygrove
committed
test: add a benchmark for the optimized PIVOT plan
CometPivotBenchmark pivots sum(v) over 4M rows on a string column with 4, 32 or 128 values, grouped into 1K, 100K or 1M groups, and compares Spark, Comet with PivotFirst disabled (so the aggregates using it fall back while the rest of the plan stays native) and Comet with PivotFirst native. The query selects every pivot column into the noop sink, and before timing each Comet run is checked against Spark's answer and for the plan its name claims.
1 parent 0c4279a commit cf2ff79

1 file changed

Lines changed: 141 additions & 0 deletions

File tree

‎spark/src/test/scala/org/apache/spark/sql/benchmark/CometPivotBenchmark.scala‎

Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,141 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.apache.spark.sql.benchmark
21+
22+
import org.apache.spark.benchmark.Benchmark
23+
import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, PivotFirst}
24+
import org.apache.spark.sql.comet.CometHashAggregateExec
25+
import org.apache.spark.sql.execution.aggregate.BaseAggregateExec
26+
27+
import org.apache.comet.CometConf
28+
29+
/**
30+
* Benchmark for the optimized `PIVOT` plan, whose second aggregation uses Spark's internal
31+
* `PivotFirst` aggregate. To run this benchmark:
32+
* {{{
33+
* SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometPivotBenchmark
34+
* }}}
35+
* Results will be written to "spark/benchmarks/CometPivotBenchmark-**results.txt".
36+
*
37+
* The optional arguments are comma-separated lists of pivot value counts and of group counts,
38+
* which default to `4,32,128` and `1000,100000,1000000`:
39+
* {{{
40+
* make benchmark-org.apache.spark.sql.benchmark.CometPivotBenchmark -- 128 1000000
41+
* }}}
42+
*
43+
* Each case pivots `sum(v)` on a string column with `numPivots` values, grouped by a column with
44+
* `numGroups` values, and compares Spark with two Comet runs: one with `PivotFirst` disabled, so
45+
* that the aggregates using it fall back to Spark while the rest of the plan stays native, and
46+
* one with `PivotFirst` native. The query selects every pivot column and writes the rows to the
47+
* noop sink, so no pivot column can be pruned. Before timing, each Comet run is checked for the
48+
* plan its name claims and for Spark's answer.
49+
*/
50+
object CometPivotBenchmark extends CometBenchmarkBase {
51+
52+
override def runCometBenchmark(mainArgs: Array[String]): Unit = {
53+
require(mainArgs.length <= 2, "Usage: CometPivotBenchmark [pivot-counts] [group-counts]")
54+
val pivotCounts = mainArgs.headOption.getOrElse("4,32,128").split(",").map(_.toInt).toSeq
55+
val groupCounts =
56+
mainArgs.lift(1).getOrElse("1000,100000,1000000").split(",").map(_.toInt).toSeq
57+
val numRows = 4L * 1024 * 1024
58+
withTempPath { dir =>
59+
withTempTable("parquetV1Table") {
60+
val columns = groupCounts.map(g => s"CAST(id % $g AS INT) AS g$g") ++
61+
pivotCounts.map(p => s"concat('k', CAST(pmod(hash(id), $p) AS STRING)) AS k$p") :+
62+
"CAST(pmod(hash(id, 17), 1000) AS INT) AS v"
63+
prepareTable(dir, spark.range(numRows).selectExpr(columns: _*))
64+
for (numPivots <- pivotCounts; numGroups <- groupCounts) {
65+
runPivotBenchmark(numRows, numPivots, numGroups)
66+
}
67+
}
68+
}
69+
}
70+
71+
private def runPivotBenchmark(numRows: Long, numPivots: Int, numGroups: Int): Unit = {
72+
val pivotValues = (0 until numPivots).map(i => s"'k$i'").mkString(", ")
73+
val query =
74+
s"""SELECT * FROM (SELECT g$numGroups AS grp, k$numPivots AS k, v FROM parquetV1Table)
75+
|PIVOT (sum(v) FOR k IN ($pivotValues))""".stripMargin
76+
val benchmark = new Benchmark(
77+
s"Pivot with $numPivots pivot values and $numGroups groups",
78+
numRows,
79+
output = output)
80+
81+
val sparkConfigs = Seq(CometConf.COMET_ENABLED.key -> "false")
82+
val nativeConfigs =
83+
Seq(CometConf.COMET_ENABLED.key -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true")
84+
val fallbackConfigs =
85+
nativeConfigs :+ (CometConf.getExprEnabledConfigKey("PivotFirst") -> "false")
86+
87+
// `withSQLConf` returns Unit on Spark 3.x, so the result leaves the block through a var.
88+
var expected = 0L
89+
withSQLConf(sparkConfigs: _*) { expected = checksum(query)._1 }
90+
verify(benchmark, query, fallbackConfigs, expected, pivotFirstNative = false)
91+
verify(benchmark, query, nativeConfigs, expected, pivotFirstNative = true)
92+
93+
benchmark.addCase("Spark") { _ =>
94+
withSQLConf(sparkConfigs: _*)(spark.sql(query).noop())
95+
}
96+
benchmark.addCase("Comet (PivotFirst in Spark)") { _ =>
97+
withSQLConf(fallbackConfigs: _*)(spark.sql(query).noop())
98+
}
99+
benchmark.addCase("Comet (PivotFirst native)") { _ =>
100+
withSQLConf(nativeConfigs: _*)(spark.sql(query).noop())
101+
}
102+
benchmark.run()
103+
}
104+
105+
/**
106+
* An order-independent checksum of every column of the query's result, along with the final
107+
* plan that computed it.
108+
*/
109+
private def checksum(query: String) = {
110+
val df = spark.sql(s"SELECT sum(hash(*)) FROM ($query)")
111+
// `collect` runs the Dataset's own query execution, so its plan is the final AQE plan.
112+
val sum = df.collect().head.getLong(0)
113+
(sum, stripAQEPlan(df.queryExecution.executedPlan))
114+
}
115+
116+
private def verify(
117+
benchmark: Benchmark,
118+
query: String,
119+
configs: Seq[(String, String)],
120+
expected: Long,
121+
pivotFirstNative: Boolean): Unit = withSQLConf(configs: _*) {
122+
val (actual, plan) = checksum(query)
123+
require(actual == expected, s"Comet's pivot result differs from Spark's for:\n$query")
124+
def hasPivotFirst(aggregates: Seq[AggregateExpression]): Boolean =
125+
aggregates.exists(_.aggregateFunction.isInstanceOf[PivotFirst])
126+
val pivotFirstRunsNatively = collect(plan) {
127+
case agg: CometHashAggregateExec if hasPivotFirst(agg.aggregateExpressions) => true
128+
case agg: BaseAggregateExec if hasPivotFirst(agg.aggregateExpressions) => false
129+
}
130+
val nativeAggregates = collect(plan) { case agg: CometHashAggregateExec => agg }
131+
if (pivotFirstRunsNatively.isEmpty || pivotFirstRunsNatively.exists(_ != pivotFirstNative) ||
132+
nativeAggregates.isEmpty) {
133+
val where = if (pivotFirstNative) "natively" else "in Spark"
134+
warn(
135+
benchmark,
136+
s"""WARNING: expected PivotFirst to run $where and the other aggregates natively, so
137+
|this case does not measure what its name says.
138+
|Comet plan:""".stripMargin + "\n" + plan.treeString)
139+
}
140+
}
141+
}

0 commit comments

Comments
 (0)