diff --git a/parquet-common/src/main/java/org/apache/parquet/bytes/CapacityByteArrayOutputStream.java b/parquet-common/src/main/java/org/apache/parquet/bytes/CapacityByteArrayOutputStream.java index 3aa2c736a3..61d92ee87c 100644 --- a/parquet-common/src/main/java/org/apache/parquet/bytes/CapacityByteArrayOutputStream.java +++ b/parquet-common/src/main/java/org/apache/parquet/bytes/CapacityByteArrayOutputStream.java @@ -269,20 +269,41 @@ private void writeToOutput(OutputStream out, ByteBuffer buf, int len) throws IOE // The OutputStream interface only takes a byte[], unfortunately this means that a ByteBuffer // not backed by a byte array must be copied to fulfil this interface byte[] copy = new byte[len]; - buf.flip(); - buf.get(copy); + // Read through a duplicate so the slab's position and limit are left unchanged + ByteBuffer content = buf.duplicate(); + content.flip(); + content.get(copy); out.write(copy); } } + /** + * Checks that the slabs hold exactly the {@link #size()} bytes that were written. The size is used for the length + * prefixes and page headers that describe this content, so serializing content of a different length would + * silently produce a corrupt file. + */ + private void checkContentSize() { + long contentSize = 0; + for (ByteBuffer slab : slabs) { + contentSize += slab.position(); + } + if (contentSize != bytesUsed) { + throw new IllegalStateException(format( + "The slabs of this stream hold %d bytes, but %d bytes were written to it (%d slabs)", + contentSize, bytesUsed, slabs.size())); + } + } + /** * Writes the complete contents of this buffer to the specified output stream argument. the output * stream's write method out.write(slab, 0, slab.length)) will be called once per slab. * * @param out the output stream to which to write the data. * @throws IOException if an I/O error occurs. + * @throws IllegalStateException if the content of the slabs does not match {@link #size()} */ public void writeTo(OutputStream out) throws IOException { + checkContentSize(); for (ByteBuffer slab : slabs) { writeToOutput(out, slab, slab.position()); } @@ -290,11 +311,16 @@ public void writeTo(OutputStream out) throws IOException { /** * It is expected that the buffer is large enough to fit the content of this. + * + * @throws IllegalStateException if the content of the slabs does not match {@link #size()} */ void writeInto(ByteBuffer buffer) { + checkContentSize(); for (ByteBuffer slab : slabs) { - slab.flip(); - buffer.put(slab); + // Read through a duplicate so the slab's position and limit are left unchanged + ByteBuffer content = slab.duplicate(); + content.flip(); + buffer.put(content); } } @@ -349,15 +375,17 @@ public long getCurrentIndex() { public void setByte(long index, byte value) { checkArgument(index < bytesUsed, "Index: %d is >= the current size of: %d", index, bytesUsed); + // A slab may be followed by a new one before it is full (see writeInt and writeLong), so the bytes written to + // each slab are counted by its position, not by its capacity or limit long seen = 0; for (int i = 0; i < slabs.size(); i++) { ByteBuffer slab = slabs.get(i); - if (index < seen + slab.limit()) { + if (index < seen + slab.position()) { // ok found index slab.put((int) (index - seen), value); break; } - seen += slab.limit(); + seen += slab.position(); } } diff --git a/parquet-common/src/test/java/org/apache/parquet/bytes/TestCapacityByteArrayOutputStream.java b/parquet-common/src/test/java/org/apache/parquet/bytes/TestCapacityByteArrayOutputStream.java new file mode 100644 index 0000000000..692c20f9d5 --- /dev/null +++ b/parquet-common/src/test/java/org/apache/parquet/bytes/TestCapacityByteArrayOutputStream.java @@ -0,0 +1,170 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.parquet.bytes; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.List; +import java.util.stream.Stream; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; + +/** + * Tests that the content serialized by {@link CapacityByteArrayOutputStream} always matches its {@link + * CapacityByteArrayOutputStream#size()}, which is used for length prefixes and page headers (GH-3712). + */ +public class TestCapacityByteArrayOutputStream { + + private TrackingByteBufferAllocator allocator; + + static Stream allocators() { + return Stream.of(new HeapByteBufferAllocator(), new DirectByteBufferAllocator()); + } + + @AfterEach + public void closeAllocator() { + if (allocator != null) { + allocator.close(); + } + } + + private CapacityByteArrayOutputStream newStream(ByteBufferAllocator baseAllocator, int initialSlabSize) { + allocator = TrackingByteBufferAllocator.wrap(baseAllocator); + return new CapacityByteArrayOutputStream(initialSlabSize, 1024 * 1024, allocator); + } + + private static byte[] writeTo(CapacityByteArrayOutputStream stream) throws IOException { + ByteArrayOutputStream out = new ByteArrayOutputStream(); + stream.writeTo(out); + return out.toByteArray(); + } + + private static byte[] writeInto(CapacityByteArrayOutputStream stream) { + ByteBuffer buffer = ByteBuffer.allocate((int) stream.size()); + stream.writeInto(buffer); + assertThat(buffer.hasRemaining()).isFalse(); + return buffer.array(); + } + + /** + * writeInt and writeLong add a new slab when the current one has fewer bytes left than they need, so a slab may be + * followed by another one before it is full. setByte, as used by the RLE encoder to patch a bit-packed run header, + * must address the written bytes, not the capacity of the slabs. + */ + @ParameterizedTest + @MethodSource("allocators") + public void testSetByteAfterPartiallyFilledSlab(ByteBufferAllocator baseAllocator) throws IOException { + try (CapacityByteArrayOutputStream stream = newStream(baseAllocator, 8)) { + stream.write(new byte[] {0, 1, 2, 3, 4, 5}, 0, 6); + // Only 2 bytes left in the first slab: the int goes to a new slab + stream.writeInt(0x09080706); + stream.write(0); + long headerIndex = stream.getCurrentIndex(); + stream.write(11); + stream.setByte(headerIndex, (byte) 10); + stream.writeLong(0x1312111009080706L); + + byte[] expected = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 6, 7, 8, 9, 0x10, 0x11, 0x12, 0x13}; + assertThat(stream.size()).isEqualTo(expected.length); + assertThat(stream.getSlabCount()).isGreaterThan(1); + assertThat(writeTo(stream)).containsExactly(expected); + assertThat(writeInto(stream)).containsExactly(expected); + } + } + + /** Serializing the content must not change the state of the slabs, so that it can be serialized again. */ + @ParameterizedTest + @MethodSource("allocators") + public void testSerializingDoesNotChangeTheContent(ByteBufferAllocator baseAllocator) throws IOException { + try (CapacityByteArrayOutputStream stream = newStream(baseAllocator, 4)) { + ByteArrayOutputStream expected = new ByteArrayOutputStream(); + for (int i = 0; i < 100; i++) { + stream.write(i); + expected.write(i); + if (i % 7 == 0) { + assertThat(writeInto(stream)).containsExactly(expected.toByteArray()); + } + if (i % 11 == 0) { + assertThat(writeTo(stream)).containsExactly(expected.toByteArray()); + } + } + assertThat(writeTo(stream)).containsExactly(expected.toByteArray()); + assertThat(writeInto(stream)).containsExactly(expected.toByteArray()); + assertThat(BytesInput.from(stream).toByteArray()).containsExactly(expected.toByteArray()); + } + } + + /** + * Reproduces the state of the corrupt pages of GH-3712: a full slab holds fewer bytes than were written to it, + * while size() is correct. The content must not be serialized, because it doesn't match the length prefixes and + * the page header computed from size(). + */ + @ParameterizedTest + @MethodSource("allocators") + public void testSerializingFailsIfTheContentDoesNotMatchTheSize(ByteBufferAllocator baseAllocator) + throws IOException { + List slabs = new ArrayList<>(); + ByteBufferAllocator recordingAllocator = new ByteBufferAllocator() { + @Override + public ByteBuffer allocate(int size) { + ByteBuffer slab = baseAllocator.allocate(size); + slabs.add(slab); + return slab; + } + + @Override + public void release(ByteBuffer b) { + baseAllocator.release(b); + } + + @Override + public boolean isDirect() { + return baseAllocator.isDirect(); + } + }; + try (CapacityByteArrayOutputStream stream = newStream(recordingAllocator, 16)) { + byte[] data = new byte[10]; + for (int i = 0; i < 10; i++) { + stream.write(data, 0, data.length); + } + assertThat(slabs).hasSizeGreaterThan(2); + ByteBuffer fullSlab = slabs.get(1); + assertThat(fullSlab.hasRemaining()).isFalse(); + fullSlab.position(fullSlab.position() - 2); + + ByteArrayOutputStream out = new ByteArrayOutputStream(); + assertThatThrownBy(() -> stream.writeTo(out)) + .isInstanceOf(IllegalStateException.class) + .hasMessageStartingWith("The slabs of this stream hold 98 bytes, but 100 bytes were written to it"); + assertThat(out.size()).isZero(); + assertThatThrownBy(() -> stream.writeInto(ByteBuffer.allocate(100))) + .isInstanceOf(IllegalStateException.class) + .hasMessageStartingWith("The slabs of this stream hold 98 bytes, but 100 bytes were written to it"); + assertThatThrownBy(() -> BytesInput.from(stream).toByteArray()) + .isInstanceOf(IllegalStateException.class) + .hasMessageStartingWith("The slabs of this stream hold 98 bytes, but 100 bytes were written to it"); + } + } +}