This is an automated email from the ASF dual-hosted git repository.
marin-ma pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 717d3e47f1 [VL] Remove gpu hash shuffle writer type (#12593)
717d3e47f1 is described below
commit 717d3e47f10988a49c2da79ed41c943bef08fe21
Author: Rong Ma <[email protected]>
AuthorDate: Wed Jul 22 15:02:07 2026 +0100
[VL] Remove gpu hash shuffle writer type (#12593)
---
.../gluten/backendsapi/velox/VeloxMetricsApi.scala | 4 ++--
.../backendsapi/velox/VeloxSparkPlanExecApi.scala | 7 +++++-
.../spark/shuffle/ColumnarShuffleWriter.scala | 15 ++----------
cpp/core/jni/JniWrapper.cc | 28 ----------------------
cpp/core/shuffle/Options.h | 21 +---------------
cpp/core/shuffle/ShuffleWriter.cc | 6 -----
cpp/velox/shuffle/VeloxShuffleReader.cc | 3 +--
cpp/velox/shuffle/VeloxShuffleWriter.cc | 1 -
cpp/velox/tests/VeloxGpuShuffleWriterTest.cc | 6 ++---
.../gluten/vectorized/ShuffleWriterJniWrapper.java | 8 -------
.../gluten/backendsapi/SparkPlanExecApi.scala | 3 ++-
.../org/apache/gluten/config/GlutenConfig.scala | 6 -----
.../org/apache/gluten/config/ReservedKeys.scala | 1 -
.../execution/ColumnarShuffleExchangeExec.scala | 18 ++++++--------
14 files changed, 24 insertions(+), 103 deletions(-)
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 421e6c3595..5d202ee153 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 0397a6dc2c..d588695cd4 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 73bb2770e5..f0d814c156 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 cadd32bb47..e4d674bd74 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<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,
diff --git a/cpp/core/shuffle/Options.h b/cpp/core/shuffle/Options.h
index ead95dfe65..ec246ce80b 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 7287181ffa..5a48757231 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 eb95f2c8ea..d10a0ca305 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>& 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 a83e3a1947..c45d941972 100644
--- a/cpp/velox/shuffle/VeloxShuffleWriter.cc
+++ b/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);
diff --git a/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc
b/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc
index d0df7ebba1..46bafbfa42 100644
--- a/cpp/velox/tests/VeloxGpuShuffleWriterTest.cc
+++ b/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;
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 2f31a529d9..edbc48a3c9 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 class ShuffleWriterJniWrapper implements RuntimeAware
{
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 b0efc00fdf..f5dc9f3f21 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 4e6dcf6f8f..98fe6a0657 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 d609a80e2a..6635a8b20a 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 3accf8e592..0fb5069458 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.
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
