style fix

apache · Nov 25, 2023 · 4724acf · 4724acf
1 parent 3032d11
commit 4724acf
Show file tree

Hide file tree

Showing 6 changed files with 7 additions and 14 deletions.
diff --git a/cpp/core/shuffle/Partitioner.cc b/cpp/core/shuffle/Partitioner.cc
@@ -23,10 +23,7 @@
 
 namespace gluten {
 
-arrow::Result<std::shared_ptr<Partitioner>> Partitioner::make(
-    Partitioning partitioning,
-    int32_t numPartitions,
-    int32_t partitionKeySeed) {
+arrow::Result<std::shared_ptr<Partitioner>> Partitioner::make(Partitioning partitioning, int32_t numPartitions, int32_t partitionKeySeed) {
   switch (partitioning) {
     case Partitioning::kHash:
       return std::make_shared<HashPartitioner>(numPartitions);

diff --git a/cpp/core/shuffle/Partitioner.h b/cpp/core/shuffle/Partitioner.h
@@ -26,10 +26,7 @@ namespace gluten {
 
 class Partitioner {
  public:
-  static arrow::Result<std::shared_ptr<Partitioner>> make(
-      Partitioning partitioning,
-      int32_t numPartitions,
-      int32_t partitionKeySeed);
+  static arrow::Result<std::shared_ptr<Partitioner>> make(Partitioning partitioning, int32_t numPartitions, int32_t partitionKeySeed);
 
   // Whether the first column is partition key.
   bool hasPid() const {

diff --git a/cpp/core/shuffle/RoundRobinPartitioner.h b/cpp/core/shuffle/RoundRobinPartitioner.h
@@ -23,8 +23,8 @@ namespace gluten {
 
 class RoundRobinPartitioner final : public Partitioner {
  public:
-  RoundRobinPartitioner(int32_t numPartitions, int32_t partitionKeySeed) 
-      : Partitioner(numPartitions, false), pidSelection_(partitionKeySeed) {}
+  RoundRobinPartitioner(int32_t numPartitions, int32_t partitionKeySeed)
+      : Partitioner(numPartitions, false), pidSelection_(partitionKeySeed % numPartitions) {}
 
   arrow::Status compute(
       const int32_t* pidArr,

diff --git a/cpp/core/tests/RoundRobinPartitionerTest.cc b/cpp/core/tests/RoundRobinPartitionerTest.cc
@@ -63,7 +63,7 @@ class RoundRobinPartitionerTest : public ::testing::Test {
 
 TEST_F(RoundRobinPartitionerTest, TestInit) {
   int numPart = 2;
-  prepareData(numPart, 1);
+  prepareData(numPart, 3);
   ASSERT_NE(partitioner_, nullptr);
   int32_t pidSelection = getPidSelection();
   ASSERT_EQ(pidSelection, 1);

diff --git a/cpp/velox/shuffle/VeloxShuffleWriter.cc b/cpp/velox/shuffle/VeloxShuffleWriter.cc
@@ -411,9 +411,7 @@ arrow::Status VeloxShuffleWriter::init() {
   VELOX_CHECK_NOT_NULL(options_.memory_pool);
 
   ARROW_ASSIGN_OR_RAISE(partitionWriter_, partitionWriterCreator_->make(this));
-  ARROW_ASSIGN_OR_RAISE(
-      partitioner_,
-      Partitioner::make(options_.partitioning, numPartitions_, options_.partition_key_seed));
+  ARROW_ASSIGN_OR_RAISE(partitioner_, Partitioner::make(options_.partitioning, numPartitions_, options_.partition_key_seed));
 
   // pre-allocated buffer size for each partition, unit is row count
   // when partitioner is SinglePart, partial variables don`t need init

diff --git a/...src/main/scala/org/apache/spark/shuffle/VeloxCelebornHashBasedColumnarShuffleWriter.scala b/...src/main/scala/org/apache/spark/shuffle/VeloxCelebornHashBasedColumnarShuffleWriter.scala
@@ -111,6 +111,7 @@ class VeloxCelebornHashBasedColumnarShuffleWriter[K, V](
               .getNativeInstanceHandle,
             handle,
             context.taskAttemptId(),
+            partitionKeySeed,
             "celeborn",
             GlutenConfig.getConf.columnarShuffleReallocThreshold
           )