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

Example 11 with TaskContext

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

the class TestSpatialJoinOperator method testDistanceQuery.

@Test
public void testDistanceQuery() {
    TaskContext taskContext = createTaskContext();
    DriverContext driverContext = taskContext.addPipelineContext(0, true, true, false).addDriverContext();
    RowPagesBuilder buildPages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR, DOUBLE)).row(stPoint(0, 0), "0_0", 1.5).row(null, "null", 1.5).row(stPoint(1, 0), "1_0", 1.5).pageBreak().row(stPoint(3, 0), "3_0", 1.5).pageBreak().row(stPoint(10, 0), "10_0", 1.5);
    PagesSpatialIndexFactory pagesSpatialIndexFactory = buildIndex(driverContext, (build, probe, r) -> build.distance(probe) <= r.getAsDouble(), Optional.of(2), Optional.empty(), buildPages);
    RowPagesBuilder probePages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR)).row(stPoint(0, 1), "0_1").row(null, "null").row(stPoint(1, 1), "1_1").pageBreak().row(stPoint(3, 1), "3_1").pageBreak().row(stPoint(10, 1), "10_1");
    OperatorFactory joinOperatorFactory = new SpatialJoinOperatorFactory(2, new PlanNodeId("test"), INNER, probePages.getTypes(), Ints.asList(1), 0, Optional.empty(), pagesSpatialIndexFactory);
    // Make sure that spatial index reference counting works with duplicate factories
    joinOperatorFactory.duplicate().noMoreOperators();
    MaterializedResult expected = resultBuilder(taskContext.getSession(), ImmutableList.of(VARCHAR, VARCHAR)).row("0_1", "0_0").row("0_1", "1_0").row("1_1", "0_0").row("1_1", "1_0").row("3_1", "3_0").row("10_1", "10_0").build();
    assertOperatorEqualsIgnoreOrder(joinOperatorFactory, driverContext, probePages.build(), expected);
}

Example 12 with TaskContext

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

the class TestSqlTaskExecution method testSimple.

