Examples with OperatorFactoryWithTypes io.trino.sql.planner.LocalExecutionPlanner.OperatorFactoryWithTypes used on opensource projects

Example 1 with OperatorFactoryWithTypes

use of io.trino.sql.planner.LocalExecutionPlanner.OperatorFactoryWithTypes in project trino by trinodb.

the class TestWorkProcessorPipelineSourceOperator method testMergePages.

@Test
public void testMergePages() {
    Transformation<Split, Page> sourceOperatorPages = split -> TransformationState.ofResult(createPage(1), false);
    Transformation<Page, Page> firstOperatorPages = page -> TransformationState.ofResult(getOnlyElement(rowPagesBuilder(BIGINT).addSequencePage(1, 0).build()));
    TestWorkProcessorSourceOperatorFactory sourceOperatorFactory = new TestWorkProcessorSourceOperatorFactory(1, new PlanNodeId("1"), sourceOperatorPages);
    TestWorkProcessorOperatorFactory firstOperatorFactory = new TestWorkProcessorOperatorFactory(2, firstOperatorPages);
    SourceOperatorFactory pipelineOperatorFactory = (SourceOperatorFactory) getOnlyElement(WorkProcessorPipelineSourceOperator.convertOperators(ImmutableList.of(new OperatorFactoryWithTypes(sourceOperatorFactory, ImmutableList.of(BIGINT)), new OperatorFactoryWithTypes(firstOperatorFactory, ImmutableList.of(BIGINT))), DataSize.ofBytes(100), 100));
    DriverContext driverContext = TestingOperatorContext.create(scheduledExecutor).getDriverContext();
    SourceOperator pipelineOperator = pipelineOperatorFactory.createOperator(driverContext);
    pipelineOperator.addSplit(createSplit());
    assertTrue(pipelineOperator.getOutput().getPositionCount() > 100);
}

Example 2 with OperatorFactoryWithTypes

use of io.trino.sql.planner.LocalExecutionPlanner.OperatorFactoryWithTypes in project trino by trinodb.

the class TestWorkProcessorPipelineSourceOperator method testWorkProcessorPipelineSourceOperator.

