Examples with OperatorContext io.prestosql.operator.OperatorContext used on opensource projects

Example 1 with OperatorContext

use of io.prestosql.operator.OperatorContext in project hetu-core by openlookeng.

the class LocalQueryRunner method executeInternal.

private MaterializedResultWithPlan executeInternal(Session session, @Language("SQL") String sql) {
    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;
        });
        TaskContext taskContext = TestingTaskContext.builder(notificationExecutor, yieldExecutor, session).setMaxSpillSize(nodeSpillConfig.getMaxSpillPerNode()).setQueryMaxSpillSize(nodeSpillConfig.getQueryMaxSpillPerNode()).build();
        Plan plan = createPlan(session, sql, WarningCollector.NOOP);
        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 2 with OperatorContext

use of io.prestosql.operator.OperatorContext in project hetu-core by openlookeng.

the class TestSingleInputSnapshotState method testStaticConstructor.

@Test
public void testStaticConstructor() {
    ScheduledExecutorService scheduler = newScheduledThreadPool(4);
    DriverContext driverContext = createTaskContext(scheduler, scheduler, TEST_SNAPSHOT_SESSION).addPipelineContext(0, true, true, false).addDriverContext();
    OperatorContext operatorContext = driverContext.addOperatorContext(1, new PlanNodeId("planNodeId"), "test");
    SingleInputSnapshotState singleInputSnapshotState = SingleInputSnapshotState.forOperator(mock(Operator.class), operatorContext);
    singleInputSnapshotState.processPage(regularPage);
}

Example 3 with OperatorContext

use of io.prestosql.operator.OperatorContext in project hetu-core by openlookeng.

the class TestMemoryRevokingScheduler method testScheduleMemoryRevoking.

@Test
public void testScheduleMemoryRevoking() throws Exception {
    SqlTask sqlTask1 = newSqlTask();
    SqlTask sqlTask2 = newSqlTask();
    TaskContext taskContext1 = sqlTask1.getQueryContext().addTaskContext(new TaskStateMachine(new TaskId("q1", 1, 1), executor), session, false, false, OptionalInt.empty(), Optional.empty(), TESTING_SERDE_FACTORY);
    PipelineContext pipelineContext11 = taskContext1.addPipelineContext(0, false, false, false);
    DriverContext driverContext111 = pipelineContext11.addDriverContext();
    OperatorContext operatorContext1 = driverContext111.addOperatorContext(1, new PlanNodeId("na"), "na");
    OperatorContext operatorContext2 = driverContext111.addOperatorContext(2, new PlanNodeId("na"), "na");
    DriverContext driverContext112 = pipelineContext11.addDriverContext();
    OperatorContext operatorContext3 = driverContext112.addOperatorContext(3, new PlanNodeId("na"), "na");
    TaskContext taskContext2 = sqlTask2.getQueryContext().addTaskContext(new TaskStateMachine(new TaskId("q2", 1, 1), executor), session, false, false, OptionalInt.empty(), Optional.empty(), TESTING_SERDE_FACTORY);
    PipelineContext pipelineContext21 = taskContext2.addPipelineContext(1, false, false, false);
    DriverContext driverContext211 = pipelineContext21.addDriverContext();
    OperatorContext operatorContext4 = driverContext211.addOperatorContext(4, new PlanNodeId("na"), "na");
    OperatorContext operatorContext5 = driverContext211.addOperatorContext(5, new PlanNodeId("na"), "na");
    Collection<SqlTask> tasks = ImmutableList.of(sqlTask1, sqlTask2);
    MemoryRevokingScheduler scheduler = new MemoryRevokingScheduler(singletonList(memoryPool), () -> tasks, executor, 1.0, 1.0, false, new DataSize(512, MEGABYTE).toBytes());
    allOperatorContexts = ImmutableSet.of(operatorContext1, operatorContext2, operatorContext3, operatorContext4, operatorContext5);
    assertMemoryRevokingNotRequested();
    requestMemoryRevoking(scheduler);
    assertEquals(10, memoryPool.getFreeBytes());
    assertMemoryRevokingNotRequested();
    LocalMemoryContext revocableMemory1 = operatorContext1.localRevocableMemoryContext();
    LocalMemoryContext revocableMemory3 = operatorContext3.localRevocableMemoryContext();
    LocalMemoryContext revocableMemory4 = operatorContext4.localRevocableMemoryContext();
    LocalMemoryContext revocableMemory5 = operatorContext5.localRevocableMemoryContext();
    revocableMemory1.setBytes(3);
    revocableMemory3.setBytes(6);
    assertEquals(1, memoryPool.getFreeBytes());
    requestMemoryRevoking(scheduler);
    // we are still good - no revoking needed
    assertMemoryRevokingNotRequested();
    revocableMemory4.setBytes(7);
    assertEquals(-6, memoryPool.getFreeBytes());
    requestMemoryRevoking(scheduler);
    // we need to revoke 3 and 6
    assertMemoryRevokingRequestedFor(operatorContext1, operatorContext3);
    // yet another revoking request should not change anything
    requestMemoryRevoking(scheduler);
    assertMemoryRevokingRequestedFor(operatorContext1, operatorContext3);
    // lets revoke some bytes
    revocableMemory1.setBytes(0);
    operatorContext1.resetMemoryRevokingRequested();
    requestMemoryRevoking(scheduler);
    assertMemoryRevokingRequestedFor(operatorContext3);
    assertEquals(-3, memoryPool.getFreeBytes());
    // and allocate some more
    revocableMemory5.setBytes(3);
    assertEquals(-6, memoryPool.getFreeBytes());
    requestMemoryRevoking(scheduler);
    // we are still good with just OC3 in process of revoking
    assertMemoryRevokingRequestedFor(operatorContext3);
    // and allocate some more
    revocableMemory5.setBytes(4);
    assertEquals(-7, memoryPool.getFreeBytes());
    requestMemoryRevoking(scheduler);
    // no we have to trigger revoking for OC4
    assertMemoryRevokingRequestedFor(operatorContext3, operatorContext4);
}

