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

Example 36 with TaskContext

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

the class TestArbitraryOutputBuffer method testMarkers.

@Test
public void testMarkers() {
    OutputBuffers outputBuffers = createInitialEmptyOutputBuffers(ARBITRARY).withBuffer(FIRST, BROADCAST_PARTITION_ID).withBuffer(SECOND, BROADCAST_PARTITION_ID);
    ArbitraryOutputBuffer buffer = createArbitraryBuffer(outputBuffers, sizeOfPages(5));
    ScheduledExecutorService scheduler = newScheduledThreadPool(4, daemonThreadsNamed("test-%s"));
    ScheduledExecutorService scheduledExecutor = newScheduledThreadPool(2, daemonThreadsNamed("test-scheduledExecutor-%s"));
    TaskContext taskContext = createTaskContext(scheduler, scheduledExecutor, TEST_SNAPSHOT_SESSION);
    buffer.setTaskContext(taskContext);
    buffer.addInputChannel("id");
    buffer.setNoMoreInputChannels();
    MarkerPage marker1 = MarkerPage.snapshotPage(1);
    MarkerPage marker2 = MarkerPage.snapshotPage(2);
    // add one item
    addPage(buffer, createPage(0));
    // broadcast 2 pages
    addPage(buffer, marker1, true);
    addPage(buffer, marker2, true);
    // target clients: ?, 1, 2, 1, 2
    // first client gets all 3 elements
    assertBufferResultEquals(TYPES, getBufferResult(buffer, FIRST, 0, sizeOfPages(3), NO_WAIT), bufferResult(0, createPage(0), marker1, marker2));
    assertEquals(buffer.getInfo().getBuffers().stream().map(BufferInfo::getBufferedPages).collect(Collectors.toList()), Arrays.asList(3, 0));
    assertEquals(buffer.getInfo().getBuffers().stream().map(BufferInfo::getPagesSent).collect(Collectors.toList()), Arrays.asList(0L, 0L));
    // acknowledge
    buffer.get(FIRST, 3, sizeOfPages(1)).cancel(true);
    assertQueueState(buffer, 2, FIRST, 0, 3);
    assertEquals(buffer.getInfo().getBuffers().stream().map(BufferInfo::getBufferedPages).collect(Collectors.toList()), Arrays.asList(0, 0));
    assertEquals(buffer.getInfo().getBuffers().stream().map(BufferInfo::getPagesSent).collect(Collectors.toList()), Arrays.asList(3L, 0L));
    // second client gets 2 elements
    assertBufferResultEquals(TYPES, getBufferResult(buffer, SECOND, 0, sizeOfPages(3), NO_WAIT), bufferResult(0, marker1, marker2));
    assertEquals(buffer.getInfo().getBuffers().stream().map(BufferInfo::getBufferedPages).collect(Collectors.toList()), Arrays.asList(0, 2));
    assertEquals(buffer.getInfo().getBuffers().stream().map(BufferInfo::getPagesSent).collect(Collectors.toList()), Arrays.asList(3L, 0L));
    // acknowledge
    buffer.get(SECOND, 2, sizeOfPages(1)).cancel(true);
    assertEquals(buffer.getInfo().getBuffers().stream().map(BufferInfo::getBufferedPages).collect(Collectors.toList()), Arrays.asList(0, 0));
    assertEquals(buffer.getInfo().getBuffers().stream().map(BufferInfo::getPagesSent).collect(Collectors.toList()), Arrays.asList(3L, 2L));
    // New buffer receives pending markers
    OutputBuffers outputBuffers1 = outputBuffers.withBuffer(THIRD, BROADCAST_PARTITION_ID);
    assertNotNull(outputBuffers1);
    buffer.setOutputBuffers(outputBuffers);
    assertBufferResultEquals(TYPES, getBufferResult(buffer, THIRD, 0, sizeOfPages(3), NO_WAIT), bufferResult(0, marker1, marker2));
    assertEquals(buffer.getInfo().getBuffers().stream().map(BufferInfo::getBufferedPages).collect(Collectors.toList()), Arrays.asList(0, 0, 2));
    assertEquals(buffer.getInfo().getBuffers().stream().map(BufferInfo::getPagesSent).collect(Collectors.toList()), Arrays.asList(3L, 2L, 0L));
    // acknowledge
    buffer.get(THIRD, 2, sizeOfPages(1)).cancel(true);
    assertEquals(buffer.getInfo().getBuffers().stream().map(BufferInfo::getBufferedPages).collect(Collectors.toList()), Arrays.asList(0, 0, 0));
    assertEquals(buffer.getInfo().getBuffers().stream().map(BufferInfo::getPagesSent).collect(Collectors.toList()), Arrays.asList(3L, 2L, 2L));
    // finish
    buffer.setNoMorePages();
    buffer.abort(FIRST);
    buffer.abort(SECOND);
    buffer.abort(THIRD);
    assertQueueClosed(buffer, 0, FIRST, 3);
    assertQueueClosed(buffer, 0, SECOND, 2);
    assertQueueClosed(buffer, 0, THIRD, 2);
    assertFinished(buffer);
}

