diff --git a/presto-orc/src/main/java/com/facebook/presto/orc/OrcOutputBuffer.java b/presto-orc/src/main/java/com/facebook/presto/orc/OrcOutputBuffer.java index 9451296aff9a..cd50133e81b3 100644 --- a/presto-orc/src/main/java/com/facebook/presto/orc/OrcOutputBuffer.java +++ b/presto-orc/src/main/java/com/facebook/presto/orc/OrcOutputBuffer.java @@ -445,7 +445,7 @@ private void writeChunkToOutputStream(byte[] chunk, int offset, int length) return; } - checkArgument(length <= buffer.length, "Write chunk length must be less than compression buffer size"); + checkArgument(length <= maxBufferSize, "Write chunk length must be less than max compression buffer size"); boolean isCompressed = false; byte[] compressionBuffer = null; @@ -497,7 +497,7 @@ private void writeDirectlyToOutputStream(byte[] bytes, int bytesOffset, int leng } while (length > 0) { - int chunkSize = Integer.min(length, buffer.length); + int chunkSize = Integer.min(length, maxBufferSize); writeChunkToOutputStream(bytes, bytesOffset, chunkSize); length -= chunkSize; bytesOffset += chunkSize; diff --git a/presto-orc/src/test/java/com/facebook/presto/orc/TestOrcOutputBuffer.java b/presto-orc/src/test/java/com/facebook/presto/orc/TestOrcOutputBuffer.java index 549a93ce8234..26fadcde6c75 100644 --- a/presto-orc/src/test/java/com/facebook/presto/orc/TestOrcOutputBuffer.java +++ b/presto-orc/src/test/java/com/facebook/presto/orc/TestOrcOutputBuffer.java @@ -20,10 +20,13 @@ import java.util.Arrays; import java.util.Optional; +import java.util.OptionalInt; import static io.airlift.slice.Slices.wrappedBuffer; import static io.airlift.units.DataSize.Unit.BYTE; +import static io.airlift.units.DataSize.Unit.KILOBYTE; import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; public class TestOrcOutputBuffer { @@ -48,6 +51,30 @@ public void testWriteHugeByteChucks() assertEquals(output.slice(), wrappedBuffer(largeByteArray, 100, size - 100)); } + @Test + public void testWriteHugeByteChucksUsesMaxCompressionBufferSizeChunks() + { + int size = 1024 * 1024; + byte[] largeByteArray = new byte[size]; + Arrays.fill(largeByteArray, (byte) 0xA); + ColumnWriterOptions columnWriterOptions = ColumnWriterOptions.builder() + .setCompressionKind(CompressionKind.ZSTD) + .setCompressionLevel(OptionalInt.of(7)) + .setCompressionMaxBufferSize(new DataSize(256, KILOBYTE)) + .build(); + OrcOutputBuffer sliceOutput = new OrcOutputBuffer(columnWriterOptions, Optional.empty()); + + // Before the fix the compressed result would be around 90KB, after the fix it went down to 117 bytes. + DynamicSliceOutput output = new DynamicSliceOutput(size); + sliceOutput.writeBytes(largeByteArray, 10, size - 10); + assertTrue(sliceOutput.writeDataTo(output) < 200); + + sliceOutput.reset(); + output.reset(); + sliceOutput.writeBytes(wrappedBuffer(largeByteArray), 100, size - 100); + assertTrue(sliceOutput.writeDataTo(output) < 200); + } + @Test public void testGrowCapacity() {