@Test(timeOut = 10_000)
public void testWorkProcessorPipelineSourceOperator() throws InterruptedException {
    Split split = createSplit();
    Page page1 = createPage(1);
    Page page2 = createPage(2);
    Page page3 = createPage(3);
    Page page4 = createPage(4);
    Page page5 = createPage(5);
    Transformation<Split, Page> sourceOperatorPages = transformationFrom(ImmutableList.of(Transform.of(Optional.of(split), TransformationState.ofResult(page1, false)), Transform.of(Optional.of(split), TransformationState.ofResult(page2, true))));
    SettableFuture<Void> firstBlockedFuture = SettableFuture.create();
    Transformation<Page, Page> firstOperatorPages = transformationFrom(ImmutableList.of(Transform.of(Optional.of(page1), TransformationState.blocked(firstBlockedFuture)), Transform.of(Optional.of(page1), TransformationState.ofResult(page3, true)), Transform.of(Optional.of(page2), TransformationState.ofResult(page4, false)), Transform.of(Optional.of(page2), TransformationState.finished())), (left, right) -> left.getPositionCount() == right.getPositionCount());
    SettableFuture<Void> secondBlockedFuture = SettableFuture.create();
    Transformation<Page, Page> secondOperatorPages = transformationFrom(ImmutableList.of(Transform.of(Optional.of(page3), TransformationState.ofResult(page5, true)), Transform.of(Optional.of(page4), TransformationState.needsMoreData()), Transform.of(Optional.empty(), TransformationState.blocked(secondBlockedFuture))), (left, right) -> left.getPositionCount() == right.getPositionCount());
    TestWorkProcessorSourceOperatorFactory sourceOperatorFactory = new TestWorkProcessorSourceOperatorFactory(1, new PlanNodeId("1"), sourceOperatorPages);
    TestWorkProcessorOperatorFactory firstOperatorFactory = new TestWorkProcessorOperatorFactory(2, firstOperatorPages);
    TestWorkProcessorOperatorFactory secondOperatorFactory = new TestWorkProcessorOperatorFactory(3, secondOperatorPages);
    SourceOperatorFactory pipelineOperatorFactory = (SourceOperatorFactory) getOnlyElement(WorkProcessorPipelineSourceOperator.convertOperators(ImmutableList.of(new OperatorFactoryWithTypes(sourceOperatorFactory, ImmutableList.of(BIGINT)), new OperatorFactoryWithTypes(firstOperatorFactory, ImmutableList.of(BIGINT)), new OperatorFactoryWithTypes(secondOperatorFactory, ImmutableList.of(BIGINT))), DataSize.ofBytes(0), 0));
    DriverContext driverContext = TestingTaskContext.builder(MoreExecutors.directExecutor(), scheduledExecutor, TEST_SESSION).build().addPipelineContext(0, true, true, false).addDriverContext();
    SourceOperator pipelineOperator = pipelineOperatorFactory.createOperator(driverContext);
    OperatorContext pipelineOperatorContext = pipelineOperator.getOperatorContext();
    // make sure WorkProcessorOperator memory is accounted for
    sourceOperatorFactory.sourceOperator.memoryTrackingContext.localUserMemoryContext().setBytes(123);
    assertEquals(driverContext.getMemoryUsage(), 123);
    assertNull(pipelineOperator.getOutput());
    assertFalse(pipelineOperator.isBlocked().isDone());
    // blocking on splits should not account for blocked time for any WorkProcessorOperator
    pipelineOperatorContext.getNestedOperatorStats().forEach(operatorStats -> assertEquals(operatorStats.getBlockedWall().toMillis(), 0));
    pipelineOperator.addSplit(split);
    assertTrue(pipelineOperator.isBlocked().isDone());
    assertNull(pipelineOperator.getOutput());
    assertFalse(pipelineOperator.isBlocked().isDone());
    Thread.sleep(100);
    firstBlockedFuture.set(null);
    assertTrue(pipelineOperator.isBlocked().isDone());
    // blocking of first WorkProcessorOperator should be accounted for in stats
    List<OperatorStats> operatorStats = pipelineOperatorContext.getNestedOperatorStats();
    assertEquals(operatorStats.get(0).getBlockedWall().toMillis(), 0);
    assertTrue(operatorStats.get(1).getBlockedWall().toMillis() > 0);
    assertEquals(operatorStats.get(2).getBlockedWall().toMillis(), 0);
    assertEquals(getTestingOperatorInfo(operatorStats.get(1)).count, 2);
    assertEquals(getTestingOperatorInfo(operatorStats.get(2)).count, 2);
    assertEquals(pipelineOperator.getOutput().getPositionCount(), page5.getPositionCount());
    // sourceOperator should yield
    driverContext.getYieldSignal().forceYieldForTesting();
    assertNull(pipelineOperator.getOutput());
    driverContext.getYieldSignal().resetYieldForTesting();
    // firstOperatorPages should finish. This should cause sourceOperator and firstOperatorPages to close.
    // secondOperatorPages should block
    assertNull(pipelineOperator.getOutput());
    assertFalse(pipelineOperator.isBlocked().isDone());
    assertTrue(sourceOperatorFactory.sourceOperator.closed);
    assertTrue(firstOperatorFactory.operator.closed);
    assertFalse(secondOperatorFactory.operator.closed);
    // first operator should return final operator info
    assertEquals(getTestingOperatorInfo(operatorStats.get(1)).count, 3);
    assertEquals(getTestingOperatorInfo(operatorStats.get(2)).count, 2);
    operatorStats = pipelineOperatorContext.getNestedOperatorStats();
    assertEquals(getTestingOperatorInfo(operatorStats.get(1)).count, 3);
    assertEquals(getTestingOperatorInfo(operatorStats.get(2)).count, 3);
    // cause early operator finish
    pipelineOperator.finish();
    // operator is still blocked on secondBlockedFuture
    assertFalse(pipelineOperator.isFinished());
    assertTrue(secondOperatorFactory.operator.closed);
    secondBlockedFuture.set(null);
    assertTrue(pipelineOperator.isBlocked().isDone());
    assertNull(pipelineOperator.getOutput());
    assertTrue(pipelineOperator.isFinished());
    // assert operator stats are correct
    operatorStats = pipelineOperatorContext.getNestedOperatorStats();
    assertEquals(operatorStats.get(0).getOutputPositions(), 3);
    assertEquals(operatorStats.get(1).getInputPositions(), 3);
    assertEquals(operatorStats.get(0).getOutputDataSize().toBytes(), 27);
    assertEquals(operatorStats.get(1).getInputDataSize().toBytes(), 27);
    assertEquals(operatorStats.get(1).getOutputPositions(), 7);
    assertEquals(operatorStats.get(2).getInputPositions(), 7);
    assertEquals(operatorStats.get(1).getOutputDataSize().toBytes(), 63);
    assertEquals(operatorStats.get(2).getInputDataSize().toBytes(), 63);
    assertEquals(operatorStats.get(2).getOutputPositions(), 5);
    assertEquals(operatorStats.get(2).getOutputDataSize().toBytes(), 45);
    assertThat(operatorStats.get(1).getMetrics().getMetrics()).hasSize(2).containsEntry("testOperatorMetric", new LongCount(1));
    // assert source operator stats are correct
    OperatorStats sourceOperatorStats = operatorStats.get(0);
    assertThat(sourceOperatorStats.getMetrics().getMetrics()).hasSize(3).containsEntry("testSourceMetric", new LongCount(1)).containsEntry("testSourceClosed", new LongCount(1));
    assertEquals(sourceOperatorStats.getConnectorMetrics().getMetrics(), ImmutableMap.of("testSourceConnectorMetric", new LongCount(2), "testSourceConnectorClosed", new LongCount(1), "Physical input read time", new DurationTiming(new Duration(7, NANOSECONDS))));
    assertEquals(sourceOperatorStats.getDynamicFilterSplitsProcessed(), 42L);
    assertEquals(sourceOperatorStats.getPhysicalInputDataSize(), DataSize.ofBytes(1));
    assertEquals(sourceOperatorStats.getPhysicalInputPositions(), 2);
    assertEquals(sourceOperatorStats.getInternalNetworkInputDataSize(), DataSize.ofBytes(3));
    assertEquals(sourceOperatorStats.getInternalNetworkInputPositions(), 4);
    assertEquals(sourceOperatorStats.getInputDataSize(), DataSize.ofBytes(5));
    assertEquals(sourceOperatorStats.getInputPositions(), 6);
    assertEquals(sourceOperatorStats.getAddInputWall(), new Duration(0, NANOSECONDS));
    // pipeline input stats should match source WorkProcessorOperator stats
    PipelineStats pipelineStats = pipelineOperator.getOperatorContext().getDriverContext().getPipelineContext().getPipelineStats();
    assertEquals(sourceOperatorStats.getPhysicalInputDataSize(), pipelineStats.getPhysicalInputDataSize());
    assertEquals(sourceOperatorStats.getPhysicalInputPositions(), pipelineStats.getPhysicalInputPositions());
    assertEquals(sourceOperatorStats.getInternalNetworkInputDataSize(), pipelineStats.getInternalNetworkInputDataSize());
    assertEquals(sourceOperatorStats.getInternalNetworkInputPositions(), pipelineStats.getInternalNetworkInputPositions());
    assertEquals(sourceOperatorStats.getInputDataSize(), pipelineStats.getProcessedInputDataSize());
    assertEquals(sourceOperatorStats.getInputPositions(), pipelineStats.getProcessedInputPositions());
    assertEquals(sourceOperatorStats.getAddInputWall(), pipelineStats.getPhysicalInputReadTime());
    // assert pipeline metrics
    List<OperatorStats> operatorSummaries = pipelineStats.getOperatorSummaries();
    assertThat(operatorSummaries.get(0).getMetrics().getMetrics()).hasSize(3).containsEntry("testSourceMetric", new LongCount(1)).containsEntry("testSourceClosed", new LongCount(1));
    assertEquals(operatorSummaries.get(0).getConnectorMetrics().getMetrics(), ImmutableMap.of("testSourceConnectorMetric", new LongCount(2), "testSourceConnectorClosed", new LongCount(1), "Physical input read time", new DurationTiming(new Duration(7, NANOSECONDS))));
    assertThat(operatorSummaries.get(1).getMetrics().getMetrics()).hasSize(2).containsEntry("testOperatorMetric", new LongCount(1));
}