Examples with PageConsumerOperator io.trino.testing.PageConsumerOperator used on opensource projects

Example 1 with PageConsumerOperator

use of io.trino.testing.PageConsumerOperator in project trino by trinodb.

the class TestMemoryBlocking method testTableScanMemoryBlocking.

@Test
public void testTableScanMemoryBlocking() {
    PlanNodeId sourceId = new PlanNodeId("source");
    List<Type> types = ImmutableList.of(VARCHAR);
    TableScanOperator source = new TableScanOperator(driverContext.addOperatorContext(1, new PlanNodeId("test"), "values"), sourceId, (session, split, table, columns, dynamicFilter) -> new FixedPageSource(rowPagesBuilder(types).addSequencePage(10, 1).addSequencePage(10, 1).addSequencePage(10, 1).addSequencePage(10, 1).addSequencePage(10, 1).build()), TEST_TABLE_HANDLE, ImmutableList.of(), DynamicFilter.EMPTY);
    PageConsumerOperator sink = createSinkOperator(types);
    Driver driver = Driver.createDriver(driverContext, source, sink);
    assertSame(driver.getDriverContext(), driverContext);
    assertFalse(driver.isFinished());
    Split testSplit = new Split(new CatalogName("test"), new TestSplit(), Lifespan.taskWide());
    driver.updateSplitAssignment(new SplitAssignment(sourceId, ImmutableSet.of(new ScheduledSplit(0, sourceId, testSplit)), true));
    ListenableFuture<Void> blocked = driver.processFor(new Duration(1, NANOSECONDS));
    // the driver shouldn't block in the first call as it will be able to move a page between source and the sink operator
    // but the operator should be blocked
    assertTrue(blocked.isDone());
    assertFalse(source.getOperatorContext().isWaitingForMemory().isDone());
    // and they should stay blocked until more memory becomes available
    for (int i = 0; i < 10; i++) {
        blocked = driver.processFor(new Duration(1, NANOSECONDS));
        assertFalse(blocked.isDone());
        assertFalse(source.getOperatorContext().isWaitingForMemory().isDone());
    }
    // free up some memory
    memoryPool.free(QUERY_ID, "test", memoryPool.getReservedBytes());
    // the operator should be unblocked
    assertTrue(source.getOperatorContext().isWaitingForMemory().isDone());
    // the driver shouldn't be blocked
    blocked = driver.processFor(new Duration(1, NANOSECONDS));
    assertTrue(blocked.isDone());
}

Example 2 with PageConsumerOperator

use of io.trino.testing.PageConsumerOperator in project trino by trinodb.

the class TestDriver method testAbruptFinish.

@Test
public void testAbruptFinish() {
    List<Type> types = ImmutableList.of(VARCHAR, BIGINT, BIGINT);
    ValuesOperator source = new ValuesOperator(driverContext.addOperatorContext(0, new PlanNodeId("test"), "values"), rowPagesBuilder(types).addSequencePage(10, 20, 30, 40).build());
    PageConsumerOperator sink = createSinkOperator(types);
    Driver driver = Driver.createDriver(driverContext, source, sink);
    assertSame(driver.getDriverContext(), driverContext);
    assertFalse(driver.isFinished());
    driver.close();
    assertTrue(driver.isFinished());
    // finish is only called in normal operations
    assertFalse(source.isFinished());
    assertFalse(sink.isFinished());
    // close is always called (values operator doesn't have a closed state)
    assertTrue(sink.isClosed());
}

Example 3 with PageConsumerOperator

use of io.trino.testing.PageConsumerOperator in project trino by trinodb.

the class TestDriver method testAddSourceFinish.

@Test
public void testAddSourceFinish() {
    PlanNodeId sourceId = new PlanNodeId("source");
    List<Type> types = ImmutableList.of(VARCHAR, BIGINT, BIGINT);
    TableScanOperator source = new TableScanOperator(driverContext.addOperatorContext(99, new PlanNodeId("test"), "values"), sourceId, (session, split, table, columns, dynamicFilter) -> new FixedPageSource(rowPagesBuilder(types).addSequencePage(10, 20, 30, 40).build()), TEST_TABLE_HANDLE, ImmutableList.of(), DynamicFilter.EMPTY);
    PageConsumerOperator sink = createSinkOperator(types);
    Driver driver = Driver.createDriver(driverContext, source, sink);
    assertSame(driver.getDriverContext(), driverContext);
    assertFalse(driver.isFinished());
    assertFalse(driver.processFor(new Duration(1, TimeUnit.MILLISECONDS)).isDone());
    assertFalse(driver.isFinished());
    driver.updateSplitAssignment(new SplitAssignment(sourceId, ImmutableSet.of(new ScheduledSplit(0, sourceId, newMockSplit())), true));
    assertFalse(driver.isFinished());
    assertTrue(driver.processFor(new Duration(1, TimeUnit.SECONDS)).isDone());
    assertTrue(driver.isFinished());
    assertTrue(sink.isFinished());
    assertTrue(source.isFinished());
}