Skip to content

Commit 87f9b6f

Browse files
maobaolongmaobaolong
authored
[MINOR]improvement(client/server): (RemoteMerger) Refactor to use mergeContext collect the arguments related (#2195)
### What changes were proposed in this pull request? Refactor to use mergeContext collect the arguments related to remote merger ### Why are the changes needed? - Make code clean and friendly to other developer who do not attention to `Remote Merger`. - Without api change while extends the `mergeContext`. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? No need, just refactor.
1 parent 3a35b0f commit 87f9b6f

17 files changed

Lines changed: 145 additions & 195 deletions

File tree

‎client-mr/core/src/main/java/org/apache/hadoop/mapreduce/v2/app/RssMRAppMaster.java‎

Lines changed: 14 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@
8080
import org.apache.uniffle.common.util.Constants;
8181
import org.apache.uniffle.common.util.RetryUtils;
8282
import org.apache.uniffle.hadoop.shim.HadoopShimImpl;
83+
import org.apache.uniffle.proto.RssProtos.MergeContext;
8384
import org.apache.uniffle.storage.util.StorageType;
8485

8586
import static org.apache.hadoop.mapreduce.RssMRConfig.RSS_REMOTE_MERGE_CLASS_LOADER;
@@ -285,17 +286,20 @@ public Thread newThread(Runnable r) {
285286
RssMRConfig.toRssConf(conf)
286287
.get(MAX_CONCURRENCY_PER_PARTITION_TO_WRITE),
287288
0,
288-
remoteMergeEnable ? conf.getMapOutputKeyClass().getName() : null,
289289
remoteMergeEnable
290-
? conf.getMapOutputValueClass().getName()
291-
: null,
292-
remoteMergeEnable
293-
? conf.getOutputKeyComparator().getClass().getName()
294-
: null,
295-
conf.getInt(
296-
RssMRConfig.RSS_MERGED_BLOCK_SZIE,
297-
RssMRConfig.RSS_MERGED_BLOCK_SZIE_DEFAULT),
298-
conf.get(RSS_REMOTE_MERGE_CLASS_LOADER)));
290+
? MergeContext.newBuilder()
291+
.setKeyClass(conf.getMapOutputKeyClass().getName())
292+
.setValueClass(conf.getMapOutputValueClass().getName())
293+
.setComparatorClass(
294+
conf.getOutputKeyComparator().getClass().getName())
295+
.setMergedBlockSize(
296+
conf.getInt(
297+
RssMRConfig.RSS_MERGED_BLOCK_SZIE,
298+
RssMRConfig.RSS_MERGED_BLOCK_SZIE_DEFAULT))
299+
.setMergeClassLoader(
300+
conf.get(RSS_REMOTE_MERGE_CLASS_LOADER, ""))
301+
.build()
302+
: null));
299303
LOG.info(
300304
"Finish register shuffle with "
301305
+ (System.currentTimeMillis() - start)

‎client-mr/core/src/test/java/org/apache/hadoop/mapred/SortWriteBufferManagerTest.java‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@
6262
import org.apache.uniffle.common.serializer.SerializerInstance;
6363
import org.apache.uniffle.common.serializer.SerializerUtils;
6464
import org.apache.uniffle.common.util.JavaUtils;
65+
import org.apache.uniffle.proto.RssProtos;
6566

6667
import static org.junit.jupiter.api.Assertions.assertEquals;
6768
import static org.junit.jupiter.api.Assertions.assertFalse;
@@ -722,11 +723,7 @@ public void registerShuffle(
722723
ShuffleDataDistributionType distributionType,
723724
int maxConcurrencyPerPartitionToWrite,
724725
int stageAttemptNumber,
725-
String keyClassName,
726-
String valueClassName,
727-
String comparatorClassName,
728-
int mergedBlockSize,
729-
String mergeClassLoader) {}
726+
RssProtos.MergeContext mergeContext) {}
730727

