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

Example 11 with Driver

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

the class LocalQueryRunner method createDrivers.

private List<Driver> createDrivers(Session session, Plan plan, OutputFactory outputFactory, TaskContext taskContext) {
    if (printPlan) {
        System.out.println(PlanPrinter.textLogicalPlan(plan.getRoot(), plan.getTypes(), plannerContext.getMetadata(), plannerContext.getFunctionManager(), plan.getStatsAndCosts(), session, 0, false));
    }
    SubPlan subplan = createSubPlans(session, plan, true);
    if (!subplan.getChildren().isEmpty()) {
        throw new AssertionError("Expected subplan to have no children");
    }
    TableExecuteContextManager tableExecuteContextManager = new TableExecuteContextManager();
    tableExecuteContextManager.registerTableExecuteContextForQuery(taskContext.getQueryContext().getQueryId());
    LocalExecutionPlanner executionPlanner = new LocalExecutionPlanner(plannerContext, new TypeAnalyzer(plannerContext, statementAnalyzerFactory), Optional.empty(), pageSourceManager, indexManager, nodePartitioningManager, pageSinkManager, null, expressionCompiler, pageFunctionCompiler, joinFilterFunctionCompiler, new IndexJoinLookupStats(), this.taskManagerConfig, spillerFactory, singleStreamSpillerFactory, partitioningSpillerFactory, new PagesIndex.TestingFactory(false), joinCompiler, operatorFactories, new OrderingCompiler(plannerContext.getTypeOperators()), new DynamicFilterConfig(), blockTypeOperators, tableExecuteContextManager, exchangeManagerRegistry);
    // plan query
    StageExecutionDescriptor stageExecutionDescriptor = subplan.getFragment().getStageExecutionDescriptor();
    LocalExecutionPlan localExecutionPlan = executionPlanner.plan(taskContext, stageExecutionDescriptor, subplan.getFragment().getRoot(), subplan.getFragment().getPartitioningScheme().getOutputLayout(), plan.getTypes(), subplan.getFragment().getPartitionedSources(), outputFactory);
    // generate splitAssignments
    List<SplitAssignment> splitAssignments = new ArrayList<>();
    long sequenceId = 0;
    for (TableScanNode tableScan : findTableScanNodes(subplan.getFragment().getRoot())) {
        TableHandle table = tableScan.getTable();
        SplitSource splitSource = splitManager.getSplits(session, table, stageExecutionDescriptor.isScanGroupedExecution(tableScan.getId()) ? GROUPED_SCHEDULING : UNGROUPED_SCHEDULING, EMPTY, alwaysTrue());
        ImmutableSet.Builder<ScheduledSplit> scheduledSplits = ImmutableSet.builder();
        while (!splitSource.isFinished()) {
            for (Split split : getNextBatch(splitSource)) {
                scheduledSplits.add(new ScheduledSplit(sequenceId++, tableScan.getId(), split));
            }
        }
        splitAssignments.add(new SplitAssignment(tableScan.getId(), scheduledSplits.build(), true));
    }
    // create drivers
    List<Driver> drivers = new ArrayList<>();
    Map<PlanNodeId, DriverFactory> driverFactoriesBySource = new HashMap<>();
    for (DriverFactory driverFactory : localExecutionPlan.getDriverFactories()) {
        for (int i = 0; i < driverFactory.getDriverInstances().orElse(1); i++) {
            if (driverFactory.getSourceId().isPresent()) {
                checkState(driverFactoriesBySource.put(driverFactory.getSourceId().get(), driverFactory) == null);
            } else {
                DriverContext driverContext = taskContext.addPipelineContext(driverFactory.getPipelineId(), driverFactory.isInputDriver(), driverFactory.isOutputDriver(), false).addDriverContext();
                Driver driver = driverFactory.createDriver(driverContext);
                drivers.add(driver);
            }
        }
    }
    // add split assignments to the drivers
    ImmutableSet<PlanNodeId> partitionedSources = ImmutableSet.copyOf(subplan.getFragment().getPartitionedSources());
    for (SplitAssignment splitAssignment : splitAssignments) {
        DriverFactory driverFactory = driverFactoriesBySource.get(splitAssignment.getPlanNodeId());
        checkState(driverFactory != null);
        boolean partitioned = partitionedSources.contains(driverFactory.getSourceId().get());
        for (ScheduledSplit split : splitAssignment.getSplits()) {
            DriverContext driverContext = taskContext.addPipelineContext(driverFactory.getPipelineId(), driverFactory.isInputDriver(), driverFactory.isOutputDriver(), partitioned).addDriverContext();
            Driver driver = driverFactory.createDriver(driverContext);
            driver.updateSplitAssignment(new SplitAssignment(split.getPlanNodeId(), ImmutableSet.of(split), true));
            drivers.add(driver);
        }
    }
    for (DriverFactory driverFactory : localExecutionPlan.getDriverFactories()) {
        driverFactory.noMoreDrivers();
    }
    return ImmutableList.copyOf(drivers);
}

