TEZ-4762: UnorderedPartitionedKVWriter: size each partition's codec stream to the partition

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,32 +373,48 @@ 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);
       }
     }
 
     /**
      * 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 {

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<WrappedBuffer> 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]++;

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.
+   * <p>
+   * 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());
   }

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<byte[][]> 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<byte[][]> 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];