Example 37 with TaskContext

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

the class TestBroadcastOutputBuffer method testMarkers.

@Test
public void testMarkers() {
    OutputBuffers outputBuffers = createInitialEmptyOutputBuffers(BROADCAST).withBuffer(FIRST, BROADCAST_PARTITION_ID).withBuffer(SECOND, BROADCAST_PARTITION_ID).withNoMoreBufferIds();
    BroadcastOutputBuffer buffer = createBroadcastBuffer(outputBuffers, sizeOfPages(10));
    ScheduledExecutorService scheduler = newScheduledThreadPool(4, daemonThreadsNamed("test-%s"));
    ScheduledExecutorService scheduledExecutor = newScheduledThreadPool(2, daemonThreadsNamed("test-scheduledExecutor-%s"));
    TaskContext taskContext = createTaskContext(scheduler, scheduledExecutor, TEST_SNAPSHOT_SESSION);
    buffer.setTaskContext(taskContext);
    buffer.addInputChannel("id");
    buffer.setNoMoreInputChannels();
    MarkerPage marker1 = MarkerPage.snapshotPage(1);
    MarkerPage marker2 = MarkerPage.snapshotPage(2);
    // add one item
    addPage(buffer, createPage(0));
    // broadcast 2 pages
    addPage(buffer, marker1, true);
    addPage(buffer, marker2, true);
    // first client gets 2 elements
    assertBufferResultEquals(TYPES, getBufferResult(buffer, FIRST, 0, sizeOfPages(3), NO_WAIT), bufferResult(0, createPage(0), marker1, marker2));
    assertBufferResultEquals(TYPES, getBufferResult(buffer, SECOND, 0, sizeOfPages(3), NO_WAIT), bufferResult(0, createPage(0), marker1, marker2));
    assertEquals(buffer.getInfo().getBuffers().stream().map(BufferInfo::getBufferedPages).collect(Collectors.toList()), Arrays.asList(3, 3));
    assertEquals(buffer.getInfo().getBuffers().stream().map(BufferInfo::getPagesSent).collect(Collectors.toList()), Arrays.asList(0L, 0L));
    // acknowledge
    buffer.get(FIRST, 3, sizeOfPages(1)).cancel(true);
    assertQueueState(buffer, FIRST, 0, 3);
    buffer.get(SECOND, 3, sizeOfPages(1)).cancel(true);
    assertQueueState(buffer, SECOND, 0, 3);
    assertEquals(buffer.getInfo().getBuffers().stream().map(BufferInfo::getBufferedPages).collect(Collectors.toList()), Arrays.asList(0, 0));
    assertEquals(buffer.getInfo().getBuffers().stream().map(BufferInfo::getPagesSent).collect(Collectors.toList()), Arrays.asList(3L, 3L));
    // finish
    buffer.setNoMorePages();
    buffer.abort(FIRST);
    buffer.abort(SECOND);
    assertQueueClosed(buffer, FIRST, 3);
    assertQueueClosed(buffer, SECOND, 3);
    assertFinished(buffer);
}

Example 38 with TaskContext

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

the class OmniLocalQueryRunner method executeInternal.

private MaterializedResultWithPlan executeInternal(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);
        List<Driver> drivers = createDrivers(session, plan, outputFactory, taskContext);
        drivers.forEach(closer::register);
        boolean done = false;
        while (!done) {
            boolean processed = false;
            for (Driver driver : drivers) {
                if (alwaysRevokeMemory) {
                    driver.getDriverContext().getOperatorContexts().stream().filter(operatorContext -> operatorContext.getOperatorStats().getRevocableMemoryReservation().getValue() > 0).forEach(OperatorContext::requestMemoryRevoking);
                }
                if (!driver.isFinished()) {
                    driver.process();
                    processed = true;
                }
            }
            done = !processed;
        }
        verify(builder.get() != null, "Output operator was not created");
        return new MaterializedResultWithPlan(builder.get().build(), plan);
    } catch (IOException e) {
        throw new UncheckedIOException(e);
    } finally {
        lock.readLock().unlock();
    }
}

Example 39 with TaskContext

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

the class VecAllocatorHelper method createOperatorLevelAllocator.

/**
 * create an operator level allocator based on driver context.
 *
 * @param driverContext diver context
 * @param limit allocator limit
 * @param jazz operator Class
 * @return operator allocator
 */
public static VecAllocator createOperatorLevelAllocator(DriverContext driverContext, long limit, Class<?> jazz) {
    TaskContext taskContext = driverContext.getPipelineContext().getTaskContext();
    VecAllocator vecAllocator = getVecAllocatorFromTaskContext(taskContext);
    return createOperatorLevelAllocator(vecAllocator, limit, taskContext.getTaskId().toString(), 0, jazz);
}