731728
@Override
732729
public boolean sendCommit(

‎client-mr/core/src/test/java/org/apache/hadoop/mapreduce/task/reduce/FetcherTest.java‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@
7777
import org.apache.uniffle.common.rpc.StatusCode;
7878
import org.apache.uniffle.common.util.JavaUtils;
7979
import org.apache.uniffle.hadoop.shim.HadoopShimImpl;
80+
import org.apache.uniffle.proto.RssProtos;
8081

8182
import static org.junit.jupiter.api.Assertions.assertEquals;
8283
import static org.mockito.Mockito.mock;
@@ -507,11 +508,7 @@ public void registerShuffle(
507508
ShuffleDataDistributionType distributionType,
508509
int maxConcurrencyPerPartitionToWrite,
509510
int stageAttemptNumber,
510-
String keyClassName,
511-
String valueClassName,
512-
String comparatorClassName,
513-
int mergedBlockSize,
514-
String mergeClassLoader) {}
511+
RssProtos.MergeContext mergeContext) {}
515512

516513
@Override
517514
public boolean sendCommit(

‎client-spark/common/src/main/java/org/apache/uniffle/shuffle/manager/RssShuffleManagerBase.java‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1028,10 +1028,6 @@ protected void registerShuffleServers(
10281028
ShuffleDataDistributionType.NORMAL,
10291029
maxConcurrencyPerPartitionToWrite,
10301030
stageAttemptNumber,
1031-
null,
1032-
null,
1033-
null,
1034-
-1,
10351031
null);
10361032
});
10371033
LOG.info(

‎client-tez/src/main/java/org/apache/tez/dag/app/TezRemoteShuffleManager.java‎

Lines changed: 18 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@
6666
import org.apache.uniffle.common.util.Constants;
6767
import org.apache.uniffle.common.util.JavaUtils;
6868
import org.apache.uniffle.common.util.RetryUtils;
69+
import org.apache.uniffle.proto.RssProtos;
6970

7071
import static org.apache.uniffle.common.config.RssClientConf.MAX_CONCURRENCY_PER_PARTITION_TO_WRITE;
7172

@@ -305,13 +306,23 @@ public ShuffleAssignmentsInfo run() throws Exception {
305306
RssTezConfig.toRssConf(conf)
306307
.get(MAX_CONCURRENCY_PER_PARTITION_TO_WRITE),
307308
0,
308-
keyClassName,
309-
valueClassName,
310-
comparatorClassName,
311-
conf.getInt(
312-
RssTezConfig.RSS_MERGED_BLOCK_SZIE,
313-
RssTezConfig.RSS_MERGED_BLOCK_SZIE_DEFAULT),
314-
conf.get(RssTezConfig.RSS_REMOTE_MERGE_CLASS_LOADER)));
309+
StringUtils.isBlank(keyClassName)
310+
? null
311+
: RssProtos.MergeContext.newBuilder()
312+
.setKeyClass(keyClassName)
313+
.setValueClass(valueClassName)
314+
.setComparatorClass(comparatorClassName)
315+
.setMergedBlockSize(
316+
conf.getInt(
317+
RssTezConfig.RSS_MERGED_BLOCK_SZIE,
318+
RssTezConfig
319+
.RSS_MERGED_BLOCK_SZIE_DEFAULT))
320+
.setMergeClassLoader(
321+
conf.get(
322+
RssTezConfig
323+
.RSS_REMOTE_MERGE_CLASS_LOADER,
324+
""))
325+
.build()));
315326
LOG.info(
316327
"Finish register shuffle with "
317328
+ (System.currentTimeMillis() - start)

‎client-tez/src/test/java/org/apache/tez/runtime/library/common/sort/buffer/WriteBufferManagerTest.java‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,7 @@
7878
import org.apache.uniffle.common.serializer.SerializerInstance;
7979
import org.apache.uniffle.common.serializer.SerializerUtils;
8080
import org.apache.uniffle.common.util.JavaUtils;
81+
import org.apache.uniffle.proto.RssProtos;
8182
import org.apache.uniffle.storage.util.StorageType;
8283

8384
import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -719,11 +720,7 @@ public void registerShuffle(
719720
ShuffleDataDistributionType dataDistributionType,
720721
int maxConcurrencyPerPartitionToWrite,
721722
int stageAttemptNumber,
722-
String keyClassName,
723-
String valueClassName,
724-
String comparatorClassName,
725-
int mergedBlockSize,
726-
String mergeClassLoader) {}
723+
RssProtos.MergeContext mergeContext) {}
727724