@Test(dataProvider = "executionStrategies", timeOut = 20_000)
public void testSimple(PipelineExecutionStrategy executionStrategy) throws Exception {
    ScheduledExecutorService taskNotificationExecutor = newScheduledThreadPool(10, threadsNamed("task-notification-%s"));
    ScheduledExecutorService driverYieldExecutor = newScheduledThreadPool(2, threadsNamed("driver-yield-%s"));
    TaskExecutor taskExecutor = new TaskExecutor(5, 10, 3, 4, TASK_FAIR, Ticker.systemTicker());
    taskExecutor.start();
    try {
        TaskStateMachine taskStateMachine = new TaskStateMachine(TASK_ID, taskNotificationExecutor);
        PartitionedOutputBuffer outputBuffer = newTestingOutputBuffer(taskNotificationExecutor);
        OutputBufferConsumer outputBufferConsumer = new OutputBufferConsumer(outputBuffer, OUTPUT_BUFFER_ID);
        // 
        // test initialization: simple task with 1 pipeline
        // 
        // pipeline 0  ... pipeline id
        // partitioned ... partitioned/unpartitioned pipeline
        // grouped   ... execution strategy (in grouped test)
        // ungrouped  ... execution strategy (in ungrouped test)
        // 
        // TaskOutput
        // |
        // Scan
        // 
        // See #testComplex for all the bahaviors that are tested. Not all of them apply here.
        TestingScanOperatorFactory testingScanOperatorFactory = new TestingScanOperatorFactory(0, TABLE_SCAN_NODE_ID, ImmutableList.of(VARCHAR));
        TaskOutputOperatorFactory taskOutputOperatorFactory = new TaskOutputOperatorFactory(1, TABLE_SCAN_NODE_ID, outputBuffer, Function.identity(), new PagesSerdeFactory(new BlockEncodingManager(), false));
        LocalExecutionPlan localExecutionPlan = new LocalExecutionPlan(ImmutableList.of(new DriverFactory(0, true, true, ImmutableList.of(testingScanOperatorFactory, taskOutputOperatorFactory), OptionalInt.empty(), executionStrategy, Optional.empty())), ImmutableList.of(TABLE_SCAN_NODE_ID), executionStrategy == GROUPED_EXECUTION ? StageExecutionDescriptor.fixedLifespanScheduleGroupedExecution(ImmutableList.of(TABLE_SCAN_NODE_ID), 8) : StageExecutionDescriptor.ungroupedExecution());
        TaskContext taskContext = newTestingTaskContext(taskNotificationExecutor, driverYieldExecutor, taskStateMachine);
        SqlTaskExecution sqlTaskExecution = SqlTaskExecution.createSqlTaskExecution(taskStateMachine, taskContext, outputBuffer, ImmutableList.of(), localExecutionPlan, taskExecutor, taskNotificationExecutor, createTestSplitMonitor());
        // 
        // test body
        assertEquals(taskStateMachine.getState(), TaskState.RUNNING);
        switch(executionStrategy) {
            case UNGROUPED_EXECUTION:
                // add source for pipeline
                sqlTaskExecution.addSources(ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(newScheduledSplit(0, TABLE_SCAN_NODE_ID, Lifespan.taskWide(), 100000, 123)), false)));
                // assert that partial task result is produced
                outputBufferConsumer.consume(123, ASSERT_WAIT_TIMEOUT);
                // pause operator execution to make sure that
                // * operatorFactory will be closed even though operator can't execute
                // * completedDriverGroups will NOT include the newly scheduled driver group while pause is in place
                testingScanOperatorFactory.getPauser().pause();
                // add source for pipeline, mark as no more splits
                sqlTaskExecution.addSources(ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(newScheduledSplit(1, TABLE_SCAN_NODE_ID, Lifespan.taskWide(), 200000, 300), newScheduledSplit(2, TABLE_SCAN_NODE_ID, Lifespan.taskWide(), 300000, 200)), true)));
                // assert that pipeline will have no more drivers
                waitUntilEquals(testingScanOperatorFactory::isOverallNoMoreOperators, true, ASSERT_WAIT_TIMEOUT);
                // assert that no DriverGroup is fully completed
                assertEquals(taskContext.getCompletedDriverGroups(), ImmutableSet.of());
                // resume operator execution
                testingScanOperatorFactory.getPauser().resume();
                // assert that task result is produced
                outputBufferConsumer.consume(300 + 200, ASSERT_WAIT_TIMEOUT);
                outputBufferConsumer.assertBufferComplete(ASSERT_WAIT_TIMEOUT);
                break;
            case GROUPED_EXECUTION:
                // add source for pipeline (driver group [1, 5]), mark driver group [1] as noMoreSplits
                sqlTaskExecution.addSources(ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(newScheduledSplit(0, TABLE_SCAN_NODE_ID, Lifespan.driverGroup(1), 0, 1), newScheduledSplit(1, TABLE_SCAN_NODE_ID, Lifespan.driverGroup(5), 100000, 10)), ImmutableSet.of(Lifespan.driverGroup(1)), false)));
                // assert that pipeline will have no more drivers for driver group [1]
                waitUntilEquals(testingScanOperatorFactory::getDriverGroupsWithNoMoreOperators, ImmutableSet.of(Lifespan.driverGroup(1)), ASSERT_WAIT_TIMEOUT);
                // assert that partial result is produced for both driver groups
                outputBufferConsumer.consume(1 + 10, ASSERT_WAIT_TIMEOUT);
                // assert that driver group [1] is fully completed
                waitUntilEquals(taskContext::getCompletedDriverGroups, ImmutableSet.of(Lifespan.driverGroup(1)), ASSERT_WAIT_TIMEOUT);
                // pause operator execution to make sure that
                // * operatorFactory will be closed even though operator can't execute
                // * completedDriverGroups will NOT include the newly scheduled driver group while pause is in place
                testingScanOperatorFactory.getPauser().pause();
                // add source for pipeline (driver group [5]), mark driver group [5] as noMoreSplits
                sqlTaskExecution.addSources(ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(newScheduledSplit(2, TABLE_SCAN_NODE_ID, Lifespan.driverGroup(5), 200000, 300)), ImmutableSet.of(Lifespan.driverGroup(5)), false)));
                // assert that pipeline will have no more drivers for driver group [1, 5]
                waitUntilEquals(testingScanOperatorFactory::getDriverGroupsWithNoMoreOperators, ImmutableSet.of(Lifespan.driverGroup(1), Lifespan.driverGroup(5)), ASSERT_WAIT_TIMEOUT);
                // assert that driver group [5] is NOT YET fully completed
                assertEquals(taskContext.getCompletedDriverGroups(), ImmutableSet.of(Lifespan.driverGroup(1)));
                // resume operator execution
                testingScanOperatorFactory.getPauser().resume();
                // assert that partial result is produced
                outputBufferConsumer.consume(300, ASSERT_WAIT_TIMEOUT);
                // assert that driver group [1, 5] is fully completed
                waitUntilEquals(taskContext::getCompletedDriverGroups, ImmutableSet.of(Lifespan.driverGroup(1), Lifespan.driverGroup(5)), ASSERT_WAIT_TIMEOUT);
                // pause operator execution to make sure that
                // * operatorFactory will be closed even though operator can't execute
                // * completedDriverGroups will NOT include the newly scheduled driver group while pause is in place
                testingScanOperatorFactory.getPauser().pause();
                // add source for pipeline (driver group [7]), mark pipeline as noMoreSplits without explicitly marking driver group 7
                sqlTaskExecution.addSources(ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(newScheduledSplit(3, TABLE_SCAN_NODE_ID, Lifespan.driverGroup(7), 300000, 45), newScheduledSplit(4, TABLE_SCAN_NODE_ID, Lifespan.driverGroup(7), 400000, 54)), ImmutableSet.of(), true)));
                // assert that pipeline will have no more drivers for driver group [1, 5, 7]
                waitUntilEquals(testingScanOperatorFactory::getDriverGroupsWithNoMoreOperators, ImmutableSet.of(Lifespan.driverGroup(1), Lifespan.driverGroup(5), Lifespan.driverGroup(7)), ASSERT_WAIT_TIMEOUT);
                // assert that pipeline will have no more drivers
                waitUntilEquals(testingScanOperatorFactory::isOverallNoMoreOperators, true, ASSERT_WAIT_TIMEOUT);
                // assert that driver group [1, 5] is fully completed
                assertEquals(taskContext.getCompletedDriverGroups(), ImmutableSet.of(Lifespan.driverGroup(1), Lifespan.driverGroup(5)));
                // resume operator execution
                testingScanOperatorFactory.getPauser().resume();
                // assert driver group [7] is not completed before output buffer is consumed
                MILLISECONDS.sleep(1000);
                assertEquals(taskContext.getCompletedDriverGroups(), ImmutableSet.of(Lifespan.driverGroup(1), Lifespan.driverGroup(5)));
                // assert that result is produced
                outputBufferConsumer.consume(45 + 54, ASSERT_WAIT_TIMEOUT);
                outputBufferConsumer.assertBufferComplete(ASSERT_WAIT_TIMEOUT);
                // assert that driver group [1, 5, 7] is fully completed
                waitUntilEquals(taskContext::getCompletedDriverGroups, ImmutableSet.of(Lifespan.driverGroup(1), Lifespan.driverGroup(5), Lifespan.driverGroup(7)), ASSERT_WAIT_TIMEOUT);
                break;
            default:
                throw new UnsupportedOperationException();
        }
        // complete the task by calling abort on it
        outputBufferConsumer.abort();
        TaskState taskState = taskStateMachine.getStateChange(TaskState.RUNNING).get(10, SECONDS);
        assertEquals(taskState, TaskState.FINISHED);
    } finally {
        taskExecutor.stop();
        taskNotificationExecutor.shutdownNow();
        driverYieldExecutor.shutdown();
    }
}