Example 12 with Driver

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

the class JoinTestUtils method buildLookupSource.

public static void buildLookupSource(ExecutorService executor, BuildSideSetup buildSideSetup) {
    requireNonNull(buildSideSetup, "buildSideSetup is null");
    LookupSourceFactory lookupSourceFactory = buildSideSetup.getLookupSourceFactoryManager().getJoinBridge(Lifespan.taskWide());
    Future<LookupSourceProvider> lookupSourceProvider = lookupSourceFactory.createLookupSourceProvider();
    List<Driver> buildDrivers = buildSideSetup.getBuildDrivers();
    while (!lookupSourceProvider.isDone()) {
        for (Driver buildDriver : buildDrivers) {
            buildDriver.process();
        }
    }
    getFutureValue(lookupSourceProvider).close();
    for (Driver buildDriver : buildDrivers) {
        runDriverInThread(executor, buildDriver);
    }
}

Example 13 with Driver

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

the class JoinTestUtils method setupBuildSide.

public static BuildSideSetup setupBuildSide(NodePartitioningManager nodePartitioningManager, boolean parallelBuild, TaskContext taskContext, RowPagesBuilder buildPages, Optional<InternalJoinFilterFunction> filterFunction, boolean spillEnabled, SingleStreamSpillerFactory singleStreamSpillerFactory) {
    Optional<JoinFilterFunctionCompiler.JoinFilterFunctionFactory> filterFunctionFactory = filterFunction.map(function -> (session, addresses, pages) -> new StandardJoinFilterFunction(function, addresses, pages));
    int partitionCount = parallelBuild ? PARTITION_COUNT : 1;
    List<Integer> hashChannels = buildPages.getHashChannels().orElseThrow();
    LocalExchange.LocalExchangeFactory localExchangeFactory = new LocalExchange.LocalExchangeFactory(nodePartitioningManager, taskContext.getSession(), FIXED_HASH_DISTRIBUTION, partitionCount, buildPages.getTypes(), hashChannels, buildPages.getHashChannel(), UNGROUPED_EXECUTION, DataSize.of(32, DataSize.Unit.MEGABYTE), TYPE_OPERATOR_FACTORY);
    LocalExchange.LocalExchangeSinkFactoryId localExchangeSinkFactoryId = localExchangeFactory.newSinkFactoryId();
    localExchangeFactory.noMoreSinkFactories();
    // collect input data into the partitioned exchange
    DriverContext collectDriverContext = taskContext.addPipelineContext(0, true, true, false).addDriverContext();
    ValuesOperator.ValuesOperatorFactory valuesOperatorFactory = new ValuesOperator.ValuesOperatorFactory(0, new PlanNodeId("values"), buildPages.build());
    LocalExchangeSinkOperator.LocalExchangeSinkOperatorFactory sinkOperatorFactory = new LocalExchangeSinkOperator.LocalExchangeSinkOperatorFactory(localExchangeFactory, 1, new PlanNodeId("sink"), localExchangeSinkFactoryId, Function.identity());
    Driver sourceDriver = Driver.createDriver(collectDriverContext, valuesOperatorFactory.createOperator(collectDriverContext), sinkOperatorFactory.createOperator(collectDriverContext));
    valuesOperatorFactory.noMoreOperators();
    sinkOperatorFactory.noMoreOperators();
    while (!sourceDriver.isFinished()) {
        sourceDriver.process();
    }
    // build side operator factories
    LocalExchangeSourceOperatorFactory sourceOperatorFactory = new LocalExchangeSourceOperatorFactory(0, new PlanNodeId("source"), localExchangeFactory);
    JoinBridgeManager<PartitionedLookupSourceFactory> lookupSourceFactoryManager = JoinBridgeManager.lookupAllAtOnce(new PartitionedLookupSourceFactory(buildPages.getTypes(), rangeList(buildPages.getTypes().size()).stream().map(buildPages.getTypes()::get).collect(toImmutableList()), hashChannels.stream().map(buildPages.getTypes()::get).collect(toImmutableList()), partitionCount, false, TYPE_OPERATOR_FACTORY));
    HashBuilderOperatorFactory buildOperatorFactory = new HashBuilderOperatorFactory(1, new PlanNodeId("build"), lookupSourceFactoryManager, rangeList(buildPages.getTypes().size()), hashChannels, buildPages.getHashChannel().map(OptionalInt::of).orElse(OptionalInt.empty()), filterFunctionFactory, Optional.empty(), ImmutableList.of(), 100, new PagesIndex.TestingFactory(false), spillEnabled, singleStreamSpillerFactory, incrementalLoadFactorHashArraySizeSupplier(taskContext.getSession()));
    return new BuildSideSetup(lookupSourceFactoryManager, buildOperatorFactory, sourceOperatorFactory, partitionCount);
}

