GH-3558: Properly close buffers

parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetFileReader.java

@@ -65,6 +65,7 @@
 import org.apache.parquet.HadoopReadOptions;
 import org.apache.parquet.ParquetReadOptions;
 import org.apache.parquet.Preconditions;
+import org.apache.parquet.bytes.ByteBufferAllocator;
 import org.apache.parquet.bytes.ByteBufferInputStream;
 import org.apache.parquet.bytes.ByteBufferReleaser;
 import org.apache.parquet.bytes.BytesInput;
@@ -1190,15 +1191,34 @@ private ColumnChunkPageReadStore internalReadRowGroup(int blockIndex) throws IOE
     }
     // actually read all the chunks
     ChunkListBuilder builder = new ChunkListBuilder(block.getRowCount());
-    readAllPartsVectoredOrNormal(allParts, builder);
-    rowGroup.setReleaser(builder.releaser);
-    for (Chunk chunk : builder.build()) {
-      readChunkPages(chunk, block, rowGroup);
-    }
+    readChunkPagesForBlock(block, allParts, builder, rowGroup);
 
     return rowGroup;
   }
 
+  /**
+   * Reads all parts into {@code builder}, transfers ownership of the acquired buffers to {@code rowGroup} and reads
+   * the pages for every chunk. If anything fails before ownership is fully transferred, the buffers already acquired
+   * by the builder are released so a partially-read row group does not leak them.
+   */
+  private void readChunkPagesForBlock(
+      BlockMetaData block,
+      List<ConsecutivePartList> allParts,
+      ChunkListBuilder builder,
+      ColumnChunkPageReadStore rowGroup)
+      throws IOException {
+    try {
+      readAllPartsVectoredOrNormal(allParts, builder);
+      rowGroup.setReleaser(builder.releaser);
+      for (Chunk chunk : builder.build()) {
+        readChunkPages(chunk, block, rowGroup);
+      }
+    } catch (RuntimeException | IOException e) {
+      builder.releaser.close();
+      throw e;
+    }
+  }
+
   /**
    * Reads all the columns requested from the specified row group. It may skip specific pages based on the column
    * indexes according to the actual filter. As the rows are not aligned among the pages of the different columns row
@@ -1377,12 +1397,23 @@ private void readVectored(List<ConsecutivePartList> allParts, ChunkListBuilder b
       totalSize += len;
     }
     LOG.debug("Reading {} bytes of data with vectored IO in {} ranges", totalSize, ranges.size());
+    // Wrap the allocator so we release the exact buffers Hadoop allocates, regardless of what the
+    // futures return. With checksum verification enabled, ChecksumFileSystem hands back a sliced view
+    // of the allocated buffer and allocates additional checksum buffers it never releases; both would
+    // otherwise leak or trip strict allocators such as TrackingByteBufferAllocator.
+    RecordingByteBufferAllocator recordingAllocator = new RecordingByteBufferAllocator(options.getAllocator());
     // Request a vectored read;
-    f.readVectored(ranges, options.getAllocator());
-    int k = 0;
-    for (ConsecutivePartList consecutivePart : allParts) {
-      ParquetFileRange currRange = ranges.get(k++);
-      consecutivePart.readFromVectoredRange(currRange, builder);
+    f.readVectored(ranges, recordingAllocator);
+    try {
+      int k = 0;
+      for (ConsecutivePartList consecutivePart : allParts) {
+        ParquetFileRange currRange = ranges.get(k++);
+        consecutivePart.readFromVectoredRange(currRange, builder);
+      }
+    } finally {
+      // Register every buffer Hadoop allocated for this read, even if a later range fails, so that a
+      // partially-completed vectored read still releases the buffers already acquired for it.
+      builder.addBuffersToRelease(recordingAllocator.getAllocatedBuffers());
     }
   }
 
@@ -1464,11 +1495,7 @@ private ColumnChunkPageReadStore internalReadFilteredRowGroup(
         }
       }
     }
-    readAllPartsVectoredOrNormal(allParts, builder);
-    rowGroup.setReleaser(builder.releaser);
-    for (Chunk chunk : builder.build()) {
-      readChunkPages(chunk, block, rowGroup);
-    }
+    readChunkPagesForBlock(block, allParts, builder, rowGroup);
 
     return rowGroup;
   }
@@ -1865,6 +1892,42 @@ public void close() throws IOException {
     }
   }
 
+  /**
+   * A {@link ByteBufferAllocator} decorator that records every buffer it hands out so the caller can release the
+   * exact allocations later. This is used for vectored IO, where the buffers returned by the read futures may be
+   * sliced views of the allocated buffers (e.g. when {@code ChecksumFileSystem} verifies checksums), and where
+   * additional buffers (such as checksum buffers) may be allocated that are not otherwise visible to the caller.
+   */
+  private static class RecordingByteBufferAllocator implements ByteBufferAllocator {
+    private final ByteBufferAllocator delegate;
+    private final List<ByteBuffer> allocated = new ArrayList<>();
+
+    RecordingByteBufferAllocator(ByteBufferAllocator delegate) {
+      this.delegate = delegate;
+    }
+
+    @Override
+    public synchronized ByteBuffer allocate(int size) {
+      ByteBuffer buffer = delegate.allocate(size);
+      allocated.add(buffer);
+      return buffer;
+    }
+
+    @Override
+    public void release(ByteBuffer b) {
+      delegate.release(b);
+    }
+
+    @Override
+    public boolean isDirect() {
+      return delegate.isDirect();
+    }
+
+    synchronized List<ByteBuffer> getAllocatedBuffers() {
+      return new ArrayList<>(allocated);
+    }
+  }
+
   /*
    * Builder to concatenate the buffers of the discontinuous parts for the same column. These parts are generated as a
    * result of the column-index based filtering when some pages might be skipped at reading.
@@ -2368,6 +2431,10 @@ public void readFromVectoredRange(ParquetFileRange currRange, ChunkListBuilder b
         LOG.error(error, e);
         throw new IOException(error, e);
       }
+      // Note: the buffers acquired from the allocator are registered for release by
+      // readVectored() via the RecordingByteBufferAllocator. We must not register {@code buffer}
+      // here: with checksum verification enabled Hadoop returns a sliced view of the allocated
+      // buffer, so releasing {@code buffer} would either fail or leak the underlying allocation.
       ByteBufferInputStream stream = ByteBufferInputStream.wrap(buffer);
       for (ChunkDescriptor descriptor : chunks) {
         builder.add(descriptor, stream.sliceBuffers(descriptor.size), f);

parquet-hadoop/src/test/java/org/apache/parquet/crypto/propertiesfactory/SchemaControlEncryptionTest.java

@@ -167,21 +167,22 @@ private String encryptParquetFile(String file, Configuration conf) throws IOExce
   }
 
   private void decryptParquetFileAndValid(String file, Configuration conf) throws IOException {
-    ParquetReader<Group> reader = ParquetReader.builder(new GroupReadSupport(), new Path(file))
+    try (ParquetReader<Group> reader = ParquetReader.builder(new GroupReadSupport(), new Path(file))
         .withConf(conf)
-        .build();
-    for (int i = 0; i < numRecord; i++) {
-      Group group = reader.read();
-      assertEquals(testData.get("Name")[i], group.getBinary("Name", 0).toStringUsingUTF8());
-      assertEquals(testData.get("Age")[i], group.getLong("Age", 0));
-
-      Group subGroup = group.getGroup("WebLinks", 0);
-      assertArrayEquals(
-          subGroup.getBinary("LinkedIn", 0).getBytes(), ((String) testData.get("LinkedIn")[i]).getBytes());
-      assertArrayEquals(
-          subGroup.getBinary("Twitter", 0).getBytes(), ((String) testData.get("Twitter")[i]).getBytes());
+        .build()) {
+      for (int i = 0; i < numRecord; i++) {
+        Group group = reader.read();
+        assertEquals(testData.get("Name")[i], group.getBinary("Name", 0).toStringUsingUTF8());
+        assertEquals(testData.get("Age")[i], group.getLong("Age", 0));
+
+        Group subGroup = group.getGroup("WebLinks", 0);
+        assertArrayEquals(
+            subGroup.getBinary("LinkedIn", 0).getBytes(),
+            ((String) testData.get("LinkedIn")[i]).getBytes());
+        assertArrayEquals(
+            subGroup.getBinary("Twitter", 0).getBytes(), ((String) testData.get("Twitter")[i]).getBytes());
+      }
     }
-    reader.close();
   }
 
   private static String createTempFile(String prefix) {

parquet-hadoop/src/test/java/org/apache/parquet/encodings/FileEncodingsIT.java

@@ -254,7 +254,7 @@ private void writeValuesToFile(
     SimpleGroupFactory message = new SimpleGroupFactory(schema);
     GroupWriteSupport.setSchema(schema, configuration);
 
-    ParquetWriter<Group> writer = ExampleParquetWriter.builder(file)
+    try (ParquetWriter<Group> writer = ExampleParquetWriter.builder(file)
         .withAllocator(allocator)
         .withCompressionCodec(compression)
         .withRowGroupSize(rowGroupSize)
@@ -263,36 +263,35 @@ private void writeValuesToFile(
         .withDictionaryEncoding(enableDictionary)
         .withWriterVersion(version)
         .withConf(configuration)
-        .build();
-
-    for (Object o : values) {
-      switch (type) {
-        case BOOLEAN:
-          writer.write(message.newGroup().append("field", (Boolean) o));
-          break;
-        case INT32:
-          writer.write(message.newGroup().append("field", (Integer) o));
-          break;
-        case INT64:
-          writer.write(message.newGroup().append("field", (Long) o));
-          break;
-        case FLOAT:
-          writer.write(message.newGroup().append("field", (Float) o));
-          break;
-        case DOUBLE:
-          writer.write(message.newGroup().append("field", (Double) o));
-          break;
-        case INT96:
-        case BINARY:
-        case FIXED_LEN_BYTE_ARRAY:
-          writer.write(message.newGroup().append("field", (Binary) o));
-          break;
-        default:
-          throw new IllegalArgumentException("Unknown type name: " + type);
+        .build()) {
+
+      for (Object o : values) {
+        switch (type) {
+          case BOOLEAN:
+            writer.write(message.newGroup().append("field", (Boolean) o));
+            break;
+          case INT32:
+            writer.write(message.newGroup().append("field", (Integer) o));
+            break;
+          case INT64:
+            writer.write(message.newGroup().append("field", (Long) o));
+            break;
+          case FLOAT:
+            writer.write(message.newGroup().append("field", (Float) o));
+            break;
+          case DOUBLE:
+            writer.write(message.newGroup().append("field", (Double) o));
+            break;
+          case INT96:
+          case BINARY:
+          case FIXED_LEN_BYTE_ARRAY:
+            writer.write(message.newGroup().append("field", (Binary) o));
+            break;
+          default:
+            throw new IllegalArgumentException("Unknown type name: " + type);
+        }
       }
     }
-
-    writer.close();
   }
 
   private List<?> generateRandomValues(PrimitiveTypeName type, int count) {
@@ -522,6 +521,8 @@ private static List<PageReadStore> readBlocksFromFile(Path file) throws IOExcept
 
       ParquetMetadata metadata =
           ParquetFileReader.readFooter(configuration, file, ParquetMetadataConverter.NO_FILTER);
+      // Not using try-with-resources here because closing the reader releases the codec factory,
+      // but the returned PageReadStore objects still hold compressed pages that need decompression later.
       ParquetFileReader fileReader = new ParquetFileReader(
           configuration,
           metadata.getFileMetaData(),

parquet-hadoop/src/test/java/org/apache/parquet/filter2/TestFiltersWithMissingColumns.java

@@ -212,17 +212,13 @@ public void testOrMissingColumnFilter() throws Exception {
   }
 
   public static long countFilteredRecords(Path path, FilterPredicate pred) throws IOException {
-    ParquetReader<Group> reader = ParquetReader.builder(new GroupReadSupport(), path)
-        .withFilter(FilterCompat.get(pred))
-        .build();
-
     long count = 0;
-    try {
+    try (ParquetReader<Group> reader = ParquetReader.builder(new GroupReadSupport(), path)
+        .withFilter(FilterCompat.get(pred))
+        .build()) {
       while (reader.read() != null) {
         count += 1;
       }
-    } finally {
-      reader.close();
     }
     return count;
   }

parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestColumnChunkPageWriteStore.java

@@ -218,36 +218,39 @@ public void test(Configuration config, ByteBufferAllocator allocator) throws Exc
 
     {
       ParquetMetadata footer = ParquetFileReader.readFooter(conf, file, NO_FILTER);
-      ParquetFileReader reader = new ParquetFileReader(
-          config, footer.getFileMetaData(), file, footer.getBlocks(), schema.getColumns());
-      PageReadStore rowGroup = reader.readNextRowGroup();
-      PageReader pageReader = rowGroup.getPageReader(col);
-      DataPageV2 page = (DataPageV2) pageReader.readPage();
-      assertEquals(rowCount, page.getRowCount());
-      assertEquals(nullCount, page.getNullCount());
-      assertEquals(valueCount, page.getValueCount());
-      assertEquals(d, intValue(page.getDefinitionLevels()));
-      assertEquals(r, intValue(page.getRepetitionLevels()));
-      assertEquals(dataEncoding, page.getDataEncoding());
-      assertEquals(v, intValue(page.getData()));
-
-      // Checking column/offset indexes for the one page
-      ColumnChunkMetaData column = footer.getBlocks().get(0).getColumns().get(0);
-      ColumnIndex columnIndex = reader.readColumnIndex(column);
-      assertArrayEquals(
-          statistics.getMinBytes(), columnIndex.getMinValues().get(0).array());
-      assertArrayEquals(
-          statistics.getMaxBytes(), columnIndex.getMaxValues().get(0).array());
-      assertEquals(
-          statistics.getNumNulls(), columnIndex.getNullCounts().get(0).longValue());
-      assertFalse(columnIndex.getNullPages().get(0));
-      OffsetIndex offsetIndex = reader.readOffsetIndex(column);
-      assertEquals(1, offsetIndex.getPageCount());
-      assertEquals(pageSize, offsetIndex.getCompressedPageSize(0));
-      assertEquals(0, offsetIndex.getFirstRowIndex(0));
-      assertEquals(pageOffset, offsetIndex.getOffset(0));
-
-      reader.close();
+      try (ParquetFileReader reader = new ParquetFileReader(
+          config, footer.getFileMetaData(), file, footer.getBlocks(), schema.getColumns())) {
+        PageReadStore rowGroup = reader.readNextRowGroup();
+        PageReader pageReader = rowGroup.getPageReader(col);
+        DataPageV2 page = (DataPageV2) pageReader.readPage();
+        assertEquals(rowCount, page.getRowCount());
+        assertEquals(nullCount, page.getNullCount());
+        assertEquals(valueCount, page.getValueCount());
+        assertEquals(d, intValue(page.getDefinitionLevels()));
+        assertEquals(r, intValue(page.getRepetitionLevels()));
+        assertEquals(dataEncoding, page.getDataEncoding());
+        assertEquals(v, intValue(page.getData()));
+
+        // Checking column/offset indexes for the one page
+        ColumnChunkMetaData column =
+            footer.getBlocks().get(0).getColumns().get(0);
+        ColumnIndex columnIndex = reader.readColumnIndex(column);
+        assertArrayEquals(
+            statistics.getMinBytes(),
+            columnIndex.getMinValues().get(0).array());
+        assertArrayEquals(
+            statistics.getMaxBytes(),
+            columnIndex.getMaxValues().get(0).array());
+        assertEquals(
+            statistics.getNumNulls(),
+            columnIndex.getNullCounts().get(0).longValue());
+        assertFalse(columnIndex.getNullPages().get(0));
+        OffsetIndex offsetIndex = reader.readOffsetIndex(column);
+        assertEquals(1, offsetIndex.getPageCount());
+        assertEquals(pageSize, offsetIndex.getCompressedPageSize(0));
+        assertEquals(0, offsetIndex.getFirstRowIndex(0));
+        assertEquals(pageOffset, offsetIndex.getOffset(0));
+      }
     }
   }