GH-2037: Keep writer size available after close

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

@@ -143,6 +143,7 @@ public void close() throws IOException, InterruptedException {
         }
         finalMetadata.putAll(finalWriteContext.getExtraMetaData());
         parquetFileWriter.end(finalMetadata);
+        lastRowGroupEndPos = parquetFileWriter.getPos();
       } catch (Throwable t) {
         parquetFileWriter.abort();
         throw t;
@@ -171,7 +172,7 @@ public void write(T value) throws IOException, InterruptedException {
    * @return the total size of data written to the file and buffered in memory
    */
   public long getDataSize() {
-    return lastRowGroupEndPos + columnStore.getBufferedSize();
+    return lastRowGroupEndPos + (columnStore == null ? 0 : columnStore.getBufferedSize());
   }
 
   private void checkBlockSizeReached() throws IOException {

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

@@ -174,6 +174,7 @@ public static enum Mode {
 
   // set when end is called
   private ParquetMetadata footer = null;
+  private long finalPosition = -1;
   private boolean aborted;
   private boolean closed;
 
@@ -1924,6 +1925,7 @@ public void end(Map<String, String> extraMetaData) throws IOException {
         this.footer =
             new ParquetMetadata(new FileMetaData(schema, extraMetaData, Version.FULL_VERSION), blocks);
         serializeFooter(footer, out, fileEncryptor, metadataConverter);
+        finalPosition = out.getPos();
       });
     } finally {
       close();
@@ -2393,11 +2395,13 @@ static ParquetMetadata mergeFooters(
   }
 
   /**
+   * Returns the cached position once the footer has been written.
+   *
    * @return the current position in the underlying file
    * @throws IOException if there is an error while getting the current stream's position
    */
   public long getPos() throws IOException {
-    return withAbortOnFailure(() -> out.getPos());
+    return finalPosition >= 0 ? finalPosition : withAbortOnFailure(() -> out.getPos());
   }
 
   public long getNextRowGroupSize() throws IOException {

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

@@ -455,6 +455,8 @@ public ParquetMetadata getFooter() {
   }
 
   /**
+   * After a successful close, this is the final file size including the footer.
+   *
    * @return the total size of data written to the file and buffered in memory
    */
   public long getDataSize() {

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

@@ -45,6 +45,10 @@
 import com.google.common.collect.ImmutableMap;
 import java.io.IOException;
 import java.lang.reflect.Field;
+import java.nio.channels.Channels;
+import java.nio.channels.FileChannel;
+import java.nio.file.Files;
+import java.nio.file.StandardOpenOption;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -95,6 +99,8 @@
 import org.apache.parquet.hadoop.metadata.ParquetMetadata;
 import org.apache.parquet.hadoop.util.HadoopInputFile;
 import org.apache.parquet.hadoop.util.HadoopOutputFile;
+import org.apache.parquet.io.DelegatingPositionOutputStream;
+import org.apache.parquet.io.LocalOutputFile;
 import org.apache.parquet.io.OutputFile;
 import org.apache.parquet.io.PositionOutputStream;
 import org.apache.parquet.io.api.Binary;
@@ -106,6 +112,8 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
 
 public class TestParquetWriter {
 
@@ -240,6 +248,53 @@ public void test() throws Exception {
   @TempDir
   private java.nio.file.Path tempDir;
 
+  @ParameterizedTest
+  @CsvSource({"0,false", "1,false", "5,false", "0,true", "1,true", "5,true"})
+  public void testDataSizeAfterClose(int records, boolean channelBacked) throws IOException {
+    java.nio.file.Path file = tempDir.resolve("size.parquet");
+    MessageType schema = parseMessageType("message test { required int32 value; }");
+    OutputFile output = channelBacked ? channelOutputFile(file) : new LocalOutputFile(file);
+    ParquetWriter<Group> writer = ExampleParquetWriter.builder(output)
+        .withType(schema)
+        .withAllocator(allocator)
+        .withRowGroupRowCountLimit(2)
+        .build();
+    try (writer) {
+      assertThat(writer.getDataSize()).isZero();
+      SimpleGroupFactory factory = new SimpleGroupFactory(schema);
+      for (int i = 0; i < records; i++) {
+        writer.write(factory.newGroup().append("value", i));
+      }
+    }
+
+    assertThat(writer.getDataSize()).isEqualTo(Files.size(file));
+    writer.close();
+    assertThat(writer.getDataSize()).isEqualTo(Files.size(file));
+
+    try (ParquetReader<Group> reader = ParquetReader.builder(new GroupReadSupport(), new Path(file.toUri()))
+        .build()) {
+      for (int i = 0; i < records; i++) {
+        assertThat(reader.read().getInteger("value", 0)).isEqualTo(i);
+      }
+      assertThat(reader.read()).isNull();
+    }
+  }
+
+  private static OutputFile channelOutputFile(java.nio.file.Path path) {
+    return new LocalOutputFile(path) {
+      @Override
+      public PositionOutputStream create(long blockSizeHint) throws IOException {
+        FileChannel channel = FileChannel.open(path, StandardOpenOption.CREATE_NEW, StandardOpenOption.WRITE);
+        return new DelegatingPositionOutputStream(Channels.newOutputStream(channel)) {
+          @Override
+          public long getPos() throws IOException {
+            return channel.position();
+          }
+        };
+      }
+    };
+  }
+
   @Test
   public void testBadWriteSchema() throws IOException {
     Path path = tempPath("test.parquet");