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

Example 6 with DriverContext

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

the class TestOrcPageSourceMemoryTracking method testTableScanOperator.

@Test
public void testTableScanOperator() {
    // Numbers used in assertions in this test may change when implementation is modified,
    // feel free to change them if they break in the future
    DriverContext driverContext = testPreparer.newDriverContext();
    SourceOperator operator = testPreparer.newTableScanOperator(driverContext);
    assertEquals(driverContext.getMemoryUsage(), 0);
    long memoryUsage = -1;
    int totalRows = 0;
    while (totalRows < 20000) {
        assertFalse(operator.isFinished());
        Page page = operator.getOutput();
        assertNotNull(page);
        page.getBlock(1);
        if (memoryUsage == -1) {
            memoryUsage = driverContext.getMemoryUsage();
            assertBetweenInclusive(memoryUsage, 460000L, 469999L);
        } else {
            assertEquals(driverContext.getMemoryUsage(), memoryUsage);
        }
        totalRows += page.getPositionCount();
    }
    memoryUsage = -1;
    while (totalRows < 40000) {
        assertFalse(operator.isFinished());
        Page page = operator.getOutput();
        assertNotNull(page);
        page.getBlock(1);
        if (memoryUsage == -1) {
            memoryUsage = driverContext.getMemoryUsage();
            assertBetweenInclusive(memoryUsage, 460000L, 469999L);
        } else {
            assertEquals(driverContext.getMemoryUsage(), memoryUsage);
        }
        totalRows += page.getPositionCount();
    }
    memoryUsage = -1;
    while (totalRows < NUM_ROWS) {
        assertFalse(operator.isFinished());
        Page page = operator.getOutput();
        assertNotNull(page);
        page.getBlock(1);
        if (memoryUsage == -1) {
            memoryUsage = driverContext.getMemoryUsage();
            assertBetweenInclusive(memoryUsage, 360000L, 369999L);
        } else {
            assertEquals(driverContext.getMemoryUsage(), memoryUsage);
        }
        totalRows += page.getPositionCount();
    }
    assertFalse(operator.isFinished());
    assertNull(operator.getOutput());
    assertTrue(operator.isFinished());
    assertEquals(driverContext.getMemoryUsage(), 0);
}

Example 7 with DriverContext

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

the class TestMemoryRevokingScheduler method testScheduleMemoryRevoking.

