Fix flat map writer's getBufferedBytes

prestodb · Jun 1, 2022 · 51f5cc1 · 51f5cc1
1 parent 3131bea
commit 51f5cc1
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/presto-orc/src/main/java/com/facebook/presto/orc/writer/MapFlatColumnWriter.java b/presto-orc/src/main/java/com/facebook/presto/orc/writer/MapFlatColumnWriter.java
@@ -378,11 +378,10 @@ public List<StreamDataOutput> getIndexStreams(Optional<List<? extends StreamChec
     @Override
     public long getBufferedBytes()
     {
-        long bufferedBytes = 0;
+        long bufferedBytes = presentStream.getBufferedBytes();
         for (MapFlatValueWriter valueWriter : valueWriters) {
-            bufferedBytes += valueWriter.getValueWriter().getBufferedBytes();
+            bufferedBytes += valueWriter.getBufferedBytes();
         }
-        // TODO Implement me
         return bufferedBytes;
     }
 

diff --git a/presto-orc/src/main/java/com/facebook/presto/orc/writer/MapFlatValueWriter.java b/presto-orc/src/main/java/com/facebook/presto/orc/writer/MapFlatValueWriter.java
@@ -131,6 +131,11 @@ public List<StreamDataOutput> getIndexStreams()
         return valueWriter.getIndexStreams(Optional.of(inMapStream.getCheckpoints()));
     }
 
+    public long getBufferedBytes()
+    {
+        return valueWriter.getBufferedBytes() + inMapStream.getBufferedBytes();
+    }
+
     public void close()
     {
         inMapStream.close();