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

Example 26 with TaskContext

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

the class TestMultiInputSnapshotState method testStaticConstructor.

@Test
public void testStaticConstructor() {
    ScheduledExecutorService scheduler = newScheduledThreadPool(4);
    TaskContext taskContext = createTaskContext(scheduler, scheduler, TEST_SNAPSHOT_SESSION);
    DriverContext driverContext = taskContext.addPipelineContext(0, true, true, false).addDriverContext();
    OperatorContext operatorContext = driverContext.addOperatorContext(1, new PlanNodeId("planNodeId"), "test");
    MultiInputSnapshotState inputSnapshotState = MultiInputSnapshotState.forOperator(mock(MultiInputRestorable.class), operatorContext);
    processPage(inputSnapshotState, source1, regularPage);
    inputSnapshotState = MultiInputSnapshotState.forTaskComponent(mock(MultiInputRestorable.class), taskContext, TestMultiInputSnapshotState::createSnapshotStateId);
    processPage(inputSnapshotState, source1, regularPage);
}

Example 27 with TaskContext

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

the class QueryContext method addTaskContext.

public TaskContext addTaskContext(String taskInstanceId, TaskStateMachine taskStateMachine, Session session, boolean perOperatorCpuTimerEnabled, boolean cpuTimerEnabled, OptionalInt totalPartitions, Optional<PlanNodeId> parent, PagesSerdeFactory serdeFactory) {
    // Task instance id has format "<resume count>-<random UUID>"
    long resumeCount = Long.valueOf(taskInstanceId.substring(0, taskInstanceId.indexOf('-')));
    TaskContext taskContext = TaskContext.createTaskContext(this, taskStateMachine, gcMonitor, notificationExecutor, yieldExecutor, session, queryMemoryContext.newMemoryTrackingContext(), perOperatorCpuTimerEnabled, cpuTimerEnabled, totalPartitions, parent.orElse(null), serdeFactory, new TaskSnapshotManager(taskStateMachine.getTaskId(), resumeCount, snapshotUtils));
    taskContexts.put(taskInstanceId, taskContext);
    return taskContext;
}

Example 28 with TaskContext

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

the class QueryContext method getTaskContext.

public TaskContext getTaskContext(String taskInstanceId) {
    TaskContext taskContext = taskContexts.get(taskInstanceId);
    verify(taskContext != null, "task does not exist");
    return taskContext;
}

Example 29 with TaskContext

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

the class IndexLoader method initializeStateIfNecessary.

private synchronized void initializeStateIfNecessary() {
    if (pipelineContext == null) {
        TaskContext taskContext = taskContextReference.get();
        checkState(taskContext != null, "Task context must be set before index can be built");
        pipelineContext = taskContext.addPipelineContext(indexBuildDriverFactoryProvider.getPipelineId(), true, true, false);
    }
    if (indexSnapshotLoader == null) {
        indexSnapshotLoader = new IndexSnapshotLoader(indexBuildDriverFactoryProvider, pipelineContext, indexSnapshotReference, lookupSourceInputChannels, keyTypes, keyOutputChannels, keyOutputHashChannel, expectedPositions, maxIndexMemorySize, pagesIndexFactory, joinCompiler);
    }
}

Example 30 with TaskContext

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

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("task-id"), TEST_EXECUTOR);
    TaskContext taskContext = queryContext.addTaskContext(taskStateMachine, TEST_SESSION, false, false, OptionalInt.empty(), Optional.empty(), TESTING_SERDE_FACTORY);
    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().get(queryId);
    assertEquals(allocations, ImmutableMap.of("test_context", 1_000L));
    queryContext.setMemoryPool(reservedPool);
    assertNull(generalPool.getTaggedMemoryAllocations().get(queryId));
    allocations = reservedPool.getTaggedMemoryAllocations().get(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);
}