[GLUTEN-10933][VL] Introduce GPU ShuffleWriterType kGpuHashShuffle

backends-clickhouse/src/main/scala/org/apache/gluten/backendsapi/clickhouse/CHSparkPlanExecApi.scala

@@ -454,8 +454,7 @@ class CHSparkPlanExecApi extends SparkPlanExecApi with Logging {
   override def createColumnarBatchSerializer(
       schema: StructType,
       metrics: Map[String, SQLMetric],
-      shuffleWriterType: ShuffleWriterType,
-      enableCudf: Boolean): Serializer = {
+      shuffleWriterType: ShuffleWriterType): Serializer = {
     val readBatchNumRows = metrics("avgReadBatchNumRows")
     val numOutputRows = metrics("numOutputRows")
     val dataSize = metrics("dataSize")

backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxSparkPlanExecApi.scala

@@ -633,8 +633,7 @@ class VeloxSparkPlanExecApi extends SparkPlanExecApi {
   override def createColumnarBatchSerializer(
       schema: StructType,
       metrics: Map[String, SQLMetric],
-      shuffleWriterType: ShuffleWriterType,
-      enableCudf: Boolean): Serializer = {
+      shuffleWriterType: ShuffleWriterType): Serializer = {
     val numOutputRows = metrics("numOutputRows")
     val deserializeTime = metrics("deserializeTime")
     val readBatchNumRows = metrics("avgReadBatchNumRows")
@@ -659,8 +658,7 @@ class VeloxSparkPlanExecApi extends SparkPlanExecApi {
           numOutputRows,
           deserializeTime,
           decompressTime,
-          shuffleWriterType,
-          enableCudf)
+          shuffleWriterType)
     }
   }
 

backends-velox/src/main/scala/org/apache/gluten/vectorized/ColumnarBatchSerializer.scala

@@ -51,8 +51,7 @@ class ColumnarBatchSerializer(
     numOutputRows: SQLMetric,
     deserializeTime: SQLMetric,
     decompressTime: SQLMetric,
-    shuffleWriterType: ShuffleWriterType,
-    enableCudf: Boolean)
+    shuffleWriterType: ShuffleWriterType)
   extends Serializer
   with Serializable {
 
@@ -64,8 +63,7 @@ class ColumnarBatchSerializer(
       numOutputRows,
       deserializeTime,
       decompressTime,
-      shuffleWriterType,
-      enableCudf)
+      shuffleWriterType)
   }
 
   override def supportsRelocationOfSerializedObjects: Boolean = true
@@ -77,8 +75,7 @@ private class ColumnarBatchSerializerInstanceImpl(
     numOutputRows: SQLMetric,
     deserializeTime: SQLMetric,
     decompressTime: SQLMetric,
-    shuffleWriterType: ShuffleWriterType,
-    enableCudf: Boolean)
+    shuffleWriterType: ShuffleWriterType)
   extends ColumnarBatchSerializerInstance
   with Logging {
 
@@ -114,8 +111,7 @@ private class ColumnarBatchSerializerInstanceImpl(
       batchSize,
       readerBufferSize,
       deserializerBufferSize,
-      shuffleWriterType.name,
-      enableCudf)
+      shuffleWriterType.name)
     // Close shuffle reader instance as lately as the end of task processing,
     // since the native reader could hold a reference to memory pool that
     // was used to create all buffers read from shuffle reader. The pool

cpp/core/jni/JniWrapper.cc

@@ -1083,8 +1083,7 @@ JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_ShuffleReaderJniWrappe
     jint batchSize,
     jlong readerBufferSize,
     jlong deserializerBufferSize,
-    jstring shuffleWriterType,
-    jboolean enableCudf) {
+    jstring shuffleWriterType) {
   JNI_METHOD_START
   auto ctx = getRuntime(env, wrapper);
 
@@ -1096,7 +1095,6 @@ JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_ShuffleReaderJniWrappe
   options.batchSize = batchSize;
   options.readerBufferSize = readerBufferSize;
   options.deserializerBufferSize = deserializerBufferSize;
-  options.enableCudf = enableCudf;
 
   options.shuffleWriterType = ShuffleWriter::stringToType(jStringToCString(env, shuffleWriterType));
   std::shared_ptr<arrow::Schema> schema =

cpp/core/shuffle/Options.h

@@ -43,7 +43,7 @@ static constexpr int64_t kDefaultDeserializerBufferSize = 1 << 20;
 static constexpr int64_t kDefaultShuffleFileBufferSize = 32 << 10;
 static constexpr bool kDefaultEnableDictionary = false;
 
-enum class ShuffleWriterType { kHashShuffle, kSortShuffle, kRssSortShuffle };
+enum class ShuffleWriterType { kHashShuffle, kSortShuffle, kRssSortShuffle, kGpuHashShuffle };
 
 enum class PartitionWriterType { kLocal, kRss };
 
@@ -62,11 +62,6 @@ struct ShuffleReaderOptions {
 
   // Buffer size when deserializing rows into columnar batches. Only used for sort-based shuffle.
   int64_t deserializerBufferSize = kDefaultDeserializerBufferSize;
-
-  // When true, convert the buffers to cudf table.
-  // Add a lock after reader produces the Vector, the next operator should be CudfFromVelox.
-  // After move the shuffle read operation to gpu, move the lock to start read.
-  bool enableCudf = false;
 };
 
 struct ShuffleWriterOptions {

cpp/core/shuffle/ShuffleWriter.cc

@@ -24,6 +24,7 @@ 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) {
@@ -36,6 +37,9 @@ 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);
 }
 
