@@ -86,14 +86,14 @@ public class RssShuffleWriter<K, V, C> extends ShuffleWriter<K, V> {
86
86
private final Map <Integer , List <ShuffleServerInfo >> partitionToServers ;
87
87
private final Set <ShuffleServerInfo > shuffleServersForData ;
88
88
private final long [] partitionLengths ;
89
- private final boolean isMemoryShuffleEnabled ;
90
89
private final Function <String , Boolean > taskFailureCallback ;
91
90
private final Set <Long > blockIds = Sets .newConcurrentHashSet ();
92
91
93
92
/** used by columnar rss shuffle writer implementation */
94
93
protected final long taskAttemptId ;
95
94
96
95
protected final ShuffleWriteMetrics shuffleWriteMetrics ;
96
+ protected final boolean isMemoryShuffleEnabled ;
97
97
98
98
private final BlockingQueue <Object > finishEventQueue = new LinkedBlockingQueue <>();
99
99
@@ -213,7 +213,7 @@ public void write(Iterator<Product2<K, V>> records) throws IOException {
213
213
}
214
214
}
215
215
216
- private void writeImpl (Iterator <Product2 <K , V >> records ) {
216
+ protected void writeImpl (Iterator <Product2 <K , V >> records ) throws IOException {
217
217
List <ShuffleBlockInfo > shuffleBlockInfos ;
218
218
boolean isCombine = shuffleDependency .mapSideCombine ();
219
219
Function1 <V , C > createCombiner = null ;
@@ -243,7 +243,7 @@ private void writeImpl(Iterator<Product2<K, V>> records) {
243
243
processShuffleBlockInfos (shuffleBlockInfos );
244
244
}
245
245
long checkStartTs = System .currentTimeMillis ();
246
- checkBlockSendResult ( blockIds );
246
+ internalCheckBlockSendResult ( );
247
247
long commitStartTs = System .currentTimeMillis ();
248
248
long checkDuration = commitStartTs - checkStartTs ;
249
249
if (!isMemoryShuffleEnabled ) {
@@ -309,6 +309,10 @@ protected List<CompletableFuture<Long>> postBlockEvent(
309
309
return futures ;
310
310
}
311
311
312
+ protected void internalCheckBlockSendResult () {
313
+ checkBlockSendResult (blockIds );
314
+ }
315
+
312
316
@ VisibleForTesting
313
317
protected void checkBlockSendResult (Set <Long > blockIds ) {
314
318
boolean interrupted = false ;
0 commit comments