728725
@Override
729726
public boolean sendCommit(

‎client/src/main/java/org/apache/uniffle/client/api/ShuffleWriteClient.java‎

Lines changed: 2 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
import org.apache.uniffle.common.ShuffleBlockInfo;
3434
import org.apache.uniffle.common.ShuffleDataDistributionType;
3535
import org.apache.uniffle.common.ShuffleServerInfo;
36+
import org.apache.uniffle.proto.RssProtos.MergeContext;
3637

3738
public interface ShuffleWriteClient {
3839

@@ -72,10 +73,6 @@ default void registerShuffle(
7273
dataDistributionType,
7374
maxConcurrencyPerPartitionToWrite,
7475
0,
75-
null,
76-
null,
77-
null,
78-
-1,
7976
null);
8077
}
8178

@@ -88,11 +85,7 @@ void registerShuffle(
8885
ShuffleDataDistributionType dataDistributionType,
8986
int maxConcurrencyPerPartitionToWrite,
9087
int stageAttemptNumber,
91-
String keyClassName,
92-
String valueClassName,
93-
String comparatorClassName,
94-
int mergedBlockSize,
95-
String mergeClassLoader);
88+
MergeContext mergeContext);
9689

9790
boolean sendCommit(
9891
Set<ShuffleServerInfo> shuffleServerInfoSet, String appId, int shuffleId, int numMaps);

‎client/src/main/java/org/apache/uniffle/client/impl/ShuffleWriteClientImpl.java‎

Lines changed: 3 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,7 @@
9595
import org.apache.uniffle.common.util.BlockIdLayout;
9696
import org.apache.uniffle.common.util.JavaUtils;
9797
import org.apache.uniffle.common.util.ThreadUtils;
98+
import org.apache.uniffle.proto.RssProtos.MergeContext;
9899

