Source:
public void testMemoryTracking(boolean buildHashEnabled, boolean spillEnabled) { RowPagesBuilder buildPages = rowPagesBuilder(buildHashEnabled, Ints.asList(0), ImmutableList.of(VARCHAR, BIGINT, BIGINT)); Random r = new Random(384013974023187L); for (int i = 0; i < 1000; i++) { buildPages.addSequencePage(1024, r.nextInt(), r.nextInt(), r.nextInt()); } HashBuilderOperator buildOperator = createTestOperator(buildPages, spillEnabled); assertEquals(buildOperator.getState(), State.CONSUMING_INPUT); List<Page> inPages = buildPages.build(); inPages.stream().forEach(buildOperator::addInput); long preciseFootprintBefore = buildOperator.getMemoryFootprint(); long predictedFootprintAfter = buildOperator.getPredictedSize(); long estimatedFootprintBefore; if (spillEnabled) { estimatedFootprintBefore = buildOperator.getOperatorContext().localRevocableMemoryContext().getBytes() - buildOperator.getPredictedSizeDelta(); } else { estimatedFootprintBefore = buildOperator.getOperatorContext().localUserMemoryContext().getBytes() - buildOperator.getPredictedSizeDelta(); } buildOperator.finish(); long preciseFootprintAfter = buildOperator.getMemoryFootprint(); long estimatedFootprintAfter; if (spillEnabled) { estimatedFootprintAfter = buildOperator.getOperatorContext().localRevocableMemoryContext().getBytes(); } else { estimatedFootprintAfter = buildOperator.getOperatorContext().localUserMemoryContext().getBytes(); } assertEquals(buildOperator.getState(), State.LOOKUP_SOURCE_BUILT); assertTrue(isCheckFootprintCloseEnough(preciseFootprintBefore, estimatedFootprintBefore, 2)); assertTrue(isCheckFootprintCloseEnough(preciseFootprintAfter, estimatedFootprintAfter, 2)); assertTrue(isCheckFootprintCloseEnough(predictedFootprintAfter, estimatedFootprintAfter, 2) || predictedFootprintAfter > estimatedFootprintAfter); assertTrue(isCheckFootprintCloseEnough(predictedFootprintAfter, preciseFootprintAfter, 2) || predictedFootprintAfter > preciseFootprintAfter); }

Target:
public void testMemoryTracking(boolean buildHashEnabled, boolean spillEnabled) { RowPagesBuilder pagesBuilder = rowPagesBuilder(buildHashEnabled, Ints.asList(0), ImmutableList.of(BIGINT, BIGINT, BIGINT)); int totalRows = NUM_PAGES * PAGE_SIZE; List<Integer> someNumbers = rangeList(totalRows + 1); TaskContext taskContext = TestingTaskContext.createTaskContext(executor, scheduledExecutor, TEST_SESSION); for (int i = 0; i < NUM_PAGES; i++) { pagesBuilder.addSequencePage(PAGE_SIZE, 1, someNumbers.get(i), someNumbers.get((i * i) % totalRows)); } List<Page> inputPages = pagesBuilder.build(); HashBuilderOperator buildOperator = createTestOperator(pagesBuilder, spillEnabled, taskContext); assertEquals(buildOperator.getState(), CONSUMING_INPUT); inputPages.stream().forEach(buildOperator::addInput); long estimatedFootprintBefore; if (spillEnabled) { estimatedFootprintBefore = taskContext.getTaskMemoryContext().getRevocableMemory(); } else { estimatedFootprintBefore = taskContext.getMemoryReservation().toBytes(); } buildOperator.finish(); long estimatedFootprintAfter; if (spillEnabled) { estimatedFootprintAfter = taskContext.getTaskMemoryContext().getRevocableMemory(); } else { estimatedFootprintAfter = taskContext.getMemoryReservation().toBytes(); } assertEquals(buildOperator.getState(), LOOKUP_SOURCE_BUILT); if (buildHashEnabled) { assertEquals(estimatedFootprintBefore, 59656384); assertEquals(estimatedFootprintAfter, 59640024); } else { assertEquals(estimatedFootprintBefore, 51396272); assertEquals(estimatedFootprintAfter, 51384016); } }