@@ -86,14 +86,14 @@ public class RssShuffleWriter<K, V, C> extends ShuffleWriter<K, V> {
8686 private final Map <Integer , List <ShuffleServerInfo >> partitionToServers ;
8787 private final Set <ShuffleServerInfo > shuffleServersForData ;
8888 private final long [] partitionLengths ;
89- private final boolean isMemoryShuffleEnabled ;
9089 private final Function <String , Boolean > taskFailureCallback ;
9190 private final Set <Long > blockIds = Sets .newConcurrentHashSet ();
9291
9392 /** used by columnar rss shuffle writer implementation */
9493 protected final long taskAttemptId ;
9594
9695 protected final ShuffleWriteMetrics shuffleWriteMetrics ;
96+ protected final boolean isMemoryShuffleEnabled ;
9797
9898 private final BlockingQueue <Object > finishEventQueue = new LinkedBlockingQueue <>();
9999
@@ -213,7 +213,7 @@ public void write(Iterator<Product2<K, V>> records) throws IOException {
213213 }
214214 }
215215
216- private void writeImpl (Iterator <Product2 <K , V >> records ) {
216+ protected void writeImpl (Iterator <Product2 <K , V >> records ) throws IOException {
217217 List <ShuffleBlockInfo > shuffleBlockInfos ;
218218 boolean isCombine = shuffleDependency .mapSideCombine ();
219219 Function1 <V , C > createCombiner = null ;
@@ -243,7 +243,7 @@ private void writeImpl(Iterator<Product2<K, V>> records) {
243243 processShuffleBlockInfos (shuffleBlockInfos );
244244 }
245245 long checkStartTs = System .currentTimeMillis ();
246- checkBlockSendResult ( blockIds );
246+ internalCheckBlockSendResult ( );
247247 long commitStartTs = System .currentTimeMillis ();
248248 long checkDuration = commitStartTs - checkStartTs ;
249249 if (!isMemoryShuffleEnabled ) {
@@ -309,6 +309,10 @@ protected List<CompletableFuture<Long>> postBlockEvent(
309309 return futures ;
310310 }
311311
312+ protected void internalCheckBlockSendResult () {
313+ checkBlockSendResult (blockIds );
314+ }
315+
312316 @ VisibleForTesting
313317 protected void checkBlockSendResult (Set <Long > blockIds ) {
314318 boolean interrupted = false ;
0 commit comments