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");
+ }
+ }
+}