diff --git a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxMetricsApi.scala b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxMetricsApi.scala index 421e6c35959..5d202ee1538 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxMetricsApi.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxMetricsApi.scala @@ -17,7 +17,7 @@ package org.apache.gluten.backendsapi.velox import org.apache.gluten.backendsapi.MetricsApi -import org.apache.gluten.config.{GpuHashShuffleWriterType, HashShuffleWriterType, RssSortShuffleWriterType, ShuffleWriterType, SortShuffleWriterType} +import org.apache.gluten.config.{HashShuffleWriterType, RssSortShuffleWriterType, ShuffleWriterType, SortShuffleWriterType} import org.apache.gluten.metrics._ import org.apache.gluten.substrait.{AggregationParams, JoinParams} @@ -430,7 +430,7 @@ class VeloxMetricsApi extends MetricsApi with Logging { "peakBytes" -> SQLMetrics.createSizeMetric(sparkContext, "peak bytes allocated") ) shuffleWriterType match { - case HashShuffleWriterType | GpuHashShuffleWriterType => + case HashShuffleWriterType => baseMetrics ++ Map( "splitTime" -> SQLMetrics.createNanoTimingMetric(sparkContext, "time to split"), "avgDictionaryFields" -> SQLMetrics diff --git a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala index 0397a6dc2c7..d588695cd40 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala @@ -607,7 +607,12 @@ class VeloxSparkPlanExecApi extends SparkPlanExecApi with Logging { /** Determine whether to use sort-based shuffle based on shuffle partitioning and output. */ override def getShuffleWriterType( partitioning: Partitioning, - output: Seq[Attribute]): ShuffleWriterType = { + output: Seq[Attribute], + executionMode: Option[StageExecutionMode] = None): ShuffleWriterType = { + if (executionMode.contains(GPUStageMode)) { + return HashShuffleWriterType + } + val conf = GlutenConfig.get // todo: remove isUseCelebornShuffleManager here if (conf.isUseCelebornShuffleManager) { diff --git a/backends-velox/src/main/scala/org/apache/spark/shuffle/ColumnarShuffleWriter.scala b/backends-velox/src/main/scala/org/apache/spark/shuffle/ColumnarShuffleWriter.scala index 73bb2770e51..f0d814c1561 100644 --- a/backends-velox/src/main/scala/org/apache/spark/shuffle/ColumnarShuffleWriter.scala +++ b/backends-velox/src/main/scala/org/apache/spark/shuffle/ColumnarShuffleWriter.scala @@ -18,7 +18,7 @@ package org.apache.spark.shuffle import org.apache.gluten.backendsapi.BackendsApiManager import org.apache.gluten.columnarbatch.ColumnarBatches -import org.apache.gluten.config.{GlutenConfig, GpuHashShuffleWriterType, HashShuffleWriterType, SortShuffleWriterType} +import org.apache.gluten.config.{GlutenConfig, HashShuffleWriterType, SortShuffleWriterType} import org.apache.gluten.memory.memtarget.{MemoryTarget, Spiller} import org.apache.gluten.runtime.Runtimes import org.apache.gluten.vectorized._ @@ -44,7 +44,7 @@ class ColumnarShuffleWriter[K, V]( private val dep = handle.dependency.asInstanceOf[ColumnarShuffleDependency[K, V, V]] dep.shuffleWriterType match { - case HashShuffleWriterType | SortShuffleWriterType | GpuHashShuffleWriterType => + case HashShuffleWriterType | SortShuffleWriterType => // Valid shuffle writer types case _ => throw new IllegalArgumentException( @@ -174,17 +174,6 @@ class ColumnarShuffleWriter[K, V]( conf.get(SHUFFLE_SORT_USE_RADIXSORT), partitionWriterHandle ) - } else if (dep.shuffleWriterType == GpuHashShuffleWriterType) { - shuffleWriterJniWrapper.createGpuHashShuffleWriter( - numPartitions, - dep.nativePartitioning.getShortName, - GlutenShuffleUtils.getStartPartitionId( - dep.nativePartitioning, - taskContext.partitionId), - nativeBufferSize, - reallocThreshold, - partitionWriterHandle - ) } else { shuffleWriterJniWrapper.createHashShuffleWriter( numPartitions, diff --git a/cpp/core/jni/JniWrapper.cc b/cpp/core/jni/JniWrapper.cc index cadd32bb47f..e4d674bd74e 100644 --- a/cpp/core/jni/JniWrapper.cc +++ b/cpp/core/jni/JniWrapper.cc @@ -1029,34 +1029,6 @@ JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_ShuffleWriterJniWrappe JNI_METHOD_END(kInvalidObjectHandle) } -JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_ShuffleWriterJniWrapper_createGpuHashShuffleWriter( - JNIEnv* env, - jobject wrapper, - jint numPartitions, - jstring partitioningNameJstr, - jint startPartitionId, - jint splitBufferSize, - jdouble splitBufferReallocThreshold, - jlong partitionWriterHandle) { - JNI_METHOD_START - - const auto ctx = getRuntime(env, wrapper); - - auto partitionWriter = ObjectStore::retrieve(partitionWriterHandle); - if (partitionWriter == nullptr) { - throw GlutenException("Partition writer handle is invalid: " + std::to_string(partitionWriterHandle)); - } - ObjectStore::release(partitionWriterHandle); - - auto shuffleWriterOptions = std::make_shared( - toPartitioning(jStringToCString(env, partitioningNameJstr)), - startPartitionId, - splitBufferSize, - splitBufferReallocThreshold); - return ctx->saveObject(ctx->createShuffleWriter(numPartitions, partitionWriter, shuffleWriterOptions)); - JNI_METHOD_END(kInvalidObjectHandle) -} - JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_ShuffleWriterJniWrapper_createSortShuffleWriter( JNIEnv* env, jobject wrapper, diff --git a/cpp/core/shuffle/Options.h b/cpp/core/shuffle/Options.h index ead95dfe657..ec246ce80bc 100644 --- a/cpp/core/shuffle/Options.h +++ b/cpp/core/shuffle/Options.h @@ -46,7 +46,7 @@ static constexpr int64_t kDefaultShuffleFileBufferSize = 32 << 10; static constexpr bool kDefaultEnableDictionary = false; static constexpr bool kDefaultEnableTypeAwareCompress = false; -enum class ShuffleWriterType { kHashShuffle, kSortShuffle, kRssSortShuffle, kGpuHashShuffle }; +enum class ShuffleWriterType { kHashShuffle, kSortShuffle, kRssSortShuffle }; enum class PartitionWriterType { kLocal, kRss }; @@ -161,25 +161,6 @@ struct RssSortShuffleWriterOptions : ShuffleWriterOptions { compressionType(compressionType) {} }; -struct GpuHashShuffleWriterOptions : HashShuffleWriterOptions { - int32_t splitBufferSize = kDefaultShuffleWriterBufferSize; - double splitBufferReallocThreshold = kDefaultSplitBufferReallocThreshold; - - GpuHashShuffleWriterOptions() : HashShuffleWriterOptions(ShuffleWriterType::kGpuHashShuffle) {} - - GpuHashShuffleWriterOptions( - Partitioning partitioning, - int32_t startPartitionId, - int32_t partitionBufferSize, - double partitionBufferReallocThreshold) - : HashShuffleWriterOptions( - ShuffleWriterType::kGpuHashShuffle, - partitioning, - startPartitionId, - partitionBufferSize, - partitionBufferReallocThreshold) {} -}; - struct LocalPartitionWriterOptions { int64_t shuffleFileBufferSize = kDefaultShuffleFileBufferSize; // spark.shuffle.file.buffer int32_t compressionBufferSize = diff --git a/cpp/core/shuffle/ShuffleWriter.cc b/cpp/core/shuffle/ShuffleWriter.cc index 7287181ffa9..5a48757231b 100644 --- a/cpp/core/shuffle/ShuffleWriter.cc +++ b/cpp/core/shuffle/ShuffleWriter.cc @@ -24,7 +24,6 @@ namespace { const std::string kHashShuffleName = "hash"; const std::string kSortShuffleName = "sort"; const std::string kRssSortShuffleName = "rss_sort"; -const std::string kGpuHashShuffleName = "gpu_hash"; } // namespace ShuffleWriterType ShuffleWriter::stringToType(const std::string& typeString) { @@ -37,9 +36,6 @@ ShuffleWriterType ShuffleWriter::stringToType(const std::string& typeString) { if (typeString == kRssSortShuffleName) { return ShuffleWriterType::kRssSortShuffle; } - if (typeString == kGpuHashShuffleName) { - return ShuffleWriterType::kGpuHashShuffle; - } throw GlutenException("Unrecognized shuffle writer type: " + typeString); } @@ -51,8 +47,6 @@ std::string ShuffleWriter::typeToString(ShuffleWriterType type) { return kSortShuffleName; case ShuffleWriterType::kRssSortShuffle: return kRssSortShuffleName; - case ShuffleWriterType::kGpuHashShuffle: - return kGpuHashShuffleName; } GLUTEN_UNREACHABLE(); } diff --git a/cpp/velox/shuffle/VeloxShuffleReader.cc b/cpp/velox/shuffle/VeloxShuffleReader.cc index eb95f2c8ea8..d10a0ca3054 100644 --- a/cpp/velox/shuffle/VeloxShuffleReader.cc +++ b/cpp/velox/shuffle/VeloxShuffleReader.cc @@ -967,8 +967,7 @@ void VeloxShuffleReader::createDeserializer( const std::shared_ptr& streamReader, const OutputType& outputType) { switch (options_->shuffleWriterType) { - case ShuffleWriterType::kHashShuffle: - case ShuffleWriterType::kGpuHashShuffle: { + case ShuffleWriterType::kHashShuffle: { if (outputType == OutputType::kCudfTable) { #ifdef GLUTEN_ENABLE_GPU VELOX_CHECK(!hasComplexType_); diff --git a/cpp/velox/shuffle/VeloxShuffleWriter.cc b/cpp/velox/shuffle/VeloxShuffleWriter.cc index a83e3a19471..c45d9419722 100644 --- a/cpp/velox/shuffle/VeloxShuffleWriter.cc +++ b/cpp/velox/shuffle/VeloxShuffleWriter.cc @@ -31,7 +31,6 @@ arrow::Result> VeloxShuffleWriter::create( std::shared_ptr shuffleWriter; switch (type) { case ShuffleWriterType::kHashShuffle: - case ShuffleWriterType::kGpuHashShuffle: return VeloxHashShuffleWriter::create(numPartitions, std::move(partitionWriter), options, memoryManager); case ShuffleWriterType::kSortShuffle: return VeloxSortShuffleWriter::create(numPartitions, std::move(partitionWriter), options, memoryManager); diff --git a/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc b/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc index d0df7ebba1b..46bafbfa42c 100644 --- a/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc +++ b/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc @@ -152,7 +152,7 @@ std::vector getTestParams() { for (const auto mergeBufferSize : mergeBufferSizes) { for (const auto enableGpuAsyncReader : {false, true}) { params.push_back(GpuShuffleTestParams{ - .shuffleWriterType = ShuffleWriterType::kGpuHashShuffle, + .shuffleWriterType = ShuffleWriterType::kHashShuffle, .partitionWriterType = PartitionWriterType::kLocal, .compressionType = compression, .compressionThreshold = compressionThreshold, @@ -163,7 +163,7 @@ std::vector getTestParams() { // Rss. params.push_back(GpuShuffleTestParams{ - .shuffleWriterType = ShuffleWriterType::kGpuHashShuffle, + .shuffleWriterType = ShuffleWriterType::kHashShuffle, .partitionWriterType = PartitionWriterType::kRss, .compressionType = compression, .compressionThreshold = compressionThreshold}); @@ -225,7 +225,7 @@ class GpuVeloxShuffleWriterTest : public ::testing::TestWithParam options; const auto& params = GetParam(); switch (params.shuffleWriterType) { - case ShuffleWriterType::kGpuHashShuffle: { + case ShuffleWriterType::kHashShuffle: { auto hashOptions = std::make_shared(); hashOptions->splitBufferSize = splitBufferSize; options = hashOptions; diff --git a/gluten-arrow/src/main/java/org/apache/gluten/vectorized/ShuffleWriterJniWrapper.java b/gluten-arrow/src/main/java/org/apache/gluten/vectorized/ShuffleWriterJniWrapper.java index 2f31a529d98..edbc48a3c97 100644 --- a/gluten-arrow/src/main/java/org/apache/gluten/vectorized/ShuffleWriterJniWrapper.java +++ b/gluten-arrow/src/main/java/org/apache/gluten/vectorized/ShuffleWriterJniWrapper.java @@ -65,14 +65,6 @@ public native long createRssSortShuffleWriter( String codec, long partitionWriterHandle); - public native long createGpuHashShuffleWriter( - int numPartitions, - String partitioningName, - int startPartitionId, - int splitBufferSize, - double splitBufferReallocThreshold, - long partitionWriterHandle); - /** * Reclaim memory from the shuffle writer instance. It will first try to shrink allocated memory, * and may trigger a spill if needed. diff --git a/gluten-substrait/src/main/scala/org/apache/gluten/backendsapi/SparkPlanExecApi.scala b/gluten-substrait/src/main/scala/org/apache/gluten/backendsapi/SparkPlanExecApi.scala index b0efc00fdfe..f5dc9f3f216 100644 --- a/gluten-substrait/src/main/scala/org/apache/gluten/backendsapi/SparkPlanExecApi.scala +++ b/gluten-substrait/src/main/scala/org/apache/gluten/backendsapi/SparkPlanExecApi.scala @@ -406,7 +406,8 @@ trait SparkPlanExecApi { /** Determine whether to use sort-based shuffle based on shuffle partitioning and output. */ def getShuffleWriterType( partitioning: Partitioning, - output: Seq[Attribute]): ShuffleWriterType = { + output: Seq[Attribute], + executionMode: Option[StageExecutionMode] = None): ShuffleWriterType = { HashShuffleWriterType } diff --git a/gluten-substrait/src/main/scala/org/apache/gluten/config/GlutenConfig.scala b/gluten-substrait/src/main/scala/org/apache/gluten/config/GlutenConfig.scala index 93fb3888e99..e5408bc8b5e 100644 --- a/gluten-substrait/src/main/scala/org/apache/gluten/config/GlutenConfig.scala +++ b/gluten-substrait/src/main/scala/org/apache/gluten/config/GlutenConfig.scala @@ -54,12 +54,6 @@ case object RssSortShuffleWriterType extends ShuffleWriterType { override val requiresResizingShuffleOutput: Boolean = false } -case object GpuHashShuffleWriterType extends ShuffleWriterType { - override val name: String = ReservedKeys.GLUTEN_GPU_HASH_SHUFFLE_WRITER - override val requiresResizingShuffleInput: Boolean = true - override val requiresResizingShuffleOutput: Boolean = true -} - /* * Note: Gluten configiguration.md is automatically generated from this code. * Make sure to run dev/gen-all-config-docs.sh after making changes to this file. diff --git a/gluten-substrait/src/main/scala/org/apache/gluten/config/ReservedKeys.scala b/gluten-substrait/src/main/scala/org/apache/gluten/config/ReservedKeys.scala index d609a80e2a7..6635a8b20ab 100644 --- a/gluten-substrait/src/main/scala/org/apache/gluten/config/ReservedKeys.scala +++ b/gluten-substrait/src/main/scala/org/apache/gluten/config/ReservedKeys.scala @@ -34,5 +34,4 @@ object ReservedKeys { val GLUTEN_HASH_SHUFFLE_WRITER = "hash" val GLUTEN_SORT_SHUFFLE_WRITER = "sort" val GLUTEN_RSS_SORT_SHUFFLE_WRITER = "rss_sort" - val GLUTEN_GPU_HASH_SHUFFLE_WRITER = "gpu_hash" } diff --git a/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ColumnarShuffleExchangeExec.scala b/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ColumnarShuffleExchangeExec.scala index 3accf8e592a..0fb50694587 100644 --- a/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ColumnarShuffleExchangeExec.scala +++ b/gluten-substrait/src/main/scala/org/apache/spark/sql/execution/ColumnarShuffleExchangeExec.scala @@ -17,8 +17,8 @@ package org.apache.spark.sql.execution import org.apache.gluten.backendsapi.BackendsApiManager -import org.apache.gluten.config.{GpuHashShuffleWriterType, ShuffleWriterType} -import org.apache.gluten.execution.{CPUStageMode, GPUStageMode, StageExecutionMode, ValidatablePlan, ValidationResult} +import org.apache.gluten.config.ShuffleWriterType +import org.apache.gluten.execution.{CPUStageMode, StageExecutionMode, ValidatablePlan, ValidationResult} import org.apache.gluten.extension.columnar.transition.Convention import org.apache.gluten.sql.shims.SparkShimLoader @@ -29,7 +29,7 @@ import org.apache.spark.serializer.Serializer import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.Attribute import org.apache.spark.sql.catalyst.plans.logical.Statistics -import org.apache.spark.sql.catalyst.plans.physical.{SinglePartition, _} +import org.apache.spark.sql.catalyst.plans.physical._ import org.apache.spark.sql.catalyst.util.truncatedString import org.apache.spark.sql.execution.exchange._ import org.apache.spark.sql.execution.metric.SQLShuffleWriteMetricsReporter @@ -129,14 +129,10 @@ case class ColumnarShuffleExchangeExec( } def getShuffleWriterType: ShuffleWriterType = { - mapperStageMode match { - case Some(GPUStageMode) => - GpuHashShuffleWriterType - case _ => - BackendsApiManager.getSparkPlanExecApiInstance.getShuffleWriterType( - outputPartitioning, - output) - } + BackendsApiManager.getSparkPlanExecApiInstance.getShuffleWriterType( + outputPartitioning, + output, + mapperStageMode) } // Required for Spark 4.0 to implement a trait method.