Skip to content

Commit ddb428b

Browse files
andygroveandygrove
committed
test: check where native shuffle puts each key type the conversion admits
Checks the partition of every row and a join against the JVM columnar shuffle for each hash key type the conversion admits, with NULL and boundary values. Adds array<string> and map<string,string> columns with NULLs to the test over several batches, and plan assertions to the test of aggregates whose partial aggregate runs on Spark.
1 parent 54defa1 commit ddb428b

1 file changed

Lines changed: 113 additions & 32 deletions

File tree

‎spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala‎

Lines changed: 113 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,17 @@
1919

2020
package org.apache.comet.exec
2121

22+
import java.sql.{Date, Timestamp}
23+
import java.time.LocalDateTime
24+
2225
import org.apache.spark.sql.{CometTestBase, DataFrame, Row}
23-
import org.apache.spark.sql.comet.{CometPlan, CometSparkToColumnarExec}
26+
import org.apache.spark.sql.comet.{CometHashAggregateExec, CometPlan, CometSparkToColumnarExec}
2427
import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec}
2528
import org.apache.spark.sql.execution.{ColumnarToRowExec, ColumnarToRowTransition, SparkPlan}
26-
import org.apache.spark.sql.functions.{array, avg, col, count, hash, length, max, min, size, sum}
29+
import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, SortAggregateExec}
30+
import org.apache.spark.sql.functions.{array, avg, col, count, hash, length, max, min, size, spark_partition_id, sum}
2731
import org.apache.spark.sql.internal.SQLConf
28-
import org.apache.spark.sql.types.{ArrayType, BinaryType, CalendarIntervalType, DataTypes, DecimalType, IntegerType, LongType, StringType, StructType}
32+
import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, CalendarIntervalType, DataType, DataTypes, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, ShortType, StringType, StructType, TimestampNTZType, TimestampType}
2933
import org.apache.spark.unsafe.types.CalendarInterval
3034

3135
import org.apache.comet.{CometConf, ExtendedExplainInfo}
@@ -124,6 +128,62 @@ class CometShuffleInputConversionSuite extends CometTestBase {
124128
assert(!reasons.exists(_.contains("ColumnarToRow")), reasons)
125129
}
126130

