use of io.prestosql.memory.MemoryPool in project hetu-core by openlookeng.
the class MemoryRevokingScheduler method requestRevoking.
private void requestRevoking(MemoryPool memoryPool, Collection<SqlTask> sqlTasks, long remainingBytesToRevoke) {
AtomicLong remainingBytesToRevokeAtomic = new AtomicLong(remainingBytesToRevoke);
Map<Long, List<SqlTask>> revocableMemoryMapping = new HashMap<>();
// This approach disregards create time
if (prioritizeLargerSpiltsMemoryRevoke) {
sqlTasks.stream().filter(task -> task.getTaskStatus().getState() == TaskState.RUNNING).filter(task -> task.getQueryContext().getMemoryPool() == memoryPool).forEach(task -> task.getQueryContext().accept(new VoidTraversingQueryContextVisitor<AtomicLong>() {
@Override
public Void visitQueryContext(QueryContext queryContext, AtomicLong remainingBytesToRevoke) {
if (remainingBytesToRevoke.get() < 0) {
// exit immediately if no work needs to be done
return null;
}
return super.visitQueryContext(queryContext, remainingBytesToRevoke);
}
@Override
public Void visitOperatorContext(OperatorContext operatorContext, AtomicLong remainingBytesToRevoke) {
if (remainingBytesToRevoke.get() <= 0) {
return null;
}
long revocableBytes = operatorContext.getRevocableMemory();
if (revocableBytes > 0) {
if (revocableBytes >= revocableMemorySelectionThreshold) {
remainingBytesToRevoke.addAndGet(-revocableBytes);
}
revocableMemoryMapping.computeIfAbsent(revocableBytes, list -> new ArrayList()).add(task);
}
return null;
}
}, remainingBytesToRevokeAtomic));
if (revocableMemoryMapping.size() > 0) {
revocableMemoryMapping.entrySet().stream().sorted(Map.Entry.comparingByKey(Comparator.reverseOrder())).flatMap(list -> list.getValue().stream()).filter(task -> task.getTaskStatus().getState() == TaskState.RUNNING).forEach(task -> task.getQueryContext().accept(new VoidTraversingQueryContextVisitor<AtomicLong>() {
@Override
public Void visitQueryContext(QueryContext queryContext, AtomicLong remainingBytesToRevoke) {
if (remainingBytesToRevoke.get() < 0) {
// exit immediately if no work needs to be done
return null;
}
return super.visitQueryContext(queryContext, remainingBytesToRevoke);
}
@Override
public Void visitOperatorContext(OperatorContext operatorContext, AtomicLong remainingBytesToRevoke) {
if (remainingBytesToRevoke.get() > 0) {
long revokedBytes = operatorContext.requestMemoryRevoking();
if (revokedBytes > 0) {
remainingBytesToRevoke.addAndGet(-revokedBytes);
log.debug("memoryPool=%s: requested revoking %s; remaining %s", memoryPool.getId(), revokedBytes, remainingBytesToRevoke.get());
}
}
return null;
}
}, remainingBytesToRevokeAtomic));
}
} else {
sqlTasks.stream().filter(task -> task.getTaskStatus().getState() == TaskState.RUNNING).filter(task -> task.getQueryContext().getMemoryPool() == memoryPool).sorted(ORDER_BY_CREATE_TIME).forEach(task -> task.getQueryContext().accept(new VoidTraversingQueryContextVisitor<AtomicLong>() {
@Override
public Void visitQueryContext(QueryContext queryContext, AtomicLong remainingBytesToRevoke) {
if (remainingBytesToRevoke.get() < 0) {
// exit immediately if no work needs to be done
return null;
}
return super.visitQueryContext(queryContext, remainingBytesToRevoke);
}
@Override
public Void visitOperatorContext(OperatorContext operatorContext, AtomicLong remainingBytesToRevoke) {
if (remainingBytesToRevoke.get() > 0) {
long revokedBytes = operatorContext.requestMemoryRevoking();
if (revokedBytes > 0) {
remainingBytesToRevoke.addAndGet(-revokedBytes);
log.debug("memoryPool=%s: requested revoking %s; remaining %s", memoryPool.getId(), revokedBytes, remainingBytesToRevoke.get());
}
}
return null;
}
}, remainingBytesToRevokeAtomic));
}
}
use of io.prestosql.memory.MemoryPool in project hetu-core by openlookeng.
the class GroupByHashYieldAssertion method finishOperatorWithYieldingGroupByHash.
/**
* @param operatorFactory creates an Operator that should directly or indirectly contain GroupByHash
* @param getHashCapacity returns the hash table capacity for the input operator
* @param additionalMemoryInBytes the memory used in addition to the GroupByHash in the operator (e.g., aggregator)
*/
public static GroupByHashYieldResult finishOperatorWithYieldingGroupByHash(List<Page> input, Type hashKeyType, OperatorFactory operatorFactory, Function<Operator, Integer> getHashCapacity, long additionalMemoryInBytes) {
assertLessThan(additionalMemoryInBytes, 1L << 21, "additionalMemoryInBytes should be a relatively small number");
List<Page> result = new LinkedList<>();
// mock an adjustable memory pool
QueryId queryId = new QueryId("test_query");
MemoryPool memoryPool = new MemoryPool(new MemoryPoolId("test"), new DataSize(1, GIGABYTE));
QueryContext queryContext = new QueryContext(queryId, new DataSize(512, MEGABYTE), new DataSize(1024, MEGABYTE), memoryPool, new TestingGcMonitor(), EXECUTOR, SCHEDULED_EXECUTOR, new DataSize(512, MEGABYTE), new SpillSpaceTracker(new DataSize(512, MEGABYTE)), NOOP_SNAPSHOT_UTILS);
DriverContext driverContext = createTaskContext(queryContext, EXECUTOR, TEST_SESSION).addPipelineContext(0, true, true, false).addDriverContext();
Operator operator = operatorFactory.createOperator(driverContext);
// run operator
int yieldCount = 0;
long expectedReservedExtraBytes = 0;
for (Page page : input) {
// unblocked
assertTrue(operator.needsInput());
// saturate the pool with a tiny memory left
long reservedMemoryInBytes = memoryPool.getFreeBytes() - additionalMemoryInBytes;
memoryPool.reserve(queryId, "test", reservedMemoryInBytes);
long oldMemoryUsage = operator.getOperatorContext().getDriverContext().getMemoryUsage();
int oldCapacity = getHashCapacity.apply(operator);
// add a page and verify different behaviors
operator.addInput(page);
// get output to consume the input
Page output = operator.getOutput();
if (output != null) {
result.add(output);
}
long newMemoryUsage = operator.getOperatorContext().getDriverContext().getMemoryUsage();
// between rehash and memory used by aggregator
if (newMemoryUsage < new DataSize(4, MEGABYTE).toBytes()) {
// free the pool for the next iteration
memoryPool.free(queryId, "test", reservedMemoryInBytes);
// this required in case input is blocked
operator.getOutput();
continue;
}
long actualIncreasedMemory = newMemoryUsage - oldMemoryUsage;
if (operator.needsInput()) {
// We have successfully added a page
// Assert we are not blocked
assertTrue(operator.getOperatorContext().isWaitingForMemory().isDone());
// assert the hash capacity is not changed; otherwise, we should have yielded
assertTrue(oldCapacity == getHashCapacity.apply(operator));
// We are not going to rehash; therefore, assert the memory increase only comes from the aggregator
assertLessThan(actualIncreasedMemory, additionalMemoryInBytes);
// free the pool for the next iteration
memoryPool.free(queryId, "test", reservedMemoryInBytes);
} else {
// We failed to finish the page processing i.e. we yielded
yieldCount++;
// Assert we are blocked
assertFalse(operator.getOperatorContext().isWaitingForMemory().isDone());
// Hash table capacity should not change
assertEquals(oldCapacity, (long) getHashCapacity.apply(operator));
// Increased memory is no smaller than the hash table size and no greater than the hash table size + the memory used by aggregator
if (hashKeyType == BIGINT) {
// groupIds and values double by hashCapacity; while valuesByGroupId double by maxFill = hashCapacity / 0.75
expectedReservedExtraBytes = oldCapacity * (long) (Long.BYTES * 1.75 + Integer.BYTES) + page.getRetainedSizeInBytes();
} else {
// groupAddressByHash, groupIdsByHash, and rawHashByHashPosition double by hashCapacity; while groupAddressByGroupId double by maxFill = hashCapacity / 0.75
expectedReservedExtraBytes = oldCapacity * (long) (Long.BYTES * 1.75 + Integer.BYTES + Byte.BYTES) + page.getRetainedSizeInBytes();
}
assertBetweenInclusive(actualIncreasedMemory, expectedReservedExtraBytes, expectedReservedExtraBytes + additionalMemoryInBytes);
// Output should be blocked as well
assertNull(operator.getOutput());
// Free the pool to unblock
memoryPool.free(queryId, "test", reservedMemoryInBytes);
// Trigger a process through getOutput() or needsInput()
output = operator.getOutput();
if (output != null) {
result.add(output);
}
assertTrue(operator.needsInput());
// Hash table capacity has increased
assertGreaterThan(getHashCapacity.apply(operator), oldCapacity);
// Assert the estimated reserved memory before rehash is very close to the one after rehash
long rehashedMemoryUsage = operator.getOperatorContext().getDriverContext().getMemoryUsage();
assertBetweenInclusive(rehashedMemoryUsage * 1.0 / newMemoryUsage, 0.99, 1.01);
// unblocked
assertTrue(operator.needsInput());
}
}
result.addAll(finishOperator(operator));
return new GroupByHashYieldResult(yieldCount, expectedReservedExtraBytes, result);
}
use of io.prestosql.memory.MemoryPool in project hetu-core by openlookeng.
the class AbstractOperatorBenchmark method runOnce.
@Override
protected Map<String, Long> runOnce() {
Session setSession = testSessionBuilder().setSystemProperty("optimizer.optimize-hash-generation", "true").build();
MemoryPool memoryPool = new MemoryPool(new MemoryPoolId("test"), new DataSize(1, GIGABYTE));
SpillSpaceTracker spillSpaceTracker = new SpillSpaceTracker(new DataSize(1, GIGABYTE));
TaskContext taskContext = new QueryContext(new QueryId("test"), new DataSize(256, MEGABYTE), new DataSize(512, MEGABYTE), memoryPool, new TestingGcMonitor(), localQueryRunner.getExecutor(), localQueryRunner.getScheduler(), new DataSize(256, MEGABYTE), spillSpaceTracker, NOOP_SNAPSHOT_UTILS).addTaskContext(new TaskStateMachine(new TaskId("query", 0, 0), localQueryRunner.getExecutor()), setSession, false, false, OptionalInt.empty(), Optional.empty(), TESTING_SERDE_FACTORY);
CpuTimer cpuTimer = new CpuTimer();
Map<String, Long> executionStats = execute(taskContext);
CpuDuration executionTime = cpuTimer.elapsedTime();
TaskStats taskStats = taskContext.getTaskStats();
long inputRows = taskStats.getRawInputPositions();
long inputBytes = taskStats.getRawInputDataSize().toBytes();
long outputRows = taskStats.getOutputPositions();
long outputBytes = taskStats.getOutputDataSize().toBytes();
double inputMegaBytes = new DataSize(inputBytes, BYTE).getValue(MEGABYTE);
return ImmutableMap.<String, Long>builder().putAll(executionStats).put("elapsed_millis", executionTime.getWall().toMillis()).put("input_rows_per_second", (long) (inputRows / executionTime.getWall().getValue(SECONDS))).put("output_rows_per_second", (long) (outputRows / executionTime.getWall().getValue(SECONDS))).put("input_megabytes", (long) inputMegaBytes).put("input_megabytes_per_second", (long) (inputMegaBytes / executionTime.getWall().getValue(SECONDS))).put("wall_nanos", executionTime.getWall().roundTo(NANOSECONDS)).put("cpu_nanos", executionTime.getCpu().roundTo(NANOSECONDS)).put("user_nanos", executionTime.getUser().roundTo(NANOSECONDS)).put("input_rows", inputRows).put("input_bytes", inputBytes).put("output_rows", outputRows).put("output_bytes", outputBytes).build();
}
use of io.prestosql.memory.MemoryPool in project hetu-core by openlookeng.
the class TestMemoryRevokingScheduler method testCountAlreadyRevokedMemoryWithinAPool.
@Test
public void testCountAlreadyRevokedMemoryWithinAPool() throws Exception {
// Given
SqlTask sqlTask1 = newSqlTask();
MemoryPool anotherMemoryPool = new MemoryPool(new MemoryPoolId("test"), new DataSize(10, BYTE));
sqlTask1.getQueryContext().setMemoryPool(anotherMemoryPool);
OperatorContext operatorContext1 = createContexts(sqlTask1);
SqlTask sqlTask2 = newSqlTask();
OperatorContext operatorContext2 = createContexts(sqlTask2);
List<SqlTask> tasks = ImmutableList.of(sqlTask1, sqlTask2);
MemoryRevokingScheduler scheduler = new MemoryRevokingScheduler(asList(memoryPool, anotherMemoryPool), () -> tasks, executor, 1.0, 1.0, false, new DataSize(512, MEGABYTE).toBytes());
allOperatorContexts = ImmutableSet.of(operatorContext1, operatorContext2);
/*
* sqlTask1 fills its pool
*/
operatorContext1.localRevocableMemoryContext().setBytes(12);
requestMemoryRevoking(scheduler);
assertMemoryRevokingRequestedFor(operatorContext1);
/*
* When sqlTask2 fills its pool
*/
operatorContext2.localRevocableMemoryContext().setBytes(12);
requestMemoryRevoking(scheduler);
/*
* Then sqlTask2 should be asked to revoke its memory too
*/
assertMemoryRevokingRequestedFor(operatorContext1, operatorContext2);
}
Aggregations