Examples with OperatorContext com.facebook.presto.operator.OperatorContext used on opensource projects

Example 11 with OperatorContext

use of com.facebook.presto.operator.OperatorContext in project presto by prestodb.

the class TestMemoryRevokingScheduler method testCountAlreadyRevokedMemoryWithinAPool.

@Test
public void testCountAlreadyRevokedMemoryWithinAPool() throws Exception {
    // Given
    MemoryPool anotherMemoryPool = new MemoryPool(new MemoryPoolId("test"), new DataSize(10, BYTE));
    SqlTask sqlTask1 = newSqlTask(new QueryId("q1"), anotherMemoryPool);
    OperatorContext operatorContext1 = createContexts(sqlTask1);
    SqlTask sqlTask2 = newSqlTask(new QueryId("q2"), memoryPool);
    OperatorContext operatorContext2 = createContexts(sqlTask2);
    List<SqlTask> tasks = ImmutableList.of(sqlTask1, sqlTask2);
    MemoryRevokingScheduler scheduler = new MemoryRevokingScheduler(asList(memoryPool, anotherMemoryPool), () -> tasks, queryContexts::get, 1.0, 1.0, ORDER_BY_CREATE_TIME, false);
    try {
        scheduler.start();
        allOperatorContexts = ImmutableSet.of(operatorContext1, operatorContext2);
        /*
             * sqlTask1 fills its pool
             */
        operatorContext1.localRevocableMemoryContext().setBytes(12);
        scheduler.awaitAsynchronousCallbacksRun();
        assertMemoryRevokingRequestedFor(operatorContext1);
        /*
             * When sqlTask2 fills its pool
             */
        operatorContext2.localRevocableMemoryContext().setBytes(12);
        scheduler.awaitAsynchronousCallbacksRun();
        /*
             * Then sqlTask2 should be asked to revoke its memory too
             */
        assertMemoryRevokingRequestedFor(operatorContext1, operatorContext2);
    } finally {
        scheduler.stop();
    }
}

Example 12 with OperatorContext

use of com.facebook.presto.operator.OperatorContext in project presto by prestodb.

the class LocalQueryRunner method executeInternal.

private MaterializedResultWithPlan executeInternal(Session session, @Language("SQL") String sql, WarningCollector warningCollector) {
    lock.readLock().lock();
    try (Closer closer = Closer.create()) {
        AtomicReference<MaterializedResult.Builder> builder = new AtomicReference<>();
        PageConsumerOutputFactory outputFactory = new PageConsumerOutputFactory(types -> {
            builder.compareAndSet(null, MaterializedResult.resultBuilder(session, types));
            return builder.get()::page;
        });
        Plan plan = createPlan(session, sql, warningCollector);
        TaskContext taskContext = TestingTaskContext.builder(notificationExecutor, yieldExecutor, session).setMaxSpillSize(nodeSpillConfig.getMaxSpillPerNode()).setQueryMaxSpillSize(nodeSpillConfig.getQueryMaxSpillPerNode()).setQueryMaxTotalMemory(getQueryMaxTotalMemoryPerNode(session)).setTaskPlan(plan.getRoot()).build();
        taskContext.getQueryContext().setVerboseExceededMemoryLimitErrorsEnabled(isVerboseExceededMemoryLimitErrorsEnabled(session));
        taskContext.getQueryContext().setHeapDumpOnExceededMemoryLimitEnabled(isHeapDumpOnExceededMemoryLimitEnabled(session));
        String heapDumpFilePath = Paths.get(getHeapDumpFileDirectory(session), format("%s_%s.hprof", session.getQueryId().getId(), taskContext.getTaskId().getStageExecutionId().getStageId().getId())).toString();
        taskContext.getQueryContext().setHeapDumpFilePath(heapDumpFilePath);
        List<Driver> drivers = createDrivers(session, plan, outputFactory, taskContext);
        drivers.forEach(closer::register);
        boolean done = false;
        while (!done) {
            boolean processed = false;
            for (Driver driver : drivers) {
                if (alwaysRevokeMemory) {
                    driver.getDriverContext().getOperatorContexts().stream().filter(operatorContext -> operatorContext.getOperatorStats().getRevocableMemoryReservation().getValue() > 0).forEach(OperatorContext::requestMemoryRevoking);
                }
                if (!driver.isFinished()) {
                    driver.process();
                    processed = true;
                }
            }
            done = !processed;
        }
        verify(builder.get() != null, "Output operator was not created");
        return new MaterializedResultWithPlan(builder.get().build(), plan);
    } catch (IOException e) {
        throw new UncheckedIOException(e);
    } finally {
        lock.readLock().unlock();
    }
}

Example 13 with OperatorContext

use of com.facebook.presto.operator.OperatorContext in project presto by prestodb.

the class TestQueryContext method testMoveTaggedAllocations.

@Test
public void testMoveTaggedAllocations() {
    MemoryPool generalPool = new MemoryPool(GENERAL_POOL, new DataSize(10_000, BYTE));
    MemoryPool reservedPool = new MemoryPool(RESERVED_POOL, new DataSize(10_000, BYTE));
    QueryId queryId = new QueryId("query");
    QueryContext queryContext = createQueryContext(queryId, generalPool);
    TaskStateMachine taskStateMachine = new TaskStateMachine(TaskId.valueOf("queryid.0.0.0"), TEST_EXECUTOR);
    TaskContext taskContext = queryContext.addTaskContext(taskStateMachine, TEST_SESSION, Optional.of(PLAN_FRAGMENT.getRoot()), false, false, false, false, false);
    DriverContext driverContext = taskContext.addPipelineContext(0, false, false, false).addDriverContext();
    OperatorContext operatorContext = driverContext.addOperatorContext(0, new PlanNodeId("test"), "test");
    // allocate some memory in the general pool
    LocalMemoryContext memoryContext = operatorContext.aggregateUserMemoryContext().newLocalMemoryContext("test_context");
    memoryContext.setBytes(1_000);
    Map<String, Long> allocations = generalPool.getTaggedMemoryAllocations(queryId);
    assertEquals(allocations, ImmutableMap.of("test_context", 1_000L));
    queryContext.setMemoryPool(reservedPool);
    assertNull(generalPool.getTaggedMemoryAllocations(queryId));
    allocations = reservedPool.getTaggedMemoryAllocations(queryId);
    assertEquals(allocations, ImmutableMap.of("test_context", 1_000L));
    assertEquals(generalPool.getFreeBytes(), 10_000);
    assertEquals(reservedPool.getFreeBytes(), 9_000);
    memoryContext.close();
    assertEquals(generalPool.getFreeBytes(), 10_000);
    assertEquals(reservedPool.getFreeBytes(), 10_000);
}

Example 14 with OperatorContext

use of com.facebook.presto.operator.OperatorContext in project presto by prestodb.

the class TestOptimizedPartitionedOutputOperator method verifyOutputSizes.

private static void verifyOutputSizes(OptimizedPartitionedOutputOperator operator, long expectedSizeInBytes, long expectedPositionCount) {
    OperatorContext operatorContext = operator.getOperatorContext();
    assertBetweenInclusive(operatorContext.getOutputDataSize().getTotalCount(), (long) (expectedSizeInBytes / OUTPUT_SIZE_ESTIMATION_ERROR_ALLOWANCE), (long) (expectedSizeInBytes * OUTPUT_SIZE_ESTIMATION_ERROR_ALLOWANCE));
    assertEquals(operatorContext.getOutputPositions().getTotalCount(), expectedPositionCount);
}