Skip to content

Commit

Permalink
Add testMemoryTracking in TestSelectiveOrcReader
Browse files Browse the repository at this point in the history
  • Loading branch information
@mbasmanova
Ying Su authored and mbasmanova committed Sep 2, 2020
1 parent 696ffca commit 3d4931b
Show file tree
Hide file tree
Showing 2 changed files with 81 additions and 4 deletions.
11 changes: 7 additions & 4 deletions presto-orc/src/test/java/com/facebook/presto/orc/OrcTester.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -891,7 +891,8 @@ public static void assertFileContentsPresto(
intermediateEncryptionKeys,
includedColumns,
outputColumns,
false)) {
false,
new TestingHiveOrcAggregatedMemoryContext())) {
assertEquals(recordReader.getReaderPosition(), 0);
assertEquals(recordReader.getFilePosition(), 0);

Expand Down Expand Up @@ -1554,7 +1555,8 @@ static OrcSelectiveRecordReader createCustomOrcSelectiveRecordReader(
ImmutableMap.of(),
ImmutableMap.of(0, type),
ImmutableList.of(0),
mapNullKeysEnabled);
mapNullKeysEnabled,
new TestingHiveOrcAggregatedMemoryContext());
}

public static OrcSelectiveRecordReader createCustomOrcSelectiveRecordReader(
Expand All @@ -1570,7 +1572,8 @@ public static OrcSelectiveRecordReader createCustomOrcSelectiveRecordReader(
Map<Integer, Slice> intermediateEncryptionKeys,
Map<Integer, Type> includedColumns,
List<Integer> outputColumns,
boolean mapNullKeysEnabled)
boolean mapNullKeysEnabled,
OrcAggregatedMemoryContext systemMemoryUsage)
throws IOException
{
OrcDataSource orcDataSource = new FileOrcDataSource(file, new DataSize(1, MEGABYTE), new DataSize(1, MEGABYTE), new DataSize(1, MEGABYTE), true);
Expand Down Expand Up @@ -1607,7 +1610,7 @@ public static OrcSelectiveRecordReader createCustomOrcSelectiveRecordReader(
orcDataSource.getSize(),
HIVE_STORAGE_TIME_ZONE,
LEGACY_MAP_SUBSCRIPT,
new TestingHiveOrcAggregatedMemoryContext(),
systemMemoryUsage,
Optional.empty(),
initialBatchSize,
intermediateEncryptionKeys);
Expand Down
74 changes: 74 additions & 0 deletions presto-orc/src/test/java/com/facebook/presto/orc/TestSelectiveOrcReader.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
package com.facebook.presto.orc;

import com.facebook.presto.common.InvalidFunctionArgumentException;
import com.facebook.presto.common.Page;
import com.facebook.presto.common.Subfield;
import com.facebook.presto.common.type.CharType;
import com.facebook.presto.common.type.DecimalType;
Expand All @@ -30,6 +31,7 @@
import com.facebook.presto.orc.TupleDomainFilter.BytesValues;
import com.facebook.presto.orc.TupleDomainFilter.DoubleRange;
import com.facebook.presto.orc.TupleDomainFilter.FloatRange;
import com.facebook.presto.orc.metadata.CompressionKind;
import com.google.common.base.Strings;
import com.google.common.collect.AbstractIterator;
import com.google.common.collect.ContiguousSet;
Expand All @@ -51,10 +53,12 @@
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Random;
import java.util.function.Function;
import java.util.stream.IntStream;

import static com.facebook.airlift.testing.Assertions.assertBetweenInclusive;
import static com.facebook.presto.common.type.BigintType.BIGINT;
import static com.facebook.presto.common.type.BooleanType.BOOLEAN;
import static com.facebook.presto.common.type.CharType.createCharType;
Expand All @@ -67,11 +71,16 @@
import static com.facebook.presto.common.type.TinyintType.TINYINT;
import static com.facebook.presto.common.type.VarbinaryType.VARBINARY;
import static com.facebook.presto.common.type.VarcharType.VARCHAR;
import static com.facebook.presto.orc.OrcReader.MAX_BATCH_SIZE;
import static com.facebook.presto.orc.OrcTester.Format.DWRF;
import static com.facebook.presto.orc.OrcTester.HIVE_STORAGE_TIME_ZONE;
import static com.facebook.presto.orc.OrcTester.arrayType;
import static com.facebook.presto.orc.OrcTester.createCustomOrcSelectiveRecordReader;
import static com.facebook.presto.orc.OrcTester.mapType;
import static com.facebook.presto.orc.OrcTester.quickSelectiveOrcTester;
import static com.facebook.presto.orc.OrcTester.rowType;
import static com.facebook.presto.orc.OrcTester.writeOrcColumnsPresto;
import static com.facebook.presto.orc.TestingOrcPredicate.createOrcPredicate;
import static com.facebook.presto.orc.TupleDomainFilter.IS_NOT_NULL;
import static com.facebook.presto.orc.TupleDomainFilter.IS_NULL;
import static com.facebook.presto.orc.TupleDomainFilterUtils.toBigintValues;
Expand Down Expand Up @@ -894,6 +903,71 @@ public void testVarBinaries()
bytesBetween(false, new byte[] {1}, new byte[] {12}));
}

