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

Example 11 with DriverContext

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

the class BenchmarkHashBuildAndJoinOperators method benchmarkJoinHash.

@Benchmark
public List<Page> benchmarkJoinHash(JoinContext joinContext) throws Exception {
    DriverContext driverContext = joinContext.createTaskContext().addPipelineContext(0, true, true, false).addDriverContext();
    Operator joinOperator = joinContext.getJoinOperatorFactory().createOperator(driverContext);
    Iterator<Page> input = joinContext.getProbePages().iterator();
    ImmutableList.Builder<Page> outputPages = ImmutableList.builder();
    boolean finishing = false;
    for (int loops = 0; !joinOperator.isFinished() && loops < 1_000_000; loops++) {
        if (joinOperator.needsInput()) {
            if (input.hasNext()) {
                Page inputPage = input.next();
                joinOperator.addInput(inputPage);
            } else if (!finishing) {
                joinOperator.finish();
                finishing = true;
            }
        }
        Page outputPage = joinOperator.getOutput();
        if (outputPage != null) {
            outputPages.add(outputPage);
        }
    }
    joinOperator.close();
    return outputPages.build();
}

Example 12 with DriverContext

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

the class JoinTestUtils method instantiateBuildDrivers.

public static void instantiateBuildDrivers(BuildSideSetup buildSideSetup, TaskContext taskContext) {
    PipelineContext buildPipeline = taskContext.addPipelineContext(1, true, true, false);
    List<Driver> buildDrivers = new ArrayList<>();
    List<HashBuilderOperator> buildOperators = new ArrayList<>();
    for (int i = 0; i < buildSideSetup.getPartitionCount(); i++) {
        DriverContext buildDriverContext = buildPipeline.addDriverContext();
        HashBuilderOperator buildOperator = buildSideSetup.getBuildOperatorFactory().createOperator(buildDriverContext);
        Driver driver = Driver.createDriver(buildDriverContext, buildSideSetup.getBuildSideSourceOperatorFactory().createOperator(buildDriverContext), buildOperator);
        buildDrivers.add(driver);
        buildOperators.add(buildOperator);
    }
    buildSideSetup.setDriversAndOperators(buildDrivers, buildOperators);
}

Example 13 with DriverContext

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

the class TestHashJoinOperator method testInnerJoinLoadsPagesInOrder.

@Test
public void testInnerJoinLoadsPagesInOrder() {
    TaskContext taskContext = createTaskContext();
    // build factory
    List<Type> buildTypes = ImmutableList.of(VARCHAR);
    RowPagesBuilder buildPages = rowPagesBuilder(false, Ints.asList(0), buildTypes);
    for (int i = 0; i < 100_000; ++i) {
        buildPages.row("a");
    }
    BuildSideSetup buildSideSetup = setupBuildSide(nodePartitioningManager, false, taskContext, buildPages, Optional.empty(), false, SINGLE_STREAM_SPILLER_FACTORY);
    JoinBridgeManager<PartitionedLookupSourceFactory> lookupSourceFactory = buildSideSetup.getLookupSourceFactoryManager();
    // probe factory
    List<Type> probeTypes = ImmutableList.of(VARCHAR, INTEGER, INTEGER);
    RowPagesBuilder probePages = rowPagesBuilder(false, Ints.asList(0), probeTypes);
    probePages.row("a", 1L, 2L);
    WorkProcessorOperatorFactory joinOperatorFactory = (WorkProcessorOperatorFactory) innerJoinOperatorFactory(operatorFactories, lookupSourceFactory, probePages, PARTITIONING_SPILLER_FACTORY, false);
    // build drivers and operators
    instantiateBuildDrivers(buildSideSetup, taskContext);
    buildLookupSource(executor, buildSideSetup);
    Page probePage = getOnlyElement(probePages.build());
    AtomicInteger totalProbePages = new AtomicInteger();
    WorkProcessor<Page> inputPages = WorkProcessor.create(() -> {
        int probePageNumber = totalProbePages.incrementAndGet();
        if (probePageNumber == 5) {
            return finished();
        }
        return ofResult(new Page(1, probePage.getBlock(0), // this block should not be loaded by join operator as it's not being used by join
        new LazyBlock(1, () -> probePage.getBlock(1)), // and outputting current probe page
        new LazyBlock(1, () -> {
            // when loaded this block should be the latest one
            assertEquals(probePageNumber, totalProbePages.get());
            return probePage.getBlock(2);
        })));
    });
    DriverContext driverContext = taskContext.addPipelineContext(0, true, true, false).addDriverContext();
    OperatorContext operatorContext = driverContext.addOperatorContext(joinOperatorFactory.getOperatorId(), joinOperatorFactory.getPlanNodeId(), joinOperatorFactory.getOperatorType());
    WorkProcessorOperator joinOperator = joinOperatorFactory.create(new ProcessorContext(taskContext.getSession(), taskContext.getTaskMemoryContext(), operatorContext), inputPages);
    WorkProcessor<Page> outputPages = joinOperator.getOutputPages();
    int totalOutputPages = 0;
    for (int i = 0; i < 1_000_000; ++i) {
        if (!outputPages.process()) {
            // allow join to progress
            driverContext.getYieldSignal().resetYieldForTesting();
            continue;
        }
        if (outputPages.isFinished()) {
            break;
        }
        Page page = outputPages.getResult();
        totalOutputPages++;
        assertFalse(page.getBlock(1).isLoaded());
        page.getBlock(2).getLoadedBlock();
        // yield to enforce more complex execution
        driverContext.getYieldSignal().forceYieldForTesting();
    }
    // make sure that multiple pages were produced for some probe pages
    assertTrue(totalOutputPages > totalProbePages.get());
    assertTrue(outputPages.isFinished());
}

Example 14 with DriverContext

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

the class TestNestedLoopBuildOperator method testNestedLoopBuild.

@Test
public void testNestedLoopBuild() throws Exception {
    TaskContext taskContext = createTaskContext();
    List<Type> buildTypes = ImmutableList.of(BIGINT);
    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, createLongSequenceBlock(11, 14));
    Page buildPageEmpty = new Page(0);
    Page buildPage2 = new Page(3000, createLongSequenceBlock(4000, 7000));
    nestedLoopBuildOperator.addInput(buildPage1);
    nestedLoopBuildOperator.addInput(buildPageEmpty);
    nestedLoopBuildOperator.addInput(buildPage2);
    nestedLoopBuildOperator.finish();
    assertTrue(nestedLoopJoinBridge.getPagesFuture().isDone());
    List<Page> buildPages = nestedLoopJoinBridge.getPagesFuture().get().getPages();
    assertEquals(buildPages.get(0), buildPage1);
    assertEquals(buildPages.get(1), buildPage2);
    assertEquals(buildPages.size(), 2);
}

Example 15 with DriverContext

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

the class TestNestedLoopBuildOperator method testNestedLoopNoBlocksMaxSizeLimit.

@Test
public void testNestedLoopNoBlocksMaxSizeLimit() 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 massivePage = new Page(PageProcessor.MAX_BATCH_SIZE + 100);
    nestedLoopBuildOperator.addInput(massivePage);
    nestedLoopBuildOperator.finish();
    assertTrue(nestedLoopJoinBridge.getPagesFuture().isDone());
    List<Page> buildPages = nestedLoopJoinBridge.getPagesFuture().get().getPages();
    assertEquals(buildPages.size(), 2);
    assertEquals(buildPages.get(0).getPositionCount(), PageProcessor.MAX_BATCH_SIZE);
    assertEquals(buildPages.get(1).getPositionCount(), 100);
}