@Test
public void testScheduleMemoryRevoking() throws Exception {
    QueryContext q1 = getOrCreateQueryContext(new QueryId("q1"));
    QueryContext q2 = getOrCreateQueryContext(new QueryId("q2"));
    SqlTask sqlTask1 = newSqlTask(q1.getQueryId());
    SqlTask sqlTask2 = newSqlTask(q2.getQueryId());
    TaskContext taskContext1 = getOrCreateTaskContext(sqlTask1);
    PipelineContext pipelineContext11 = taskContext1.addPipelineContext(0, false, false, false);
    DriverContext driverContext111 = pipelineContext11.addDriverContext();
    OperatorContext operatorContext1 = driverContext111.addOperatorContext(1, new PlanNodeId("na"), "na");
    OperatorContext operatorContext2 = driverContext111.addOperatorContext(2, new PlanNodeId("na"), "na");
    DriverContext driverContext112 = pipelineContext11.addDriverContext();
    OperatorContext operatorContext3 = driverContext112.addOperatorContext(3, new PlanNodeId("na"), "na");
    TaskContext taskContext2 = getOrCreateTaskContext(sqlTask2);
    PipelineContext pipelineContext21 = taskContext2.addPipelineContext(1, false, false, false);
    DriverContext driverContext211 = pipelineContext21.addDriverContext();
    OperatorContext operatorContext4 = driverContext211.addOperatorContext(4, new PlanNodeId("na"), "na");
    OperatorContext operatorContext5 = driverContext211.addOperatorContext(5, new PlanNodeId("na"), "na");
    Collection<SqlTask> tasks = ImmutableList.of(sqlTask1, sqlTask2);
    MemoryRevokingScheduler scheduler = new MemoryRevokingScheduler(memoryPool, () -> tasks, executor, 1.0, 1.0);
    allOperatorContexts = ImmutableSet.of(operatorContext1, operatorContext2, operatorContext3, operatorContext4, operatorContext5);
    assertMemoryRevokingNotRequested();
    requestMemoryRevoking(scheduler);
    assertEquals(10, memoryPool.getFreeBytes());
    assertMemoryRevokingNotRequested();
    LocalMemoryContext revocableMemory1 = operatorContext1.localRevocableMemoryContext();
    LocalMemoryContext revocableMemory3 = operatorContext3.localRevocableMemoryContext();
    LocalMemoryContext revocableMemory4 = operatorContext4.localRevocableMemoryContext();
    LocalMemoryContext revocableMemory5 = operatorContext5.localRevocableMemoryContext();
    revocableMemory1.setBytes(3);
    revocableMemory3.setBytes(6);
    assertEquals(1, memoryPool.getFreeBytes());
    requestMemoryRevoking(scheduler);
    // we are still good - no revoking needed
    assertMemoryRevokingNotRequested();
    revocableMemory4.setBytes(7);
    assertEquals(-6, memoryPool.getFreeBytes());
    requestMemoryRevoking(scheduler);
    // we need to revoke 3 and 6
    assertMemoryRevokingRequestedFor(operatorContext1, operatorContext3);
    // yet another revoking request should not change anything
    requestMemoryRevoking(scheduler);
    assertMemoryRevokingRequestedFor(operatorContext1, operatorContext3);
    // lets revoke some bytes
    revocableMemory1.setBytes(0);
    operatorContext1.resetMemoryRevokingRequested();
    requestMemoryRevoking(scheduler);
    assertMemoryRevokingRequestedFor(operatorContext3);
    assertEquals(-3, memoryPool.getFreeBytes());
    // and allocate some more
    revocableMemory5.setBytes(3);
    assertEquals(-6, memoryPool.getFreeBytes());
    requestMemoryRevoking(scheduler);
    // we are still good with just OC3 in process of revoking
    assertMemoryRevokingRequestedFor(operatorContext3);
    // and allocate some more
    revocableMemory5.setBytes(4);
    assertEquals(-7, memoryPool.getFreeBytes());
    requestMemoryRevoking(scheduler);
    // no we have to trigger revoking for OC4
    assertMemoryRevokingRequestedFor(operatorContext3, operatorContext4);
}

Example 8 with DriverContext

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

the class TestMemoryRevokingScheduler method createContexts.

private OperatorContext createContexts(SqlTask sqlTask) {
    TaskContext taskContext = getOrCreateTaskContext(sqlTask);
    PipelineContext pipelineContext = taskContext.addPipelineContext(0, false, false, false);
    DriverContext driverContext = pipelineContext.addDriverContext();
    OperatorContext operatorContext = driverContext.addOperatorContext(1, new PlanNodeId("na"), "na");
    return operatorContext;
}

Example 9 with DriverContext

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

the class TestHashJoinOperator method testInnerJoinWithBlockingLookupSource.

@Test(dataProvider = "hashJoinTestValues")
public void testInnerJoinWithBlockingLookupSource(boolean parallelBuild, boolean probeHashEnabled, boolean buildHashEnabled) throws Exception {
    RowPagesBuilder probePages = rowPagesBuilder(probeHashEnabled, Ints.asList(0), ImmutableList.of(VARCHAR));
    Page probePage = getOnlyElement(probePages.addSequencePage(1, 0).build());
    // 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());
        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());
        assertNull(joinOperator.getOutput());
        // join needs input page
        assertTrue(joinOperator.isBlocked().isDone());
        assertFalse(joinOperator.isFinished());
        joinOperator.addInput(probePage);
        assertNull(joinOperator.getOutput());
        // lookup join operator got blocked waiting for build side
        assertFalse(joinOperator.isBlocked().isDone());
        assertFalse(joinOperator.isFinished());
    }
}

Example 10 with DriverContext

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

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(new StageId("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(false, Ints.asList(0), ImmutableList.of(VARCHAR, BIGINT)).addSequencePage(4, 20, 200).addSequencePage(4, 20, 200).addSequencePage(4, 30, 300).addSequencePage(4, 40, 400);
    BuildSideSetup buildSideSetup = setupBuildSide(nodePartitioningManager, true, taskContext, 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(operatorFactories, lookupSourceFactoryManager, probePages, joinSpillerFactory, true);
    // 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();
    }
}