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

Example 16 with TaskContext

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

the class TestArbitraryOutputBuffer method testSnapshotState.

@Test(dataProvider = "bufferTypes")
public void testSnapshotState(OutputBuffers.BufferType bufferType) {
    OutputBuffer buffer = createOutputBuffer(bufferType);
    TaskContext taskContext = BufferTestUtils.newTestingTaskContext();
    buffer.setTaskContext(taskContext);
    buffer.addInputChannel("1");
    buffer.addInputChannel("2");
    buffer.setNoMoreInputChannels();
    MarkerPage markerPage = MarkerPage.snapshotPage(1);
    SerializedPage marker = SerializedPage.forMarker(markerPage);
    buffer.enqueue(Collections.singletonList(marker), "1");
    List<SerializedPage> pages = getBufferResult(buffer, FIRST, 0, sizeOfPages(3), NO_WAIT).getSerializedPages();
    assertEquals(pages.size(), 1);
    assertTrue(pages.get(0).isMarkerPage());
    // Acknowledge
    buffer.get(FIRST, 1, sizeOfPages(1)).cancel(true);
    buffer.enqueue(Collections.singletonList(marker), "2");
    pages = getBufferResult(buffer, FIRST, 0, sizeOfPages(3), NO_WAIT).getSerializedPages();
    assertEquals(pages.size(), 0);
}

Example 17 with TaskContext

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

the class TestArbitraryOutputBuffer method testSetMultipleTaskContext.

@Test(expectedExceptions = IllegalStateException.class, dataProvider = "bufferTypes")
public void testSetMultipleTaskContext(OutputBuffers.BufferType bufferType) {
    OutputBuffer buffer = createOutputBuffer(bufferType);
    TaskContext taskContext = BufferTestUtils.newTestingTaskContext();
    buffer.setTaskContext(taskContext);
    buffer.setTaskContext(taskContext);
}

Example 18 with TaskContext

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

the class BufferTestUtils method newTestingTaskContext.

static TaskContext newTestingTaskContext() {
    ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();
    QueryContext queryContext = new QueryContext(new QueryId("query_id"), new DataSize(1, MEGABYTE), new DataSize(2, MEGABYTE), new MemoryPool(new MemoryPoolId("test"), new DataSize(1, GIGABYTE)), new TestingGcMonitor(), executor, executor, new DataSize(1, MEGABYTE), new SpillSpaceTracker(new DataSize(1, GIGABYTE)), NOOP_SNAPSHOT_UTILS);
    TaskId taskId = TaskId.valueOf("query_id.1.2");
    TaskContext taskContext = queryContext.addTaskContext(new TaskStateMachine(taskId, executor), TEST_SNAPSHOT_SESSION, false, false, OptionalInt.empty(), Optional.empty(), TESTING_SERDE_FACTORY);
    taskContext.getSnapshotManager().setTotalComponents(100);
    return taskContext;
}

Example 19 with TaskContext

use of io.prestosql.operator.TaskContext in project boostkit-bigdata by kunpengcompute.

the class OmniLocalQueryRunner method executeInternalWithPlanOnly.

private MaterializedResultWithPlan executeInternalWithPlanOnly(Session session, @Language("SQL") String sql) {
    lock.readLock().lock();
    try (Closer closer = Closer.create()) {
        AtomicReference<MaterializedResult.Builder> builder = new AtomicReference<>();
        PageConsumerOutputFactory outputFactory = new PageConsumerOutputFactory(types -> {
            builder.compareAndSet(null, MaterializedResult.resultBuilder(session, types));
            return builder.get()::page;
        });
        TaskContext taskContext = TestingTaskContext.builder(notificationExecutor, yieldExecutor, session).setMaxSpillSize(nodeSpillConfig.getMaxSpillPerNode()).setQueryMaxSpillSize(nodeSpillConfig.getQueryMaxSpillPerNode()).build();
        Plan plan = createPlan(session, sql, WarningCollector.NOOP);
        createDriversWithPlanOnly(session, plan, outputFactory, taskContext);
    } catch (IOException e) {
        throw new UncheckedIOException(e);
    } finally {
        lock.readLock().unlock();
    }
    return null;
}

Example 20 with TaskContext

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

the class TestSpatialJoinOperator method testSpatialJoin.

@Test
public void testSpatialJoin() {
    TaskContext taskContext = createTaskContext();
    RowPagesBuilder buildPages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR)).row(POLYGON_A, "A").row(null, "null").pageBreak().row(POLYGON_B, "B");
    RowPagesBuilder probePages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR)).row(POINT_X, "x").row(null, "null").row(POINT_Y, "y").pageBreak().row(POINT_Z, "z").pageBreak().row(POINT_W, "w");
    MaterializedResult expected = resultBuilder(taskContext.getSession(), ImmutableList.of(VARCHAR, VARCHAR)).row("x", "A").row("y", "A").row("y", "B").row("z", "B").build();
    assertSpatialJoin(taskContext, INNER, buildPages, probePages, expected);
}