[VL] Remove gpu hash shuffle writer type (#12593)

Author: marin-ma

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

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) {

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,

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<PartitionWriter>(partitionWriterHandle);
-  if (partitionWriter == nullptr) {
-    throw GlutenException("Partition writer handle is invalid: " + std::to_string(partitionWriterHandle));
-  }
-  ObjectStore::release(partitionWriterHandle);
-
-  auto shuffleWriterOptions = std::make_shared<HashShuffleWriterOptions>(
-      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,

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 =

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();
 }

cpp/velox/shuffle/VeloxShuffleReader.cc

@@ -967,8 +967,7 @@ void VeloxShuffleReader::createDeserializer(
     const std::shared_ptr<StreamReader>& 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_);

cpp/velox/shuffle/VeloxShuffleWriter.cc

@@ -31,7 +31,6 @@ arrow::Result<std::shared_ptr<VeloxShuffleWriter>> VeloxShuffleWriter::create(
   std::shared_ptr<VeloxShuffleWriter> 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);

cpp/velox/tests/VeloxGpuShuffleWriterTest.cc

@@ -152,7 +152,7 @@ std::vector<GpuShuffleTestParams> 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<GpuShuffleTestParams> 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<GpuShuffleTest
     std::shared_ptr<ShuffleWriterOptions> options;
     const auto& params = GetParam();
     switch (params.shuffleWriterType) {
-      case ShuffleWriterType::kGpuHashShuffle: {
+      case ShuffleWriterType::kHashShuffle: {
         auto hashOptions = std::make_shared<HashShuffleWriterOptions>();
         hashOptions->splitBufferSize = splitBufferSize;
         options = hashOptions;

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.