Example 13 with TaskContext

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

the class SqlTask method createTaskStatus.

private TaskStatus createTaskStatus(TaskHolder taskHolder) {
    long taskStatusAgeInMilis = System.currentTimeMillis() - creationTimeInMillis;
    // Always return a new TaskInfo with a larger version number;
    // otherwise a client will not accept the update
    long versionNumber = nextTaskInfoVersion.getAndIncrement();
    TaskState state = taskStateMachine.getState();
    List<ExecutionFailureInfo> failures = ImmutableList.of();
    if (state == FAILED) {
        failures = toFailures(taskStateMachine.getFailureCauses());
    }
    int queuedPartitionedDrivers = 0;
    long queuedPartitionedSplitsWeight = 0L;
    int runningPartitionedDrivers = 0;
    long runningPartitionedSplitsWeight = 0L;
    long physicalWrittenDataSizeInBytes = 0L;
    long userMemoryReservationInBytes = 0L;
    long systemMemoryReservationInBytes = 0L;
    // TODO: add a mechanism to avoid sending the whole completedDriverGroups set over the wire for every task status reply
    Set<Lifespan> completedDriverGroups = ImmutableSet.of();
    long fullGcCount = 0;
    long fullGcTimeInMillis = 0L;
    long totalCpuTimeInNanos = 0L;
    if (taskHolder.getFinalTaskInfo() != null) {
        TaskStats taskStats = taskHolder.getFinalTaskInfo().getStats();
        queuedPartitionedDrivers = taskStats.getQueuedPartitionedDrivers();
        queuedPartitionedSplitsWeight = taskStats.getQueuedPartitionedSplitsWeight();
        runningPartitionedDrivers = taskStats.getRunningPartitionedDrivers();
        runningPartitionedSplitsWeight = taskStats.getRunningPartitionedSplitsWeight();
        physicalWrittenDataSizeInBytes = taskStats.getPhysicalWrittenDataSizeInBytes();
        userMemoryReservationInBytes = taskStats.getUserMemoryReservationInBytes();
        systemMemoryReservationInBytes = taskStats.getSystemMemoryReservationInBytes();
        fullGcCount = taskStats.getFullGcCount();
        fullGcTimeInMillis = taskStats.getFullGcTimeInMillis();
        totalCpuTimeInNanos = taskStats.getTotalCpuTimeInNanos();
    } else if (taskHolder.getTaskExecution() != null) {
        long physicalWrittenBytes = 0;
        TaskContext taskContext = taskHolder.getTaskExecution().getTaskContext();
        for (PipelineContext pipelineContext : taskContext.getPipelineContexts()) {
            PipelineStatus pipelineStatus = pipelineContext.getPipelineStatus();
            queuedPartitionedDrivers += pipelineStatus.getQueuedPartitionedDrivers();
            queuedPartitionedSplitsWeight += pipelineStatus.getQueuedPartitionedSplitsWeight();
            runningPartitionedDrivers += pipelineStatus.getRunningPartitionedDrivers();
            runningPartitionedSplitsWeight += pipelineStatus.getRunningPartitionedSplitsWeight();
            physicalWrittenBytes += pipelineContext.getPhysicalWrittenDataSize();
            totalCpuTimeInNanos += pipelineContext.getPipelineStats().getTotalCpuTimeInNanos();
        }
        physicalWrittenDataSizeInBytes = physicalWrittenBytes;
        userMemoryReservationInBytes = taskContext.getMemoryReservation().toBytes();
        systemMemoryReservationInBytes = taskContext.getSystemMemoryReservation().toBytes();
        completedDriverGroups = taskContext.getCompletedDriverGroups();
        fullGcCount = taskContext.getFullGcCount();
        fullGcTimeInMillis = taskContext.getFullGcTime().toMillis();
    }
    return new TaskStatus(taskInstanceId.getUuidLeastSignificantBits(), taskInstanceId.getUuidMostSignificantBits(), versionNumber, state, location, completedDriverGroups, failures, queuedPartitionedDrivers, runningPartitionedDrivers, outputBuffer.getUtilization(), isOutputBufferOverutilized(), physicalWrittenDataSizeInBytes, userMemoryReservationInBytes, systemMemoryReservationInBytes, queryContext.getPeakNodeTotalMemory(), fullGcCount, fullGcTimeInMillis, totalCpuTimeInNanos, taskStatusAgeInMilis, queuedPartitionedSplitsWeight, runningPartitionedSplitsWeight);
}

