Examples with PageBuffer com.facebook.presto.operator.index.PageBuffer used on opensource projects

Example 1 with PageBuffer

use of com.facebook.presto.operator.index.PageBuffer in project presto by prestodb.

the class TestHashJoinOperator method testInnerJoinWithSpillWithEarlyTermination.

@Test(timeOut = 60000)
public void testInnerJoinWithSpillWithEarlyTermination() {
    TaskStateMachine taskStateMachine = new TaskStateMachine(new TaskId("query", 0, 0, 0), executor);
    TaskContext taskContext = TestingTaskContext.createTaskContext(executor, scheduledExecutor, TEST_SESSION, taskStateMachine);
    PipelineContext joinPipelineContext = taskContext.addPipelineContext(2, true, true, false);
    DriverContext joinDriverContext1 = joinPipelineContext.addDriverContext();
    DriverContext joinDriverContext2 = joinPipelineContext.addDriverContext();
    DriverContext joinDriverContext3 = joinPipelineContext.addDriverContext();
    // build factory
    RowPagesBuilder buildPages = rowPagesBuilder(ImmutableList.of(VARCHAR, BIGINT)).addSequencePage(4, 20, 200).addSequencePage(4, 20, 200).addSequencePage(4, 30, 300).addSequencePage(4, 40, 400);
    // force a yield for every match in LookupJoinOperator, set called to true after first
    AtomicBoolean called = new AtomicBoolean(false);
    InternalJoinFilterFunction filterFunction = new TestInternalJoinFilterFunction((leftPosition, leftPage, rightPosition, rightPage) -> {
        called.set(true);
        return true;
    });
    BuildSideSetup buildSideSetup = setupBuildSide(true, taskContext, Ints.asList(0), buildPages, Optional.of(filterFunction), true, SINGLE_STREAM_SPILLER_FACTORY);
    JoinBridgeManager<PartitionedLookupSourceFactory> lookupSourceFactoryManager = buildSideSetup.getLookupSourceFactoryManager();
    // probe factory
    RowPagesBuilder probe1Pages = rowPagesBuilder(true, Ints.asList(0), ImmutableList.of(VARCHAR, BIGINT)).row("no_match_1", 123_000L).row("no_match_2", 123_000L);
    RowPagesBuilder probe2Pages = rowPagesBuilder(true, Ints.asList(0), ImmutableList.of(VARCHAR, BIGINT)).row("20", 123_000L).row("20", 123_000L).pageBreak().addSequencePage(20, 0, 123_000).addSequencePage(10, 30, 123_000);
    OperatorFactory joinOperatorFactory = innerJoinOperatorFactory(lookupSourceFactoryManager, probe2Pages, PARTITIONING_SPILLER_FACTORY, OptionalInt.of(3));
    // build drivers and operators
    instantiateBuildDrivers(buildSideSetup, taskContext);
    List<Driver> buildDrivers = buildSideSetup.getBuildDrivers();
    int buildOperatorCount = buildDrivers.size();
    LookupSourceFactory lookupSourceFactory = lookupSourceFactoryManager.getJoinBridge(Lifespan.taskWide());
    Operator lookupOperator1 = joinOperatorFactory.createOperator(joinDriverContext1);
    Operator lookupOperator2 = joinOperatorFactory.createOperator(joinDriverContext2);
    Operator lookupOperator3 = joinOperatorFactory.createOperator(joinDriverContext3);
    joinOperatorFactory.noMoreOperators();
    ListenableFuture<LookupSourceProvider> lookupSourceProvider = lookupSourceFactory.createLookupSourceProvider();
    while (!lookupSourceProvider.isDone()) {
        for (Driver buildDriver : buildDrivers) {
            checkErrors(taskStateMachine);
            buildDriver.process();
        }
    }
    getFutureValue(lookupSourceProvider).close();
    for (int i = 0; i < buildOperatorCount; i++) {
        revokeMemory(buildSideSetup.getBuildOperators().get(i));
    }
    for (Driver buildDriver : buildDrivers) {
        runDriverInThread(executor, buildDriver);
    }
    ValuesOperatorFactory valuesOperatorFactory1 = new ValuesOperatorFactory(17, new PlanNodeId("values1"), probe1Pages.build());
    ValuesOperatorFactory valuesOperatorFactory2 = new ValuesOperatorFactory(18, new PlanNodeId("values2"), probe2Pages.build());
    ValuesOperatorFactory valuesOperatorFactory3 = new ValuesOperatorFactory(18, new PlanNodeId("values3"), ImmutableList.of());
    PageBuffer pageBuffer = new PageBuffer(10);
    PageBufferOperatorFactory pageBufferOperatorFactory = new PageBufferOperatorFactory(19, new PlanNodeId("pageBuffer"), pageBuffer);
    Driver joinDriver1 = Driver.createDriver(joinDriverContext1, valuesOperatorFactory1.createOperator(joinDriverContext1), lookupOperator1, pageBufferOperatorFactory.createOperator(joinDriverContext1));
    Driver joinDriver2 = Driver.createDriver(joinDriverContext2, valuesOperatorFactory2.createOperator(joinDriverContext2), lookupOperator2, pageBufferOperatorFactory.createOperator(joinDriverContext2));
    Driver joinDriver3 = Driver.createDriver(joinDriverContext3, valuesOperatorFactory3.createOperator(joinDriverContext3), lookupOperator3, pageBufferOperatorFactory.createOperator(joinDriverContext3));
    joinDriver3.close();
    joinDriver3.process();
    while (!called.get()) {
        checkErrors(taskStateMachine);
        processRow(joinDriver1, taskStateMachine);
        processRow(joinDriver2, taskStateMachine);
    }
    joinDriver1.close();
    joinDriver1.process();
    while (!joinDriver2.isFinished()) {
        processRow(joinDriver2, taskStateMachine);
    }
    checkErrors(taskStateMachine);
    List<Page> pages = getPages(pageBuffer);
    MaterializedResult expected = MaterializedResult.resultBuilder(taskContext.getSession(), concat(probe2Pages.getTypesWithoutHash(), buildPages.getTypesWithoutHash())).row("20", 123_000L, "20", 200L).row("20", 123_000L, "20", 200L).row("20", 123_000L, "20", 200L).row("20", 123_000L, "20", 200L).row("30", 123_000L, "30", 300L).row("31", 123_001L, "31", 301L).row("32", 123_002L, "32", 302L).row("33", 123_003L, "33", 303L).build();
    assertEqualsIgnoreOrder(getProperColumns(lookupOperator1, concat(probe2Pages.getTypes(), buildPages.getTypes()), probe2Pages, pages).getMaterializedRows(), expected.getMaterializedRows());
}

