diff --git a/parquet-format-structures/src/main/java/org/apache/parquet/format/Util.java b/parquet-format-structures/src/main/java/org/apache/parquet/format/Util.java index 419d757702..6e50d7ed00 100644 --- a/parquet-format-structures/src/main/java/org/apache/parquet/format/Util.java +++ b/parquet-format-structures/src/main/java/org/apache/parquet/format/Util.java @@ -158,6 +158,14 @@ public static FileMetaData readFileMetaData(InputStream from, BlockCipher.Decryp return read(from, new FileMetaData(), decryptor, AAD); } + /** + * Reads file metadata using a custom Thrift max message size. + * + * @param from the stream to read the metadata from + * @param maxMessageSize maximum Thrift message size in bytes; {@code -1} uses the default (100 MB) + * @return the resulting metadata + * @throws IOException if any I/O error occurs during the reading + */ public static FileMetaData readFileMetaData(InputStream from, int maxMessageSize) throws IOException { return readFileMetaData(from, null, null, maxMessageSize); } diff --git a/parquet-format-structures/src/test/java/org/apache/parquet/format/TestUtil.java b/parquet-format-structures/src/test/java/org/apache/parquet/format/TestUtil.java index d78023d9eb..179fcf633d 100644 --- a/parquet-format-structures/src/test/java/org/apache/parquet/format/TestUtil.java +++ b/parquet-format-structures/src/test/java/org/apache/parquet/format/TestUtil.java @@ -70,6 +70,36 @@ public void testInvalidPageHeader() throws IOException { .hasMessageContaining("Compressed page size"); } + @Test + public void testReadFileMetaDataAcceptsMinusOneAsDefaultMaxMessageSize() throws Exception { + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + FileMetaData md = sampleFileMetaData(); + writeFileMetaData(md, baos); + + assertThat(readFileMetaData(in(baos), -1)).isEqualTo(md); + } + + @Test + public void testReadFileMetaDataRejectsZeroMaxMessageSize() { + assertRejectsNonPositiveMaxMessageSize(0); + } + + @Test + public void testReadFileMetaDataRejectsInvalidNegativeMaxMessageSize() { + assertRejectsNonPositiveMaxMessageSize(-5); + } + + private static FileMetaData sampleFileMetaData() { + return new FileMetaData( + 1, asList(new SchemaElement("foo")), 10, asList(new RowGroup(asList(new ColumnChunk(0)), 10, 5))); + } + + private static void assertRejectsNonPositiveMaxMessageSize(int maxMessageSize) { + assertThatThrownBy(() -> readFileMetaData(new ByteArrayInputStream(new byte[0]), maxMessageSize)) + .isInstanceOf(NumberFormatException.class) + .hasMessage("Max message size must be positive: " + maxMessageSize); + } + private ByteArrayInputStream in(ByteArrayOutputStream baos) { return new ByteArrayInputStream(baos.toByteArray()); }