From 9c4c863277e64c104eaa252c1fad78a4129226b0 Mon Sep 17 00:00:00 2001 From: Denys Kuzmenko Date: Wed, 4 Nov 2026 20:50:01 +0200 Subject: [PATCH] TEZ-4762: UnorderedPartitionedKVWriter: size each partition's codec stream to the partition - A spill opens a codec stream per partition: SnappyCodec and Lz4Codec parse the buffer size from the Configuration and allocate and zero a 256 KiB buffer for each, and IFile.Writer a 64 KiB write buffer, while a partition holds a few KB (Hive LLAP, TPC-DS 10 TB q75 / q64: writer setup in spills is 3.2 / 3.5 % of LLAP CPU, Integer.parseInt 1.5 / 1.5 % and these buffers 1.0 / 1.4 % of it) - CodecUtils reads the buffer size once per output and builds these codecs' BlockCompressorStream itself: a stream known to fit one block gets a block of its length and a buffer of its compression bound, which writes the same bytes; IFile.Writer sizes its write buffer to the stream too - The spill passes each partition's raw length bound from the buffers' record and byte counts, the final merge the sum of the spilled raw lengths; other codecs and writers keep codec.createOutputStream - JMH, one spill of 1,009 partitions (Snappy, 22 B / 2.2 KB / 22 KB per partition): -76 % / -48 % / -10 %, 336 MB allocated -> 5 / 11 / 62 MB; LZ4 -59 % / -41 % / -12 % --- .../library/common/sort/impl/IFile.java | 56 ++++++-- .../writers/UnorderedPartitionedKVWriter.java | 40 +++++- .../tez/runtime/library/utils/CodecUtils.java | 66 +++++++++ .../library/common/sort/impl/TestIFile.java | 64 +++++++++ .../TestUnorderedPartitionedKVWriter.java | 133 ++++++++++++++++++ .../runtime/library/utils/TestCodecUtils.java | 127 +++++++++++++++++ 6 files changed, 473 insertions(+), 13 deletions(-) diff --git a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/sort/impl/IFile.java b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/sort/impl/IFile.java index 90d3cd3c67..5e239ec6b6 100644 --- a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/sort/impl/IFile.java +++ b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/sort/impl/IFile.java @@ -180,7 +180,7 @@ private void resetToFileBasedWriter() throws IOException { this.rawOut = newRawOut; this.ownOutputStream = true; - setupOutputStream(fileCodec); + setupOutputStream(fileCodec, -1, -1); // Write header to file headerWritten = false; @@ -332,13 +332,33 @@ public Writer(Serialization keySerialization, Serialization valSerialization, FS Class keyClass, Class valueClass, CompressionCodec codec, TezCounter writesCounter, TezCounter serializedBytesCounter, boolean rle) throws IOException { + this(keySerialization, valSerialization, outputStream, keyClass, valueClass, codec, -1, -1, + writesCounter, serializedBytesCounter, rle); + } + + /** + * A writer that sizes its stream buffers to the data when the caller knows how much it appends. + * + * @param codecBufferSize {@link CodecUtils#getOutputBufferSize(CompressionCodec)} of the codec + * @param maxRawLength upper bound of {@link #getRawLength()}, or -1 if unknown + */ + public Writer(Serialization keySerialization, Serialization valSerialization, FSDataOutputStream outputStream, + Class keyClass, Class valueClass, CompressionCodec codec, int codecBufferSize, long maxRawLength, + TezCounter writesCounter, TezCounter serializedBytesCounter) throws IOException { + this(keySerialization, valSerialization, outputStream, keyClass, valueClass, codec, codecBufferSize, + maxRawLength, writesCounter, serializedBytesCounter, false); + } + + private Writer(Serialization keySerialization, Serialization valSerialization, FSDataOutputStream outputStream, + Class keyClass, Class valueClass, CompressionCodec codec, int codecBufferSize, long maxRawLength, + TezCounter writesCounter, TezCounter serializedBytesCounter, boolean rle) throws IOException { this.rawOut = outputStream; this.writtenRecordsCounter = writesCounter; this.serializedUncompressedBytes = serializedBytesCounter; this.start = this.rawOut.getPos(); this.rle = rle; - setupOutputStream(codec); + setupOutputStream(codec, codecBufferSize, maxRawLength < 0 ? -1 : maxRawLength - HEADER.length); writeHeader(outputStream); @@ -353,21 +373,26 @@ public Writer(Serialization keySerialization, Serialization valSerialization, FS } } - void setupOutputStream(CompressionCodec codec) throws IOException { + /** + * @param streamLength upper bound of the bytes written to the codec or checksum stream (the + * header bypasses them), or -1 if unknown + */ + void setupOutputStream(CompressionCodec codec, int codecBufferSize, long streamLength) throws IOException { this.checksumOut = new IFileOutputStream(this.rawOut); if (codec != null) { this.compressor = CodecUtils.getCompressor(codec); if (this.compressor != null) { this.compressor.reset(); - this.compressedOut = CodecUtils.createOutputStream(codec, checksumOut, compressor); - this.out = bufferedOutput(this.compressedOut); + this.compressedOut = CodecUtils.createOutputStream(codec, checksumOut, compressor, codecBufferSize, + streamLength); + this.out = bufferedOutput(this.compressedOut, streamLength); this.compressOutput = true; } else { LOG.warn("Could not obtain compressor from CodecPool"); - this.out = bufferedOutput(checksumOut); + this.out = bufferedOutput(checksumOut, streamLength); } } else { - this.out = bufferedOutput(checksumOut); + this.out = bufferedOutput(checksumOut, streamLength); } } @@ -375,10 +400,21 @@ void setupOutputStream(CompressionCodec codec) throws IOException { * The writer emits several small writes per record (vint lengths, key, value). Unbuffered, * each one is a full call into the codec stream (synchronized Compressor.setInput and * friends) or the checksum stream. The buffer is unsynchronized: a Writer is confined to - * one thread. + * one thread. A stream known to be shorter than the buffer gets a buffer of its length. + */ + private static DataOutputStream bufferedOutput(OutputStream stream, long streamLength) { + int size = streamLength > 0 ? (int) Math.min(streamLength, WRITE_BUFFER_SIZE) : WRITE_BUFFER_SIZE; + return new NonSyncDataOutputStream(new NonSyncBufferedOutputStream(stream, size)); + } + + /** + * Upper bound of {@link #getRawLength()} of a writer without RLE that appends numRecords + * records of keyValueBytes key and value bytes in total: no record length needs a longer + * vint than the total. */ - private static DataOutputStream bufferedOutput(OutputStream stream) { - return new NonSyncDataOutputStream(new NonSyncBufferedOutputStream(stream, WRITE_BUFFER_SIZE)); + public static long maxRawLength(long numRecords, long keyValueBytes) { + return HEADER.length + keyValueBytes + numRecords * 2L * WritableUtils.getVIntSize(keyValueBytes) + + 2L * WritableUtils.getVIntSize(EOF_MARKER); } public Writer(Serialization keySerialization, Serialization valSerialization, FileSystem fs, Path file) throws IOException { diff --git a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/writers/UnorderedPartitionedKVWriter.java b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/writers/UnorderedPartitionedKVWriter.java index 7b5e8f922b..e178b00c8e 100644 --- a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/writers/UnorderedPartitionedKVWriter.java +++ b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/common/writers/UnorderedPartitionedKVWriter.java @@ -76,6 +76,7 @@ import org.apache.tez.runtime.library.common.sort.impl.TezSpillRecord; import org.apache.tez.runtime.library.shuffle.impl.ShuffleUserPayloads; import org.apache.tez.runtime.library.shuffle.impl.ShuffleUserPayloads.DataMovementEventPayloadProto; +import org.apache.tez.runtime.library.utils.CodecUtils; import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.Lists; @@ -202,6 +203,9 @@ public Deflater get() { // When enabled, uses in-mem ifile writer private final boolean useCachedStream; + // The codec's output stream buffer size, read once instead of per partition stream + private final int codecBufferSize; + public UnorderedPartitionedKVWriter(OutputContext outputContext, Configuration conf, int numOutputs, long availableMemoryBytes) throws IOException { super(outputContext, conf, numOutputs); @@ -236,6 +240,7 @@ public UnorderedPartitionedKVWriter(OutputContext outputContext, Configuration c this.useCachedStream = useCachedStreamConfig && (this.dataViaEventsEnabled && (numPartitions == 1) && !pipelinedShuffle); + this.codecBufferSize = CodecUtils.getOutputBufferSize(codec); if (availableMemoryBytes == 0) { Preconditions.checkArgument(((numPartitions == 1) && !pipelinedShuffle), "availableMemory " @@ -645,7 +650,8 @@ protected SpillResult callInternal() throws IOException { continue; } if (writer == null) { - writer = new Writer(keySerialization, valSerialization, out, keyClass, valClass, codec, null, null); + writer = new Writer(keySerialization, valSerialization, out, keyClass, valClass, codec, + codecBufferSize, maxRawLength(filledBuffers, i), null, null); } numRecords += writePartition(buffer.partitionPositions[i], buffer, writer, key, val); } @@ -686,6 +692,17 @@ protected SpillResult callInternal() throws IOException { } } + /** Upper bound of the raw length of a partition spilled from the buffers. */ + private static long maxRawLength(List buffers, int partition) { + long numRecords = 0; + long keyValueBytes = 0; + for (WrappedBuffer buffer : buffers) { + numRecords += buffer.recordsPerPartition[partition]; + keyValueBytes += buffer.sizePerPartition[partition]; + } + return Writer.maxRawLength(numRecords, keyValueBytes); + } + private long writePartition(int pos, WrappedBuffer wrappedBuffer, Writer writer, DataInputBuffer keyBuffer, DataInputBuffer valBuffer) throws IOException { long numRecords = 0; @@ -1094,7 +1111,8 @@ private void mergeAll() throws IOException { sourceDestNameTrimmed + ": " + "Skipping partition: " + i + " in final merge since it has no records"); continue; } - writer = new Writer(keySerialization, valSerialization, out, keyClass, valClass, codec, null, null); + writer = new Writer(keySerialization, valSerialization, out, keyClass, valClass, codec, + codecBufferSize, mergedMaxRawLength(i), null, null); try { if (currentBuffer.nextPosition != 0 && currentBuffer.partitionPositions[i] != WrappedBuffer.PARTITION_ABSENT_POSITION) { @@ -1148,6 +1166,21 @@ private void mergeAll() throws IOException { LOG.info(sourceDestNameTrimmed + ": " + "Finished final spill after merging : " + numSpills.get() + " spills"); } + /** + * Upper bound of the raw length of a partition merged from the current buffer and the spills: + * each part counts the header and EOF markers the merged stream writes once. + */ + private long mergedMaxRawLength(int partition) { + long maxRawLength = Writer.maxRawLength(currentBuffer.recordsPerPartition[partition], + currentBuffer.sizePerPartition[partition]); + synchronized (spillInfoList) { + for (SpillInfo spillInfo : spillInfoList) { + maxRawLength += spillInfo.spillRecord.getIndex(partition).getRawLength(); + } + } + return maxRawLength; + } + private void deleteIntermediateSpills() { // Delete the intermediate spill files synchronized (spillInfoList) { @@ -1185,7 +1218,8 @@ private void writeLargeRecord(final Object key, final Object value, final int pa spilledRecordsCounter.increment(1); Writer writer = null; try { - writer = new IFile.Writer(keySerialization, valSerialization, out, keyClass, valClass, codec, null, null); + writer = new IFile.Writer(keySerialization, valSerialization, out, keyClass, valClass, codec, + codecBufferSize, -1, null, null); writer.append(key, value); outputLargeRecordsCounter.increment(1); numRecordsPerPartition[i]++; diff --git a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/utils/CodecUtils.java b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/utils/CodecUtils.java index 674de15612..c8f212b0ae 100644 --- a/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/utils/CodecUtils.java +++ b/tez-runtime-library/src/main/java/org/apache/tez/runtime/library/utils/CodecUtils.java @@ -26,6 +26,7 @@ import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.CommonConfigurationKeys; import org.apache.hadoop.fs.CommonConfigurationKeysPublic; +import org.apache.hadoop.io.compress.BlockCompressorStream; import org.apache.hadoop.io.compress.CodecPool; import org.apache.hadoop.io.compress.CompressionCodec; import org.apache.hadoop.io.compress.CompressionInputStream; @@ -160,6 +161,71 @@ public static CompressionOutputStream createOutputStream(CompressionCodec codec, } } + /** + * The buffer size of the codec's output streams. The codec reads it from its Configuration for + * every stream it creates; a writer of many streams reads it once and passes it to + * {@link #createOutputStream(CompressionCodec, OutputStream, Compressor, int, long)}. + * + * @return the buffer size, or -1 if the codec creates its streams itself + */ + public static int getOutputBufferSize(CompressionCodec codec) { + if (codec == null || getBlockCompressionOverhead(codec.getClass().getName(), 0) < 0) { + return -1; + } + Configuration conf = ((Configurable) codec).getConf(); + synchronized (conf) { + return conf.getInt(getBufferSizeProperty(codec), getDefaultBufferSize(conf, codec)); + } + } + + /** + * Creates the output stream of a Snappy or LZ4 codec without reading its Configuration, sized to + * the data when that fits in one block; leaves any other codec to create its own stream. + *

+ * These codecs write a {@link BlockCompressorStream}: the input in blocks of at most + * bufferSize - overhead bytes, each compressed into a buffer of bufferSize bytes. Any block size + * of at least dataLength writes a stream of at most dataLength bytes as a single block, so a + * block of dataLength bytes, compressed into a buffer of its compression bound (the overhead + * formulas are the codecs' bounds), writes the same bytes as the configured block size. A small + * per-partition stream of a spill then no longer allocates and zeroes the configured buffer + * (256 KiB by default). A block of 0 bytes would never take input, hence a positive dataLength. + * + * @param bufferSize the codec's {@link #getOutputBufferSize(CompressionCodec)} + * @param dataLength upper bound of the bytes written to the stream, or -1 if unknown + */ + public static CompressionOutputStream createOutputStream(CompressionCodec codec, + OutputStream checksumOut, Compressor compressor, int bufferSize, long dataLength) + throws IOException { + if (bufferSize < 0) { + return createOutputStream(codec, checksumOut, compressor); + } + String codecClassName = codec.getClass().getName(); + int overhead = getBlockCompressionOverhead(codecClassName, bufferSize); + if (dataLength > 0 && dataLength < bufferSize - overhead) { + int blockSize = (int) dataLength; + int blockOverhead = getBlockCompressionOverhead(codecClassName, blockSize); + return new BlockCompressorStream(checksumOut, compressor, blockSize + blockOverhead, + blockOverhead); + } + return new BlockCompressorStream(checksumOut, compressor, bufferSize, overhead); + } + + /** + * The overhead SnappyCodec and Lz4Codec reserve in their BlockCompressorStream for a buffer of + * the given size, or -1 for any other codec. + */ + @VisibleForTesting + static int getBlockCompressionOverhead(String codecClassName, int bufferSize) { + switch (codecClassName) { + case "org.apache.hadoop.io.compress.SnappyCodec": + return bufferSize / 6 + 32; + case "org.apache.hadoop.io.compress.Lz4Codec": + return bufferSize / 255 + 16; + default: + return -1; + } + } + public static String getBufferSizeProperty(CompressionCodec codec) { return getBufferSizeProperty(codec.getClass().getName()); } diff --git a/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/sort/impl/TestIFile.java b/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/sort/impl/TestIFile.java index 29704e27fa..f79187847c 100644 --- a/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/sort/impl/TestIFile.java +++ b/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/sort/impl/TestIFile.java @@ -18,6 +18,7 @@ */ package org.apache.tez.runtime.library.common.sort.impl; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertInstanceOf; @@ -53,11 +54,13 @@ import org.apache.hadoop.io.Text; import org.apache.hadoop.io.compress.CompressionCodec; import org.apache.hadoop.io.compress.CompressionCodecFactory; +import org.apache.hadoop.io.compress.SnappyCodec; import org.apache.hadoop.io.compress.lz4.Lz4Compressor; import org.apache.hadoop.io.serializer.Deserializer; import org.apache.hadoop.io.serializer.SerializationFactory; import org.apache.hadoop.io.serializer.WritableSerialization; import org.apache.hadoop.util.NativeCodeLoader; +import org.apache.hadoop.util.ReflectionUtils; import org.apache.tez.common.TezRuntimeFrameworkConfigs; import org.apache.tez.runtime.library.common.InputAttemptIdentifier; import org.apache.tez.runtime.library.common.shuffle.orderedgrouped.InMemoryReader; @@ -821,6 +824,67 @@ public void testLz4CompressedDataIsLargerThanOriginal() throws IOException { comp.compress(buf, 0, 32); } + /** + * A writer given an upper bound of its raw length writes the bytes of an unsized writer, from a + * record to several codec blocks, and with values longer than the write buffer. + */ + @Test + public void testSizedWriterWritesUnsizedBytes() throws IOException { + CompressionCodec snappy = ReflectionUtils.newInstance(SnappyCodec.class, new Configuration(false)); + int codecBufferSize = CodecUtils.getOutputBufferSize(snappy); + Random random = new Random(0); + // {records, value length bound}: one record, one block, several blocks, values above the write buffer + for (int[] shape : new int[][] {{1, 50}, {100, 50}, {10_000, 50}, {5, 200_000}}) { + List records = new ArrayList<>(); + long keyValueBytes = 0; + for (int i = 0; i < shape[0]; i++) { + byte[] key = new byte[1 + random.nextInt(20)]; + byte[] value = new byte[random.nextInt(shape[1])]; + random.nextBytes(key); + random.nextBytes(value); + records.add(new byte[][] {key, value}); + keyValueBytes += key.length + value.length; + } + long maxRawLength = Writer.maxRawLength(records.size(), keyValueBytes); + + ByteArrayOutputStream unsized = new ByteArrayOutputStream(); + ByteArrayOutputStream sized = new ByteArrayOutputStream(); + writeRecords(new Writer(null, null, new FSDataOutputStream(unsized, null), null, null, snappy, null, null), + records); + Writer writer = writeRecords(new Writer(null, null, new FSDataOutputStream(sized, null), null, null, snappy, + codecBufferSize, maxRawLength, null, null), records); + + String what = Arrays.toString(shape); + assertTrue(writer.getRawLength() <= maxRawLength, what); + assertArrayEquals(unsized.toByteArray(), sized.toByteArray(), what); + Reader reader = new Reader(new ByteArrayInputStream(sized.toByteArray()), sized.size(), snappy, null, null, + false, 0, -1); + DataInputBuffer key = new DataInputBuffer(); + DataInputBuffer value = new DataInputBuffer(); + for (byte[][] record : records) { + assertTrue(reader.nextRawKey(key), what); + reader.nextRawValue(value); + assertArrayEquals(record[0], Arrays.copyOfRange(key.getData(), key.getPosition(), key.getLength()), what); + assertArrayEquals(record[1], Arrays.copyOfRange(value.getData(), value.getPosition(), value.getLength()), + what); + } + assertFalse(reader.nextRawKey(key), what); + reader.close(); + } + } + + private static Writer writeRecords(Writer writer, List records) throws IOException { + DataInputBuffer key = new DataInputBuffer(); + DataInputBuffer value = new DataInputBuffer(); + for (byte[][] record : records) { + key.reset(record[0], record[0].length); + value.reset(record[1], record[1].length); + writer.append(key, value); + } + writer.close(); + return writer; + } + private void initBufWithNumbers(byte[] buf, int... args) { for (int i = 0; i < args.length; i++) { buf[i] = (byte) args[i]; diff --git a/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/writers/TestUnorderedPartitionedKVWriter.java b/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/writers/TestUnorderedPartitionedKVWriter.java index 8fb045bae8..d684e6f2c5 100644 --- a/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/writers/TestUnorderedPartitionedKVWriter.java +++ b/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/common/writers/TestUnorderedPartitionedKVWriter.java @@ -18,6 +18,7 @@ */ package org.apache.tez.runtime.library.common.writers; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertInstanceOf; @@ -37,11 +38,14 @@ import static org.mockito.Mockito.verify; import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; import java.io.DataInputStream; import java.io.IOException; import java.io.InputStream; +import java.lang.management.ManagementFactory; import java.nio.ByteBuffer; import java.util.ArrayList; +import java.util.Arrays; import java.util.BitSet; import java.util.HashMap; import java.util.Iterator; @@ -51,6 +55,7 @@ import java.util.Set; import java.util.UUID; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; import java.util.regex.Matcher; import java.util.regex.Pattern; import java.util.stream.Stream; @@ -59,6 +64,7 @@ import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.CommonConfigurationKeys; import org.apache.hadoop.fs.FSDataInputStream; +import org.apache.hadoop.fs.FSDataOutputStream; import org.apache.hadoop.fs.FileSystem; import org.apache.hadoop.fs.Path; import org.apache.hadoop.fs.permission.FsAction; @@ -67,8 +73,10 @@ import org.apache.hadoop.io.LongWritable; import org.apache.hadoop.io.Text; import org.apache.hadoop.io.Writable; +import org.apache.hadoop.io.compress.CodecPool; import org.apache.hadoop.io.compress.CompressionCodec; import org.apache.hadoop.io.compress.DefaultCodec; +import org.apache.hadoop.io.compress.SnappyCodec; import org.apache.hadoop.yarn.api.records.ApplicationId; import org.apache.tez.common.TezCommonUtils; import org.apache.tez.common.TezRuntimeFrameworkConfigs; @@ -105,6 +113,7 @@ import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; import org.junit.jupiter.api.Timeout; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; @@ -1601,6 +1610,130 @@ private static String createRandomString(int size) { return sb.toString(); } + /** + * A spill sizes each partition's stream to the partition without reading the codec's buffer size + * per stream, and writes the bytes of an unsized writer: a final spill of many partitions, spills + * of several buffers each, and a merge of many spills. + */ + @Test + public void testSpillStreamsSizedToPartitions() throws IOException, InterruptedException { + setupInit(true, ReportPartitionStats.DISABLED); + int numPartitions = 1009; + UnorderedPartitionedKVWriter kvWriter = createSnappyWriter(numPartitions, 16 << 20, -1, true); + write(kvWriter, 50_000); + ReadCountingConfiguration conf = (ReadCountingConfiguration) kvWriter.conf; + // The final spill runs on this thread; it borrows the pooled compressor instead of creating one. + CodecPool.returnCompressor(CodecPool.getCompressor(kvWriter.codec)); + int reads = conf.reads.get(); + long allocated = getCurrentThreadAllocatedBytes(); + kvWriter.close(); + allocated = getCurrentThreadAllocatedBytes() - allocated; + assertEquals(1, kvWriter.numSpills.get()); + assertEquals(reads, conf.reads.get()); + // The configured buffers would take 320 KiB per partition: 256 KiB codec + 64 KiB write buffer. + assertTrue(allocated < numPartitions * 32L * 1024, "allocated " + allocated); + assertPartitionsWrittenUnsized(kvWriter, kvWriter.finalOutPath, kvWriter.finalIndexPath, numPartitions); + + numPartitions = 10; + kvWriter = createSnappyWriter(numPartitions, 10 << 10, 1 << 10, false); + write(kvWriter, 5000); + kvWriter.close(); + assertTrue(kvWriter.spillLimit > 1); + for (int i = 0; i < kvWriter.numSpills.get(); i++) { + assertPartitionsWrittenUnsized(kvWriter, kvWriter.outputFileHandler.getSpillFileForWrite(i, 0), + kvWriter.outputFileHandler.getSpillIndexFileForWrite(i, 0), numPartitions); + } + + kvWriter = createSnappyWriter(numPartitions, 2048, -1, true); + write(kvWriter, 5000); + kvWriter.close(); + assertTrue(kvWriter.numSpills.get() > 100); + assertPartitionsWrittenUnsized(kvWriter, kvWriter.finalOutPath, kvWriter.finalIndexPath, numPartitions); + } + + /** Spills of a writer without final merge take all its buffers. */ + private UnorderedPartitionedKVWriter createSnappyWriter(int numPartitions, long availableMemory, + int maxSingleBufferSizeBytes, boolean finalMerge) throws IOException { + OutputContext outputContext = createMockOutputContext(new TezCounters(), ApplicationId.newInstance(10000000, 1), + UUID.randomUUID().toString(), defaultConf.get(TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID, + TezConfiguration.TEZ_AM_SHUFFLE_AUXILIARY_SERVICE_ID_DEFAULT)); + Configuration conf = new ReadCountingConfiguration( + createConfiguration(outputContext, IntWritable.class, LongWritable.class, true, maxSingleBufferSizeBytes), + CommonConfigurationKeys.IO_COMPRESSION_CODEC_SNAPPY_BUFFERSIZE_KEY); + conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_COMPRESS_CODEC, SnappyCodec.class.getName()); + conf.setBoolean(TezRuntimeConfiguration.TEZ_RUNTIME_ENABLE_FINAL_MERGE_IN_OUTPUT, finalMerge); + if (!finalMerge) { + conf.setInt(TezRuntimeConfiguration.TEZ_RUNTIME_UNORDERED_PARTITIONED_KVWRITER_BUFFER_MERGE_PERCENT, 100); + } + return new UnorderedPartitionedKVWriterForTest(outputContext, conf, numPartitions, availableMemory); + } + + private static void write(UnorderedPartitionedKVWriter kvWriter, int numRecords) throws IOException { + IntWritable key = new IntWritable(); + LongWritable value = new LongWritable(); + for (int i = 0; i < numRecords; i++) { + key.set(i); + value.set(i); + kvWriter.write(key, value); + } + } + + private static long getCurrentThreadAllocatedBytes() { + return ((com.sun.management.ThreadMXBean) ManagementFactory.getThreadMXBean()) + .getCurrentThreadAllocatedBytes(); + } + + /** Each partition of the output holds the bytes an unsized writer writes for its records. */ + private static void assertPartitionsWrittenUnsized(UnorderedPartitionedKVWriter kvWriter, Path outputPath, + Path indexPath, int numPartitions) throws IOException { + byte[] output = new byte[(int) localFs.getFileStatus(outputPath).getLen()]; + try (FSDataInputStream in = localFs.open(outputPath)) { + in.readFully(output); + } + TezSpillRecord spillRecord = new TezSpillRecord(indexPath, kvWriter.conf); + DataInputBuffer key = new DataInputBuffer(); + DataInputBuffer value = new DataInputBuffer(); + for (int i = 0; i < numPartitions; i++) { + TezIndexRecord indexRecord = spillRecord.getIndex(i); + if (indexRecord.getPartLength() == 0) { + continue; + } + int start = (int) indexRecord.getStartOffset(); + byte[] segment = Arrays.copyOfRange(output, start, start + (int) indexRecord.getPartLength()); + Reader reader = new Reader(new ByteArrayInputStream(segment), segment.length, kvWriter.codec, null, null, + false, 0, -1); + ByteArrayOutputStream unsized = new ByteArrayOutputStream(); + IFile.Writer writer = new IFile.Writer(null, null, new FSDataOutputStream(unsized, null), null, null, + kvWriter.codec, null, null); + while (reader.nextRawKey(key)) { + reader.nextRawValue(value); + writer.append(key, value); + } + reader.close(); + writer.close(); + assertArrayEquals(unsized.toByteArray(), segment, "partition " + i); + } + } + + /** Counts the reads of one property. */ + private static class ReadCountingConfiguration extends Configuration { + private final String property; + private final AtomicInteger reads = new AtomicInteger(); + + ReadCountingConfiguration(Configuration conf, String property) { + super(conf); + this.property = property; + } + + @Override + public String get(String name) { + if (property.equals(name)) { + reads.incrementAndGet(); + } + return super.get(name); + } + } + private OutputContext createMockOutputContext(TezCounters counters, ApplicationId appId, String uniqueId, String auxiliaryService) { OutputContext outputContext = mock(OutputContext.class); diff --git a/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/utils/TestCodecUtils.java b/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/utils/TestCodecUtils.java index 5be628819a..0b6717b6c9 100644 --- a/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/utils/TestCodecUtils.java +++ b/tez-runtime-library/src/test/java/org/apache/tez/runtime/library/utils/TestCodecUtils.java @@ -18,13 +18,18 @@ */ package org.apache.tez.runtime.library.utils; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assumptions.assumeTrue; import static org.mockito.Mockito.mock; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; import java.lang.reflect.Field; +import java.util.Random; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; @@ -35,12 +40,14 @@ import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.CommonConfigurationKeys; import org.apache.hadoop.fs.CommonConfigurationKeysPublic; +import org.apache.hadoop.io.IOUtils; import org.apache.hadoop.io.compress.BZip2Codec; import org.apache.hadoop.io.compress.CodecPool; import org.apache.hadoop.io.compress.CompressionCodec; import org.apache.hadoop.io.compress.CompressionInputStream; import org.apache.hadoop.io.compress.CompressionOutputStream; import org.apache.hadoop.io.compress.Compressor; +import org.apache.hadoop.io.compress.CompressorStream; import org.apache.hadoop.io.compress.Decompressor; import org.apache.hadoop.io.compress.DecompressorStream; import org.apache.hadoop.io.compress.DefaultCodec; @@ -48,11 +55,14 @@ import org.apache.hadoop.io.compress.Lz4Codec; import org.apache.hadoop.io.compress.SnappyCodec; import org.apache.hadoop.io.compress.ZStandardCodec; +import org.apache.hadoop.util.ReflectionUtils; import org.apache.tez.runtime.library.api.TezRuntimeConfiguration; import org.apache.tez.runtime.library.common.shuffle.orderedgrouped.DummyCompressionCodec; import org.apache.tez.runtime.library.common.sort.impl.IFileInputStream; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; public class TestCodecUtils { @@ -203,6 +213,123 @@ public void testDefaultBufferSize() { CodecUtils.getDefaultBufferSize(conf, new Lz4Codec())); } + /** + * A stream sized to its data writes the bytes of the configured stream, from one byte to several + * configured blocks of compressible and of incompressible data, and gets a buffer of the + * compression bound of the data when the data fits one configured block. + */ + @ParameterizedTest + @ValueSource(classes = {SnappyCodec.class, Lz4Codec.class}) + public void testOutputStreamSizedToData(Class codecClass) throws IOException { + Configuration conf = new Configuration(false); + conf.setInt(CodecUtils.getBufferSizeProperty(codecClass.getName()), 64 * 1024); + CompressionCodec codec = newBlockCodec(codecClass, conf); + int bufferSize = CodecUtils.getOutputBufferSize(codec); + assertEquals(64 * 1024, bufferSize); + int maxBlockSize = bufferSize - CodecUtils.getBlockCompressionOverhead(codecClass.getName(), bufferSize); + for (boolean compressible : new boolean[] {true, false}) { + for (int length : new int[] {1, 1000, maxBlockSize - 1, maxBlockSize, maxBlockSize + 1, 5 * bufferSize}) { + byte[] data = randomBytes(length, compressible); + ByteArrayOutputStream configured = new ByteArrayOutputStream(); + ByteArrayOutputStream sized = new ByteArrayOutputStream(); + compress(codec, data, configured, -1, -1); + int sizedBuffer = compress(codec, data, sized, bufferSize, length); + + String what = "length " + length + ", compressible " + compressible; + assertEquals(length < maxBlockSize + ? length + CodecUtils.getBlockCompressionOverhead(codecClass.getName(), length) : bufferSize, + sizedBuffer, what); + assertArrayEquals(configured.toByteArray(), sized.toByteArray(), what); + assertArrayEquals(data, decompress(codec, configured.toByteArray(), length), what); + assertArrayEquals(data, decompress(codec, sized.toByteArray(), length), what); + } + } + } + + /** A length that underestimates the data costs more blocks, not the data. */ + @ParameterizedTest + @ValueSource(classes = {SnappyCodec.class, Lz4Codec.class}) + public void testOutputStreamSizedToUnderestimatedLength(Class codecClass) + throws IOException { + CompressionCodec codec = newBlockCodec(codecClass, new Configuration(false)); + byte[] data = randomBytes(100_000, false); + ByteArrayOutputStream sized = new ByteArrayOutputStream(); + compress(codec, data, sized, CodecUtils.getOutputBufferSize(codec), 1000); + assertArrayEquals(data, decompress(codec, sized.toByteArray(), data.length)); + } + + /** Lz4Codec needs lz4-java, which Hadoop declares provided. */ + private static CompressionCodec newBlockCodec(Class codecClass, + Configuration conf) { + if (codecClass == Lz4Codec.class) { + try { + Class.forName("net.jpountz.lz4.LZ4Factory"); + } catch (ClassNotFoundException e) { + assumeTrue(false, "lz4-java is not on the classpath"); + } + } + return ReflectionUtils.newInstance(codecClass, conf); + } + + private static byte[] randomBytes(int length, boolean compressible) { + byte[] data = new byte[length]; + Random random = new Random(length); + random.nextBytes(data); + if (compressible) { + for (int i = 0; i < length; i++) { + data[i] &= 0x0f; + } + } + return data; + } + + /** + * Writes the data in 64 KiB chunks as IFile.Writer does, through the codec's own stream when + * bufferSize is -1. + * + * @return the length of the stream's buffer + */ + private static int compress(CompressionCodec codec, byte[] data, OutputStream out, int bufferSize, + long dataLength) throws IOException { + Compressor compressor = CodecUtils.getCompressor(codec); + try { + CompressionOutputStream stream = + CodecUtils.createOutputStream(codec, out, compressor, bufferSize, dataLength); + for (int off = 0; off < data.length; off += 64 * 1024) { + stream.write(data, off, Math.min(64 * 1024, data.length - off)); + } + stream.finish(); + return getCompressorStreamBufferSize(stream); + } finally { + CodecPool.returnCompressor(compressor); + } + } + + private static byte[] decompress(CompressionCodec codec, byte[] compressed, int length) + throws IOException { + Decompressor decompressor = CodecUtils.getDecompressor(codec); + try { + InputStream in = CodecUtils.createInputStream(codec, new ByteArrayInputStream(compressed), + decompressor); + byte[] data = new byte[length]; + IOUtils.readFully(in, data, 0, length); + assertEquals(-1, in.read()); + return data; + } finally { + CodecPool.returnDecompressor(decompressor); + } + } + + private static int getCompressorStreamBufferSize(CompressionOutputStream stream) { + try { + Field field = CompressorStream.class.getDeclaredField("buffer"); + field.setAccessible(true); + return ((byte[]) field.get(stream)).length; + } catch (ReflectiveOperationException e) { + throw new RuntimeException(e); + } + } + private void waitForLatch(CountDownLatch latch) { try { latch.await();