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

Example 6 with DriverContext

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

the class TestSpatialJoinOperator method assertSpatialJoinSnapshot.

private void assertSpatialJoinSnapshot(TaskContext taskContext, Type joinType, RowPagesBuilder buildPages, RowPagesBuilder probePages, MaterializedResult expected) {
    DriverContext driverContext = taskContext.addPipelineContext(0, true, true, false).addDriverContext();
    PagesSpatialIndexFactory pagesSpatialIndexFactory = buildIndexSnapshot(driverContext, (build, probe, r) -> build.contains(probe), Optional.empty(), Optional.empty(), buildPages);
    OperatorFactory joinOperatorFactory = new SpatialJoinOperatorFactory(2, new PlanNodeId("test"), joinType, probePages.getTypes(), Ints.asList(1), 0, Optional.empty(), pagesSpatialIndexFactory);
    assetResultEqualsSnapshot(joinOperatorFactory, driverContext, probePages.build(), expected);
}

Example 7 with DriverContext

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

the class TestSpatialJoinOperator method testDistanceQuery.

@Test
public void testDistanceQuery() {
    TaskContext taskContext = createTaskContext();
    DriverContext driverContext = taskContext.addPipelineContext(0, true, true, false).addDriverContext();
    RowPagesBuilder buildPages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR, DOUBLE)).row(stPoint(0, 0), "0_0", 1.5).row(null, "null", 1.5).row(stPoint(1, 0), "1_0", 1.5).pageBreak().row(stPoint(3, 0), "3_0", 1.5).pageBreak().row(stPoint(10, 0), "10_0", 1.5);
    PagesSpatialIndexFactory pagesSpatialIndexFactory = buildIndex(driverContext, (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").row(null, "null").row(stPoint(1, 1), "1_1").pageBreak().row(stPoint(3, 1), "3_1").pageBreak().row(stPoint(10, 1), "10_1");
    OperatorFactory joinOperatorFactory = new SpatialJoinOperatorFactory(2, new PlanNodeId("test"), INNER, probePages.getTypes(), Ints.asList(1), 0, Optional.empty(), pagesSpatialIndexFactory);
    MaterializedResult expected = resultBuilder(taskContext.getSession(), ImmutableList.of(VARCHAR, VARCHAR)).row("0_1", "0_0").row("0_1", "1_0").row("1_1", "0_0").row("1_1", "1_0").row("3_1", "3_0").row("10_1", "10_0").build();
    assertOperatorEquals(joinOperatorFactory, driverContext, probePages.build(), expected);
}

Example 8 with DriverContext

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

the class TestSpatialJoinOperator method testDistributedSpatialJoin.

@Test
public void testDistributedSpatialJoin() {
    TaskContext taskContext = createTaskContext();
    DriverContext driverContext = taskContext.addPipelineContext(0, true, true, true).addDriverContext();
    RowPagesBuilder buildPages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR, INTEGER)).row(POLYGON_A, "A", 1).row(POLYGON_A, "A", 2).row(null, "null", null).pageBreak().row(POLYGON_B, "B", 0).row(POLYGON_B, "B", 2);
    RowPagesBuilder probePages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR, INTEGER)).row(POINT_X, "x", 2).row(null, "null", null).row(POINT_Y, "y", 2).pageBreak().row(POINT_Z, "z", 0);
    MaterializedResult expected = resultBuilder(taskContext.getSession(), ImmutableList.of(VARCHAR, VARCHAR)).row("x", "A").row("y", "A").row("y", "B").row("z", "B").build();
    PagesSpatialIndexFactory pagesSpatialIndexFactory = buildIndex(driverContext, (build, probe, r) -> build.contains(probe), Optional.empty(), Optional.of(2), Optional.of(KDB_TREE_JSON), Optional.empty(), buildPages);
    OperatorFactory joinOperatorFactory = new SpatialJoinOperatorFactory(2, new PlanNodeId("test"), INNER, probePages.getTypes(), Ints.asList(1), 0, Optional.of(2), pagesSpatialIndexFactory);
    assertOperatorEquals(joinOperatorFactory, driverContext, probePages.build(), expected);
}

Example 9 with DriverContext

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

the class SnapshotStateId method forDriverComponent.

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

Example 10 with DriverContext

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

the class TestMemoryRevokingScheduler method testScheduleMemoryRevoking.

@Test
public void testScheduleMemoryRevoking() throws Exception {
    SqlTask sqlTask1 = newSqlTask();
    SqlTask sqlTask2 = newSqlTask();
    TaskContext taskContext1 = sqlTask1.getQueryContext().addTaskContext(new TaskStateMachine(new TaskId("q1", 1, 1), executor), session, false, false, OptionalInt.empty(), Optional.empty(), TESTING_SERDE_FACTORY);
    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 = sqlTask2.getQueryContext().addTaskContext(new TaskStateMachine(new TaskId("q2", 1, 1), executor), session, false, false, OptionalInt.empty(), Optional.empty(), TESTING_SERDE_FACTORY);
    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(singletonList(memoryPool), () -> tasks, executor, 1.0, 1.0, false, new DataSize(512, MEGABYTE).toBytes());
    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);
}