You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ag...@apache.org on 2026/07/21 13:00:21 UTC
(datafusion-comet) branch main updated: refactor: remove untested async columnar shuffle (#4985)
This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 32e2d7c4a8 refactor: remove untested async columnar shuffle (#4985)
32e2d7c4a8 is described below
commit 32e2d7c4a82fafd70ba023dd848c6076609f3c9a
Author: Andy Grove <ag...@apache.org>
AuthorDate: Tue Jul 21 07:00:15 2026 -0600
refactor: remove untested async columnar shuffle (#4985)
* feat: add withAlternative alias mechanism for CometConf, rename replaceSortMergeJoin
Introduces `.withAlternative(oldKey, ...)` on `ConfigBuilder`. Reads check the primary
key first, then any registered alternatives in order; reading a value from an alternative
logs a one-time deprecation warning per JVM naming the current key. Threaded through
`TypedConfigBuilder` into `ConfigEntryWithDefault` and `OptionalConfigEntry` so the same
mechanism covers configs with defaults, env-var overrides, and optional configs.
Pilots the mechanism by renaming `spark.comet.exec.replaceSortMergeJoin` to
`spark.comet.exec.forceShuffleHashJoin.enabled` (Scala symbol `COMET_FORCE_SHJ`), keeping
the old key working as an alias.
Adds `CometConfSuite` with tests for primary-wins, alternative fallback, default,
multiple alternatives, OptionalConfigEntry, and type conversion through the alias path.
Adds a Configuration Conventions contributor-guide page that documents the key/symbol
naming rules and the rename workflow.
Broader rename backlog tracked in #4978, per the 1.0.0 tracker #4082.
* fixup: forceShuffleHashJoin → forceShuffledHashJoin, drop .enabled suffix
Rename correction: use the past-participle form matching Spark's `ShuffledHashJoin`
class. Also drops the `.enabled` suffix — the `force` verb already encodes the
boolean action, so `.enabled` is redundant for this key.
Softens the Configuration Conventions doc's boolean-suffix rule accordingly:
descriptor-form flags (debug.enabled, metrics.enabled) still end in `.enabled`;
action-verb flags (force..., allow...) can omit it.
* ci: register CometConfSuite in linux and macos PR workflows
* refactor: remove untested async columnar shuffle
The async shuffle path (spark.comet.columnar.shuffle.async.enabled) has been
disabled by default and lacked meaningful test coverage. Remove it along with
its supporting classes and configs:
- Delete CometShuffleExternalSorterAsync and ShuffleThreadPool
- Merge CometShuffleExternalSorterSync back into CometShuffleExternalSorter
as a single concrete class (was an interface + factory)
- Strip async branches from CometBypassMergeSortShuffleWriter and
CometDiskBlockWriter
- Delete spark.comet.columnar.shuffle.async.{enabled,thread.num} configs
- Rename spark.comet.columnar.shuffle.async.max.thread.num to
spark.comet.columnar.shuffle.max.writers.per.executor (the cap is actually
on hash-shuffle writers, not threads); keep the old key as an alias via
withAlternative
- Remove CometAsyncShuffleSuite, DisableAQECometAsyncShuffleSuite, and the
asyncShuffleEnable machinery from CometColumnarShuffleSuite; drop the
async benchmark case
- Drop the async config table entry from the JVM shuffle contributor doc
Net: -773 lines.
---
docs/source/contributor-guide/jvm_shuffle.md | 13 +-
.../shuffle/sort/CometShuffleExternalSorter.java | 317 +++++++++++++---
.../sort/CometShuffleExternalSorterAsync.java | 417 ---------------------
.../sort/CometShuffleExternalSorterSync.java | 346 -----------------
.../shuffle/CometBypassMergeSortShuffleWriter.java | 21 --
.../execution/shuffle/CometDiskBlockWriter.java | 91 +----
.../shuffle/CometUnsafeShuffleWriter.java | 2 +-
.../comet/execution/shuffle/ShuffleThreadPool.java | 71 ----
.../main/scala/org/apache/comet/CometConf.scala | 32 +-
.../execution/shuffle/CometShuffleManager.scala | 5 +-
.../comet/exec/CometColumnarShuffleSuite.scala | 52 +--
.../spark/sql/benchmark/CometExecBenchmark.scala | 1 -
.../sql/benchmark/CometShuffleBenchmark.scala | 33 +-
.../apache/comet/exec/CometShuffle4_0Suite.scala | 2 -
14 files changed, 315 insertions(+), 1088 deletions(-)
diff --git a/docs/source/contributor-guide/jvm_shuffle.md b/docs/source/contributor-guide/jvm_shuffle.md
index 2145c82eba..e67f79a524 100644
--- a/docs/source/contributor-guide/jvm_shuffle.md
+++ b/docs/source/contributor-guide/jvm_shuffle.md
@@ -184,14 +184,11 @@ Selection logic in `CometShuffleManager.shouldBypassMergeSort()`:
- `CometShuffleMemoryAllocator`: Custom allocator for off-heap memory pages
- Memory is allocated in pages; when allocation fails, writers spill to disk
- `CometDiskBlockWriter` coordinates spilling across all partition writers (largest first)
-- Async spilling is supported via `ShuffleThreadPool`
## Configuration
-| Config | Description |
-| ----------------------------------------------- | ----------------------------------- |
-| `spark.comet.columnar.shuffle.async.enabled` | Enable async spill writes |
-| `spark.comet.columnar.shuffle.async.thread.num` | Threads per writer for async |
-| `spark.comet.columnar.shuffle.batch.size` | Rows per Arrow batch |
-| `spark.comet.columnar.shuffle.spill.threshold` | Row count threshold for spill |
-| `spark.comet.exec.shuffle.compression.codec` | Compression codec (zstd, lz4, etc.) |
+| Config | Description |
+| ---------------------------------------------- | ----------------------------------- |
+| `spark.comet.columnar.shuffle.batch.size` | Rows per Arrow batch |
+| `spark.comet.columnar.shuffle.spill.threshold` | Row count threshold for spill |
+| `spark.comet.exec.shuffle.compression.codec` | Compression codec (zstd, lz4, etc.) |
diff --git a/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorter.java b/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorter.java
index 9266309019..e8bb611d89 100644
--- a/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorter.java
+++ b/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorter.java
@@ -19,20 +19,35 @@
package org.apache.spark.shuffle.sort;
+import java.io.File;
import java.io.IOException;
+import java.util.LinkedList;
+import java.util.concurrent.atomic.AtomicLong;
+
+import scala.Tuple2;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
import org.apache.spark.SparkConf;
import org.apache.spark.TaskContext;
+import org.apache.spark.memory.SparkOutOfMemoryError;
import org.apache.spark.shuffle.ShuffleWriteMetricsReporter;
+import org.apache.spark.shuffle.comet.CometShuffleChecksumSupport;
import org.apache.spark.shuffle.comet.CometShuffleMemoryAllocatorTrait;
+import org.apache.spark.shuffle.comet.TooLargePageException;
import org.apache.spark.sql.comet.execution.shuffle.SpillInfo;
import org.apache.spark.sql.types.StructType;
import org.apache.spark.storage.BlockManager;
+import org.apache.spark.storage.TempShuffleBlockId;
+import org.apache.spark.unsafe.UnsafeAlignedOffset;
+import org.apache.spark.unsafe.array.LongArray;
+import org.apache.spark.util.Utils;
import org.apache.comet.CometConf$;
/**
- * An external sorter interface specialized for sort-based shuffle.
+ * External sorter specialized for sort-based shuffle.
*
* <p>Incoming records are appended to data pages. When all records have been inserted (or when the
* current thread's shuffle memory limit is reached), the in-memory records are sorted according to
@@ -44,32 +59,261 @@ import org.apache.comet.CometConf$;
* org.apache.spark.sql.comet.execution.shuffle.CometUnsafeShuffleWriter}, which uses a specialized
* merge procedure that avoids extra serialization/deserialization.
*/
-public interface CometShuffleExternalSorter {
+public final class CometShuffleExternalSorter implements CometShuffleChecksumSupport {
- int MAXIMUM_PAGE_SIZE_BYTES = PackedRecordPointer.MAXIMUM_PAGE_SIZE_BYTES;
+ private static final Logger logger = LoggerFactory.getLogger(CometShuffleExternalSorter.class);
- /** Returns the checksums for each partition. */
- long[] getChecksums();
+ public static final int MAXIMUM_PAGE_SIZE_BYTES = PackedRecordPointer.MAXIMUM_PAGE_SIZE_BYTES;
+
+ private final int numPartitions;
+ private final BlockManager blockManager;
+ private final TaskContext taskContext;
+ private final ShuffleWriteMetricsReporter writeMetrics;
+
+ private final StructType schema;
+
+ /** Force this sorter to spill when there are this many elements in memory. */
+ private final int numElementsForSpillThreshold;
+
+ // When this external sorter allocates memory of `sorterArray`, we need to keep its
+ // assigned initial size. After spilling, we will reset the array to its initial size.
+ // See `sorterArray` comment for more details.
+ private int initialSize;
+
+ private SpillSorter activeSpillSorter;
+
+ private final LinkedList<SpillInfo> spills = new LinkedList<>();
+
+ /** Shared accumulator for encode + compression time across all SpillSorter instances. */
+ private final AtomicLong encodeNanos = new AtomicLong(0);
/** Returns total encode + compression time in nanoseconds accumulated across all spills. */
- long getEncodeNanos();
+ public long getEncodeNanos() {
+ return encodeNanos.get();
+ }
+
+ /** Peak memory used by this sorter so far, in bytes. */
+ private long peakMemoryUsedBytes;
+
+ // Checksum calculator for each partition. Empty when shuffle checksum disabled.
+ private final long[] partitionChecksums;
+
+ private final String checksumAlgorithm;
+ private final String compressionCodec;
+ private final int compressionLevel;
+
+ // The memory allocator for this sorter. It is used to allocate/free memory pages for this sorter.
+ // Because we need to allocate off-heap memory regardless of configured Spark memory mode
+ // (on-heap/off-heap), we need a separate memory allocator.
+ private final CometShuffleMemoryAllocatorTrait allocator;
+
+ private boolean spilling = false;
+
+ private final int uaoSize = UnsafeAlignedOffset.getUaoSize();
+ private final double preferDictionaryRatio;
+ private final boolean tracingEnabled;
+
+ public CometShuffleExternalSorter(
+ CometShuffleMemoryAllocatorTrait allocator,
+ BlockManager blockManager,
+ TaskContext taskContext,
+ int initialSize,
+ int numPartitions,
+ SparkConf conf,
+ ShuffleWriteMetricsReporter writeMetrics,
+ StructType schema) {
+ this.allocator = allocator;
+ this.blockManager = blockManager;
+ this.taskContext = taskContext;
+ this.numPartitions = numPartitions;
+ this.schema = schema;
+ this.numElementsForSpillThreshold =
+ (int) CometConf$.MODULE$.COMET_COLUMNAR_SHUFFLE_SPILL_THRESHOLD().get();
+ this.writeMetrics = writeMetrics;
+
+ this.peakMemoryUsedBytes = getMemoryUsage();
+ this.partitionChecksums = createPartitionChecksums(numPartitions, conf);
+ this.checksumAlgorithm = getChecksumAlgorithm(conf);
+ this.compressionCodec = CometConf$.MODULE$.COMET_EXEC_SHUFFLE_COMPRESSION_CODEC().get();
+ this.compressionLevel =
+ (int) CometConf$.MODULE$.COMET_EXEC_SHUFFLE_COMPRESSION_ZSTD_LEVEL().get();
+
+ this.initialSize = initialSize;
+ this.tracingEnabled = (boolean) CometConf$.MODULE$.COMET_TRACING_ENABLED().get();
+
+ this.preferDictionaryRatio =
+ (double) CometConf$.MODULE$.COMET_SHUFFLE_PREFER_DICTIONARY_RATIO().get();
+
+ this.activeSpillSorter = createSpillSorter();
+ }
+
+ /** Creates a new SpillSorter with all required dependencies. */
+ private SpillSorter createSpillSorter() {
+ SpillSorter sorter =
+ new SpillSorter(
+ allocator,
+ initialSize,
+ schema,
+ uaoSize,
+ preferDictionaryRatio,
+ compressionCodec,
+ compressionLevel,
+ checksumAlgorithm,
+ partitionChecksums,
+ writeMetrics,
+ taskContext,
+ spills,
+ this::spill);
+ sorter.setEncodeNanosAccumulator(encodeNanos);
+ return sorter;
+ }
+
+ /** Returns the checksums for each partition. */
+ public long[] getChecksums() {
+ return partitionChecksums;
+ }
/** Sort and spill the current records in response to memory pressure. */
- void spill() throws IOException;
+ public void spill() throws IOException {
+ if (spilling || activeSpillSorter == null || activeSpillSorter.numRecords() == 0) {
+ return;
+ }
+
+ spilling = true;
+
+ logger.info(
+ "Thread {} spilling sort data of {} to disk ({} {} so far)",
+ Thread.currentThread().getId(),
+ Utils.bytesToString(getMemoryUsage()),
+ spills.size(),
+ spills.size() > 1 ? " times" : " time");
+
+ final Tuple2<TempShuffleBlockId, File> spilledFileInfo =
+ blockManager.diskBlockManager().createTempShuffleBlock();
+ final File file = spilledFileInfo._2();
+ final TempShuffleBlockId blockId = spilledFileInfo._1();
+ final SpillInfo spillInfo = new SpillInfo(numPartitions, file, blockId);
+
+ activeSpillSorter.setSpillInfo(spillInfo);
+
+ activeSpillSorter.writeSortedFileNative(false, tracingEnabled);
+ final long spillSize = activeSpillSorter.freeMemory();
+ activeSpillSorter.reset();
+
+ // Reset the in-memory sorter's pointer array only after freeing up the memory pages holding
+ // the records. Otherwise, if the task is over allocated memory, then without freeing the
+ // memory pages, we might not be able to get memory for the pointer array.
+ taskContext.taskMetrics().incMemoryBytesSpilled(spillSize);
+
+ spilling = false;
+ }
+
+ private long getMemoryUsage() {
+ if (activeSpillSorter != null) {
+ return activeSpillSorter.getMemoryUsage();
+ }
+ return 0;
+ }
+
+ private void updatePeakMemoryUsed() {
+ long mem = getMemoryUsage();
+ if (mem > peakMemoryUsedBytes) {
+ peakMemoryUsedBytes = mem;
+ }
+ }
/** Return the peak memory used so far, in bytes. */
- long getPeakMemoryUsedBytes();
+ public long getPeakMemoryUsedBytes() {
+ updatePeakMemoryUsed();
+ return peakMemoryUsedBytes;
+ }
+
+ private long freeMemory() {
+ updatePeakMemoryUsed();
+ long memoryFreed = activeSpillSorter.freeMemory();
+ activeSpillSorter.freeArray();
+ return memoryFreed;
+ }
/** Force all memory and spill files to be deleted; called by shuffle error-handling code. */
- void cleanupResources();
+ public void cleanupResources() {
+ freeMemory();
+
+ for (SpillInfo spill : spills) {
+ if (spill.file.exists() && !spill.file.delete()) {
+ logger.error("Unable to delete spill file {}", spill.file.getPath());
+ }
+ }
+ }
+
+ /**
+ * Checks whether there is enough space to insert an additional record in to the sort pointer
+ * array and grows the array if additional space is required. If the required space cannot be
+ * obtained, then the in-memory data will be spilled to disk.
+ */
+ private void growPointerArrayIfNecessary() throws IOException {
+ assert (activeSpillSorter != null);
+ if (!activeSpillSorter.hasSpaceForAnotherRecord()) {
+ long used = activeSpillSorter.getMemoryUsage();
+ LongArray array;
+ try {
+ // could trigger spilling
+ array = allocator.allocateArray(used / 8 * 2);
+ } catch (TooLargePageException e) {
+ // The pointer array is too big to fix in a single page, spill.
+ spill();
+ return;
+ } catch (SparkOutOfMemoryError e) {
+ // Cannot allocate enough memory, spill and reset pointer array.
+ try {
+ spill();
+ } catch (SparkOutOfMemoryError e2) {
+ // Cannot allocate memory even after spilling, throw the error.
+ if (!activeSpillSorter.hasSpaceForAnotherRecord()) {
+ logger.error("Unable to grow the pointer array");
+ throw e2;
+ }
+ }
+ return;
+ }
+ // check if spilling is triggered or not
+ if (activeSpillSorter.hasSpaceForAnotherRecord()) {
+ allocator.freeArray(array);
+ } else {
+ activeSpillSorter.expandPointerArray(array);
+ }
+ }
+ }
/**
* Writes a record to the shuffle sorter. This copies the record data into this external sorter's
* managed memory, which may trigger spilling if the copy would exceed the memory limit. It
* inserts a pointer for the record and record's partition id into the in-memory sorter.
*/
- void insertRecord(Object recordBase, long recordOffset, int length, int partitionId)
- throws IOException;
+ public void insertRecord(Object recordBase, long recordOffset, int length, int partitionId)
+ throws IOException {
+
+ assert (activeSpillSorter != null);
+ int threshold = numElementsForSpillThreshold;
+ if (activeSpillSorter.numRecords() >= threshold) {
+ logger.info(
+ "Spilling data because number of spilledRecords crossed the threshold " + threshold);
+ spill();
+ }
+
+ growPointerArrayIfNecessary();
+
+ // Need 4 or 8 bytes to store the record length.
+ final int required = length + uaoSize;
+ // Acquire enough memory to store the record.
+ // If we cannot acquire enough memory, we will spill current writers.
+ if (!activeSpillSorter.acquireNewPageIfNecessary(required)) {
+ // Spilling is happened, initiate new memory page for new writer.
+ activeSpillSorter.initialCurrentPage(required);
+ }
+
+ activeSpillSorter.insertRecord(recordBase, recordOffset, length, partitionId);
+ }
/**
* Close the sorter, causing any buffered data to be sorted and written out to disk.
@@ -77,46 +321,21 @@ public interface CometShuffleExternalSorter {
* @return metadata for the spill files written by this sorter. If no records were ever inserted
* into this sorter, then this will return an empty array.
*/
- SpillInfo[] closeAndGetSpills() throws IOException;
+ public SpillInfo[] closeAndGetSpills() throws IOException {
+ if (activeSpillSorter != null) {
+ // Do not count the final file towards the spill count.
+ final Tuple2<TempShuffleBlockId, File> spilledFileInfo =
+ blockManager.diskBlockManager().createTempShuffleBlock();
+ final File file = spilledFileInfo._2();
+ final TempShuffleBlockId blockId = spilledFileInfo._1();
+ final SpillInfo spillInfo = new SpillInfo(numPartitions, file, blockId);
- /**
- * Factory method to create the appropriate sorter implementation based on configuration.
- *
- * @return either a sync or async implementation based on the COMET_COLUMNAR_SHUFFLE_ASYNC_ENABLED
- * configuration
- */
- static CometShuffleExternalSorter create(
- CometShuffleMemoryAllocatorTrait allocator,
- BlockManager blockManager,
- TaskContext taskContext,
- int initialSize,
- int numPartitions,
- SparkConf conf,
- ShuffleWriteMetricsReporter writeMetrics,
- StructType schema) {
+ activeSpillSorter.setSpillInfo(spillInfo);
+ activeSpillSorter.writeSortedFileNative(true, tracingEnabled);
- boolean isAsync = (boolean) CometConf$.MODULE$.COMET_COLUMNAR_SHUFFLE_ASYNC_ENABLED().get();
-
- if (isAsync) {
- return new CometShuffleExternalSorterAsync(
- allocator,
- blockManager,
- taskContext,
- initialSize,
- numPartitions,
- conf,
- writeMetrics,
- schema);
- } else {
- return new CometShuffleExternalSorterSync(
- allocator,
- blockManager,
- taskContext,
- initialSize,
- numPartitions,
- conf,
- writeMetrics,
- schema);
+ freeMemory();
}
+
+ return spills.toArray(new SpillInfo[spills.size()]);
}
}
diff --git a/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorterAsync.java b/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorterAsync.java
deleted file mode 100644
index f128fdf33c..0000000000
--- a/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorterAsync.java
+++ /dev/null
@@ -1,417 +0,0 @@
-/*
- * 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.shuffle.sort;
-
-import java.io.File;
-import java.io.IOException;
-import java.util.LinkedList;
-import java.util.concurrent.*;
-import java.util.concurrent.atomic.AtomicLong;
-
-import scala.Tuple2;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import org.apache.spark.SparkConf;
-import org.apache.spark.TaskContext;
-import org.apache.spark.memory.SparkOutOfMemoryError;
-import org.apache.spark.shuffle.ShuffleWriteMetricsReporter;
-import org.apache.spark.shuffle.comet.CometShuffleChecksumSupport;
-import org.apache.spark.shuffle.comet.CometShuffleMemoryAllocatorTrait;
-import org.apache.spark.shuffle.comet.TooLargePageException;
-import org.apache.spark.sql.comet.execution.shuffle.ShuffleThreadPool;
-import org.apache.spark.sql.comet.execution.shuffle.SpillInfo;
-import org.apache.spark.sql.types.StructType;
-import org.apache.spark.storage.BlockManager;
-import org.apache.spark.storage.TempShuffleBlockId;
-import org.apache.spark.unsafe.UnsafeAlignedOffset;
-import org.apache.spark.unsafe.array.LongArray;
-import org.apache.spark.util.Utils;
-
-import org.apache.comet.CometConf$;
-
-/**
- * Asynchronous implementation of the external sorter for sort-based shuffle.
- *
- * <p>Incoming records are appended to data pages. When all records have been inserted (or when the
- * current thread's shuffle memory limit is reached), the in-memory records are sorted according to
- * their partition ids using native sorter. The sorted records are then written to a single output
- * file (or multiple files, if we've spilled).
- *
- * <p>Unlike {@link org.apache.spark.util.collection.ExternalSorter}, this sorter does not merge its
- * spill files. Instead, this merging is performed in {@link
- * org.apache.spark.sql.comet.execution.shuffle.CometUnsafeShuffleWriter}, which uses a specialized
- * merge procedure that avoids extra serialization/deserialization.
- *
- * <p>This sorter provides async spilling write mode. When spilling, it will submit a task to thread
- * pool to write shuffle spilling file. After submitting the task, it will continue to buffer, sort
- * incoming records and submit another spilling task once spilling threshold reached again or memory
- * is not enough to buffer incoming records. Each spilling task will write a shuffle spilling file
- * separately. After all records have been sorted and spilled, all spill files will be merged by
- * {@link org.apache.spark.sql.comet.execution.shuffle.CometUnsafeShuffleWriter}.
- */
-public final class CometShuffleExternalSorterAsync
- implements CometShuffleExternalSorter, CometShuffleChecksumSupport {
-
- private static final Logger logger =
- LoggerFactory.getLogger(CometShuffleExternalSorterAsync.class);
-
- private final int numPartitions;
- private final BlockManager blockManager;
- private final TaskContext taskContext;
- private final ShuffleWriteMetricsReporter writeMetrics;
-
- private final StructType schema;
-
- /** Force this sorter to spill when there are this many elements in memory. */
- private final int numElementsForSpillThreshold;
-
- // When this external sorter allocates memory of `sorterArray`, we need to keep its
- // assigned initial size. After spilling, we will reset the array to its initial size.
- // See `sorterArray` comment for more details.
- private int initialSize;
-
- /** All sorters with memory pages used by the sorters. */
- private final ConcurrentLinkedQueue<SpillSorter> spillingSorters = new ConcurrentLinkedQueue<>();
-
- private SpillSorter activeSpillSorter;
-
- private final LinkedList<SpillInfo> spills = new LinkedList<>();
-
- /** Shared accumulator for encode + compression time across all SpillSorter instances. */
- private final AtomicLong encodeNanos = new AtomicLong(0);
-
- @Override
- public long getEncodeNanos() {
- return encodeNanos.get();
- }
-
- /** Peak memory used by this sorter so far, in bytes. */
- private long peakMemoryUsedBytes;
-
- // Checksum calculator for each partition. Empty when shuffle checksum disabled.
- private final long[] partitionChecksums;
-
- private final String checksumAlgorithm;
- private final String compressionCodec;
- private final int compressionLevel;
-
- // The memory allocator for this sorter. It is used to allocate/free memory pages for this sorter.
- // Because we need to allocate off-heap memory regardless of configured Spark memory mode
- // (on-heap/off-heap), we need a separate memory allocator.
- private final CometShuffleMemoryAllocatorTrait allocator;
-
- /** Thread pool shared for async spilling write */
- private final ExecutorService threadPool;
-
- private final int threadNum;
-
- private ConcurrentLinkedQueue<Future<Void>> asyncSpillTasks = new ConcurrentLinkedQueue<>();
-
- private boolean spilling = false;
-
- private final int uaoSize = UnsafeAlignedOffset.getUaoSize();
- private final double preferDictionaryRatio;
- private final boolean tracingEnabled;
-
- public CometShuffleExternalSorterAsync(
- CometShuffleMemoryAllocatorTrait allocator,
- BlockManager blockManager,
- TaskContext taskContext,
- int initialSize,
- int numPartitions,
- SparkConf conf,
- ShuffleWriteMetricsReporter writeMetrics,
- StructType schema) {
- this.allocator = allocator;
- this.blockManager = blockManager;
- this.taskContext = taskContext;
- this.numPartitions = numPartitions;
- this.schema = schema;
- this.numElementsForSpillThreshold =
- (int) CometConf$.MODULE$.COMET_COLUMNAR_SHUFFLE_SPILL_THRESHOLD().get();
- this.writeMetrics = writeMetrics;
-
- this.peakMemoryUsedBytes = getMemoryUsage();
- this.partitionChecksums = createPartitionChecksums(numPartitions, conf);
- this.checksumAlgorithm = getChecksumAlgorithm(conf);
- this.compressionCodec = CometConf$.MODULE$.COMET_EXEC_SHUFFLE_COMPRESSION_CODEC().get();
- this.compressionLevel =
- (int) CometConf$.MODULE$.COMET_EXEC_SHUFFLE_COMPRESSION_ZSTD_LEVEL().get();
-
- this.initialSize = initialSize;
- this.tracingEnabled = (boolean) CometConf$.MODULE$.COMET_TRACING_ENABLED().get();
-
- this.threadNum = (int) CometConf$.MODULE$.COMET_COLUMNAR_SHUFFLE_ASYNC_THREAD_NUM().get();
- assert (this.threadNum > 0);
- this.threadPool = ShuffleThreadPool.getThreadPool();
-
- this.preferDictionaryRatio =
- (double) CometConf$.MODULE$.COMET_SHUFFLE_PREFER_DICTIONARY_RATIO().get();
-
- this.activeSpillSorter = createSpillSorter();
- }
-
- /** Creates a new SpillSorter with all required dependencies. */
- private SpillSorter createSpillSorter() {
- SpillSorter sorter =
- new SpillSorter(
- allocator,
- initialSize,
- schema,
- uaoSize,
- preferDictionaryRatio,
- compressionCodec,
- compressionLevel,
- checksumAlgorithm,
- partitionChecksums,
- writeMetrics,
- taskContext,
- spills,
- this::spill);
- sorter.setEncodeNanosAccumulator(encodeNanos);
- return sorter;
- }
-
- @Override
- public long[] getChecksums() {
- return partitionChecksums;
- }
-
- /** Sort and spill the current records in response to memory pressure. */
- @Override
- public void spill() throws IOException {
- if (spilling || activeSpillSorter == null || activeSpillSorter.numRecords() == 0) {
- return;
- }
-
- // In async mode, if new in-memory sorter cannot allocate required array, it triggers spill
- // here. This method will initiate new sorter following normal spill logic and cause stack
- // overflow eventually. So we need to avoid triggering spilling again while spilling. But
- // we cannot make this as "synchronized" because it will block the caller thread.
- spilling = true;
-
- logger.info(
- "Thread {} spilling sort data of {} to disk ({} {} so far)",
- Thread.currentThread().getId(),
- Utils.bytesToString(getMemoryUsage()),
- spills.size(),
- spills.size() > 1 ? " times" : " time");
-
- final Tuple2<TempShuffleBlockId, File> spilledFileInfo =
- blockManager.diskBlockManager().createTempShuffleBlock();
- final File file = spilledFileInfo._2();
- final TempShuffleBlockId blockId = spilledFileInfo._1();
- final SpillInfo spillInfo = new SpillInfo(numPartitions, file, blockId);
-
- activeSpillSorter.setSpillInfo(spillInfo);
-
- SpillSorter spillingSorter = activeSpillSorter;
- Callable<Void> task =
- () -> {
- spillingSorter.writeSortedFileNative(false, tracingEnabled);
- final long spillSize = spillingSorter.freeMemory();
- spillingSorter.freeArray();
- spillingSorters.remove(spillingSorter);
-
- // Reset the in-memory sorter's pointer array only after freeing up the memory pages
- // holding the records. Otherwise, if the task is over allocated memory, then without
- // freeing the memory pages, we might not be able to get memory for the pointer array.
- synchronized (CometShuffleExternalSorterAsync.this) {
- taskContext.taskMetrics().incMemoryBytesSpilled(spillSize);
- }
-
- return null;
- };
-
- spillingSorters.add(spillingSorter);
- asyncSpillTasks.add(threadPool.submit(task));
-
- while (asyncSpillTasks.size() == threadNum) {
- for (Future<Void> spillingTask : asyncSpillTasks) {
- if (spillingTask.isDone()) {
- asyncSpillTasks.remove(spillingTask);
- break;
- }
- }
- }
-
- activeSpillSorter = createSpillSorter();
-
- spilling = false;
- }
-
- private long getMemoryUsage() {
- long totalPageSize = 0;
- for (SpillSorter sorter : spillingSorters) {
- totalPageSize += sorter.getMemoryUsage();
- }
- if (activeSpillSorter != null) {
- totalPageSize += activeSpillSorter.getMemoryUsage();
- }
- return totalPageSize;
- }
-
- private void updatePeakMemoryUsed() {
- long mem = getMemoryUsage();
- if (mem > peakMemoryUsedBytes) {
- peakMemoryUsedBytes = mem;
- }
- }
-
- /** Return the peak memory used so far, in bytes. */
- @Override
- public long getPeakMemoryUsedBytes() {
- updatePeakMemoryUsed();
- return peakMemoryUsedBytes;
- }
-
- private long freeMemory() {
- updatePeakMemoryUsed();
- long memoryFreed = 0;
- for (SpillSorter sorter : spillingSorters) {
- memoryFreed += sorter.freeMemory();
- sorter.freeArray();
- }
- memoryFreed += activeSpillSorter.freeMemory();
- activeSpillSorter.freeArray();
-
- return memoryFreed;
- }
-
- /** Force all memory and spill files to be deleted; called by shuffle error-handling code. */
- @Override
- public void cleanupResources() {
- freeMemory();
-
- for (SpillInfo spill : spills) {
- if (spill.file.exists() && !spill.file.delete()) {
- logger.error("Unable to delete spill file {}", spill.file.getPath());
- }
- }
- }
-
- /**
- * Checks whether there is enough space to insert an additional record in to the sort pointer
- * array and grows the array if additional space is required. If the required space cannot be
- * obtained, then the in-memory data will be spilled to disk.
- */
- private void growPointerArrayIfNecessary() throws IOException {
- assert (activeSpillSorter != null);
- if (!activeSpillSorter.hasSpaceForAnotherRecord()) {
- long used = activeSpillSorter.getMemoryUsage();
- LongArray array;
- try {
- // could trigger spilling
- array = allocator.allocateArray(used / 8 * 2);
- } catch (TooLargePageException e) {
- // The pointer array is too big to fix in a single page, spill.
- spill();
- return;
- } catch (SparkOutOfMemoryError e) {
- // Cannot allocate enough memory, spill and reset pointer array.
- try {
- spill();
- } catch (SparkOutOfMemoryError e2) {
- // Cannot allocate memory even after spilling, throw the error.
- if (!activeSpillSorter.hasSpaceForAnotherRecord()) {
- logger.error("Unable to grow the pointer array");
- throw e2;
- }
- }
- return;
- }
- // check if spilling is triggered or not
- if (activeSpillSorter.hasSpaceForAnotherRecord()) {
- allocator.freeArray(array);
- } else {
- activeSpillSorter.expandPointerArray(array);
- }
- }
- }
-
- /**
- * Writes a record to the shuffle sorter. This copies the record data into this external sorter's
- * managed memory, which may trigger spilling if the copy would exceed the memory limit. It
- * inserts a pointer for the record and record's partition id into the in-memory sorter.
- */
- @Override
- public void insertRecord(Object recordBase, long recordOffset, int length, int partitionId)
- throws IOException {
-
- assert (activeSpillSorter != null);
- int threshold = numElementsForSpillThreshold;
- if (activeSpillSorter.numRecords() >= threshold) {
- logger.info(
- "Spilling data because number of spilledRecords crossed the threshold " + threshold);
- spill();
- }
-
- growPointerArrayIfNecessary();
-
- // Need 4 or 8 bytes to store the record length.
- final int required = length + uaoSize;
- // Acquire enough memory to store the record.
- // If we cannot acquire enough memory, we will spill current writers.
- if (!activeSpillSorter.acquireNewPageIfNecessary(required)) {
- // Spilling is happened, initiate new memory page for new writer.
- activeSpillSorter.initialCurrentPage(required);
- }
-
- activeSpillSorter.insertRecord(recordBase, recordOffset, length, partitionId);
- }
-
- /**
- * Close the sorter, causing any buffered data to be sorted and written out to disk.
- *
- * @return metadata for the spill files written by this sorter. If no records were ever inserted
- * into this sorter, then this will return an empty array.
- */
- @Override
- public SpillInfo[] closeAndGetSpills() throws IOException {
- if (activeSpillSorter != null) {
- // Do not count the final file towards the spill count.
- final Tuple2<TempShuffleBlockId, File> spilledFileInfo =
- blockManager.diskBlockManager().createTempShuffleBlock();
- final File file = spilledFileInfo._2();
- final TempShuffleBlockId blockId = spilledFileInfo._1();
- final SpillInfo spillInfo = new SpillInfo(numPartitions, file, blockId);
-
- // Waits for all async tasks to finish.
- for (Future<Void> task : asyncSpillTasks) {
- try {
- task.get();
- } catch (Exception e) {
- throw new IOException(e);
- }
- }
-
- asyncSpillTasks.clear();
-
- activeSpillSorter.setSpillInfo(spillInfo);
- activeSpillSorter.writeSortedFileNative(true, tracingEnabled);
-
- freeMemory();
- }
-
- return spills.toArray(new SpillInfo[spills.size()]);
- }
-}
diff --git a/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorterSync.java b/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorterSync.java
deleted file mode 100644
index 0f0db86284..0000000000
--- a/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorterSync.java
+++ /dev/null
@@ -1,346 +0,0 @@
-/*
- * 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.shuffle.sort;
-
-import java.io.File;
-import java.io.IOException;
-import java.util.LinkedList;
-import java.util.concurrent.atomic.AtomicLong;
-
-import scala.Tuple2;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import org.apache.spark.SparkConf;
-import org.apache.spark.TaskContext;
-import org.apache.spark.memory.SparkOutOfMemoryError;
-import org.apache.spark.shuffle.ShuffleWriteMetricsReporter;
-import org.apache.spark.shuffle.comet.CometShuffleChecksumSupport;
-import org.apache.spark.shuffle.comet.CometShuffleMemoryAllocatorTrait;
-import org.apache.spark.shuffle.comet.TooLargePageException;
-import org.apache.spark.sql.comet.execution.shuffle.SpillInfo;
-import org.apache.spark.sql.types.StructType;
-import org.apache.spark.storage.BlockManager;
-import org.apache.spark.storage.TempShuffleBlockId;
-import org.apache.spark.unsafe.UnsafeAlignedOffset;
-import org.apache.spark.unsafe.array.LongArray;
-import org.apache.spark.util.Utils;
-
-import org.apache.comet.CometConf$;
-
-/**
- * Synchronous implementation of the external sorter for sort-based shuffle.
- *
- * <p>Incoming records are appended to data pages. When all records have been inserted (or when the
- * current thread's shuffle memory limit is reached), the in-memory records are sorted according to
- * their partition ids using native sorter. The sorted records are then written to a single output
- * file (or multiple files, if we've spilled).
- *
- * <p>Unlike {@link org.apache.spark.util.collection.ExternalSorter}, this sorter does not merge its
- * spill files. Instead, this merging is performed in {@link
- * org.apache.spark.sql.comet.execution.shuffle.CometUnsafeShuffleWriter}, which uses a specialized
- * merge procedure that avoids extra serialization/deserialization.
- */
-public final class CometShuffleExternalSorterSync
- implements CometShuffleExternalSorter, CometShuffleChecksumSupport {
-
- private static final Logger logger =
- LoggerFactory.getLogger(CometShuffleExternalSorterSync.class);
-
- private final int numPartitions;
- private final BlockManager blockManager;
- private final TaskContext taskContext;
- private final ShuffleWriteMetricsReporter writeMetrics;
-
- private final StructType schema;
-
- /** Force this sorter to spill when there are this many elements in memory. */
- private final int numElementsForSpillThreshold;
-
- // When this external sorter allocates memory of `sorterArray`, we need to keep its
- // assigned initial size. After spilling, we will reset the array to its initial size.
- // See `sorterArray` comment for more details.
- private int initialSize;
-
- private SpillSorter activeSpillSorter;
-
- private final LinkedList<SpillInfo> spills = new LinkedList<>();
-
- /** Shared accumulator for encode + compression time across all SpillSorter instances. */
- private final AtomicLong encodeNanos = new AtomicLong(0);
-
- @Override
- public long getEncodeNanos() {
- return encodeNanos.get();
- }
-
- /** Peak memory used by this sorter so far, in bytes. */
- private long peakMemoryUsedBytes;
-
- // Checksum calculator for each partition. Empty when shuffle checksum disabled.
- private final long[] partitionChecksums;
-
- private final String checksumAlgorithm;
- private final String compressionCodec;
- private final int compressionLevel;
-
- // The memory allocator for this sorter. It is used to allocate/free memory pages for this sorter.
- // Because we need to allocate off-heap memory regardless of configured Spark memory mode
- // (on-heap/off-heap), we need a separate memory allocator.
- private final CometShuffleMemoryAllocatorTrait allocator;
-
- private boolean spilling = false;
-
- private final int uaoSize = UnsafeAlignedOffset.getUaoSize();
- private final double preferDictionaryRatio;
- private final boolean tracingEnabled;
-
- public CometShuffleExternalSorterSync(
- CometShuffleMemoryAllocatorTrait allocator,
- BlockManager blockManager,
- TaskContext taskContext,
- int initialSize,
- int numPartitions,
- SparkConf conf,
- ShuffleWriteMetricsReporter writeMetrics,
- StructType schema) {
- this.allocator = allocator;
- this.blockManager = blockManager;
- this.taskContext = taskContext;
- this.numPartitions = numPartitions;
- this.schema = schema;
- this.numElementsForSpillThreshold =
- (int) CometConf$.MODULE$.COMET_COLUMNAR_SHUFFLE_SPILL_THRESHOLD().get();
- this.writeMetrics = writeMetrics;
-
- this.peakMemoryUsedBytes = getMemoryUsage();
- this.partitionChecksums = createPartitionChecksums(numPartitions, conf);
- this.checksumAlgorithm = getChecksumAlgorithm(conf);
- this.compressionCodec = CometConf$.MODULE$.COMET_EXEC_SHUFFLE_COMPRESSION_CODEC().get();
- this.compressionLevel =
- (int) CometConf$.MODULE$.COMET_EXEC_SHUFFLE_COMPRESSION_ZSTD_LEVEL().get();
-
- this.initialSize = initialSize;
- this.tracingEnabled = (boolean) CometConf$.MODULE$.COMET_TRACING_ENABLED().get();
-
- this.preferDictionaryRatio =
- (double) CometConf$.MODULE$.COMET_SHUFFLE_PREFER_DICTIONARY_RATIO().get();
-
- this.activeSpillSorter = createSpillSorter();
- }
-
- /** Creates a new SpillSorter with all required dependencies. */
- private SpillSorter createSpillSorter() {
- SpillSorter sorter =
- new SpillSorter(
- allocator,
- initialSize,
- schema,
- uaoSize,
- preferDictionaryRatio,
- compressionCodec,
- compressionLevel,
- checksumAlgorithm,
- partitionChecksums,
- writeMetrics,
- taskContext,
- spills,
- this::spill);
- sorter.setEncodeNanosAccumulator(encodeNanos);
- return sorter;
- }
-
- @Override
- public long[] getChecksums() {
- return partitionChecksums;
- }
-
- /** Sort and spill the current records in response to memory pressure. */
- @Override
- public void spill() throws IOException {
- if (spilling || activeSpillSorter == null || activeSpillSorter.numRecords() == 0) {
- return;
- }
-
- spilling = true;
-
- logger.info(
- "Thread {} spilling sort data of {} to disk ({} {} so far)",
- Thread.currentThread().getId(),
- Utils.bytesToString(getMemoryUsage()),
- spills.size(),
- spills.size() > 1 ? " times" : " time");
-
- final Tuple2<TempShuffleBlockId, File> spilledFileInfo =
- blockManager.diskBlockManager().createTempShuffleBlock();
- final File file = spilledFileInfo._2();
- final TempShuffleBlockId blockId = spilledFileInfo._1();
- final SpillInfo spillInfo = new SpillInfo(numPartitions, file, blockId);
-
- activeSpillSorter.setSpillInfo(spillInfo);
-
- activeSpillSorter.writeSortedFileNative(false, tracingEnabled);
- final long spillSize = activeSpillSorter.freeMemory();
- activeSpillSorter.reset();
-
- // Reset the in-memory sorter's pointer array only after freeing up the memory pages holding
- // the records. Otherwise, if the task is over allocated memory, then without freeing the
- // memory pages, we might not be able to get memory for the pointer array.
- taskContext.taskMetrics().incMemoryBytesSpilled(spillSize);
-
- spilling = false;
- }
-
- private long getMemoryUsage() {
- if (activeSpillSorter != null) {
- return activeSpillSorter.getMemoryUsage();
- }
- return 0;
- }
-
- private void updatePeakMemoryUsed() {
- long mem = getMemoryUsage();
- if (mem > peakMemoryUsedBytes) {
- peakMemoryUsedBytes = mem;
- }
- }
-
- /** Return the peak memory used so far, in bytes. */
- @Override
- public long getPeakMemoryUsedBytes() {
- updatePeakMemoryUsed();
- return peakMemoryUsedBytes;
- }
-
- private long freeMemory() {
- updatePeakMemoryUsed();
- long memoryFreed = activeSpillSorter.freeMemory();
- activeSpillSorter.freeArray();
- return memoryFreed;
- }
-
- /** Force all memory and spill files to be deleted; called by shuffle error-handling code. */
- @Override
- public void cleanupResources() {
- freeMemory();
-
- for (SpillInfo spill : spills) {
- if (spill.file.exists() && !spill.file.delete()) {
- logger.error("Unable to delete spill file {}", spill.file.getPath());
- }
- }
- }
-
- /**
- * Checks whether there is enough space to insert an additional record in to the sort pointer
- * array and grows the array if additional space is required. If the required space cannot be
- * obtained, then the in-memory data will be spilled to disk.
- */
- private void growPointerArrayIfNecessary() throws IOException {
- assert (activeSpillSorter != null);
- if (!activeSpillSorter.hasSpaceForAnotherRecord()) {
- long used = activeSpillSorter.getMemoryUsage();
- LongArray array;
- try {
- // could trigger spilling
- array = allocator.allocateArray(used / 8 * 2);
- } catch (TooLargePageException e) {
- // The pointer array is too big to fix in a single page, spill.
- spill();
- return;
- } catch (SparkOutOfMemoryError e) {
- // Cannot allocate enough memory, spill and reset pointer array.
- try {
- spill();
- } catch (SparkOutOfMemoryError e2) {
- // Cannot allocate memory even after spilling, throw the error.
- if (!activeSpillSorter.hasSpaceForAnotherRecord()) {
- logger.error("Unable to grow the pointer array");
- throw e2;
- }
- }
- return;
- }
- // check if spilling is triggered or not
- if (activeSpillSorter.hasSpaceForAnotherRecord()) {
- allocator.freeArray(array);
- } else {
- activeSpillSorter.expandPointerArray(array);
- }
- }
- }
-
- /**
- * Writes a record to the shuffle sorter. This copies the record data into this external sorter's
- * managed memory, which may trigger spilling if the copy would exceed the memory limit. It
- * inserts a pointer for the record and record's partition id into the in-memory sorter.
- */
- @Override
- public void insertRecord(Object recordBase, long recordOffset, int length, int partitionId)
- throws IOException {
-
- assert (activeSpillSorter != null);
- int threshold = numElementsForSpillThreshold;
- if (activeSpillSorter.numRecords() >= threshold) {
- logger.info(
- "Spilling data because number of spilledRecords crossed the threshold " + threshold);
- spill();
- }
-
- growPointerArrayIfNecessary();
-
- // Need 4 or 8 bytes to store the record length.
- final int required = length + uaoSize;
- // Acquire enough memory to store the record.
- // If we cannot acquire enough memory, we will spill current writers.
- if (!activeSpillSorter.acquireNewPageIfNecessary(required)) {
- // Spilling is happened, initiate new memory page for new writer.
- activeSpillSorter.initialCurrentPage(required);
- }
-
- activeSpillSorter.insertRecord(recordBase, recordOffset, length, partitionId);
- }
-
- /**
- * Close the sorter, causing any buffered data to be sorted and written out to disk.
- *
- * @return metadata for the spill files written by this sorter. If no records were ever inserted
- * into this sorter, then this will return an empty array.
- */
- @Override
- public SpillInfo[] closeAndGetSpills() throws IOException {
- if (activeSpillSorter != null) {
- // Do not count the final file towards the spill count.
- final Tuple2<TempShuffleBlockId, File> spilledFileInfo =
- blockManager.diskBlockManager().createTempShuffleBlock();
- final File file = spilledFileInfo._2();
- final TempShuffleBlockId blockId = spilledFileInfo._1();
- final SpillInfo spillInfo = new SpillInfo(numPartitions, file, blockId);
-
- activeSpillSorter.setSpillInfo(spillInfo);
- activeSpillSorter.writeSortedFileNative(true, tracingEnabled);
-
- freeMemory();
- }
-
- return spills.toArray(new SpillInfo[spills.size()]);
- }
-}
diff --git a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometBypassMergeSortShuffleWriter.java b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometBypassMergeSortShuffleWriter.java
index a364d34ed2..07da649f4a 100644
--- a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometBypassMergeSortShuffleWriter.java
+++ b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometBypassMergeSortShuffleWriter.java
@@ -25,7 +25,6 @@ import java.io.IOException;
import java.io.OutputStream;
import java.nio.channels.FileChannel;
import java.util.Optional;
-import java.util.concurrent.ExecutorService;
import javax.annotation.Nullable;
import scala.*;
@@ -103,13 +102,6 @@ final class CometBypassMergeSortShuffleWriter<K, V> extends ShuffleWriter<K, V>
/** Checksum calculator for each partition. Empty when shuffle checksum disabled. */
private final long[] partitionChecksums;
- private final boolean isAsync;
-
- private final int asyncThreadNum;
-
- /** Thread pool shared across all partition writers, for async write batch */
- private final ExecutorService threadPool;
-
/**
* Are we in the process of stopping? Because map tasks can call stop() with success = true and
* then call stop() with success = false if they get an exception, we want to make sure we don't
@@ -152,17 +144,7 @@ final class CometBypassMergeSortShuffleWriter<K, V> extends ShuffleWriter<K, V>
this.schema = ((CometShuffleDependency<?, ?, ?>) dep).schema().get();
this.partitionChecksums = createPartitionChecksums(numPartitions, conf);
- this.isAsync = (boolean) CometConf$.MODULE$.COMET_COLUMNAR_SHUFFLE_ASYNC_ENABLED().get();
- this.asyncThreadNum = (int) CometConf$.MODULE$.COMET_COLUMNAR_SHUFFLE_ASYNC_THREAD_NUM().get();
this.tracingEnabled = (boolean) CometConf$.MODULE$.COMET_TRACING_ENABLED().get();
-
- if (isAsync) {
- logger.info("Async shuffle writer enabled");
- this.threadPool = ShuffleThreadPool.getThreadPool();
- } else {
- logger.info("Async shuffle writer disabled");
- this.threadPool = null;
- }
this.encodeTimeMetric = encodeTimeMetric;
}
@@ -208,9 +190,6 @@ final class CometBypassMergeSortShuffleWriter<K, V> extends ShuffleWriter<K, V>
schema,
writeMetrics,
conf,
- isAsync,
- asyncThreadNum,
- threadPool,
tracingEnabled);
if (partitionChecksums.length > 0) {
writer.setChecksum(partitionChecksums[i]);
diff --git a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriter.java b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriter.java
index fd871a7ef9..c47953b051 100644
--- a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriter.java
+++ b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriter.java
@@ -23,9 +23,6 @@ import java.io.*;
import java.util.Collections;
import java.util.Comparator;
import java.util.LinkedList;
-import java.util.concurrent.ConcurrentLinkedQueue;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Future;
import java.util.concurrent.atomic.AtomicLong;
import scala.reflect.ClassTag;
@@ -60,11 +57,6 @@ import org.apache.comet.Native;
* stream the data to disk. However, for Comet, we need to buffer rows in memory page and then write
* them to disk as batches in Arrow IPC format. So, we need to extend `MemoryConsumer` to be able to
* spill the buffered rows to disk when memory pressure is high.
- *
- * <p>Similar to `CometShuffleExternalSorter`, this class also provides asynchronous spill
- * mechanism. But different from `CometShuffleExternalSorter`, as a writer class for Spark
- * hash-based shuffle, it writes all the rows for a partition into a single file, instead of each
- * file for each spill.
*/
public final class CometDiskBlockWriter {
private static final Logger logger = LoggerFactory.getLogger(CometDiskBlockWriter.class);
@@ -73,9 +65,6 @@ public final class CometDiskBlockWriter {
/** List of all `NativeDiskBlockArrowIPCWriter`s of same shuffle task. */
private static final LinkedList<CometDiskBlockWriter> currentWriters = new LinkedList<>();
- /** Queue of pending asynchronous spill tasks. */
- private ConcurrentLinkedQueue<Future<Void>> asyncSpillTasks = new ConcurrentLinkedQueue<>();
-
/** List of `ArrowIPCWriter`s which are spilling. */
private final LinkedList<ArrowIPCWriter> spillingWriters = new LinkedList<>();
@@ -104,10 +93,7 @@ public final class CometDiskBlockWriter {
private final int columnarBatchSize;
private final String compressionCodec;
private final int compressionLevel;
- private final boolean isAsync;
private final boolean tracingEnabled;
- private final int asyncThreadNum;
- private final ExecutorService threadPool;
private final int numElementsForSpillThreshold;
private final double preferDictionaryRatio;
@@ -145,9 +131,6 @@ public final class CometDiskBlockWriter {
StructType schema,
ShuffleWriteMetricsReporter writeMetrics,
SparkConf conf,
- boolean isAsync,
- int asyncThreadNum,
- ExecutorService threadPool,
boolean tracingEnabled) {
this.nativeLib = new Native();
this.allocator = allocator;
@@ -156,10 +139,7 @@ public final class CometDiskBlockWriter {
this.schema = schema;
this.writeMetrics = writeMetrics;
this.file = file;
- this.isAsync = isAsync;
this.tracingEnabled = tracingEnabled;
- this.asyncThreadNum = asyncThreadNum;
- this.threadPool = threadPool;
this.columnarBatchSize = (int) CometConf$.MODULE$.COMET_COLUMNAR_SHUFFLE_BATCH_SIZE().get();
this.compressionCodec = CometConf$.MODULE$.COMET_EXEC_SHUFFLE_COMPRESSION_CODEC().get();
@@ -191,7 +171,7 @@ public final class CometDiskBlockWriter {
return this.activeWriter.getChecksum();
}
- private void doSpill(boolean forceSync) throws IOException {
+ private void doSpill() throws IOException {
// We only allow spilling request from `NativeDiskBlockArrowIPCWriter`.
if (spilling || activeWriter.numRecords() == 0) {
return;
@@ -200,50 +180,11 @@ public final class CometDiskBlockWriter {
// Set this into spilling state first, so it cannot recursively trigger another spill on itself.
spilling = true;
- if (isAsync && !forceSync) {
- // Although we can continue to submit spill tasks to thread pool, buffering more rows in
- // memory page will increase memory usage. So, we need to wait for at least one spilling
- // task to finish.
- while (asyncSpillTasks.size() == asyncThreadNum) {
- for (Future<Void> task : asyncSpillTasks) {
- if (task.isDone()) {
- asyncSpillTasks.remove(task);
- break;
- }
- }
- }
-
- final ArrowIPCWriter spillingWriter = activeWriter;
- activeWriter = new ArrowIPCWriter();
-
- spillingWriters.add(spillingWriter);
-
- asyncSpillTasks.add(
- threadPool.submit(
- new Runnable() {
- @Override
- public void run() {
- try {
- long written = spillingWriter.doSpilling(false);
- totalWritten += written;
- } catch (IOException e) {
- throw new RuntimeException(e);
- } finally {
- spillingWriter.freeMemory();
- spillingWriters.remove(spillingWriter);
- }
- }
- },
- null));
-
- } else {
- // Spill in a synchronous way.
- // This spill could be triggered by other thread (i.e., other `CometDiskBlockWriter`),
- // so we need to synchronize it.
- synchronized (CometDiskBlockWriter.this) {
- totalWritten += activeWriter.doSpilling(false);
- activeWriter.freeMemory();
- }
+ // This spill could be triggered by other thread (i.e., other `CometDiskBlockWriter`),
+ // so we need to synchronize it.
+ synchronized (CometDiskBlockWriter.this) {
+ totalWritten += activeWriter.doSpilling(false);
+ activeWriter.freeMemory();
}
spilling = false;
@@ -281,7 +222,7 @@ public final class CometDiskBlockWriter {
logger.info(
"Spilling data because number of spilledRecords crossed the threshold " + threshold);
// Spill the current writer
- doSpill(false);
+ doSpill();
if (activeWriter.numRecords() != 0) {
throw new RuntimeException(
"activeWriter.numRecords()(" + activeWriter.numRecords() + ") != 0");
@@ -302,16 +243,6 @@ public final class CometDiskBlockWriter {
}
FileSegment close() throws IOException {
- if (isAsync) {
- for (Future<Void> task : asyncSpillTasks) {
- try {
- task.get();
- } catch (Exception e) {
- throw new RuntimeException(e);
- }
- }
- }
-
totalWritten += activeWriter.doSpilling(true);
if (outputRecords != insertRecords) {
@@ -367,8 +298,8 @@ public final class CometDiskBlockWriter {
this.nativeLib = CometDiskBlockWriter.this.nativeLib;
this.dataTypes = serializeSchema(schema);
- // Share the writer-level AtomicLong so all ArrowIPCWriter instances
- // (including async spilling ones) accumulate into the same counter.
+ // Share the writer-level AtomicLong so all ArrowIPCWriter instances (active and spilling)
+ // accumulate into the same counter.
this.setEncodeNanosAccumulator(CometDiskBlockWriter.this.encodeNanos);
}
@@ -462,11 +393,9 @@ public final class CometDiskBlockWriter {
long totalFreed = 0;
for (CometDiskBlockWriter writer : currentWriters) {
- // Force to spill the writer in a synchronous way, otherwise, we may not be able to
- // acquire enough memory.
long used = writer.getActiveMemoryUsage();
- writer.doSpill(true);
+ writer.doSpill();
totalFreed += used;
diff --git a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometUnsafeShuffleWriter.java b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometUnsafeShuffleWriter.java
index 392846a56e..7997f60532 100644
--- a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometUnsafeShuffleWriter.java
+++ b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometUnsafeShuffleWriter.java
@@ -263,7 +263,7 @@ public class CometUnsafeShuffleWriter<K, V> extends ShuffleWriter<K, V> {
Math.min(
CometShuffleExternalSorter.MAXIMUM_PAGE_SIZE_BYTES, memoryManager.pageSizeBytes()));
sorter =
- CometShuffleExternalSorter.create(
+ new CometShuffleExternalSorter(
allocator,
blockManager,
taskContext,
diff --git a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/ShuffleThreadPool.java b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/ShuffleThreadPool.java
deleted file mode 100644
index 86d7d7dc21..0000000000
--- a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/ShuffleThreadPool.java
+++ /dev/null
@@ -1,71 +0,0 @@
-/*
- * 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.comet.execution.shuffle;
-
-import java.util.concurrent.*;
-
-import com.google.common.util.concurrent.ThreadFactoryBuilder;
-
-import org.apache.comet.CometConf$;
-
-public class ShuffleThreadPool {
- private static ThreadPoolExecutor INSTANCE;
-
- /** Get the thread pool instance for shuffle writer. */
- public static synchronized ExecutorService getThreadPool() {
- if (INSTANCE == null) {
- boolean isAsync = (boolean) CometConf$.MODULE$.COMET_COLUMNAR_SHUFFLE_ASYNC_ENABLED().get();
-
- if (isAsync) {
- ThreadFactory factory =
- new ThreadFactoryBuilder().setNameFormat("async-shuffle-writer-%d").build();
-
- int threadNum =
- (int) CometConf$.MODULE$.COMET_COLUMNAR_SHUFFLE_ASYNC_MAX_THREAD_NUM().get();
- INSTANCE =
- new ThreadPoolExecutor(
- 0, threadNum, 1L, TimeUnit.SECONDS, new ThreadPoolQueue(threadNum), factory);
- }
- }
-
- return INSTANCE;
- }
-}
-
-/**
- * A blocking queue for thread pool. This will block new task submission until there is space in the
- * queue.
- */
-final class ThreadPoolQueue extends ArrayBlockingQueue<Runnable> {
- public ThreadPoolQueue(int capacity) {
- super(capacity);
- }
-
- @Override
- public boolean offer(Runnable e) {
- try {
- put(e);
- } catch (InterruptedException e1) {
- Thread.currentThread().interrupt();
- return false;
- }
- return true;
- }
-}
diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala
index b1fbeea5cb..41c99d2724 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/spark/src/main/scala/org/apache/comet/CometConf.scala
@@ -461,33 +461,15 @@ object CometConf extends ShimCometConf {
.intConf
.createWithDefault(1)
- val COMET_COLUMNAR_SHUFFLE_ASYNC_ENABLED: ConfigEntry[Boolean] =
- conf("spark.comet.columnar.shuffle.async.enabled")
- .category(CATEGORY_SHUFFLE)
- .doc("Whether to enable asynchronous shuffle for Arrow-based shuffle.")
- .booleanConf
- .createWithDefault(false)
-
- val COMET_COLUMNAR_SHUFFLE_ASYNC_THREAD_NUM: ConfigEntry[Int] =
- conf("spark.comet.columnar.shuffle.async.thread.num")
+ val COMET_COLUMNAR_SHUFFLE_MAX_WRITERS_PER_EXECUTOR: ConfigEntry[Int] = {
+ conf("spark.comet.columnar.shuffle.max.writers.per.executor")
+ .withAlternative("spark.comet.columnar.shuffle.async.max.thread.num")
.category(CATEGORY_SHUFFLE)
.doc(
- "Number of threads used for Comet async columnar shuffle per shuffle task. " +
- "Note that more threads means more memory requirement to " +
- "buffer shuffle data before flushing to disk. Also, more threads may not always " +
- "improve performance, and should be set based on the number of cores available.")
- .intConf
- .createWithDefault(3)
-
- val COMET_COLUMNAR_SHUFFLE_ASYNC_MAX_THREAD_NUM: ConfigEntry[Int] = {
- conf("spark.comet.columnar.shuffle.async.max.thread.num")
- .category(CATEGORY_SHUFFLE)
- .doc("Maximum number of threads on an executor used for Comet async columnar shuffle. " +
- "This is the upper bound of total number of shuffle " +
- "threads per executor. In other words, if the number of cores * the number of shuffle " +
- "threads per task `spark.comet.columnar.shuffle.async.thread.num` is larger than " +
- "this config. Comet will use this config as the number of shuffle threads per " +
- "executor instead.")
+ "Maximum number of concurrent hash-based shuffle writers per executor. Comet's " +
+ "hash-based (bypass merge sort) columnar shuffle allocates one writer per partition, " +
+ "so `executor cores * numPartitions` writers can be active concurrently. When that " +
+ "product exceeds this cap, Comet falls back to sort-based shuffle to avoid OOM.")
.intConf
.createWithDefault(100)
}
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleManager.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleManager.scala
index bd69e91898..06f2ff02d3 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleManager.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleManager.scala
@@ -303,9 +303,10 @@ object CometShuffleManager extends Logging {
val partitionCond = SortShuffleWriter.shouldBypassMergeSort(conf, dep)
// Bypass merge sort if we have partition * cores fewer than
- // `spark.comet.columnar.shuffle.async.max.thread.num`
+ // `spark.comet.columnar.shuffle.max.writers.per.executor`
val executorCores = conf.get(config.EXECUTOR_CORES)
- val maxThreads = CometConf.COMET_COLUMNAR_SHUFFLE_ASYNC_MAX_THREAD_NUM.get(SQLConf.get)
+ val maxThreads =
+ CometConf.COMET_COLUMNAR_SHUFFLE_MAX_WRITERS_PER_EXECUTOR.get(SQLConf.get)
val threadCond = dep.partitioner.numPartitions * executorCores <= maxThreads
// Comet columnar shuffle buffers rows in memory. If too many cores are used with
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala
index bc484b6599..1cbb414c04 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala
@@ -50,13 +50,10 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar
.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, adaptiveExecutionEnabled.toString)
}
- protected val asyncShuffleEnable: Boolean
-
override protected def test(testName: String, testTags: Tag*)(testFun: => Any)(implicit
pos: Position): Unit = {
super.test(testName, testTags: _*) {
withSQLConf(
- CometConf.COMET_COLUMNAR_SHUFFLE_ASYNC_ENABLED.key -> asyncShuffleEnable.toString,
CometConf.COMET_COLUMNAR_SHUFFLE_SPILL_THRESHOLD.key -> numElementsForceSpillThreshold.toString,
CometConf.COMET_EXEC_ENABLED.key -> "true",
CometConf.COMET_SHUFFLE_MODE.key -> "jvm",
@@ -801,15 +798,7 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar
}
}
-class CometAsyncShuffleSuite extends CometColumnarShuffleSuite {
- override protected val asyncShuffleEnable: Boolean = true
-
- protected val adaptiveExecutionEnabled: Boolean = true
-}
-
class CometShuffleSuite extends CometColumnarShuffleSuite {
- override protected val asyncShuffleEnable: Boolean = false
-
protected val adaptiveExecutionEnabled: Boolean = true
import testImplicits._
@@ -861,8 +850,6 @@ class CometShuffleSuite extends CometColumnarShuffleSuite {
}
class DisableAQECometShuffleSuite extends CometColumnarShuffleSuite {
- override protected val asyncShuffleEnable: Boolean = false
-
protected val adaptiveExecutionEnabled: Boolean = false
import testImplicits._
@@ -892,12 +879,6 @@ class DisableAQECometShuffleSuite extends CometColumnarShuffleSuite {
}
}
-class DisableAQECometAsyncShuffleSuite extends CometColumnarShuffleSuite {
- override protected val asyncShuffleEnable: Boolean = true
-
- protected val adaptiveExecutionEnabled: Boolean = false
-}
-
class CometShuffleEncryptionSuite extends CometTestBase {
import testImplicits._
@@ -910,23 +891,20 @@ class CometShuffleEncryptionSuite extends CometTestBase {
test("comet columnar shuffle with encryption") {
Seq(10, 201).foreach { numPartitions =>
Seq(true, false).foreach { dictionaryEnabled =>
- Seq(true, false).foreach { asyncEnabled =>
- withTempDir { dir =>
- val path = new Path(dir.toURI.toString, "test.parquet")
- makeParquetFileAllPrimitiveTypes(path, dictionaryEnabled = dictionaryEnabled, 1000)
-
- (1 until 10).map(i => $"_$i").foreach { col =>
- withSQLConf(
- CometConf.COMET_EXEC_ENABLED.key -> "true",
- CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
- CometConf.COMET_SHUFFLE_MODE.key -> "jvm",
- CometConf.COMET_COLUMNAR_SHUFFLE_ASYNC_ENABLED.key -> asyncEnabled.toString) {
- readParquetFile(path.toString) { df =>
- val shuffled = df
- .select($"_1")
- .repartition(numPartitions, col)
- checkSparkAnswer(shuffled)
- }
+ withTempDir { dir =>
+ val path = new Path(dir.toURI.toString, "test.parquet")
+ makeParquetFileAllPrimitiveTypes(path, dictionaryEnabled = dictionaryEnabled, 1000)
+
+ (1 until 10).map(i => $"_$i").foreach { col =>
+ withSQLConf(
+ CometConf.COMET_EXEC_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
+ CometConf.COMET_SHUFFLE_MODE.key -> "jvm") {
+ readParquetFile(path.toString) { df =>
+ val shuffled = df
+ .select($"_1")
+ .repartition(numPartitions, col)
+ checkSparkAnswer(shuffled)
}
}
}
@@ -939,7 +917,7 @@ class CometShuffleEncryptionSuite extends CometTestBase {
class CometShuffleManagerSuite extends CometTestBase {
test("should not bypass merge sort if executor cores are too high") {
- withSQLConf(CometConf.COMET_COLUMNAR_SHUFFLE_ASYNC_MAX_THREAD_NUM.key -> "100") {
+ withSQLConf(CometConf.COMET_COLUMNAR_SHUFFLE_MAX_WRITERS_PER_EXECUTOR.key -> "100") {
val conf = new SparkConf()
conf.set("spark.executor.cores", "1")
diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometExecBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometExecBenchmark.scala
index 740707aefd..f458d84a3f 100644
--- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometExecBenchmark.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometExecBenchmark.scala
@@ -46,7 +46,6 @@ object CometExecBenchmark extends CometBenchmarkBase {
.set(
"spark.shuffle.manager",
"org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager")
- .set("spark.comet.columnar.shuffle.async.thread.num", "7")
.set("spark.comet.columnar.shuffle.spill.threshold", "30000")
val sparkSession = SparkSession.builder
diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleBenchmark.scala
index 6751959105..a8961457ed 100644
--- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleBenchmark.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleBenchmark.scala
@@ -53,7 +53,6 @@ object CometShuffleBenchmark extends CometBenchmarkBase {
.set(
"spark.shuffle.manager",
"org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager")
- .set("spark.comet.columnar.shuffle.async.thread.num", "7")
.set("spark.comet.columnar.shuffle.spill.threshold", "30000")
val sparkSession = SparkSession.builder
@@ -109,8 +108,7 @@ object CometShuffleBenchmark extends CometBenchmarkBase {
CometConf.COMET_EXEC_ENABLED.key -> "true",
CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
CometConf.COMET_SHUFFLE_PREFER_DICTIONARY_RATIO.key -> "1.0",
- CometConf.COMET_SHUFFLE_MODE.key -> "jvm",
- CometConf.COMET_COLUMNAR_SHUFFLE_ASYNC_ENABLED.key -> "false") {
+ CometConf.COMET_SHUFFLE_MODE.key -> "jvm") {
spark
.sql(
s"SELECT ARRAY_REPEAT(CAST(1 AS ${dataType.sql}), 10) AS c1 FROM parquetV1Table")
@@ -166,8 +164,7 @@ object CometShuffleBenchmark extends CometBenchmarkBase {
CometConf.COMET_EXEC_ENABLED.key -> "true",
CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
CometConf.COMET_SHUFFLE_PREFER_DICTIONARY_RATIO.key -> "1.0",
- CometConf.COMET_SHUFFLE_MODE.key -> "jvm",
- CometConf.COMET_COLUMNAR_SHUFFLE_ASYNC_ENABLED.key -> "false") {
+ CometConf.COMET_SHUFFLE_MODE.key -> "jvm") {
spark
.sql(
s"SELECT STRUCT(CAST(c1 AS ${dataType.sql})," +
@@ -221,8 +218,7 @@ object CometShuffleBenchmark extends CometBenchmarkBase {
CometConf.COMET_EXEC_ENABLED.key -> "true",
CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
CometConf.COMET_SHUFFLE_PREFER_DICTIONARY_RATIO.key -> "1.0",
- CometConf.COMET_SHUFFLE_MODE.key -> "jvm",
- CometConf.COMET_COLUMNAR_SHUFFLE_ASYNC_ENABLED.key -> "false") {
+ CometConf.COMET_SHUFFLE_MODE.key -> "jvm") {
spark
.sql("select c1 from parquetV1Table")
.repartition(partitionNum, Column("c1"))
@@ -236,8 +232,7 @@ object CometShuffleBenchmark extends CometBenchmarkBase {
CometConf.COMET_EXEC_ENABLED.key -> "true",
CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
CometConf.COMET_SHUFFLE_PREFER_DICTIONARY_RATIO.key -> "2.0",
- CometConf.COMET_SHUFFLE_MODE.key -> "jvm",
- CometConf.COMET_COLUMNAR_SHUFFLE_ASYNC_ENABLED.key -> "false") {
+ CometConf.COMET_SHUFFLE_MODE.key -> "jvm") {
spark
.sql("select c1 from parquetV1Table")
.repartition(partitionNum, Column("c1"))
@@ -251,8 +246,7 @@ object CometShuffleBenchmark extends CometBenchmarkBase {
CometConf.COMET_EXEC_ENABLED.key -> "true",
CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
CometConf.COMET_SHUFFLE_PREFER_DICTIONARY_RATIO.key -> "1000000000.0",
- CometConf.COMET_SHUFFLE_MODE.key -> "jvm",
- CometConf.COMET_COLUMNAR_SHUFFLE_ASYNC_ENABLED.key -> "false") {
+ CometConf.COMET_SHUFFLE_MODE.key -> "jvm") {
spark
.sql("select c1 from parquetV1Table")
.repartition(partitionNum, Column("c1"))
@@ -316,22 +310,7 @@ object CometShuffleBenchmark extends CometBenchmarkBase {
CometConf.COMET_ENABLED.key -> "true",
CometConf.COMET_EXEC_ENABLED.key -> "true",
CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
- CometConf.COMET_SHUFFLE_MODE.key -> "jvm",
- CometConf.COMET_COLUMNAR_SHUFFLE_ASYNC_ENABLED.key -> "false") {
- spark
- .sql("select c1 from parquetV1Table")
- .repartition(partitionNum, Column("c1"))
- .noop()
- }
- }
-
- benchmark.addCase("SQL Parquet - Comet (Comet Async JVM Shuffle)") { _ =>
- withSQLConf(
- CometConf.COMET_ENABLED.key -> "true",
- CometConf.COMET_EXEC_ENABLED.key -> "true",
- CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
- CometConf.COMET_SHUFFLE_MODE.key -> "jvm",
- CometConf.COMET_COLUMNAR_SHUFFLE_ASYNC_ENABLED.key -> "true") {
+ CometConf.COMET_SHUFFLE_MODE.key -> "jvm") {
spark
.sql("select c1 from parquetV1Table")
.repartition(partitionNum, Column("c1"))
diff --git a/spark/src/test/spark-4.x/org/apache/comet/exec/CometShuffle4_0Suite.scala b/spark/src/test/spark-4.x/org/apache/comet/exec/CometShuffle4_0Suite.scala
index 3d0ec5006d..8629d1a78c 100644
--- a/spark/src/test/spark-4.x/org/apache/comet/exec/CometShuffle4_0Suite.scala
+++ b/spark/src/test/spark-4.x/org/apache/comet/exec/CometShuffle4_0Suite.scala
@@ -29,8 +29,6 @@ import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.{FloatType, LongType, StringType, TimestampType}
class CometShuffle4_0Suite extends CometColumnarShuffleSuite {
- override protected val asyncShuffleEnable: Boolean = false
-
protected val adaptiveExecutionEnabled: Boolean = true
override def beforeAll(): Unit = {
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org