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

Example 16 with DriverContext

use of io.prestosql.operator.DriverContext in project hetu-core by openlookeng.

the class TestSpatialJoinOperator method testDuplicateProbeFactory.

@Test(dataProvider = "testDuplicateProbeFactoryDataProvider")
public void testDuplicateProbeFactory(boolean createSecondaryOperators) throws Exception {
    TaskContext taskContext = createTaskContext();
    PipelineContext pipelineContext = taskContext.addPipelineContext(0, true, true, false);
    DriverContext probeDriver = pipelineContext.addDriverContext();
    DriverContext buildDriver = pipelineContext.addDriverContext();
    RowPagesBuilder buildPages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR, DOUBLE)).row(stPoint(0, 0), "0_0", 1.5);
    PagesSpatialIndexFactory pagesSpatialIndexFactory = buildIndex(buildDriver, (build, probe, r) -> build.distance(probe) <= r.getAsDouble(), Optional.of(2), Optional.empty(), buildPages);
    RowPagesBuilder probePages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR)).row(stPoint(0, 1), "0_1");
    OperatorFactory firstFactory = new SpatialJoinOperatorFactory(2, new PlanNodeId("test"), INNER, probePages.getTypes(), Ints.asList(1), 0, Optional.empty(), pagesSpatialIndexFactory);
    for (int i = 0; i < 3; i++) {
        DriverContext secondDriver = pipelineContext.addDriverContext();
        OperatorFactory secondFactory = firstFactory.duplicate();
        if (createSecondaryOperators) {
            try (Operator secondOperator = secondFactory.createOperator(secondDriver)) {
                assertEquals(toPages(secondOperator, emptyIterator()), ImmutableList.of());
            }
        }
        secondFactory.noMoreOperators();
    }
    MaterializedResult expected = resultBuilder(taskContext.getSession(), ImmutableList.of(VARCHAR, VARCHAR)).row("0_1", "0_0").build();
    assertOperatorEquals(firstFactory, probeDriver, probePages.build(), expected);
}

Example 17 with DriverContext

use of io.prestosql.operator.DriverContext in project hetu-core by openlookeng.

the class TestOrcPageSourceMemoryTracking method testScanFilterAndProjectOperator.

@Test
public void testScanFilterAndProjectOperator() {
    // 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.newScanFilterAndProjectOperator(driverContext);
    assertEquals(driverContext.getSystemMemoryUsage(), 0);
    int totalRows = 0;
    while (totalRows < NUM_ROWS) {
        assertFalse(operator.isFinished());
        Page page = operator.getOutput();
        assertNotNull(page);
        assertBetweenInclusive(driverContext.getSystemMemoryUsage(), 90_000L, 499_999L);
        totalRows += page.getPositionCount();
    }
    // done... in the current implementation finish is not set until output returns a null page
    assertNull(operator.getOutput());
    assertTrue(operator.isFinished());
    assertBetweenInclusive(driverContext.getSystemMemoryUsage(), 0L, 500L);
}

Example 18 with DriverContext

use of io.prestosql.operator.DriverContext in project hetu-core by openlookeng.

the class TestMultiInputSnapshotState method testStaticConstructor.

@Test
public void testStaticConstructor() {
    ScheduledExecutorService scheduler = newScheduledThreadPool(4);
    TaskContext taskContext = createTaskContext(scheduler, scheduler, TEST_SNAPSHOT_SESSION);
    DriverContext driverContext = taskContext.addPipelineContext(0, true, true, false).addDriverContext();
    OperatorContext operatorContext = driverContext.addOperatorContext(1, new PlanNodeId("planNodeId"), "test");
    MultiInputSnapshotState inputSnapshotState = MultiInputSnapshotState.forOperator(mock(MultiInputRestorable.class), operatorContext);
    processPage(inputSnapshotState, source1, regularPage);
    inputSnapshotState = MultiInputSnapshotState.forTaskComponent(mock(MultiInputRestorable.class), taskContext, TestMultiInputSnapshotState::createSnapshotStateId);
    processPage(inputSnapshotState, source1, regularPage);
}

Example 19 with DriverContext

use of io.prestosql.operator.DriverContext in project hetu-core by openlookeng.

the class TestCrossRegionDynamicFilterOperator method createBloomFilterOperator.

private CrossRegionDynamicFilterOperator createBloomFilterOperator(String queryId, DynamicFilterCacheManager dynamicFilterCacheManager) {
    int filterOperatorId = operatorId.getAndIncrement();
    CrossRegionDynamicFilterOperator.CrossRegionDynamicFilterOperatorFactory factory = new CrossRegionDynamicFilterOperator.CrossRegionDynamicFilterOperatorFactory(filterOperatorId, planNodeId, queryId, createSymbolList(), createTypeProvider(), dynamicFilterCacheManager, createColumns(), createSymbolList());
    DriverContext driverContext = createTaskContext(executor, executor, TEST_SESSION).addPipelineContext(0, true, true, false).addDriverContext();
    return (CrossRegionDynamicFilterOperator) factory.createOperator(driverContext);
}

Example 20 with DriverContext

use of io.prestosql.operator.DriverContext in project hetu-core by openlookeng.

the class SnapshotStateId method forOperator.

public static SnapshotStateId forOperator(long snapshotId, OperatorContext operatorContext) {
    DriverContext driverContext = operatorContext.getDriverContext();
    TaskId taskIdGen = driverContext.getTaskId();
    int pipelineId = driverContext.getPipelineContext().getPipelineId();
    int driverId = driverContext.getDriverId();
    return new SnapshotStateId(snapshotId, taskIdGen, pipelineId, driverId, operatorContext.getOperatorId());
}