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

Example 11 with DriverContext

use of com.facebook.presto.operator.DriverContext 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 12 with DriverContext

use of com.facebook.presto.operator.DriverContext 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();
    }
}

Example 13 with DriverContext

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

the class TestOrcPageSourceMemoryTracking method testScanFilterAndProjectOperator.

@Test
public void testScanFilterAndProjectOperator() throws Exception {
    // Numbers used in assertions in this test may change when implementation is modified,
    // feel free to change them if they break in the future
    DriverContext driverContext = testPreparer.newDriverContext();
    SourceOperator operator = testPreparer.newScanFilterAndProjectOperator(driverContext);
    assertEquals(driverContext.getSystemMemoryUsage(), 0);
    for (int i = 0; i < 52; i++) {
        assertFalse(operator.isFinished());
        operator.getOutput();
        assertBetweenInclusive(driverContext.getSystemMemoryUsage(), 550000L, 639999L);
    }
    for (int i = 52; i < 65; i++) {
        assertFalse(operator.isFinished());
        operator.getOutput();
        assertBetweenInclusive(driverContext.getSystemMemoryUsage(), 450000L, 539999L);
    }
    // Page source is over, but data still exist in buffer of ScanFilterProjectOperator
    assertFalse(operator.isFinished());
    assertNull(operator.getOutput());
    assertBetweenInclusive(driverContext.getSystemMemoryUsage(), 100000L, 109999L);
    assertFalse(operator.isFinished());
    Page lastPage = operator.getOutput();
    assertNotNull(lastPage);
    // No data is left
    assertTrue(operator.isFinished());
    // an empty page builder of two variable width block builders is left in ScanFilterAndProjectOperator
    PageBuilder pageBuilder = new PageBuilder(ImmutableList.of(createUnboundedVarcharType(), createUnboundedVarcharType()));
    for (int i = 0; i < lastPage.getPositionCount(); i++) {
        pageBuilder.declarePosition();
        createUnboundedVarcharType().appendTo(lastPage.getBlock(0), i, pageBuilder.getBlockBuilder(0));
        createUnboundedVarcharType().appendTo(lastPage.getBlock(1), i, pageBuilder.getBlockBuilder(1));
    }
    pageBuilder.reset();
    assertEquals(driverContext.getSystemMemoryUsage(), pageBuilder.getRetainedSizeInBytes());
}

Example 14 with DriverContext

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

the class SqlTaskExecution method enqueueDrivers.

private synchronized void enqueueDrivers(boolean forceRunSplit, List<DriverSplitRunner> runners) {
    // schedule driver to be executed
    List<ListenableFuture<?>> finishedFutures = taskExecutor.enqueueSplits(taskHandle, forceRunSplit, runners);
    checkState(finishedFutures.size() == runners.size(), "Expected %s futures but got %s", runners.size(), finishedFutures.size());
    // record new driver
    remainingDrivers.addAndGet(finishedFutures.size());
    // when driver completes, update state and fire events
    for (int i = 0; i < finishedFutures.size(); i++) {
        ListenableFuture<?> finishedFuture = finishedFutures.get(i);
        final DriverSplitRunner splitRunner = runners.get(i);
        Futures.addCallback(finishedFuture, new FutureCallback<Object>() {

            @Override
            public void onSuccess(Object result) {
                try (SetThreadName ignored = new SetThreadName("Task-%s", taskId)) {
                    // record driver is finished
                    remainingDrivers.decrementAndGet();
                    checkTaskCompletion();
                    queryMonitor.splitCompletedEvent(taskId, getDriverStats());
                }
            }

            @Override
            public void onFailure(Throwable cause) {
                try (SetThreadName ignored = new SetThreadName("Task-%s", taskId)) {
                    taskStateMachine.failed(cause);
                    // record driver is finished
                    remainingDrivers.decrementAndGet();
                    // fire failed event with cause
                    queryMonitor.splitFailedEvent(taskId, getDriverStats(), cause);
                }
            }

            private DriverStats getDriverStats() {
                DriverContext driverContext = splitRunner.getDriverContext();
                DriverStats driverStats;
                if (driverContext != null) {
                    driverStats = driverContext.getDriverStats();
                } else {
                    // split runner did not start successfully
                    driverStats = new DriverStats();
                }
                return driverStats;
            }
        }, notificationExecutor);
    }
}

Example 15 with DriverContext

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

the class AbstractOperatorBenchmark method createTableScanOperator.

protected final OperatorFactory createTableScanOperator(int operatorId, PlanNodeId planNodeId, String tableName, String... columnNames) {
    checkArgument(session.getCatalog().isPresent(), "catalog not set");
    checkArgument(session.getSchema().isPresent(), "schema not set");
    // look up the table
    Metadata metadata = localQueryRunner.getMetadata();
    QualifiedObjectName qualifiedTableName = new QualifiedObjectName(session.getCatalog().get(), session.getSchema().get(), tableName);
    TableHandle tableHandle = metadata.getTableHandle(session, qualifiedTableName).orElse(null);
    checkArgument(tableHandle != null, "Table %s does not exist", qualifiedTableName);
    // lookup the columns
    Map<String, ColumnHandle> allColumnHandles = metadata.getColumnHandles(session, tableHandle);
    ImmutableList.Builder<ColumnHandle> columnHandlesBuilder = ImmutableList.builder();
    for (String columnName : columnNames) {
        ColumnHandle columnHandle = allColumnHandles.get(columnName);
        checkArgument(columnHandle != null, "Table %s does not have a column %s", tableName, columnName);
        columnHandlesBuilder.add(columnHandle);
    }
    List<ColumnHandle> columnHandles = columnHandlesBuilder.build();
    Split split = getLocalQuerySplit(session, tableHandle);
    return new OperatorFactory() {

        @Override
        public Operator createOperator(DriverContext driverContext) {
            OperatorContext operatorContext = driverContext.addOperatorContext(operatorId, planNodeId, "BenchmarkSource");
            ConnectorPageSource pageSource = localQueryRunner.getPageSourceManager().createPageSource(session, split, tableHandle.withDynamicFilter(TupleDomain::all), columnHandles);
            return new PageSourceOperator(pageSource, operatorContext);
        }

        @Override
        public void noMoreOperators() {
        }

        @Override
        public OperatorFactory duplicate() {
            throw new UnsupportedOperationException();
        }
    };
}