99100
public class ShuffleWriteClientImpl implements ShuffleWriteClient {
100101

@@ -564,11 +565,7 @@ public void registerShuffle(
564565
ShuffleDataDistributionType dataDistributionType,
565566
int maxConcurrencyPerPartitionToWrite,
566567
int stageAttemptNumber,
567-
String keyClassName,
568-
String valueClassName,
569-
String comparatorClassName,
570-
int mergedBlockSize,
571-
String mergeClassLoader) {
568+
MergeContext mergeContext) {
572569
String user = null;
573570
try {
574571
user = UserGroupInformation.getCurrentUser().getShortUserName();
@@ -586,11 +583,7 @@ public void registerShuffle(
586583
dataDistributionType,
587584
maxConcurrencyPerPartitionToWrite,
588585
stageAttemptNumber,
589-
keyClassName,
590-
valueClassName,
591-
comparatorClassName,
592-
mergedBlockSize,
593-
mergeClassLoader);
586+
mergeContext);
594587
RssRegisterShuffleResponse response =
595588
getShuffleServerClient(shuffleServerInfo).registerShuffle(request);
596589

‎client/src/test/java/org/apache/uniffle/client/record/reader/MockedShuffleWriteClient.java‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
import org.apache.uniffle.common.ShuffleBlockInfo;
3535
import org.apache.uniffle.common.ShuffleDataDistributionType;
3636
import org.apache.uniffle.common.ShuffleServerInfo;
37+
import org.apache.uniffle.proto.RssProtos;
3738

3839
public class MockedShuffleWriteClient implements ShuffleWriteClient {
3940

@@ -63,11 +64,7 @@ public void registerShuffle(
6364
ShuffleDataDistributionType dataDistributionType,
6465
int maxConcurrencyPerPartitionToWrite,
6566
int stageAttemptNumber,
66-
String keyClassName,
67-
String valueClassName,
68-
String comparatorClassName,
69-
int mergedBlockSize,
70-
String mergeClassLoader) {}
67+
RssProtos.MergeContext mergeContext) {}
7168

7269
@Override
7370
public boolean sendCommit(

‎integration-test/common/src/test/java/org/apache/uniffle/test/RemoteMergeShuffleWithRssClientTest.java‎

Lines changed: 29 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@
6464
import org.apache.uniffle.common.util.BlockIdLayout;
6565
import org.apache.uniffle.common.util.ChecksumUtils;
6666
import org.apache.uniffle.coordinator.CoordinatorConf;
67+
import org.apache.uniffle.proto.RssProtos;
6768
import org.apache.uniffle.server.ShuffleServerConf;
6869
import org.apache.uniffle.server.buffer.ShuffleBufferType;
6970
import org.apache.uniffle.storage.util.StorageType;
@@ -174,11 +175,13 @@ public void remoteMergeWriteReadTest(String classes) throws Exception {
174175
ShuffleDataDistributionType.NORMAL,
175176
0,
176177
-1,
177-
keyClass.getName(),
178-
valueClass.getName(),
179-
comparator.getClass().getName(),
180-
-1,
181-
null);
178+
RssProtos.MergeContext.newBuilder()
179+
.setKeyClass(keyClass.getName())
180+
.setValueClass(valueClass.getName())
181+
.setComparatorClass(comparator.getClass().getName())
182+
.setMergedBlockSize(-1)
183+
.setMergeClassLoader("")
184+
.build());
182185

183186
// 3 report shuffle result
184187
// task 0 attempt 0 generate three blocks
@@ -337,11 +340,13 @@ public void remoteMergeWriteReadTestWithCombine(String classes) throws Exception
337340
ShuffleDataDistributionType.NORMAL,
338341
0,
339342
-1,
340-
keyClass.getName(),
341-
valueClass.getName(),
342-
comparator.getClass().getName(),
343-
-1,
344-
null);
343+
RssProtos.MergeContext.newBuilder()
344+
.setKeyClass(keyClass.getName())
345+
.setValueClass(valueClass.getName())
346+
.setComparatorClass(comparator.getClass().getName())
347+
.setMergedBlockSize(-1)
348+
.setMergeClassLoader("")
349+
.build());
345350

346351
// 3 report shuffle result
347352
// task 0 attempt 0 generate three blocks
@@ -508,11 +513,13 @@ public void remoteMergeWriteReadTestMultiPartition(String classes) throws Except
508513
ShuffleDataDistributionType.NORMAL,
509514
0,
510515
-1,
511-
keyClass.getName(),
512-
valueClass.getName(),
513-
comparator.getClass().getName(),
514-
-1,
515-
null);
516+
RssProtos.MergeContext.newBuilder()
517+
.setKeyClass(keyClass.getName())
518+
.setValueClass(valueClass.getName())
519+
.setComparatorClass(comparator.getClass().getName())
520+
.setMergedBlockSize(-1)
521+
.setMergeClassLoader("")
522+
.build());
516523

517524
// 3 report shuffle result
518525
// this shuffle have three partition, which is hash by key index mode 3
@@ -714,11 +721,13 @@ public void remoteMergeWriteReadTestMultiPartitionWithCombine(String classes) th
714721
ShuffleDataDistributionType.NORMAL,
715722
0,
716723
-1,
717-
keyClass.getName(),
718-
valueClass.getName(),
719-
comparator.getClass().getName(),
720-
-1,
721-
null);
724+
RssProtos.MergeContext.newBuilder()
725+
.setKeyClass(keyClass.getName())
726+
.setValueClass(valueClass.getName())
727+
.setComparatorClass(comparator.getClass().getName())
728+
.setMergedBlockSize(-1)
729+
.setMergeClassLoader("")
730+
.build());
722731

723732
// 3 report shuffle result
724733
// this shuffle have three partition, which is hash by key index mode 3

0 commit comments

Comments
 (0)