andygrove commented on code in PR #6361:
URL: https://github.com/apache/datafusion-comet/pull/6361#discussion_r4550995750
##########
spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala:
##########
@@ -262,6 +268,43 @@ class SpillSorterSuite extends AnyFunSuite with
BeforeAndAfterEach {
}
}
+ test("write sorted file across partitions with shuffle checksums disabled") {
+ // With spark.shuffle.checksum.enabled=false the sorter gets an empty
checksum array and no
+ // checksum algorithm.
+ val spills = new java.util.LinkedList[SpillInfo]()
Review Comment:
Moot now. I merged `main`, where #7065 replaced the sorter code this change
guarded and added its own `SpillSorterSuite` tests, including one with
checksums disabled. I took `main`'s `SpillSorterSuite`, so the `SpillInfo`
collision is gone with the test that had it, and `main`'s version already
imports it as `CometSpillInfo` (79040689d7).
##########
spark/src/test/scala/org/apache/spark/shuffle/sort/SpillSorterSuite.scala:
##########
@@ -267,6 +273,43 @@ class SpillSorterSuite extends AnyFunSuite with
BeforeAndAfterEach {
}
}
+ test("write sorted file across partitions with shuffle checksums disabled") {
+ // With spark.shuffle.checksum.enabled=false the sorter gets an empty
checksum array and no
+ // checksum algorithm.
+ val spills = new java.util.LinkedList[SpillInfo]()
+ val sorter = createSpillSorter(
+ spills = spills,
+ partitionChecksums = Array.emptyLongArray,
+ checksumAlgorithm = null)
+ val file = File.createTempFile("spill-sorter", ".data")
+ try {
+ val numPartitions = 3
+ val rowSize = UnsafeRow.calculateBitSetWidthInBytes(1) + 8
+ // Each record is a 4-byte key followed by an UnsafeRow with one int
column.
+ val record = new Array[Byte](4 + rowSize)
+ val row = new UnsafeRow(1)
+ row.pointTo(record, Platform.BYTE_ARRAY_OFFSET + 4, rowSize)
Review Comment:
Gone with the same test, since `main`'s `SpillSorterSuite` replaced it.
`test-compile -Pspark-3.5 -Pstrict-warnings` passes on the merged tree.
##########
spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriterSuite.scala:
##########
@@ -212,6 +214,32 @@ class CometDiskBlockWriterSuite extends AnyFunSuite {
new JLinkedList[CometDiskBlockWriter]())
}
+ test("a writer computes no checksum when shuffle checksums are disabled") {
+ // With spark.shuffle.checksum.enabled=false the bypass merge sort writer
never calls
+ // setChecksum, so every write, including the ones after the first spill,
must skip it.
+ val conf = new SparkConf()
+ .set("spark.memory.offHeap.enabled", "true")
+ .set("spark.memory.offHeap.size", "1g")
+ val tmm = new TaskMemoryManager(new TestMemoryManager(conf), 0L)
+ val allocator = CometShuffleMemoryAllocator.getInstance(tmm, pageSize)
+ val tempDir = Utils.createTempDir()
+ // Spill every two rows so that the writer writes to its file several
times.
+ SQLConf.get.setConfString(CometConf.COMET_SHUFFLE_JVM_SPILL_THRESHOLD.key,
"2")
+ try {
+ val writer =
+ newWriter(new File(tempDir, "partition0"), allocator,
newTaskContext(tmm, 0L), conf)
+ val toUnsafe = UnsafeProjection.create(schema)
+ (0 until 5).foreach(_ => writer.insertRow(toUnsafe(InternalRow(new
Array[Byte](16))), 0))
+ assert(writer.close().length > 0)
+ assert(writer.getOutputRecords == 5)
+ assert(writer.getChecksum == -1)
Review Comment:
Done. The test runs with checksums disabled and enabled. The enabled run
sets the checksum to start from the sentinel with `adler32`, and checks that
`getChecksum()` equals a one-shot `Adler32` over the partition file after
several spills. The disabled run fails when I revert the `SpillWriter` change.
##########
spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala:
##########
@@ -959,6 +959,34 @@ class CometShuffleEncryptionSuite extends CometTestBase {
}
}
+class CometShuffleChecksumDisabledSuite extends CometTestBase {
+
+ override protected def sparkConf: SparkConf = {
+ val conf = super.sparkConf
+ conf.set("spark.shuffle.checksum.enabled", "false")
+ }
+
+ test("comet columnar shuffle with shuffle checksums disabled") {
+ // 10 partitions use the bypass merge sort writer, while 300 partitions
exceed
+ // spark.shuffle.sort.bypassMergeThreshold and use the sort-based writer.
The lower spill
+ // threshold makes every map task spill, so the sort-based writer also
merges spill files.
+ Seq(10, 300).foreach { numPartitions =>
+ Seq(Int.MaxValue, 2000).foreach { spillThreshold =>
+ withSQLConf(
+ CometConf.COMET_SHUFFLE_MODE.key -> "jvm",
+ CometConf.COMET_SHUFFLE_JVM_SPILL_THRESHOLD.key ->
spillThreshold.toString) {
+ val df = spark
+ .range(0, 100000, 1, 4)
+ .selectExpr("id", "cast(id as string) as s")
+ .repartition(numPartitions, col("id"))
+ checkCometExchange(df, 1, false)
Review Comment:
Added. The test asserts that the shuffle handle is the
`CometSerializedShuffleHandle` exactly when the partition count is over 200,
with your assertion, and it passes for 10 and 300 partitions.
##########
spark/src/main/java/org/apache/spark/shuffle/sort/SpillSorter.java:
##########
@@ -287,7 +287,9 @@ public void writeSortedFileNative(boolean isLastFile,
boolean tracingEnabled) th
spillInfo.partitionLengths[currentPartition] = written;
// Store the checksum for the current partition.
- partitionChecksums[currentPartition] = getChecksum();
+ if (partitionChecksums.length > 0) {
Review Comment:
No longer applies. #7065 rewrote `writeSortedFileNative` to write every
partition of a spill in one native call, so the two near-copies of that block
are gone, and the checksum store now happens in native code.
--
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]
