Expose counters for ORC writer memory usage

trinodb · Mar 13, 2018 · 1ae1910 · 1ae1910
1 parent 475fbce
commit 1ae1910
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 3 deletions.
diff --git a/presto-orc/src/main/java/com/facebook/presto/orc/OrcWriter.java b/presto-orc/src/main/java/com/facebook/presto/orc/OrcWriter.java
@@ -111,6 +111,7 @@ public class OrcWriter
     private long columnWritersRetainedBytes;
     private long closedStripesRetainedBytes;
     private long validatorRetainedBytes;
+    private long previouslyRecordedSizeInBytes;
     private boolean closed;
 
     @Nullable
@@ -195,6 +196,9 @@ public OrcWriter(
         }
 
         this.validatorRetainedBytes = validationBuilder == null ? 0 : VALIDATOR_INSTANCE_SIZE;
+
+        this.previouslyRecordedSizeInBytes = getRetainedBytes();
+        stats.updateSizeInBytes(previouslyRecordedSizeInBytes);
     }
 
     public int getBufferedBytes()
@@ -235,6 +239,10 @@ public void write(Page page)
             }
             writeChunk(chunk);
         }
+
+        long recordedSizeInBytes = getRetainedBytes();
+        stats.updateSizeInBytes(recordedSizeInBytes - previouslyRecordedSizeInBytes);
+        previouslyRecordedSizeInBytes = recordedSizeInBytes;
     }
 
     private void writeChunk(Page chunk)
@@ -383,6 +391,8 @@ public void close()
             return;
         }
         closed = true;
+        stats.updateSizeInBytes(-previouslyRecordedSizeInBytes);
+        previouslyRecordedSizeInBytes = 0;
 
         writeStripe(CLOSED);
 

diff --git a/presto-orc/src/main/java/com/facebook/presto/orc/OrcWriterStats.java b/presto-orc/src/main/java/com/facebook/presto/orc/OrcWriterStats.java
@@ -16,6 +16,8 @@
 import org.weakref.jmx.Managed;
 import org.weakref.jmx.Nested;
 
+import java.util.concurrent.atomic.AtomicLong;
+
 import static com.facebook.presto.orc.OrcWriterStats.FlushReason.CLOSED;
 import static com.facebook.presto.orc.OrcWriterStats.FlushReason.DICTIONARY_FULL;
 import static com.facebook.presto.orc.OrcWriterStats.FlushReason.MAX_BYTES;
@@ -34,13 +36,19 @@ public enum FlushReason
     private final OrcWriterFlushStats maxBytesFlush = new OrcWriterFlushStats(MAX_ROWS.name());
     private final OrcWriterFlushStats dictionaryFullFlush = new OrcWriterFlushStats(DICTIONARY_FULL.name());
     private final OrcWriterFlushStats closedFlush = new OrcWriterFlushStats(CLOSED.name());
+    private final AtomicLong writerSizeInBytes = new AtomicLong();
 
     public void recordStripeWritten(FlushReason flushReason, long stripeBytes, int stripeRows, int dictionaryBytes)
     {
         getFlushStats(flushReason).recordStripeWritten(stripeBytes, stripeRows, dictionaryBytes);
         allFlush.recordStripeWritten(stripeBytes, stripeRows, dictionaryBytes);
     }
 
+    public void updateSizeInBytes(long deltaInBytes)
+    {
+        writerSizeInBytes.addAndGet(deltaInBytes);
+    }
+
     @Managed
     @Nested
     public OrcWriterFlushStats getAllFlush()
@@ -76,6 +84,12 @@ public OrcWriterFlushStats getClosedFlush()
         return closedFlush;
     }
 
+    @Managed
+    public long getWriterSizeInBytes()
+    {
+        return writerSizeInBytes.get();
+    }
+
     private OrcWriterFlushStats getFlushStats(FlushReason flushReason)
     {
         switch (flushReason) {
@@ -101,6 +115,7 @@ public String toString()
                 .add("maxBytesFlush", maxBytesFlush)
                 .add("dictionaryFullFlush", dictionaryFullFlush)
                 .add("closedFlush", closedFlush)
+                .add("writerSizeInBytes", writerSizeInBytes.get())
                 .toString();
     }
 }
diff --git a/presto-orc/src/test/java/com/facebook/presto/orc/OrcTester.java b/presto-orc/src/test/java/com/facebook/presto/orc/OrcTester.java
@@ -439,6 +439,7 @@ public void assertRoundTrip(Type type, List<?> readValues)
     public void assertRoundTrip(Type type, List<?> readValues, boolean verifyWithHiveReader)
             throws Exception
     {
+        OrcWriterStats stats = new OrcWriterStats();
         for (Format format : formats) {
             if (!format.supportsType(type)) {
                 return;
@@ -458,7 +459,7 @@ public void assertRoundTrip(Type type, List<?> readValues, boolean verifyWithHiv
 
                 // write Presto, read Hive and Presto
                 try (TempFile tempFile = new TempFile()) {
-                    writeOrcColumnPresto(tempFile.getFile(), format, compression, type, readValues.iterator());
+                    writeOrcColumnPresto(tempFile.getFile(), format, compression, type, readValues.iterator(), stats);
 
                     if (verifyWithHiveReader && hiveSupported) {
                         assertFileContentsHive(type, tempFile, format, readValues);
@@ -476,6 +477,8 @@ public void assertRoundTrip(Type type, List<?> readValues, boolean verifyWithHiv
                 }
             }
         }
+
+        assertEquals(stats.getWriterSizeInBytes(), 0);
     }
 
     private static void assertFileContentsPresto(
@@ -614,7 +617,7 @@ static OrcRecordReader createCustomOrcRecordReader(TempFile tempFile, OrcEncodin
         return orcReader.createRecordReader(ImmutableMap.of(0, type), predicate, HIVE_STORAGE_TIME_ZONE, newSimpleAggregatedMemoryContext());
     }
 
-    private static void writeOrcColumnPresto(File outputFile, Format format, CompressionKind compression, Type type, Iterator<?> values)
+    private static void writeOrcColumnPresto(File outputFile, Format format, CompressionKind compression, Type type, Iterator<?> values, OrcWriterStats stats)
             throws Exception
     {
         ImmutableMap.Builder<String, String> metadata = ImmutableMap.builder();
@@ -632,7 +635,7 @@ private static void writeOrcColumnPresto(File outputFile, Format format, Compres
                 ImmutableMap.of(),
                 HIVE_STORAGE_TIME_ZONE,
                 true,
-                new OrcWriterStats());
+                stats);
 
         BlockBuilder blockBuilder = type.createBlockBuilder(new BlockBuilderStatus(), 1024);
         while (values.hasNext()) {