@@ -47,6 +51,8 @@ std::string ShuffleWriter::typeToString(ShuffleWriterType type) {
       return kSortShuffleName;
     case ShuffleWriterType::kRssSortShuffle:
       return kRssSortShuffleName;
+    case ShuffleWriterType::kGpuShuffle:
+      return kGpuShuffleName;
   }
   GLUTEN_UNREACHABLE();
 }

cpp/velox/shuffle/GpuShuffleReader.h

@@ -27,6 +27,9 @@
 
 namespace gluten {
 
+/// Convert the buffers to cudf table.
+/// Add a lock after reader produces the Vector, relase the lock after the thread processes all the batches.
+/// After move the shuffle read operation to gpu, move the lock to start read.
 class GpuHashShuffleReaderDeserializer final : public ColumnarBatchIterator {
  public:
   GpuHashShuffleReaderDeserializer(

cpp/velox/shuffle/VeloxShuffleReader.cc

@@ -803,10 +803,8 @@ VeloxShuffleReaderDeserializerFactory::VeloxShuffleReaderDeserializerFactory(
     int64_t readerBufferSize,
     int64_t deserializerBufferSize,
     VeloxMemoryManager* memoryManager,
-    ShuffleWriterType shuffleWriterType,
-    bool enableCudf)
-    : enableCudf_(enableCudf),
-      schema_(schema),
+    ShuffleWriterType shuffleWriterType)
+    : schema_(schema),
       codec_(codec),
       veloxCompressionType_(veloxCompressionType),
       rowType_(rowType),
@@ -821,10 +819,9 @@ VeloxShuffleReaderDeserializerFactory::VeloxShuffleReaderDeserializerFactory(
 std::unique_ptr<ColumnarBatchIterator> VeloxShuffleReaderDeserializerFactory::createDeserializer(
     const std::shared_ptr<StreamReader>& streamReader) {
   switch (shuffleWriterType_) {
-    case ShuffleWriterType::kHashShuffle:
- #ifdef GLUTEN_ENABLE_GPU       
-      if (enableCudf_) {
-        return std::make_unique<GpuHashShuffleReaderDeserializer>(
+    case ShuffleWriterType::kGpuHashShuffle:
+#ifdef GLUTEN_ENABLE_GPU
+      return std::make_unique<GpuHashShuffleReaderDeserializer>(
           streamReader,
           schema_,
           codec_,
@@ -836,7 +833,6 @@ std::unique_ptr<ColumnarBatchIterator> VeloxShuffleReaderDeserializerFactory::cr
           hasComplexType_,
           deserializeTime_,
           decompressTime_);
-      }
 #endif
       return std::make_unique<VeloxHashShuffleReaderDeserializer>(
           streamReader,

cpp/velox/shuffle/VeloxShuffleReader.h

@@ -169,8 +169,7 @@ class VeloxShuffleReaderDeserializerFactory {
       int64_t readerBufferSize,
       int64_t deserializerBufferSize,
       VeloxMemoryManager* memoryManager,
-      ShuffleWriterType shuffleWriterType,
-      bool enableCudf);
+      ShuffleWriterType shuffleWriterType);
 
   std::unique_ptr<ColumnarBatchIterator> createDeserializer(const std::shared_ptr<StreamReader>& streamReader);
 
@@ -181,7 +180,6 @@ class VeloxShuffleReaderDeserializerFactory {
  private:
   void initFromSchema();
 
-  const bool enableCudf_;
   std::shared_ptr<arrow::Schema> schema_;
   std::shared_ptr<arrow::util::Codec> codec_;
   facebook::velox::common::CompressionKind veloxCompressionType_;

gluten-arrow/src/main/java/org/apache/gluten/vectorized/ShuffleReaderJniWrapper.java

@@ -42,8 +42,7 @@ public native long make(
       int batchSize,
       long readerBufferSize,
       long deserializerBufferSize,
-      String shuffleWriterType,
-      boolean enableCudf);
+      String shuffleWriterType);
 
   public native long read(long shuffleReaderHandle, ShuffleStreamReader streamReader);
 

gluten-substrait/src/main/scala/org/apache/gluten/backendsapi/SparkPlanExecApi.scala

@@ -398,8 +398,7 @@ trait SparkPlanExecApi {
   def createColumnarBatchSerializer(
       schema: StructType,
       metrics: Map[String, SQLMetric],
-      shuffleWriterType: ShuffleWriterType,
-      enableCudf: Boolean = false): Serializer
+      shuffleWriterType: ShuffleWriterType): Serializer
 
   /** Create broadcast relation for BroadcastExchangeExec */
   def createBroadcastRelation(

gluten-substrait/src/main/scala/org/apache/gluten/config/GlutenConfig.scala

@@ -45,6 +45,10 @@ case object RssSortShuffleWriterType extends ShuffleWriterType {
   override val name: String = ReservedKeys.GLUTEN_RSS_SORT_SHUFFLE_WRITER
 }
 
+case object GpuHashShuffleWriterType extends ShuffleWriterType {
+  override val name: String = ReservedKeys.GLUTEN_GPU_HASH_SHUFFLE_WRITER
+}
+
 /*
  * 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.

gluten-substrait/src/main/scala/org/apache/gluten/config/ReservedKeys.scala

@@ -34,4 +34,5 @@ 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"
 }