Example 14 with Driver

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

the class TestNestedLoopJoinOperator method newJoinOperatorFactoryWithCompletedBuild.

private static NestedLoopJoinOperatorFactory newJoinOperatorFactoryWithCompletedBuild(TaskContext taskContext, RowPagesBuilder buildPages, List<Integer> probeChannels, List<Integer> buildChannels) {
    DriverContext driverContext = taskContext.addPipelineContext(0, true, true, false).addDriverContext();
    ValuesOperatorFactory valuesOperatorFactory = new ValuesOperatorFactory(0, new PlanNodeId("test"), buildPages.build());
    JoinBridgeManager<NestedLoopJoinBridge> nestedLoopJoinBridgeManager = new JoinBridgeManager<>(false, PipelineExecutionStrategy.UNGROUPED_EXECUTION, PipelineExecutionStrategy.UNGROUPED_EXECUTION, lifespan -> new NestedLoopJoinPagesSupplier(), buildPages.getTypes());
    NestedLoopBuildOperatorFactory nestedLoopBuildOperatorFactory = new NestedLoopBuildOperatorFactory(1, new PlanNodeId("test"), nestedLoopJoinBridgeManager);
    NestedLoopJoinOperatorFactory joinOperatorFactory = new NestedLoopJoinOperatorFactory(3, new PlanNodeId("test"), nestedLoopJoinBridgeManager, probeChannels, buildChannels);
    Operator valuesOperator = valuesOperatorFactory.createOperator(driverContext);
    Operator nestedLoopBuildOperator = nestedLoopBuildOperatorFactory.createOperator(driverContext);
    Driver driver = Driver.createDriver(driverContext, valuesOperator, nestedLoopBuildOperator);
    valuesOperatorFactory.noMoreOperators();
    nestedLoopBuildOperatorFactory.noMoreOperators();
    while (nestedLoopBuildOperator.isBlocked().isDone()) {
        driver.process();
    }
    return joinOperatorFactory;
}

Example 15 with Driver

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

the class IndexLoader method streamIndexDataForSingleKey.

public IndexedData streamIndexDataForSingleKey(UpdateRequest updateRequest) {
    Page indexKeyTuple = updateRequest.getPage().getRegion(0, 1);
    PageBuffer pageBuffer = new PageBuffer(100);
    DriverFactory driverFactory = indexBuildDriverFactoryProvider.createStreaming(pageBuffer, indexKeyTuple);
    Driver driver = driverFactory.createDriver(pipelineContext.addDriverContext());
    PageRecordSet pageRecordSet = new PageRecordSet(keyTypes, indexKeyTuple);
    PlanNodeId planNodeId = driverFactory.getSourceId().get();
    ScheduledSplit split = new ScheduledSplit(0, planNodeId, new Split(INDEX_CONNECTOR_ID, new IndexSplit(pageRecordSet), Lifespan.taskWide()));
    driver.updateSplitAssignment(new SplitAssignment(planNodeId, ImmutableSet.of(split), true));
    return new StreamingIndexedData(outputTypes, keyEqualOperators, indexKeyTuple, pageBuffer, driver);
}