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

Example 26 with DriverContext

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

the class TestHashJoinOperator method testInnerJoinWithBlockingLookupSourceAndEmptyProbe.

@Test(dataProvider = "hashJoinTestValues")
public void testInnerJoinWithBlockingLookupSourceAndEmptyProbe(boolean parallelBuild, boolean probeHashEnabled, boolean buildHashEnabled) throws Exception {
    // join that waits for build side to be collected
    TaskContext taskContext = createTaskContext();
    OperatorFactory joinOperatorFactory = createJoinOperatorFactoryWithBlockingLookupSource(taskContext, parallelBuild, probeHashEnabled, buildHashEnabled, true);
    DriverContext driverContext = taskContext.addPipelineContext(0, true, true, false).addDriverContext();
    try (Operator joinOperator = joinOperatorFactory.createOperator(driverContext)) {
        joinOperatorFactory.noMoreOperators();
        assertFalse(joinOperator.needsInput());
        joinOperator.finish();
        assertNull(joinOperator.getOutput());
        // lookup join operator got blocked waiting for build side
        assertFalse(joinOperator.isBlocked().isDone());
        assertFalse(joinOperator.isFinished());
    }
    // join that doesn't wait for build side to be collected
    taskContext = createTaskContext();
    joinOperatorFactory = createJoinOperatorFactoryWithBlockingLookupSource(taskContext, parallelBuild, probeHashEnabled, buildHashEnabled, false);
    driverContext = taskContext.addPipelineContext(0, true, true, false).addDriverContext();
    try (Operator joinOperator = joinOperatorFactory.createOperator(driverContext)) {
        joinOperatorFactory.noMoreOperators();
        assertTrue(joinOperator.needsInput());
        joinOperator.finish();
        assertNull(joinOperator.getOutput());
        // lookup join operator will yield once before finishing
        assertNull(joinOperator.getOutput());
        assertTrue(joinOperator.isBlocked().isDone());
        assertTrue(joinOperator.isFinished());
    }
}

Example 27 with DriverContext

use of io.trino.operator.DriverContext 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 28 with DriverContext

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

the class TestNestedLoopBuildOperator method testNestedLoopBuildNoBlock.

@Test
public void testNestedLoopBuildNoBlock() throws Exception {
    TaskContext taskContext = createTaskContext();
    List<Type> buildTypes = ImmutableList.of();
    JoinBridgeManager<NestedLoopJoinBridge> nestedLoopJoinBridgeManager = new JoinBridgeManager<>(false, PipelineExecutionStrategy.UNGROUPED_EXECUTION, PipelineExecutionStrategy.UNGROUPED_EXECUTION, lifespan -> new NestedLoopJoinPagesSupplier(), buildTypes);
    NestedLoopBuildOperatorFactory nestedLoopBuildOperatorFactory = new NestedLoopBuildOperatorFactory(3, new PlanNodeId("test"), nestedLoopJoinBridgeManager);
    DriverContext driverContext = taskContext.addPipelineContext(0, true, true, false).addDriverContext();
    NestedLoopBuildOperator nestedLoopBuildOperator = (NestedLoopBuildOperator) nestedLoopBuildOperatorFactory.createOperator(driverContext);
    NestedLoopJoinBridge nestedLoopJoinBridge = nestedLoopJoinBridgeManager.getJoinBridge(Lifespan.taskWide());
    assertFalse(nestedLoopJoinBridge.getPagesFuture().isDone());
    // build pages
    Page buildPage1 = new Page(3);
    Page buildPageEmpty = new Page(0);
    Page buildPage2 = new Page(3000);
    nestedLoopBuildOperator.addInput(buildPage1);
    nestedLoopBuildOperator.addInput(buildPageEmpty);
    nestedLoopBuildOperator.addInput(buildPage2);
    nestedLoopBuildOperator.finish();
    assertTrue(nestedLoopJoinBridge.getPagesFuture().isDone());
    List<Page> buildPages = nestedLoopJoinBridge.getPagesFuture().get().getPages();
    assertEquals(buildPages.size(), 1);
    assertEquals(buildPages.get(0).getPositionCount(), 3003);
}

Example 29 with DriverContext

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

the class HashJoinBenchmark method createDrivers.

/*
    select orderkey, quantity, totalprice
    from lineitem join orders using (orderkey)
     */
@Override
protected List<Driver> createDrivers(TaskContext taskContext) {
    if (probeDriverFactory == null) {
        List<Type> ordersTypes = getColumnTypes("orders", "orderkey", "totalprice");
        OperatorFactory ordersTableScan = createTableScanOperator(0, new PlanNodeId("test"), "orders", "orderkey", "totalprice");
        BlockTypeOperators blockTypeOperators = new BlockTypeOperators(new TypeOperators());
        JoinBridgeManager<PartitionedLookupSourceFactory> lookupSourceFactoryManager = JoinBridgeManager.lookupAllAtOnce(new PartitionedLookupSourceFactory(ordersTypes, ImmutableList.of(0, 1).stream().map(ordersTypes::get).collect(toImmutableList()), Ints.asList(0).stream().map(ordersTypes::get).collect(toImmutableList()), 1, false, blockTypeOperators));
        HashBuilderOperatorFactory hashBuilder = new HashBuilderOperatorFactory(1, new PlanNodeId("test"), lookupSourceFactoryManager, ImmutableList.of(0, 1), Ints.asList(0), OptionalInt.empty(), Optional.empty(), Optional.empty(), ImmutableList.of(), 1_500_000, new PagesIndex.TestingFactory(false), false, SingleStreamSpillerFactory.unsupportedSingleStreamSpillerFactory(), incrementalLoadFactorHashArraySizeSupplier(session));
        DriverContext driverContext = taskContext.addPipelineContext(0, false, false, false).addDriverContext();
        DriverFactory buildDriverFactory = new DriverFactory(0, false, false, ImmutableList.of(ordersTableScan, hashBuilder), OptionalInt.empty(), UNGROUPED_EXECUTION);
        List<Type> lineItemTypes = getColumnTypes("lineitem", "orderkey", "quantity");
        OperatorFactory lineItemTableScan = createTableScanOperator(0, new PlanNodeId("test"), "lineitem", "orderkey", "quantity");
        OperatorFactory joinOperator = operatorFactories.innerJoin(1, new PlanNodeId("test"), lookupSourceFactoryManager, false, false, false, lineItemTypes, Ints.asList(0), OptionalInt.empty(), Optional.empty(), OptionalInt.empty(), unsupportedPartitioningSpillerFactory(), blockTypeOperators);
        NullOutputOperatorFactory output = new NullOutputOperatorFactory(2, new PlanNodeId("test"));
        this.probeDriverFactory = new DriverFactory(1, true, true, ImmutableList.of(lineItemTableScan, joinOperator, output), OptionalInt.empty(), UNGROUPED_EXECUTION);
        Driver driver = buildDriverFactory.createDriver(driverContext);
        Future<LookupSourceProvider> lookupSourceProvider = lookupSourceFactoryManager.getJoinBridge(Lifespan.taskWide()).createLookupSourceProvider();
        while (!lookupSourceProvider.isDone()) {
            driver.process();
        }
        getFutureValue(lookupSourceProvider).close();
    }
    DriverContext driverContext = taskContext.addPipelineContext(1, true, true, false).addDriverContext();
    Driver driver = probeDriverFactory.createDriver(driverContext);
    return ImmutableList.of(driver);
}