131+
/** Runs `f` with each join planned as a sort-merge join of two shuffles into 10 partitions. */
132+
private def withShuffledJoins(f: => Unit): Unit =
133+
withSQLConf(
134+
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
135+
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
136+
SQLConf.SHUFFLE_PARTITIONS.key -> "10")(f)
137+
138+
/**
139+
* Values of each type the conversion admits as a hash key, with boundary values. The floating
140+
* point values include NaN and both zeros, which a join normalizes.
141+
*/
142+
private val hashKeys: Seq[(DataType, Seq[Any])] = Seq(
143+
BooleanType -> Seq(true, false),
144+
ByteType -> Seq(0.toByte, -1.toByte, Byte.MinValue, Byte.MaxValue),
145+
ShortType -> Seq(0.toShort, -1.toShort, Short.MinValue, Short.MaxValue),
146+
IntegerType -> Seq(0, -1, Int.MinValue, Int.MaxValue),
147+
LongType -> Seq(0L, -1L, Long.MinValue, Long.MaxValue),
148+
FloatType -> Seq(
149+
0.0f,
150+
-0.0f,
151+
1.5f,
152+
Float.NaN,
153+
Float.MinPositiveValue,
154+
Float.MaxValue,
155+
Float.NegativeInfinity),
156+
DoubleType -> Seq(
157+
0.0d,
158+
-0.0d,
159+
1.5d,
160+
Double.NaN,
161+
Double.MinPositiveValue,
162+
Double.MaxValue,
163+
Double.NegativeInfinity),
164+
DecimalType(18, 2) -> Seq("0.00", "-0.01", "9999999999999999.99", "-9999999999999999.99")
165+
.map(new java.math.BigDecimal(_)),
166+
DateType -> Seq("1970-01-01", "1969-12-31", "0001-01-01", "9999-12-31").map(Date.valueOf),
167+
TimestampType -> Seq(
168+
new Timestamp(0L),
169+
new Timestamp(-1L),
170+
Timestamp.valueOf("0001-01-01 00:00:00"),
171+
Timestamp.valueOf("9999-12-31 23:59:59.999999")),
172+
TimestampNTZType -> Seq(
173+
LocalDateTime.of(1970, 1, 1, 0, 0),
174+
LocalDateTime.of(1969, 12, 31, 23, 59, 59, 999999000),
175+
LocalDateTime.of(1, 1, 1, 0, 0),
176+
LocalDateTime.of(9999, 12, 31, 23, 59, 59, 999999000)),
177+
BinaryType -> Seq(Array.emptyByteArray, Array[Byte](0), Array[Byte](-1, -128, 127)))
178+
179+
/** Four rows of each value and of a NULL, as column `k`, over four input partitions. */
180+
private def keysDf(keyType: DataType, values: Seq[Any]): DataFrame = {
181+
val keys = values :+ null
182+
val data = (0 until keys.length * 4).map(i => Row(keys(i % keys.length), i.toLong))
183+
val schema = new StructType().add("k", keyType).add("v", LongType)
184+
spark.createDataFrame(spark.sparkContext.parallelize(data, 4), schema)
185+
}
186+
127187
convertTest("the shuffle of a Spark operator's rows runs as native shuffle") {
128188
Seq("true", "false").foreach { aqe =>
129189
withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) {
@@ -167,10 +227,24 @@ class CometShuffleInputConversionSuite extends CometTestBase {
167227
}
168228

169229
convertTest("aggregates whose partial aggregate runs on Spark") {
170-
checkSparkAnswer(
230+
// The partial aggregate reads Spark rows, so it runs on Spark, and native shuffle carries its
231+
// buffers to the final aggregate. With these aggregates the final one runs natively.
232+
val (_, split) = checkSparkAnswer(
233+
rowsDf()
234+
.groupBy((col("k") % 7).as("g"))
235+
.agg(sum("l"), avg("d"), max("l"), min("l")))
236+
assert(convertedShuffles(split).length == 1, split)
237+
assert(collect(split) { case a: HashAggregateExec => a }.length == 1, split)
238+
assert(collect(split) { case a: CometHashAggregateExec => a }.length == 1, split)
239+
// The min of a string makes Spark plan sort aggregates, which stay on Spark, so the final
240+
// aggregate reads the buffers back as rows. The shuffle stays native: the revert in the next
241+
// test covers hash aggregates only.
242+
val (_, sorted) = checkSparkAnswer(
171243
rowsDf()
172244
.groupBy((col("k") % 7).as("g"))
173245
.agg(sum("l"), count("s"), avg("d"), sum("m"), max("l"), min("s")))
246+
assert(convertedShuffles(sorted).length == 1, sorted)
247+
assert(collect(sorted) { case a: SortAggregateExec => a }.length == 2, sorted)
174248
}
175249

176250
convertTest("a shuffle between two Spark aggregates goes back to Spark's shuffle") {
@@ -209,30 +283,30 @@ class CometShuffleInputConversionSuite extends CometTestBase {
209283
}
210284
}
211285

212-
convertTest("a join with an input that stays on the JVM columnar shuffle") {
213-
// The left input converts and the right one, with its array<int> column, does not. Native
214-
// shuffle hashes an int key as Spark does, so matching keys still meet.
215-
withSQLConf(
216-
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
217-
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
218-
SQLConf.SHUFFLE_PARTITIONS.key -> "10") {
219-
val df = rowsDf(100)
220-
.select(col("k").as("lk"), col("l"))
221-
.join(arraysDf(), col("lk") === col("k"))
222-
val (_, plan) = checkSparkAnswer(df)
223-
assert(convertedShuffles(plan).length == 1, plan)
224-
assert(cometShuffles(plan).count(_.shuffleType == CometColumnarShuffle) == 1, plan)
286+
hashKeys.foreach { case (keyType, values) =>
287+
convertTest(s"native shuffle puts ${keyType.simpleString} keys where Spark does") {
288+
withShuffledJoins {
289+
val keys = keysDf(keyType, values)
290+
// The partition of each row, against Spark's partitioner, NULL keys included.
291+
val (_, placed) = checkSparkAnswer(
292+
keys.repartition(10, col("k")).select(col("k"), col("v"), spark_partition_id()))
293+
assert(convertedShuffles(placed).length == 1, placed)
294+
// The left input converts and the right one, with its array<bigint> column, does not, so
295+
// matching keys meet only if both shuffles put them in the same partition.
296+
val left = keys.select(col("k").as("lk"), col("v"))
297+
val right = keys.select(col("k").as("rk"), array(col("v")).as("xs"))
298+
val (_, joined) = checkSparkAnswer(left.join(right, col("lk") === col("rk")))
299+
assert(convertedShuffles(joined).length == 1, joined)
300+
assert(cometShuffles(joined).count(_.shuffleType == CometColumnarShuffle) == 1, joined)
301+
}
225302
}
226303
}
227304

228305
convertTest("a shuffle that hashes a wide decimal stays on the JVM columnar shuffle") {
229306
// Native shuffle hashes a decimal wider than 18 digits differently from Spark's partitioner
230307
// (#5994). The right input, with its array<int> column, stays on the JVM columnar shuffle,
231308
// so the left one has to as well, or matching keys land in different partitions.
232-
withSQLConf(
233-
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
234-
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
235-
SQLConf.SHUFFLE_PARTITIONS.key -> "10") {
309+
withShuffledJoins {
236310
val left = rowsDf(100).select(col("k").cast(DecimalType(38, 0)).as("lk"), col("l"))
237311
val right = arraysDf().select(col("k").cast(DecimalType(38, 0)).as("rk"), col("xs"))
238312
val df = left.join(right, col("lk") === col("rk"))
@@ -251,10 +325,7 @@ class CometShuffleInputConversionSuite extends CometTestBase {
251325
// replaced string. Both shuffles write the rows with invalid UTF-8 replaced, and the join
252326
// compares what they wrote, so the two keys stay apart only because each key expression puts
253327
// them in different partitions.
254-
withSQLConf(
255-
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
256-
SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
257-
SQLConf.SHUFFLE_PARTITIONS.key -> "10") {
328+
withShuffledJoins {
258329
val schema = new StructType().add("b", BinaryType).add("v", LongType)
259330
val data = (0 until 40).map { i =>
260331
Row(Array((if (i % 2 == 0) 0x80 else 0x83).toByte), i.toLong)
@@ -273,17 +344,27 @@ class CometShuffleInputConversionSuite extends CometTestBase {
273344
}
274345
}
275346

276-
convertTest("struct columns over several batches") {
277-
// The rows of each input partition fill several batches. With the conversion of RDD scans
278-
// on, native shuffle reads the converted scan the same way.
347+
convertTest("nested columns over several batches") {
348+
// The rows of each input partition fill several batches, and the conversion reuses its
349+
// vectors from one batch to the next. With the conversion of RDD scans on, native shuffle
350+
// reads the converted scan the same way.
351+
val schema = new StructType()
352+
.add("k", IntegerType)
353+
.add("payload", new StructType().add("v", LongType).add("s", StringType))
354+
.add("xs", ArrayType(StringType))
355+
.add("m", MapType(StringType, StringType))
356+
val data = (0 until 24).map { i =>
357+
Row(
358+
i,
359+
if (i % 5 == 4) null else Row(i.toLong, if (i % 3 == 0) null else s"s$i"),
360+
if (i % 6 == 5) null else Seq.tabulate(i % 4)(j => if (j == 1) null else s"x$i-$j"),
361+
if (i % 7 == 6) null
362+
else Seq.tabulate(i % 3)(j => s"k$j" -> (if (j == 1) null else s"v$i-$j")).toMap)
363+
}
279364
Seq("false", "true").foreach { leafConversion =>
280365
withSQLConf(
281366
CometConf.COMET_BATCH_SIZE.key -> "7",
282367
CometConf.COMET_CONVERT_FROM_RDD_ENABLED.key -> leafConversion) {
283-
val schema = new StructType()
284-
.add("k", IntegerType)
285-
.add("payload", new StructType().add("v", LongType).add("s", StringType))
286-
val data = (0 until 24).map(i => Row(i, Row(i.toLong, s"s$i")))
287368
val df = spark.createDataFrame(spark.sparkContext.parallelize(data, 1), schema)
288369
Seq(df.repartition(3, col("k")), df.repartitionByRange(3, col("k"))).foreach { shuffled =>
289370
val (_, plan) = checkSparkAnswer(shuffled)

0 commit comments

Comments
 (0)