Example 2 with PageBuffer

use of com.facebook.presto.operator.index.PageBuffer in project presto by prestodb.

the class TestHashJoinOperator method innerJoinWithSpill.

private void innerJoinWithSpill(boolean probeHashEnabled, List<WhenSpill> whenSpill, SingleStreamSpillerFactory buildSpillerFactory, PartitioningSpillerFactory joinSpillerFactory) throws Exception {
    TaskStateMachine taskStateMachine = new TaskStateMachine(new TaskId("query", 0, 0, 0), executor);
    TaskContext taskContext = TestingTaskContext.createTaskContext(executor, scheduledExecutor, TEST_SESSION, taskStateMachine);
    DriverContext joinDriverContext = taskContext.addPipelineContext(2, true, true, false).addDriverContext();
    // force a yield for every match in LookupJoinOperator, set called to true after first
    AtomicBoolean called = new AtomicBoolean(false);
    InternalJoinFilterFunction filterFunction = new TestInternalJoinFilterFunction((leftPosition, leftPage, rightPosition, rightPage) -> {
        called.set(true);
        joinDriverContext.getYieldSignal().forceYieldForTesting();
        return true;
    });
    // build factory
    RowPagesBuilder buildPages = rowPagesBuilder(ImmutableList.of(VARCHAR, BIGINT)).addSequencePage(4, 20, 200).addSequencePage(4, 20, 200).addSequencePage(4, 30, 300).addSequencePage(4, 40, 400);
    BuildSideSetup buildSideSetup = setupBuildSide(true, taskContext, Ints.asList(0), buildPages, Optional.of(filterFunction), true, buildSpillerFactory);
    JoinBridgeManager<PartitionedLookupSourceFactory> lookupSourceFactoryManager = buildSideSetup.getLookupSourceFactoryManager();
    // probe factory
    RowPagesBuilder probePages = rowPagesBuilder(probeHashEnabled, Ints.asList(0), ImmutableList.of(VARCHAR, BIGINT)).row("20", 123_000L).row("20", 123_000L).pageBreak().addSequencePage(20, 0, 123_000).addSequencePage(10, 30, 123_000);
    OperatorFactory joinOperatorFactory = innerJoinOperatorFactory(lookupSourceFactoryManager, probePages, joinSpillerFactory);
    // build drivers and operators
    instantiateBuildDrivers(buildSideSetup, taskContext);
    List<Driver> buildDrivers = buildSideSetup.getBuildDrivers();
    int buildOperatorCount = buildDrivers.size();
    checkState(buildOperatorCount == whenSpill.size());
    LookupSourceFactory lookupSourceFactory = lookupSourceFactoryManager.getJoinBridge(Lifespan.taskWide());
    try (Operator joinOperator = joinOperatorFactory.createOperator(joinDriverContext)) {
        // build lookup source
        ListenableFuture<LookupSourceProvider> lookupSourceProvider = lookupSourceFactory.createLookupSourceProvider();
        List<Boolean> revoked = new ArrayList<>(nCopies(buildOperatorCount, false));
        while (!lookupSourceProvider.isDone()) {
            for (int i = 0; i < buildOperatorCount; i++) {
                checkErrors(taskStateMachine);
                buildDrivers.get(i).process();
                HashBuilderOperator buildOperator = buildSideSetup.getBuildOperators().get(i);
                if (whenSpill.get(i) == WhenSpill.DURING_BUILD && buildOperator.getOperatorContext().getReservedRevocableBytes() > 0) {
                    checkState(!lookupSourceProvider.isDone(), "Too late, LookupSource already done");
                    revokeMemory(buildOperator);
                    revoked.set(i, true);
                }
            }
        }
        getFutureValue(lookupSourceProvider).close();
        assertEquals(revoked, whenSpill.stream().map(WhenSpill.DURING_BUILD::equals).collect(toImmutableList()), "Some operators not spilled before LookupSource built");
        for (int i = 0; i < buildOperatorCount; i++) {
            if (whenSpill.get(i) == WhenSpill.AFTER_BUILD) {
                revokeMemory(buildSideSetup.getBuildOperators().get(i));
            }
        }
        for (Driver buildDriver : buildDrivers) {
            runDriverInThread(executor, buildDriver);
        }
        ValuesOperatorFactory valuesOperatorFactory = new ValuesOperatorFactory(17, new PlanNodeId("values"), probePages.build());
        PageBuffer pageBuffer = new PageBuffer(10);
        PageBufferOperatorFactory pageBufferOperatorFactory = new PageBufferOperatorFactory(18, new PlanNodeId("pageBuffer"), pageBuffer);
        Driver joinDriver = Driver.createDriver(joinDriverContext, valuesOperatorFactory.createOperator(joinDriverContext), joinOperator, pageBufferOperatorFactory.createOperator(joinDriverContext));
        while (!called.get()) {
            // process first row of first page of LookupJoinOperator
            processRow(joinDriver, taskStateMachine);
        }
        for (int i = 0; i < buildOperatorCount; i++) {
            if (whenSpill.get(i) == WhenSpill.DURING_USAGE) {
                triggerMemoryRevokingAndWait(buildSideSetup.getBuildOperators().get(i), taskStateMachine);
            }
        }
        // process remaining LookupJoinOperator pages
        while (!joinDriver.isFinished()) {
            checkErrors(taskStateMachine);
            processRow(joinDriver, taskStateMachine);
        }
        checkErrors(taskStateMachine);
        List<Page> actualPages = getPages(pageBuffer);
        MaterializedResult expected = MaterializedResult.resultBuilder(taskContext.getSession(), concat(probePages.getTypesWithoutHash(), buildPages.getTypesWithoutHash())).row("20", 123_000L, "20", 200L).row("20", 123_000L, "20", 200L).row("20", 123_000L, "20", 200L).row("20", 123_000L, "20", 200L).row("30", 123_000L, "30", 300L).row("31", 123_001L, "31", 301L).row("32", 123_002L, "32", 302L).row("33", 123_003L, "33", 303L).build();
        assertEqualsIgnoreOrder(getProperColumns(joinOperator, concat(probePages.getTypes(), buildPages.getTypes()), probePages, actualPages).getMaterializedRows(), expected.getMaterializedRows());
    } finally {
        joinOperatorFactory.noMoreOperators();
    }
}