From 03cdc8c8df758475609157a34d656a0f5781ea93 Mon Sep 17 00:00:00 2001 From: Paul Meng Date: Wed, 8 Mar 2023 16:12:58 -0500 Subject: [PATCH] Improve the memory accounting for different BatchStreamReaders --- .../presto/orc/reader/BooleanBatchStreamReader.java | 4 ++-- .../facebook/presto/orc/reader/ByteBatchStreamReader.java | 4 ++-- .../presto/orc/reader/LongDirectBatchStreamReader.java | 8 ++++---- .../orc/reader/SliceDictionaryBatchStreamReader.java | 6 +++--- .../com/facebook/presto/orc/TestOrcReaderMemoryUsage.java | 8 ++++---- 5 files changed, 15 insertions(+), 15 deletions(-) diff --git a/presto-orc/src/main/java/com/facebook/presto/orc/reader/BooleanBatchStreamReader.java b/presto-orc/src/main/java/com/facebook/presto/orc/reader/BooleanBatchStreamReader.java index 02b4d879fdca0..1e8663dc61cfb 100644 --- a/presto-orc/src/main/java/com/facebook/presto/orc/reader/BooleanBatchStreamReader.java +++ b/presto-orc/src/main/java/com/facebook/presto/orc/reader/BooleanBatchStreamReader.java @@ -159,7 +159,7 @@ private Block readNullBlock(boolean[] isNull, int nonNullCount) int minNonNullValueSize = minNonNullValueSize(nonNullCount); if (nonNullValueTemp.length < minNonNullValueSize) { nonNullValueTemp = new byte[minNonNullValueSize]; - systemMemoryContext.setBytes(sizeOf(nonNullValueTemp)); + systemMemoryContext.setBytes(getRetainedSizeInBytes()); } dataStream.getSetBits(nonNullCount, nonNullValueTemp); @@ -225,6 +225,6 @@ public void close() @Override public long getRetainedSizeInBytes() { - return INSTANCE_SIZE; + return INSTANCE_SIZE + sizeOf(nonNullValueTemp); } } diff --git a/presto-orc/src/main/java/com/facebook/presto/orc/reader/ByteBatchStreamReader.java b/presto-orc/src/main/java/com/facebook/presto/orc/reader/ByteBatchStreamReader.java index f70024711549c..e0aef629c1339 100644 --- a/presto-orc/src/main/java/com/facebook/presto/orc/reader/ByteBatchStreamReader.java +++ b/presto-orc/src/main/java/com/facebook/presto/orc/reader/ByteBatchStreamReader.java @@ -153,7 +153,7 @@ private Block readNullBlock(boolean[] isNull, int nonNullCount) int minNonNullValueSize = minNonNullValueSize(nonNullCount); if (nonNullValueTemp.length < minNonNullValueSize) { nonNullValueTemp = new byte[minNonNullValueSize]; - systemMemoryContext.setBytes(sizeOf(nonNullValueTemp)); + systemMemoryContext.setBytes(getRetainedSizeInBytes()); } dataStream.next(nonNullValueTemp, nonNullCount); @@ -220,6 +220,6 @@ public void close() @Override public long getRetainedSizeInBytes() { - return INSTANCE_SIZE; + return INSTANCE_SIZE + sizeOf(nonNullValueTemp); } } diff --git a/presto-orc/src/main/java/com/facebook/presto/orc/reader/LongDirectBatchStreamReader.java b/presto-orc/src/main/java/com/facebook/presto/orc/reader/LongDirectBatchStreamReader.java index 0bcbb1c15c045..7d57fd794d236 100644 --- a/presto-orc/src/main/java/com/facebook/presto/orc/reader/LongDirectBatchStreamReader.java +++ b/presto-orc/src/main/java/com/facebook/presto/orc/reader/LongDirectBatchStreamReader.java @@ -193,7 +193,7 @@ private Block longReadNullBlock(boolean[] isNull, int nonNullCount) int minNonNullValueSize = minNonNullValueSize(nonNullCount); if (longNonNullValueTemp.length < minNonNullValueSize) { longNonNullValueTemp = new long[minNonNullValueSize]; - systemMemoryContext.setBytes(sizeOf(longNonNullValueTemp)); + systemMemoryContext.setBytes(getRetainedSizeInBytes()); } dataStream.next(longNonNullValueTemp, nonNullCount); @@ -210,7 +210,7 @@ private Block intReadNullBlock(boolean[] isNull, int nonNullCount) int minNonNullValueSize = minNonNullValueSize(nonNullCount); if (intNonNullValueTemp.length < minNonNullValueSize) { intNonNullValueTemp = new int[minNonNullValueSize]; - systemMemoryContext.setBytes(sizeOf(intNonNullValueTemp)); + systemMemoryContext.setBytes(getRetainedSizeInBytes()); } dataStream.next(intNonNullValueTemp, nonNullCount); @@ -227,7 +227,7 @@ private Block shortReadNullBlock(boolean[] isNull, int nonNullCount) int minNonNullValueSize = minNonNullValueSize(nonNullCount); if (shortNonNullValueTemp.length < minNonNullValueSize) { shortNonNullValueTemp = new short[minNonNullValueSize]; - systemMemoryContext.setBytes(sizeOf(shortNonNullValueTemp)); + systemMemoryContext.setBytes(getRetainedSizeInBytes()); } dataStream.next(shortNonNullValueTemp, nonNullCount); @@ -296,6 +296,6 @@ public void close() @Override public long getRetainedSizeInBytes() { - return INSTANCE_SIZE; + return INSTANCE_SIZE + sizeOf(shortNonNullValueTemp) + sizeOf(longNonNullValueTemp) + sizeOf(intNonNullValueTemp); } } diff --git a/presto-orc/src/main/java/com/facebook/presto/orc/reader/SliceDictionaryBatchStreamReader.java b/presto-orc/src/main/java/com/facebook/presto/orc/reader/SliceDictionaryBatchStreamReader.java index 8778bcc82ab72..8fb6254152d72 100644 --- a/presto-orc/src/main/java/com/facebook/presto/orc/reader/SliceDictionaryBatchStreamReader.java +++ b/presto-orc/src/main/java/com/facebook/presto/orc/reader/SliceDictionaryBatchStreamReader.java @@ -219,7 +219,7 @@ private void openRowGroup() // resize the dictionary lengths array if necessary if (stripeDictionaryLength.length < stripeDictionarySize) { stripeDictionaryLength = new int[stripeDictionarySize]; - systemMemoryContext.setBytes(sizeOf(stripeDictionaryLength)); + systemMemoryContext.setBytes(getRetainedSizeInBytes()); } // read the lengths @@ -236,11 +236,11 @@ private void openRowGroup() // we must always create a new dictionary array because the previous dictionary may still be referenced stripeDictionaryData = new byte[toIntExact(dataLength)]; - systemMemoryContext.setBytes(sizeOf(stripeDictionaryData)); + systemMemoryContext.setBytes(getRetainedSizeInBytes()); // add one extra entry for null stripeDictionaryOffsetVector = new int[stripeDictionarySize + 2]; - systemMemoryContext.setBytes(sizeOf(stripeDictionaryOffsetVector)); + systemMemoryContext.setBytes(getRetainedSizeInBytes()); // read dictionary values ByteArrayInputStream dictionaryDataStream = stripeDictionaryDataStreamSource.openStream(); diff --git a/presto-orc/src/test/java/com/facebook/presto/orc/TestOrcReaderMemoryUsage.java b/presto-orc/src/test/java/com/facebook/presto/orc/TestOrcReaderMemoryUsage.java index 12e2dc1012eea..a66fe0925550d 100644 --- a/presto-orc/src/test/java/com/facebook/presto/orc/TestOrcReaderMemoryUsage.java +++ b/presto-orc/src/test/java/com/facebook/presto/orc/TestOrcReaderMemoryUsage.java @@ -128,8 +128,8 @@ public void testBigIntTypeWithNulls() // StripeReader memory should increase after reading a block. assertGreaterThan(reader.getCurrentStripeRetainedSizeInBytes(), stripeReaderRetainedSize); - // There are no local buffers needed. - assertEquals(reader.getStreamReaderRetainedSizeInBytes() - streamReaderRetainedSize, 0L); + // There are local buffers needed. + assertGreaterThan(reader.getStreamReaderRetainedSizeInBytes() - streamReaderRetainedSize, 0L); // The total retained size and system memory usage should be strictly larger than 0L because of the instance sizes. assertGreaterThan(reader.getRetainedSizeInBytes() - readerRetainedSize, 0L); assertGreaterThan(reader.getSystemMemoryUsage() - readerSystemMemoryUsage, 0L); @@ -189,8 +189,8 @@ public void testMapTypeWithNulls() // StripeReader memory should increase after reading a block. assertGreaterThan(reader.getCurrentStripeRetainedSizeInBytes(), stripeReaderRetainedSize); - // There are no local buffers needed. - assertEquals(reader.getStreamReaderRetainedSizeInBytes() - streamReaderRetainedSize, 0L); + // There are local buffers needed. + assertGreaterThan(reader.getStreamReaderRetainedSizeInBytes() - streamReaderRetainedSize, 0L); // The total retained size and system memory usage should be strictly larger than 0L because of the instance sizes. assertGreaterThan(reader.getRetainedSizeInBytes() - readerRetainedSize, 0L); assertGreaterThan(reader.getSystemMemoryUsage() - readerSystemMemoryUsage, 0L);