viirya commented on code in PR #7065:
URL: https://github.com/apache/datafusion-comet/pull/7065#discussion_r4304034154
##########
spark/src/main/java/org/apache/spark/shuffle/sort/SpillSorter.java:
##########
@@ -235,135 +238,83 @@ public int numRecords() {
return this.inMemSorter.numRecords();
}
+ /**
+ * Sorts the buffered records by partition and writes them to the spill file
set by {@link
+ * #setSpillInfo}, every partition in one native call, recording each
partition's length and
+ * checksum.
+ */
public void writeSortedFileNative(boolean isLastFile, boolean
tracingEnabled) throws IOException {
Review Comment:
This method no longer calls `Native.writeSortedFileNative`. It calls
`writeSortedPartitionsNative`, and the bypass writer, through
`SpillWriter.doSpilling`, is now the only caller of the old native method. Two
methods with the same name doing different things is why `jvm_shuffle.md` needs
a paragraph to tell them apart. Would you consider renaming this one, for
example to `writeSortedPartitions`? Then a grep for `writeSortedFileNative`
would only find the per-partition path.
##########
native/shuffle/src/spark_unsafe/row.rs:
##########
@@ -1357,6 +1356,170 @@ fn make_builders(
Ok(builder)
}
+/// Bits of Spark's `PackedRecordPointer`: the partition id in the upper 24
bits, then the page
+/// number in 13 bits and the offset in the page in the lower 27.
+const PACKED_PARTITION_ID_SHIFT: u32 = 40;
+const PACKED_PAGE_NUMBER_SHIFT: u32 = 27;
+const PACKED_PAGE_NUMBER_MASK: u64 = (1 << 13) - 1;
+const PACKED_OFFSET_MASK: u64 = (1 << 27) - 1;
+
+/// Bytes `UnsafeRowSerializer` writes before each row for the row's size.
+const SERIALIZED_ROW_SIZE_BYTES: i64 = 4;
+
+/// Encoded blocks are buffered until they reach this many bytes, so that a
spill of many small
+/// partitions is not written one partition at a time.
+const OUTPUT_BUFFER_BYTES: usize = 1 << 20;
Review Comment:
The description says this buffer, like the builders, is native memory that
no pool reserves, and that it holds at most 1 MiB plus one block. Could the doc
comment here say the same? This is new buffering on this path, and someone
reading the code later will want to know whether it is reserved without going
back to the PR.
##########
spark/src/test/scala/org/apache/spark/sql/benchmark/CometColumnarShuffleWriteBenchmark.scala:
##########
@@ -0,0 +1,119 @@
+/*
+ * 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.benchmark
+
+import org.apache.spark.benchmark.Benchmark
+import org.apache.spark.sql.DataFrame
+import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle,
CometShuffleExchangeExec}
+import org.apache.spark.sql.functions.col
+import org.apache.spark.sql.internal.SQLConf
+
+import org.apache.comet.CometConf
+
+/**
+ * Measures Comet's JVM columnar shuffle at partition counts above Spark's
+ * `spark.shuffle.sort.bypassMergeThreshold` (200), where the sort-based
writer is used. That
+ * writer sorts its buffered rows by partition and writes every partition
holding a row, once per
+ * spill and once more at the end, so its cost per partition is paid
`partitions x (spills + 1)`
+ * times. Each case hash-partitions the same rows and discards the result.
+ *
+ * The partition count is set through `spark.sql.shuffle.partitions` rather
than `repartition(n,
+ * ...)`, so that AQE coalesces the read side into a few tasks and the
measurement is not
+ * dominated by scheduling one reduce task per partition.
+ *
+ * `spark.comet.shuffle.jvm.spillThreshold` makes Comet spill a fixed number
of rows at a time. It
+ * does not apply to Spark's shuffle, which does not spill in these cases.
+ *
+ * To run this benchmark:
+ * {{{
+ * SPARK_GENERATE_BENCHMARK_FILES=1 make
benchmark-org.apache.spark.sql.benchmark.CometColumnarShuffleWriteBenchmark
+ * }}}
+ * Results will be written to
"spark/benchmarks/CometColumnarShuffleWriteBenchmark-**results.txt".
+ */
+object CometColumnarShuffleWriteBenchmark extends CometBenchmarkBase {
Review Comment:
`CometShuffleBenchmark` already targets the sort-based writer on purpose,
through `manyPartitions = 201` and
`spark.comet.shuffle.jvm.spillThreshold=30000`. Is there a reason the 1,000 and
10,000 partition cases need their own benchmark object rather than a group in
`CometShuffleBenchmark`? If it is the different session setup
(`CometShuffleBenchmark` runs with `spark.comet.exec.onHeap.enabled=true`), a
sentence in the scaladoc would explain it. Otherwise, keeping all JVM shuffle
write benchmarks in one place would make them easier to maintain.
--
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]