Example 14 with TaskContext

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

the class TestMemoryRevokingScheduler method testMemoryPoolRevoking.

@Test
public void testMemoryPoolRevoking() throws Exception {
    QueryContext q1 = getOrCreateQueryContext(new QueryId("q1"), memoryPool);
    QueryContext q2 = getOrCreateQueryContext(new QueryId("q2"), memoryPool);
    SqlTask sqlTask1 = newSqlTask(q1.getQueryId(), memoryPool);
    SqlTask sqlTask2 = newSqlTask(q2.getQueryId(), memoryPool);
    TaskContext taskContext1 = getOrCreateTaskContext(sqlTask1);
    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 = getOrCreateTaskContext(sqlTask2);
    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");
    List<SqlTask> tasks = ImmutableList.of(sqlTask1, sqlTask2);
    MemoryRevokingScheduler scheduler = new MemoryRevokingScheduler(singletonList(memoryPool), () -> tasks, queryContexts::get, 1.0, 1.0, ORDER_BY_CREATE_TIME, false);
    try {
        scheduler.start();
        allOperatorContexts = ImmutableSet.of(operatorContext1, operatorContext2, operatorContext3, operatorContext4, operatorContext5);
        assertMemoryRevokingNotRequested();
        assertEquals(10, memoryPool.getFreeBytes());
        scheduler.awaitAsynchronousCallbacksRun();
        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());
        scheduler.awaitAsynchronousCallbacksRun();
        // we are still good - no revoking needed
        assertMemoryRevokingNotRequested();
        revocableMemory4.setBytes(7);
        assertEquals(-6, memoryPool.getFreeBytes());
        scheduler.awaitAsynchronousCallbacksRun();
        // we need to revoke 3 and 6
        assertMemoryRevokingRequestedFor(operatorContext1, operatorContext3);
        // lets revoke some bytes
        revocableMemory1.setBytes(0);
        operatorContext1.resetMemoryRevokingRequested();
        scheduler.awaitAsynchronousCallbacksRun();
        assertMemoryRevokingRequestedFor(operatorContext3);
        assertEquals(-3, memoryPool.getFreeBytes());
        // and allocate some more
        revocableMemory5.setBytes(3);
        assertEquals(-6, memoryPool.getFreeBytes());
        scheduler.awaitAsynchronousCallbacksRun();
        // we are still good with just OC3 in process of revoking
        assertMemoryRevokingRequestedFor(operatorContext3);
        // and allocate some more
        revocableMemory5.setBytes(4);
        assertEquals(-7, memoryPool.getFreeBytes());
        scheduler.awaitAsynchronousCallbacksRun();
        // now we have to trigger revoking for OC4
        assertMemoryRevokingRequestedFor(operatorContext3, operatorContext4);
    } finally {
        scheduler.stop();
    }
}