Example 4 with OperatorContext

use of io.prestosql.operator.OperatorContext in project hetu-core by openlookeng.

the class TestMemoryRevokingScheduler method testImmediateMemoryRevoking.

/**
 * Test that when a {@link MemoryPool} is over-allocated, revocable memory is revoked without delay (although asynchronously).
 */
@Test
public void testImmediateMemoryRevoking() throws Exception {
    // Given
    SqlTask sqlTask = newSqlTask();
    OperatorContext operatorContext = createContexts(sqlTask);
    allOperatorContexts = ImmutableSet.of(operatorContext);
    List<SqlTask> tasks = ImmutableList.of(sqlTask);
    MemoryRevokingScheduler scheduler = new MemoryRevokingScheduler(singletonList(memoryPool), () -> tasks, executor, 1.0, 1.0, false, new DataSize(512, MEGABYTE).toBytes());
    // no periodic check initiated
    scheduler.registerPoolListeners();
    // When
    operatorContext.localRevocableMemoryContext().setBytes(12);
    awaitAsynchronousCallbacksRun();
    // Then
    assertMemoryRevokingRequestedFor(operatorContext);
}

Example 5 with OperatorContext

use of io.prestosql.operator.OperatorContext in project hetu-core by openlookeng.

the class TestMemoryRevokingScheduler method createContexts.

private OperatorContext createContexts(SqlTask sqlTask) {
    TaskContext taskContext = sqlTask.getQueryContext().addTaskContext(new TaskStateMachine(new TaskId("q", 1, 1), executor), session, false, false, OptionalInt.empty(), Optional.empty(), TESTING_SERDE_FACTORY);
    PipelineContext pipelineContext = taskContext.addPipelineContext(0, false, false, false);
    DriverContext driverContext = pipelineContext.addDriverContext();
    OperatorContext operatorContext = driverContext.addOperatorContext(1, new PlanNodeId("na"), "na");
    return operatorContext;
}