MINOR: Close output stream when ParquetFileWriter construction fails on encryption validation

The shared ParquetFileWriter constructor opens the output stream before validating that every encrypted column exists in the file schema. If validation fails, the half-constructed writer is never returned and the stream cannot be closed by the caller.

Close the output stream before rethrowing encryption setup failures, matching ParquetFileReader. Add a regression test that verifies the stream is closed when an encrypted column is missing from the schema.

Signed-off-by: Anas Khan <83116240+anxkhn@users.noreply.github.com>

Author: anxkhn

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

@@ -519,36 +519,44 @@ private ParquetFileWriter(
       return;
     }
 
-    if (null == encryptionProperties) {
-      encryptionProperties = encryptor.getEncryptionProperties();
-    }
+    try {
+      if (null == encryptionProperties) {
+        encryptionProperties = encryptor.getEncryptionProperties();
+      }
 
-    // Verify that every encrypted column is in file schema
-    Map<ColumnPath, ColumnEncryptionProperties> columnEncryptionProperties =
-        encryptionProperties.getEncryptedColumns();
-    if (null != columnEncryptionProperties) { // if null, every column in file schema will be encrypted with footer
-      // key
-      for (Map.Entry<ColumnPath, ColumnEncryptionProperties> entry : columnEncryptionProperties.entrySet()) {
-        String[] path = entry.getKey().toArray();
-        if (!schema.containsPath(path)) {
-          StringBuilder columnList = new StringBuilder();
-          columnList.append("[");
-          for (String[] columnPath : schema.getPaths()) {
-            columnList
-                .append(ColumnPath.get(columnPath).toDotString())
-                .append("], [");
+      // Verify that every encrypted column is in file schema
+      Map<ColumnPath, ColumnEncryptionProperties> columnEncryptionProperties =
+          encryptionProperties.getEncryptedColumns();
+      if (null != columnEncryptionProperties) { // if null, every column in file schema will be encrypted with
+        // footer
+        // key
+        for (Map.Entry<ColumnPath, ColumnEncryptionProperties> entry : columnEncryptionProperties.entrySet()) {
+          String[] path = entry.getKey().toArray();
+          if (!schema.containsPath(path)) {
+            StringBuilder columnList = new StringBuilder();
+            columnList.append("[");
+            for (String[] columnPath : schema.getPaths()) {
+              columnList
+                  .append(ColumnPath.get(columnPath).toDotString())
+                  .append("], [");
+            }
+            throw new ParquetCryptoRuntimeException("Encrypted column ["
+                + entry.getKey().toDotString() + "] not in file schema column list: "
+                + columnList.substring(0, columnList.length() - 3));
           }
-          throw new ParquetCryptoRuntimeException(
-              "Encrypted column [" + entry.getKey().toDotString() + "] not in file schema column list: "
-                  + columnList.substring(0, columnList.length() - 3));
         }
       }
-    }
 
-    if (null == encryptor) {
-      this.fileEncryptor = new InternalFileEncryptor(encryptionProperties);
-    } else {
-      this.fileEncryptor = encryptor;
+      if (null == encryptor) {
+        this.fileEncryptor = new InternalFileEncryptor(encryptionProperties);
+      } else {
+        this.fileEncryptor = encryptor;
+      }
+    } catch (Exception e) {
+      // If encryption setup throws in the constructor, the output stream opened above should be
+      // closed. Otherwise, there's no way to close it outside since the object is never returned.
+      out.close();
+      throw e;
     }
   }
 

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

@@ -48,6 +48,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.FSDataInputStream;
 import org.apache.hadoop.fs.FileStatus;
@@ -76,6 +77,9 @@
 import org.apache.parquet.column.statistics.LongStatistics;
 import org.apache.parquet.column.values.bloomfilter.BlockSplitBloomFilter;
 import org.apache.parquet.column.values.bloomfilter.BloomFilter;
+import org.apache.parquet.crypto.ColumnEncryptionProperties;
+import org.apache.parquet.crypto.FileEncryptionProperties;
+import org.apache.parquet.crypto.ParquetCryptoRuntimeException;
 import org.apache.parquet.example.data.Group;
 import org.apache.parquet.example.data.simple.SimpleGroup;
 import org.apache.parquet.format.Statistics;
@@ -84,6 +88,7 @@
 import org.apache.parquet.hadoop.example.GroupWriteSupport;
 import org.apache.parquet.hadoop.metadata.BlockMetaData;
 import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData;
+import org.apache.parquet.hadoop.metadata.ColumnPath;
 import org.apache.parquet.hadoop.metadata.CompressionCodecName;
 import org.apache.parquet.hadoop.metadata.ConcatenatingKeyValueMetadataMergeStrategy;
 import org.apache.parquet.hadoop.metadata.FileMetaData;
@@ -98,7 +103,9 @@
 import org.apache.parquet.internal.column.columnindex.BoundaryOrder;
 import org.apache.parquet.internal.column.columnindex.ColumnIndex;
 import org.apache.parquet.internal.column.columnindex.OffsetIndex;
+import org.apache.parquet.io.OutputFile;
 import org.apache.parquet.io.ParquetEncodingException;
+import org.apache.parquet.io.PositionOutputStream;
 import org.apache.parquet.io.api.Binary;
 import org.apache.parquet.schema.MessageType;
 import org.apache.parquet.schema.MessageTypeParser;
@@ -1443,6 +1450,89 @@ public void testMergeMetadataWithNoConflictingKeyValues(boolean vectoredRead) {
     assertThat(mergedValues.get("c")).isEqualTo("d");
   }
 
+  @Test
+  public void testConstructorClosesStreamWhenEncryptedColumnMissing() throws Exception {
+    ColumnEncryptionProperties missingColumn = ColumnEncryptionProperties.builder("not_in_schema")
+        .withKey("0123456789012345".getBytes(StandardCharsets.UTF_8))
+        .build();
+    Map<ColumnPath, ColumnEncryptionProperties> encryptedColumns = new HashMap<>();
+    encryptedColumns.put(missingColumn.getPath(), missingColumn);
+    FileEncryptionProperties encryptionProperties = FileEncryptionProperties.builder(
+            "0123456789012345".getBytes(StandardCharsets.UTF_8))
+        .withEncryptedColumns(encryptedColumns)
+        .build();
+
+    RecordingOutputFile file = new RecordingOutputFile();
+    assertThatThrownBy(() -> new ParquetFileWriter(
+            file,
+            SCHEMA,
+            CREATE,
+            DEFAULT_BLOCK_SIZE,
+            MAX_PADDING_SIZE_DEFAULT,
+            ParquetProperties.DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH,
+            ParquetProperties.DEFAULT_STATISTICS_TRUNCATE_LENGTH,
+            ParquetProperties.DEFAULT_PAGE_WRITE_CHECKSUM_ENABLED,
+            encryptionProperties))
+        .isInstanceOf(ParquetCryptoRuntimeException.class)
+        .hasMessage("Encrypted column [not_in_schema] not in file schema column list: [a.b], [c.d]");
+
+    assertThat(file.isStreamClosed()).isTrue();
+  }
+
+  /**
+   * An {@link OutputFile} whose {@link PositionOutputStream} records whether it was closed, used to
+   * assert that a failed {@link ParquetFileWriter} construction does not leak the open stream.
+   */
+  private static class RecordingOutputFile implements OutputFile {
+
+    private final AtomicBoolean streamClosed = new AtomicBoolean(false);
+
+    boolean isStreamClosed() {
+      return streamClosed.get();
+    }
+
+    private PositionOutputStream newRecordingStream() {
+      return new PositionOutputStream() {
+        private long pos = 0;
+
+        @Override
+        public long getPos() {
+          return pos;
+        }
+
+        @Override
+        public void write(int b) {
+          pos++;
+        }
+
+        @Override
+        public void close() {
+          streamClosed.set(true);
+        }
+      };
+    }
+
+    @Override
+    public PositionOutputStream create(long blockSizeHint) {
+      return newRecordingStream();
+    }
+
+    @Override
+    public PositionOutputStream createOrOverwrite(long blockSizeHint) {
+      return newRecordingStream();
+    }
+
+    @Override
+    public boolean supportsBlockSize() {
+      return false;
+    }
+
+    @Override
+    public long defaultBlockSize() {
+      return 0;
+    }
+  }
+
   private org.apache.parquet.column.statistics.Statistics<?> statsC1(Binary... values) {
     org.apache.parquet.column.statistics.Statistics<?> stats =
         org.apache.parquet.column.statistics.Statistics.createStats(C1.getPrimitiveType());