Example 15 with TaskContext

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

the class TestMemoryRevokingScheduler method testRevokesPoolWhenFullBeforeQueryLimit.

@Test
public void testRevokesPoolWhenFullBeforeQueryLimit() throws Exception {
    QueryContext q1 = getOrCreateQueryContext(new QueryId("q1"), memoryPool);
    QueryContext q2 = getOrCreateQueryContext(new QueryId("q2"), memoryPool);
    SqlTask sqlTask1 = newSqlTask(q1.getQueryId(), memoryPool);
    SqlTask sqlTask2 = newSqlTask(q2.getQueryId(), memoryPool);
    TaskContext taskContext1 = getOrCreateTaskContext(sqlTask1);
    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 = getOrCreateTaskContext(sqlTask2);
    PipelineContext pipelineContext21 = taskContext2.addPipelineContext(1, false, false, false);
    DriverContext driverContext211 = pipelineContext21.addDriverContext();
    OperatorContext operatorContext4 = driverContext211.addOperatorContext(4, new PlanNodeId("na"), "na");
    List<SqlTask> tasks = ImmutableList.of(sqlTask1, sqlTask2);
    MemoryRevokingScheduler scheduler = new MemoryRevokingScheduler(singletonList(memoryPool), () -> tasks, queryContexts::get, 1.0, 1.0, ORDER_BY_CREATE_TIME, true);
    try {
        scheduler.start();
        allOperatorContexts = ImmutableSet.of(operatorContext1, operatorContext2, operatorContext3, operatorContext4);
        assertMemoryRevokingNotRequested();
        assertEquals(10, memoryPool.getFreeBytes());
        scheduler.awaitAsynchronousCallbacksRun();
        assertMemoryRevokingNotRequested();
        LocalMemoryContext revocableMemory1 = operatorContext1.localRevocableMemoryContext();
        LocalMemoryContext revocableMemory3 = operatorContext3.localRevocableMemoryContext();
        LocalMemoryContext revocableMemory4 = operatorContext4.localRevocableMemoryContext();
        revocableMemory1.setBytes(3);
        revocableMemory3.setBytes(6);
        assertEquals(1, memoryPool.getFreeBytes());
        scheduler.awaitAsynchronousCallbacksRun();
        // we are still good - no revoking needed
        assertMemoryRevokingNotRequested();
        revocableMemory4.setBytes(7);
        assertEquals(-6, memoryPool.getFreeBytes());
        scheduler.awaitAsynchronousCallbacksRun();
        // we need to revoke 3 and 6
        assertMemoryRevokingRequestedFor(operatorContext1, operatorContext3);
    } finally {
        scheduler.stop();
    }
}