Skip to content

Commit 1ef56c7

Browse files
JoshRosenJoshRosen
committed
Revise compression codec support in merger; test cross product of configurations.
1 parent b57c17f commit 1ef56c7

File tree

4 files changed

+151
-43
lines changed

4 files changed

+151
-43
lines changed

‎core/src/main/java/org/apache/spark/shuffle/unsafe/SpillInfo.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818
package org.apache.spark.shuffle.unsafe;
1919

20-
import org.apache.spark.storage.BlockId;
20+
import org.apache.spark.storage.TempShuffleBlockId;
2121

2222
import java.io.File;
2323

@@ -27,9 +27,9 @@
2727
final class SpillInfo {
2828
final long[] partitionLengths;
2929
final File file;
30-
final BlockId blockId;
30+
final TempShuffleBlockId blockId;
3131

32-
public SpillInfo(int numPartitions, File file, BlockId blockId) {
32+
public SpillInfo(int numPartitions, File file, TempShuffleBlockId blockId) {
3333
this.partitionLengths = new long[numPartitions];
3434
this.file = file;
3535
this.blockId = blockId;

‎core/src/main/java/org/apache/spark/shuffle/unsafe/UnsafeShuffleExternalSorter.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -153,7 +153,7 @@ private SpillInfo writeSpillFile() throws IOException {
153153
final Tuple2<TempShuffleBlockId, File> spilledFileInfo =
154154
blockManager.diskBlockManager().createTempShuffleBlock();
155155
final File file = spilledFileInfo._2();
156-
final BlockId blockId = spilledFileInfo._1();
156+
final TempShuffleBlockId blockId = spilledFileInfo._1();
157157
final SpillInfo spillInfo = new SpillInfo(numPartitions, file, blockId);
158158

159159
// Unfortunately, we need a serializer instance in order to construct a DiskBlockObjectWriter.
@@ -320,7 +320,7 @@ private void allocateSpaceForRecord(int requiredSpace) throws IOException {
320320
}
321321
}
322322
if (requiredSpace > freeSpaceInCurrentPage) {
323-
logger.debug("Required space {} is less than free space in current page ({}}", requiredSpace,
323+
logger.trace("Required space {} is less than free space in current page ({})", requiredSpace,
324324
freeSpaceInCurrentPage);
325325
// TODO: we should track metrics on the amount of space wasted when we roll over to a new page
326326
// without using the free space at the end of the current page. We should also do this for

‎core/src/main/java/org/apache/spark/shuffle/unsafe/UnsafeShuffleWriter.java‎

Lines changed: 46 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import java.io.*;
2121
import java.nio.channels.FileChannel;
2222
import java.util.Iterator;
23+
import javax.annotation.Nullable;
2324

2425
import scala.Option;
2526
import scala.Product2;
@@ -35,6 +36,9 @@
3536
import org.slf4j.LoggerFactory;
3637

3738
import org.apache.spark.*;
39+
import org.apache.spark.io.CompressionCodec;
40+
import org.apache.spark.io.CompressionCodec$;
41+
import org.apache.spark.io.LZFCompressionCodec;
3842
import org.apache.spark.executor.ShuffleWriteMetrics;
3943
import org.apache.spark.network.util.LimitedInputStream;
4044
import org.apache.spark.scheduler.MapStatus;
@@ -53,8 +57,6 @@ public class UnsafeShuffleWriter<K, V> extends ShuffleWriter<K, V> {
5357

5458
private final Logger logger = LoggerFactory.getLogger(UnsafeShuffleWriter.class);
5559

56-
@VisibleForTesting
57-
static final int MAXIMUM_RECORD_SIZE = 1024 * 1024 * 64; // 64 megabytes
5860
private static final ClassTag<Object> OBJECT_CLASS_TAG = ClassTag$.MODULE$.Object();
5961

6062
private final BlockManager blockManager;
@@ -201,6 +203,12 @@ void forceSorterToSpill() throws IOException {
201203

202204
private long[] mergeSpills(SpillInfo[] spills) throws IOException {
203205
final File outputFile = shuffleBlockResolver.getDataFile(shuffleId, mapId);
206+
final boolean compressionEnabled = sparkConf.getBoolean("spark.shuffle.compress", true);
207+
final CompressionCodec compressionCodec = CompressionCodec$.MODULE$.createCodec(sparkConf);
208+
final boolean fastMergeEnabled =
209+
sparkConf.getBoolean("spark.shuffle.unsafe.fastMergeEnabled", true);
210+
final boolean fastMergeIsSupported =
211+
!compressionEnabled || compressionCodec instanceof LZFCompressionCodec;
204212
try {
205213
if (spills.length == 0) {
206214
new FileOutputStream(outputFile).close(); // Create an empty file
@@ -215,11 +223,20 @@ private long[] mergeSpills(SpillInfo[] spills) throws IOException {
215223
Files.move(spills[0].file, outputFile);
216224
return spills[0].partitionLengths;
217225
} else {
218-
// Need to merge multiple spills.
219-
if (transferToEnabled) {
220-
return mergeSpillsWithTransferTo(spills, outputFile);
226+
if (fastMergeEnabled && fastMergeIsSupported) {
227+
// Compression is disabled or we are using an IO compression codec that supports
228+
// decompression of concatenated compressed streams, so we can perform a fast spill merge
229+
// that doesn't need to interpret the spilled bytes.
230+
if (transferToEnabled) {
231+
logger.debug("Using transferTo-based fast merge");
232+
return mergeSpillsWithTransferTo(spills, outputFile);
233+
} else {
234+
logger.debug("Using fileStream-based fast merge");
235+
return mergeSpillsWithFileStream(spills, outputFile, null);
236+
}
221237
} else {
222-
return mergeSpillsWithFileStream(spills, outputFile);
238+
logger.debug("Using slow merge");
239+
return mergeSpillsWithFileStream(spills, outputFile, compressionCodec);
223240
}
224241
}
225242
} catch (IOException e) {
@@ -230,27 +247,40 @@ private long[] mergeSpills(SpillInfo[] spills) throws IOException {
230247
}
231248
}
232249

233-
private long[] mergeSpillsWithFileStream(SpillInfo[] spills, File outputFile) throws IOException {
250+
private long[] mergeSpillsWithFileStream(
251+
SpillInfo[] spills,
252+
File outputFile,
253+
@Nullable CompressionCodec compressionCodec) throws IOException {
234254
final int numPartitions = partitioner.numPartitions();
235255
final long[] partitionLengths = new long[numPartitions];
236-
final FileInputStream[] spillInputStreams = new FileInputStream[spills.length];
237-
FileOutputStream mergedFileOutputStream = null;
256+
final InputStream[] spillInputStreams = new FileInputStream[spills.length];
257+
OutputStream mergedFileOutputStream = null;
238258

239259
try {
240260
for (int i = 0; i < spills.length; i++) {
241261
spillInputStreams[i] = new FileInputStream(spills[i].file);
242262
}
243-
mergedFileOutputStream = new FileOutputStream(outputFile);
244-
245263
for (int partition = 0; partition < numPartitions; partition++) {
264+
final long initialFileLength = outputFile.length();
265+
mergedFileOutputStream = new FileOutputStream(outputFile, true);
266+
if (compressionCodec != null) {
267+
mergedFileOutputStream = compressionCodec.compressedOutputStream(mergedFileOutputStream);
268+
}
269+
246270
for (int i = 0; i < spills.length; i++) {
247271
final long partitionLengthInSpill = spills[i].partitionLengths[partition];
248-
final FileInputStream spillInputStream = spillInputStreams[i];
249-
ByteStreams.copy
250-
(new LimitedInputStream(spillInputStream, partitionLengthInSpill),
251-
mergedFileOutputStream);
252-
partitionLengths[partition] += partitionLengthInSpill;
272+
if (partitionLengthInSpill > 0) {
273+
InputStream partitionInputStream =
274+
new LimitedInputStream(spillInputStreams[i], partitionLengthInSpill);
275+
if (compressionCodec != null) {
276+
partitionInputStream = compressionCodec.compressedInputStream(partitionInputStream);
277+
}
278+
ByteStreams.copy(partitionInputStream, mergedFileOutputStream);
279+
}
253280
}
281+
mergedFileOutputStream.flush();
282+
mergedFileOutputStream.close();
283+
partitionLengths[partition] = (outputFile.length() - initialFileLength);
254284
}
255285
} finally {
256286
for (InputStream stream : spillInputStreams) {

‎core/src/test/java/org/apache/spark/shuffle/unsafe/UnsafeShuffleWriterSuite.java‎

Lines changed: 100 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -37,11 +37,14 @@
3737
import org.mockito.invocation.InvocationOnMock;
3838
import org.mockito.stubbing.Answer;
3939
import static org.mockito.AdditionalAnswers.returnsFirstArg;
40-
import static org.mockito.AdditionalAnswers.returnsSecondArg;
4140
import static org.mockito.Answers.RETURNS_SMART_NULLS;
4241
import static org.mockito.Mockito.*;
4342

4443
import org.apache.spark.*;
44+
import org.apache.spark.io.CompressionCodec$;
45+
import org.apache.spark.io.LZ4CompressionCodec;
46+
import org.apache.spark.io.LZFCompressionCodec;
47+
import org.apache.spark.io.SnappyCompressionCodec;
4548
import org.apache.spark.executor.ShuffleWriteMetrics;
4649
import org.apache.spark.executor.TaskMetrics;
4750
import org.apache.spark.network.util.LimitedInputStream;
@@ -65,6 +68,7 @@ public class UnsafeShuffleWriterSuite {
6568
File tempDir;
6669
long[] partitionSizesInMergedFile;
6770
final LinkedList<File> spillFilesCreated = new LinkedList<File>();
71+
SparkConf conf;
6872
final Serializer serializer = new KryoSerializer(new SparkConf());
6973

7074
@Mock(answer = RETURNS_SMART_NULLS) ShuffleMemoryManager shuffleMemoryManager;
@@ -74,10 +78,14 @@ public class UnsafeShuffleWriterSuite {
7478
@Mock(answer = RETURNS_SMART_NULLS) TaskContext taskContext;
7579
@Mock(answer = RETURNS_SMART_NULLS) ShuffleDependency<Object, Object, Object> shuffleDep;
7680

77-
private static final class CompressStream extends AbstractFunction1<OutputStream, OutputStream> {
81+
private final class CompressStream extends AbstractFunction1<OutputStream, OutputStream> {
7882
@Override
7983
public OutputStream apply(OutputStream stream) {
80-
return stream;
84+
if (conf.getBoolean("spark.shuffle.compress", true)) {
85+
return CompressionCodec$.MODULE$.createCodec(conf).compressedOutputStream(stream);
86+
} else {
87+
return stream;
88+
}
8189
}
8290
}
8391

@@ -98,6 +106,7 @@ public void setUp() throws IOException {
98106
mergedOutputFile = File.createTempFile("mergedoutput", "", tempDir);
99107
partitionSizesInMergedFile = null;
100108
spillFilesCreated.clear();
109+
conf = new SparkConf();
101110

102111
when(shuffleMemoryManager.tryToAcquire(anyLong())).then(returnsFirstArg());
103112

@@ -123,8 +132,35 @@ public DiskBlockObjectWriter answer(InvocationOnMock invocationOnMock) throws Th
123132
);
124133
}
125134
});
126-
when(blockManager.wrapForCompression(any(BlockId.class), any(InputStream.class)))
127-
.then(returnsSecondArg());
135+
when(blockManager.wrapForCompression(any(BlockId.class), any(InputStream.class))).thenAnswer(
136+
new Answer<InputStream>() {
137+
@Override
138+
public InputStream answer(InvocationOnMock invocation) throws Throwable {
139+
assert (invocation.getArguments()[0] instanceof TempShuffleBlockId);
140+
InputStream is = (InputStream) invocation.getArguments()[1];
141+
if (conf.getBoolean("spark.shuffle.compress", true)) {
142+
return CompressionCodec$.MODULE$.createCodec(conf).compressedInputStream(is);
143+
} else {
144+
return is;
145+
}
146+
}
147+
}
148+
);
149+
150+
when(blockManager.wrapForCompression(any(BlockId.class), any(OutputStream.class))).thenAnswer(
151+
new Answer<OutputStream>() {
152+
@Override
153+
public OutputStream answer(InvocationOnMock invocation) throws Throwable {
154+
assert (invocation.getArguments()[0] instanceof TempShuffleBlockId);
155+
OutputStream os = (OutputStream) invocation.getArguments()[1];
156+
if (conf.getBoolean("spark.shuffle.compress", true)) {
157+
return CompressionCodec$.MODULE$.createCodec(conf).compressedOutputStream(os);
158+
} else {
159+
return os;
160+
}
161+
}
162+
}
163+
);
128164

129165
when(shuffleBlockResolver.getDataFile(anyInt(), anyInt())).thenReturn(mergedOutputFile);
130166
doAnswer(new Answer<Void>() {
@@ -136,11 +172,11 @@ public Void answer(InvocationOnMock invocationOnMock) throws Throwable {
136172
}).when(shuffleBlockResolver).writeIndexFile(anyInt(), anyInt(), any(long[].class));
137173

138174
when(diskBlockManager.createTempShuffleBlock()).thenAnswer(
139-
new Answer<Tuple2<TempLocalBlockId, File>>() {
175+
new Answer<Tuple2<TempShuffleBlockId, File>>() {
140176
@Override
141-
public Tuple2<TempLocalBlockId, File> answer(
177+
public Tuple2<TempShuffleBlockId, File> answer(
142178
InvocationOnMock invocationOnMock) throws Throwable {
143-
TempLocalBlockId blockId = new TempLocalBlockId(UUID.randomUUID());
179+
TempShuffleBlockId blockId = new TempShuffleBlockId(UUID.randomUUID());
144180
File file = File.createTempFile("spillFile", ".spill", tempDir);
145181
spillFilesCreated.add(file);
146182
return Tuple2$.MODULE$.apply(blockId, file);
@@ -154,7 +190,6 @@ public Tuple2<TempLocalBlockId, File> answer(
154190
}
155191

156192
private UnsafeShuffleWriter<Object, Object> createWriter(boolean transferToEnabled) {
157-
SparkConf conf = new SparkConf();
158193
conf.set("spark.file.transferTo", String.valueOf(transferToEnabled));
159194
return new UnsafeShuffleWriter<Object, Object>(
160195
blockManager,
@@ -164,7 +199,7 @@ private UnsafeShuffleWriter<Object, Object> createWriter(boolean transferToEnabl
164199
new UnsafeShuffleHandle<Object, Object>(0, 1, shuffleDep),
165200
0, // map id
166201
taskContext,
167-
new SparkConf()
202+
conf
168203
);
169204
}
170205

@@ -183,8 +218,11 @@ private List<Tuple2<Object, Object>> readRecordsFromFile() throws IOException {
183218
if (partitionSize > 0) {
184219
InputStream in = new FileInputStream(mergedOutputFile);
185220
ByteStreams.skipFully(in, startOffset);
186-
DeserializationStream recordsStream = serializer.newInstance().deserializeStream(
187-
new LimitedInputStream(in, partitionSize));
221+
in = new LimitedInputStream(in, partitionSize);
222+
if (conf.getBoolean("spark.shuffle.compress", true)) {
223+
in = CompressionCodec$.MODULE$.createCodec(conf).compressedInputStream(in);
224+
}
225+
DeserializationStream recordsStream = serializer.newInstance().deserializeStream(in);
188226
Iterator<Tuple2<Object, Object>> records = recordsStream.asKeyValueIterator();
189227
while (records.hasNext()) {
190228
Tuple2<Object, Object> record = records.next();
@@ -245,7 +283,15 @@ public void writeWithoutSpilling() throws Exception {
245283
assertSpillFilesWereCleanedUp();
246284
}
247285

248-
private void testMergingSpills(boolean transferToEnabled) throws IOException {
286+
private void testMergingSpills(
287+
boolean transferToEnabled,
288+
String compressionCodecName) throws IOException {
289+
if (compressionCodecName != null) {
290+
conf.set("spark.shuffle.compress", "true");
291+
conf.set("spark.io.compression.codec", compressionCodecName);
292+
} else {
293+
conf.set("spark.shuffle.compress", "false");
294+
}
249295
final UnsafeShuffleWriter<Object, Object> writer = createWriter(transferToEnabled);
250296
final ArrayList<Product2<Object, Object>> dataToWrite =
251297
new ArrayList<Product2<Object, Object>>();
@@ -265,25 +311,57 @@ private void testMergingSpills(boolean transferToEnabled) throws IOException {
265311
Assert.assertTrue(mergedOutputFile.exists());
266312
Assert.assertEquals(2, spillFilesCreated.size());
267313

268-
long sumOfPartitionSizes = 0;
269-
for (long size: partitionSizesInMergedFile) {
270-
sumOfPartitionSizes += size;
271-
}
272-
Assert.assertEquals(mergedOutputFile.length(), sumOfPartitionSizes);
314+
// This assertion only holds for the fast merging path:
315+
// long sumOfPartitionSizes = 0;
316+
// for (long size: partitionSizesInMergedFile) {
317+
// sumOfPartitionSizes += size;
318+
// }
319+
// Assert.assertEquals(sumOfPartitionSizes, mergedOutputFile.length());
320+
Assert.assertTrue(mergedOutputFile.length() > 0);
273321
Assert.assertEquals(
274322
HashMultiset.create(dataToWrite),
275323
HashMultiset.create(readRecordsFromFile()));
276324
assertSpillFilesWereCleanedUp();
277325
}
278326

279327
@Test
280-
public void mergeSpillsWithTransferTo() throws Exception {
281-
testMergingSpills(true);
328+
public void mergeSpillsWithTransferToAndLZF() throws Exception {
329+
testMergingSpills(true, LZFCompressionCodec.class.getName());
330+
}
331+
332+
@Test
333+
public void mergeSpillsWithFileStreamAndLZF() throws Exception {
334+
testMergingSpills(false, LZFCompressionCodec.class.getName());
335+
}
336+
337+
@Test
338+
public void mergeSpillsWithTransferToAndLZ4() throws Exception {
339+
testMergingSpills(true, LZ4CompressionCodec.class.getName());
340+
}
341+
342+
@Test
343+
public void mergeSpillsWithFileStreamAndLZ4() throws Exception {
344+
testMergingSpills(false, LZ4CompressionCodec.class.getName());
345+
}
346+
347+
@Test
348+
public void mergeSpillsWithTransferToAndSnappy() throws Exception {
349+
testMergingSpills(true, SnappyCompressionCodec.class.getName());
350+
}
351+
352+
@Test
353+
public void mergeSpillsWithFileStreamAndSnappy() throws Exception {
354+
testMergingSpills(false, SnappyCompressionCodec.class.getName());
355+
}
356+
357+
@Test
358+
public void mergeSpillsWithTransferToAndNoCompression() throws Exception {
359+
testMergingSpills(true, null);
282360
}
283361

284362
@Test
285-
public void mergeSpillsWithFileStream() throws Exception {
286-
testMergingSpills(false);
363+
public void mergeSpillsWithFileStreamAndNoCompression() throws Exception {
364+
testMergingSpills(false, null);
287365
}
288366

289367
@Test

0 commit comments

Comments
 (0)