diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ColumnChunkPageReadStore.java b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ColumnChunkPageReadStore.java index dfdd6c5d82..d36c4f4c24 100644 --- a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ColumnChunkPageReadStore.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ColumnChunkPageReadStore.java @@ -177,7 +177,8 @@ public DataPage visit(DataPageV1 dataPageV1) { bytes = BytesInput.from(blockDecryptor.decrypt(bytes.toByteArray(), dataPageAAD)); } long start = System.nanoTime(); - decompressed = decompressor.decompress(bytes, dataPageV1.getUncompressedSize()); + decompressed = + BytesInput.copy(decompressor.decompress(bytes, dataPageV1.getUncompressedSize())); setDecompressMetrics(bytes, start); } @@ -257,7 +258,7 @@ public DataPage visit(DataPageV2 dataPageV2) { - dataPageV2.getDefinitionLevels().size() - dataPageV2.getRepetitionLevels().size()); long start = System.nanoTime(); - pageBytes = decompressor.decompress(pageBytes, uncompressedSize); + pageBytes = BytesInput.copy(decompressor.decompress(pageBytes, uncompressedSize)); setDecompressMetrics(pageBytes, start); } } @@ -326,7 +327,9 @@ public DictionaryPage readDictionaryPage() { long start = System.nanoTime(); setDecompressMetrics(bytes, start); DictionaryPage decompressedPage = new DictionaryPage( - decompressor.decompress(bytes, compressedDictionaryPage.getUncompressedSize()), + decompressor + .decompress(bytes, compressedDictionaryPage.getUncompressedSize()) + .copy(releaser), compressedDictionaryPage.getDictionarySize(), compressedDictionaryPage.getEncoding()); if (compressedDictionaryPage.getCrc().isPresent()) { diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestColumnChunkPageReadStore.java b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestColumnChunkPageReadStore.java index a0fa4377b4..d83bae46bd 100644 --- a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestColumnChunkPageReadStore.java +++ b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestColumnChunkPageReadStore.java @@ -20,7 +20,12 @@ import static org.apache.parquet.column.Encoding.PLAIN; import static org.apache.parquet.column.Encoding.RLE; +import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.verify; import java.nio.ByteBuffer; import java.util.Collections; @@ -58,6 +63,23 @@ public void decompress(ByteBuffer input, int compressedSize, ByteBuffer output, public void release() {} }; + @Test + public void heapDataPagesAreNotRetainedUntilRowGroupClose() throws Exception { + ByteBufferAllocator allocator = spy(new HeapByteBufferAllocator()); + ParquetReadOptions options = + ParquetReadOptions.builder().withAllocator(allocator).build(); + byte[] bytes = {1, 2, 3, 4}; + DataPageV1 page = new DataPageV1(BytesInput.from(bytes), 1, bytes.length, null, RLE, RLE, PLAIN); + ColumnChunkPageReader reader = new ColumnChunkPageReader( + NOOP_DECOMPRESSOR, Collections.singletonList(page), null, null, 1L, null, null, 0, 0, options); + try (ColumnChunkPageReadStore store = new ColumnChunkPageReadStore(1L)) { + store.addColumn(COLUMN, reader); + assertThat(((DataPageV1) reader.readPage()).getBytes().toByteArray()) + .isEqualTo(bytes); + verify(allocator, never()).allocate(anyInt()); + } + } + @Test public void closeWithoutSetReleaserDoesNotThrow() { try (TrackingByteBufferAllocator allocator = TrackingByteBufferAllocator.wrap(new HeapByteBufferAllocator())) { diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestParquetReader.java b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestParquetReader.java index 87c5630ec7..4aee0d0b9d 100644 --- a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestParquetReader.java +++ b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestParquetReader.java @@ -35,24 +35,35 @@ import java.util.stream.Stream; import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.Path; +import org.apache.parquet.ParquetReadOptions; +import org.apache.parquet.bytes.DirectByteBufferAllocator; import org.apache.parquet.bytes.HeapByteBufferAllocator; import org.apache.parquet.bytes.TrackingByteBufferAllocator; import org.apache.parquet.column.ParquetProperties; +import org.apache.parquet.column.ParquetProperties.WriterVersion; +import org.apache.parquet.column.page.DictionaryPage; +import org.apache.parquet.column.page.PageReadStore; import org.apache.parquet.example.data.Group; +import org.apache.parquet.example.data.simple.SimpleGroupFactory; import org.apache.parquet.filter2.compat.FilterCompat; import org.apache.parquet.filter2.recordlevel.PhoneBookWriter; import org.apache.parquet.hadoop.ParquetReader.Builder; import org.apache.parquet.hadoop.api.ReadSupport; import org.apache.parquet.hadoop.example.ExampleParquetWriter; import org.apache.parquet.hadoop.example.GroupReadSupport; +import org.apache.parquet.hadoop.metadata.CompressionCodecName; import org.apache.parquet.hadoop.util.HadoopInputFile; import org.apache.parquet.io.InputFile; +import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.MessageTypeParser; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.io.TempDir; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.CsvSource; import org.junit.jupiter.params.provider.MethodSource; public class TestParquetReader { @@ -65,6 +76,9 @@ public class TestParquetReader { private TrackingByteBufferAllocator allocator; + @TempDir + private java.nio.file.Path tempDir; + static Stream data() { return Stream.of(Arguments.of(FILE_V1), Arguments.of(FILE_V2), Arguments.of(STATIC_FILE_WITHOUT_COL_INDEXES)); } @@ -174,6 +188,84 @@ public void closeAllocator() { allocator.close(); } + @ParameterizedTest + @CsvSource({ + "SNAPPY, PARQUET_1_0, false", + "SNAPPY, PARQUET_2_0, false", + "ZSTD, PARQUET_1_0, false", + "ZSTD, PARQUET_2_0, false", + "SNAPPY, PARQUET_1_0, true", + "SNAPPY, PARQUET_2_0, true", + "ZSTD, PARQUET_1_0, true", + "ZSTD, PARQUET_2_0, true" + }) + public void testSharedDirectCodecPreservesColumnValues( + CompressionCodecName codec, WriterVersion version, boolean dictionaryEnabled) throws IOException { + Path path = new Path(tempDir.resolve("compressed.parquet").toUri()); + Configuration conf = new Configuration(); + MessageType schema = MessageTypeParser.parseMessageType( + "message records { optional binary key (STRING); optional binary value (STRING); }"); + SimpleGroupFactory groups = new SimpleGroupFactory(schema); + try (ParquetWriter writer = ExampleParquetWriter.builder(path) + .withConf(conf) + .withType(schema) + .withCompressionCodec(codec) + .withWriterVersion(version) + .withDictionaryEncoding(dictionaryEnabled) + .withPageRowCountLimit(100) + .build()) { + for (int i = 1; i <= 1000; i++) { + int valueIndex = dictionaryEnabled ? i % 17 : i; + writer.write( + groups.newGroup().append("key", "key_" + valueIndex).append("value", "value_" + valueIndex)); + } + } + + for (boolean direct : new boolean[] {false, true}) { + try (TrackingByteBufferAllocator codecAllocator = + TrackingByteBufferAllocator.wrap(new DirectByteBufferAllocator())) { + ParquetReader.Builder builder = ParquetReader.builder(new GroupReadSupport(), path) + .withConf(conf) + .withAllocator(allocator); + if (direct) { + builder.withCodecFactory(CodecFactory.createDirectCodecFactory( + conf, codecAllocator, ParquetProperties.DEFAULT_PAGE_SIZE)); + } + try (ParquetReader reader = builder.build()) { + for (int i = 1; i <= 1000; i++) { + int valueIndex = dictionaryEnabled ? i % 17 : i; + Group record = reader.read(); + assertThat(record).isNotNull(); + assertThat(record.getBinary("key", 0).toStringUsingUTF8()) + .as("key at row %s with direct codec %s", i, direct) + .isEqualTo("key_" + valueIndex); + assertThat(record.getBinary("value", 0).toStringUsingUTF8()) + .as("value at row %s with direct codec %s", i, direct) + .isEqualTo("value_" + valueIndex); + } + assertThat(reader.read()).isNull(); + } + if (direct && dictionaryEnabled) { + ParquetReadOptions options = ParquetReadOptions.builder() + .withAllocator(allocator) + .withCodecFactory(CodecFactory.createDirectCodecFactory( + conf, codecAllocator, ParquetProperties.DEFAULT_PAGE_SIZE)) + .build(); + try (ParquetFileReader reader = + ParquetFileReader.open(HadoopInputFile.fromPath(path, conf), options); + PageReadStore pages = reader.readNextRowGroup()) { + DictionaryPage values = + pages.getPageReader(schema.getColumns().get(1)).readDictionaryPage(); + assertThat(values).isNotNull(); + byte[] expected = values.getBytes().toByteArray(); + pages.getPageReader(schema.getColumns().get(0)).readDictionaryPage(); + assertThat(values.getBytes().toByteArray()).isEqualTo(expected); + } + } + } + } + } + @ParameterizedTest @MethodSource("data") public void testCurrentRowIndex(Path file) throws Exception {