GH-3150: Keep column values intact with direct codecs

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()) {

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())) {

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<Arguments> 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<Group> 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<Group> builder = ParquetReader.builder(new GroupReadSupport(), path)
+            .withConf(conf)
+            .withAllocator(allocator);
+        if (direct) {
+          builder.withCodecFactory(CodecFactory.createDirectCodecFactory(
+              conf, codecAllocator, ParquetProperties.DEFAULT_PAGE_SIZE));
+        }
+        try (ParquetReader<Group> 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 {