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

Example 1 with WorkProcessorOperatorFactory

use of io.trino.operator.WorkProcessorOperatorFactory 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());
}