Examples with DriverContext io.trino.operator.DriverContext used on opensource projects

Example 16 with DriverContext

use of io.trino.operator.DriverContext in project trino by trinodb.

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();
    // get the split for this table
    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, columnHandles, DynamicFilter.EMPTY);
            return new PageSourceOperator(pageSource, operatorContext);
        }

        @Override
        public void noMoreOperators() {
        }

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

Example 17 with DriverContext

use of io.trino.operator.DriverContext in project trino by trinodb.

the class AbstractSimpleOperatorBenchmark method createDrivers.

@Override
protected List<Driver> createDrivers(TaskContext taskContext) {
    DriverFactory driverFactory = createDriverFactory();
    DriverContext driverContext = taskContext.addPipelineContext(0, true, true, false).addDriverContext();
    Driver driver = driverFactory.createDriver(driverContext);
    return ImmutableList.of(driver);
}

Example 18 with DriverContext

use of io.trino.operator.DriverContext in project trino by trinodb.

the class SqlTaskExecution method enqueueDriverSplitRunner.

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

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

            @Override
            public void onFailure(Throwable cause) {
                try (SetThreadName ignored = new SetThreadName("Task-%s", taskId)) {
                    taskStateMachine.failed(cause);
                    // record driver is finished
                    status.decrementRemainingDriver(splitRunner.getLifespan());
                    // fire failed event with cause
                    splitMonitor.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 19 with DriverContext

use of io.trino.operator.DriverContext in project trino by trinodb.

the class TestSpatialJoinOperator method testDistributedSpatialJoin.

@Test
public void testDistributedSpatialJoin() {
    TaskContext taskContext = createTaskContext();
    DriverContext driverContext = taskContext.addPipelineContext(0, true, true, true).addDriverContext();
    RowPagesBuilder buildPages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR, INTEGER)).row(POLYGON_A, "A", 1).row(POLYGON_A, "A", 2).row(null, "null", null).pageBreak().row(POLYGON_B, "B", 0).row(POLYGON_B, "B", 2);
    RowPagesBuilder probePages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR, INTEGER)).row(POINT_X, "x", 2).row(null, "null", null).row(POINT_Y, "y", 2).pageBreak().row(POINT_Z, "z", 0);
    MaterializedResult expected = resultBuilder(taskContext.getSession(), ImmutableList.of(VARCHAR, VARCHAR)).row("x", "A").row("y", "A").row("y", "B").row("z", "B").build();
    PagesSpatialIndexFactory pagesSpatialIndexFactory = buildIndex(driverContext, (build, probe, r) -> build.contains(probe), Optional.empty(), Optional.of(2), Optional.of(KDB_TREE_JSON), Optional.empty(), buildPages);
    OperatorFactory joinOperatorFactory = new SpatialJoinOperatorFactory(2, new PlanNodeId("test"), INNER, probePages.getTypes(), Ints.asList(1), 0, Optional.of(2), pagesSpatialIndexFactory);
    assertOperatorEquals(joinOperatorFactory, driverContext, probePages.build(), expected);
}

Example 20 with DriverContext

use of io.trino.operator.DriverContext in project trino by trinodb.

the class TestSpatialJoinOperator method testDistributedSpatialSelfJoin.

@Test
public void testDistributedSpatialSelfJoin() {
    TaskContext taskContext = createTaskContext();
    DriverContext driverContext = taskContext.addPipelineContext(0, true, true, true).addDriverContext();
    RowPagesBuilder pages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR, INTEGER)).row(POLYGON_A, "A", 1).row(POLYGON_A, "A", 2).row(null, "null", null).pageBreak().row(POLYGON_B, "B", 0).row(POLYGON_B, "B", 2);
    MaterializedResult expected = resultBuilder(taskContext.getSession(), ImmutableList.of(VARCHAR, VARCHAR)).row("A", "A").row("A", "B").row("B", "A").row("B", "B").build();
    PagesSpatialIndexFactory pagesSpatialIndexFactory = buildIndex(driverContext, (build, probe, r) -> build.intersects(probe), Optional.empty(), Optional.of(2), Optional.of(KDB_TREE_JSON), Optional.empty(), pages);
    OperatorFactory joinOperatorFactory = new SpatialJoinOperatorFactory(2, new PlanNodeId("test"), INNER, pages.getTypes(), Ints.asList(1), 0, Optional.of(2), pagesSpatialIndexFactory);
    assertOperatorEquals(joinOperatorFactory, driverContext, pages.build(), expected);
}