@Test
public void testMemoryTracking()
throws Exception
{
List<Type> types = ImmutableList.of(INTEGER, VARCHAR);
TempFile tempFile = new TempFile();
List<Integer> intValues = newArrayList(limit(
cycle(concat(
ImmutableList.of(1), nCopies(9999, 123),
ImmutableList.of(2), nCopies(9999, 123),
ImmutableList.of(3), nCopies(9999, 123),
nCopies(1_000_000, null))),
NUM_ROWS));
List<String> varcharValues = newArrayList(limit(cycle(ImmutableList.of("A", "B", "C")), NUM_ROWS));

writeOrcColumnsPresto(tempFile.getFile(), DWRF, CompressionKind.NONE, Optional.empty(), types, ImmutableList.of(intValues, varcharValues), new OrcWriterStats());

OrcPredicate orcPredicate = createOrcPredicate(types, ImmutableList.of(intValues, varcharValues), DWRF, false);
Map<Integer, Type> includedColumns = IntStream.range(0, types.size())
.boxed()
.collect(toImmutableMap(Function.identity(), types::get));
List<Integer> outputColumns = IntStream.range(0, types.size())
.boxed()
.collect(toImmutableList());
OrcAggregatedMemoryContext systemMemoryUsage = new TestingHiveOrcAggregatedMemoryContext();
try (OrcSelectiveRecordReader recordReader = createCustomOrcSelectiveRecordReader(
tempFile.getFile(),
DWRF.getOrcEncoding(),
orcPredicate,
types,
MAX_BATCH_SIZE,
ImmutableMap.of(),
OrcReaderSettings.builder().build().getFilterFunctions(),
OrcReaderSettings.builder().build().getFilterFunctionInputMapping(),
OrcReaderSettings.builder().build().getRequiredSubfields(),
ImmutableMap.of(),
includedColumns,
outputColumns,
false,
systemMemoryUsage)) {
assertEquals(recordReader.getReaderPosition(), 0);
assertEquals(recordReader.getFilePosition(), 0);

int rowsProcessed = 0;
while (true) {
Page page = recordReader.getNextPage();
if (page == null) {
break;
}

int positionCount = page.getPositionCount();
if (positionCount == 0) {
continue;
}

page.getLoadedPage();

assertBetweenInclusive(systemMemoryUsage.getBytes(), 110000L, 130000L);

rowsProcessed += positionCount;
}
assertEquals(rowsProcessed, NUM_ROWS);
}
}

private void testRoundTripNumeric(Iterable<? extends Number> values, TupleDomainFilter filter)
throws Exception
{
Expand Down

0 comments on commit 3d4931b

Please sign in to comment.