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

Example 1 with TaskContext

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

the class TestDynamicFiltersCollector method TestCollectingGlobalDynamicFilters.

@Test
public void TestCollectingGlobalDynamicFilters() throws InterruptedException {
    final QueryId queryId = new QueryId("test_query");
    final String filterId = "1";
    final String columnName = "column";
    final TestingColumnHandle columnHandle = new TestingColumnHandle(columnName);
    final Set<String> valueSet = ImmutableSet.of("1", "2", "3");
    TaskContext taskContext = mock(TaskContext.class);
    Session session = testSessionBuilder().setQueryId(queryId).setSystemProperty(ENABLE_DYNAMIC_FILTERING, "true").setSystemProperty(DYNAMIC_FILTERING_DATA_TYPE, "HASHSET").build();
    when(taskContext.getSession()).thenReturn(session);
    // set up state store and merged dynamic filters map
    Map mockMap = new HashMap<>();
    StateStoreProvider stateStoreProvider = mock(StateStoreProvider.class);
    StateStore stateStore = mock(StateStore.class);
    StateMap stateMap = new MockStateMap<>("test-map", mockMap);
    when(stateStoreProvider.getStateStore()).thenReturn(stateStore);
    when(stateStore.getStateCollection(any())).thenReturn(stateMap);
    when(stateStore.createStateMap(any())).thenReturn(stateMap);
    when(stateStore.getOrCreateStateCollection(any(), any())).thenReturn(stateMap);
    // set up state store listener and dynamic filter cache
    StateStoreListenerManager stateStoreListenerManager = new StateStoreListenerManager(stateStoreProvider);
    DynamicFilterCacheManager dynamicFilterCacheManager = new DynamicFilterCacheManager();
    stateStoreListenerManager.addStateStoreListener(new DynamicFilterListener(dynamicFilterCacheManager), MERGED_DYNAMIC_FILTERS);
    LocalDynamicFiltersCollector collector = new LocalDynamicFiltersCollector(taskContext, Optional.empty(), dynamicFilterCacheManager);
    TableScanNode tableScan = mock(TableScanNode.class);
    when(tableScan.getAssignments()).thenReturn(ImmutableMap.of(new Symbol(columnName), columnHandle));
    List<DynamicFilters.Descriptor> dynamicFilterDescriptors = ImmutableList.of(new DynamicFilters.Descriptor(filterId, new VariableReferenceExpression(columnName, BIGINT)));
    collector.initContext(ImmutableList.of(dynamicFilterDescriptors), SymbolUtils.toLayOut(tableScan.getOutputSymbols()));
    assertTrue(collector.getDynamicFilters(tableScan).isEmpty(), "there should be no dynamic filter available");
    // put some values in state store as a new dynamic filter
    // and wait for the listener to process the event
    stateMap.put(createKey(DynamicFilterUtils.FILTERPREFIX, filterId, queryId.getId()), valueSet);
    TimeUnit.MILLISECONDS.sleep(100);
    // get available dynamic filter and verify it
    List<Map<ColumnHandle, DynamicFilter>> dynamicFilters = collector.getDynamicFilters(tableScan);
    assertEquals(dynamicFilters.size(), 1, "there should be a new dynamic filter");
    assertEquals(dynamicFilters.size(), 1);
    DynamicFilter dynamicFilter = dynamicFilters.get(0).get(columnHandle);
    assertTrue(dynamicFilter instanceof HashSetDynamicFilter, "new dynamic filter should be hashset");
    assertEquals(dynamicFilter.getSize(), valueSet.size(), "new dynamic filter should have correct size");
    for (String value : valueSet) {
        assertTrue(dynamicFilter.contains(value), "new dynamic filter should contain correct values");
    }
    // clean up when task finishes
    collector.removeDynamicFilter(true);
    DynamicFilter cachedFilter = dynamicFilterCacheManager.getDynamicFilter(DynamicFilterCacheManager.createCacheKey(filterId, queryId.getId()));
    assertNull(cachedFilter, "cached dynamic filter should have been removed");
}

Example 2 with TaskContext

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

the class LocalQueryRunner 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 3 with TaskContext

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

the class TestingTaskContext method createTaskContext.

private static TaskContext createTaskContext(QueryContext queryContext, Session session, TaskStateMachine taskStateMachine) {
    TaskContext taskContext = queryContext.addTaskContext(taskStateMachine, session, true, true, OptionalInt.empty(), Optional.empty(), TESTING_SERDE_FACTORY);
    taskContext.getSnapshotManager().setTotalComponents(100);
    return taskContext;
}

Example 4 with TaskContext

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

the class MemoryLocalQueryRunner method execute.

public List<Page> execute(@Language("SQL") String query) {
    MemoryPool memoryPool = new MemoryPool(new MemoryPoolId("test"), new DataSize(2, GIGABYTE));
    SpillSpaceTracker spillSpaceTracker = new SpillSpaceTracker(new DataSize(1, GIGABYTE));
    QueryContext queryContext = new QueryContext(new QueryId("test"), new DataSize(1, GIGABYTE), new DataSize(2, GIGABYTE), memoryPool, new TestingGcMonitor(), localQueryRunner.getExecutor(), localQueryRunner.getScheduler(), new DataSize(4, GIGABYTE), spillSpaceTracker, NOOP_SNAPSHOT_UTILS);
    TaskContext taskContext = queryContext.addTaskContext(new TaskStateMachine(new TaskId("query", 0, 0), localQueryRunner.getExecutor()), localQueryRunner.getDefaultSession(), false, false, OptionalInt.empty(), Optional.empty(), TESTING_SERDE_FACTORY);
    // Use NullOutputFactory to avoid coping out results to avoid affecting benchmark results
    ImmutableList.Builder<Page> output = ImmutableList.builder();
    List<Driver> drivers = localQueryRunner.createDrivers(query, new PageConsumerOperator.PageConsumerOutputFactory(types -> output::add), taskContext);
    boolean done = false;
    while (!done) {
        boolean processed = false;
        for (Driver driver : drivers) {
            if (!driver.isFinished()) {
                driver.process();
                processed = true;
            }
        }
        done = !processed;
    }
    return output.build();
}

Example 5 with TaskContext

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

the class TestPartitionedOutputBuffer method testMarkers.

@Test
public void testMarkers() {
    int firstPartition = 0;
    int secondPartition = 1;
    PartitionedOutputBuffer buffer = createPartitionedBuffer(createInitialEmptyOutputBuffers(PARTITIONED).withBuffer(FIRST, firstPartition).withBuffer(SECOND, secondPartition).withNoMoreBufferIds(), sizeOfPages(6));
    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 1 item to each buffer
    addPage(buffer, createPage(0), firstPartition);
    addPage(buffer, createPage(0), secondPartition);
    // broadcast 2 pages
    addPage(buffer, marker1, true);
    addPage(buffer, marker2, true);
    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));
    // try to add one more page, which should block
    ListenableFuture<?> future = enqueuePage(buffer, createPage(3), firstPartition);
    assertFalse(future.isDone());
    assertEquals(buffer.getInfo().getBuffers().stream().map(BufferInfo::getBufferedPages).collect(Collectors.toList()), Arrays.asList(4, 3));
    assertEquals(buffer.getInfo().getBuffers().stream().map(BufferInfo::getPagesSent).collect(Collectors.toList()), Arrays.asList(0L, 0L));
    // get most 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(4, 3));
    assertEquals(buffer.getInfo().getBuffers().stream().map(BufferInfo::getPagesSent).collect(Collectors.toList()), Arrays.asList(0L, 0L));
    // we should still be blocked
    assertFalse(future.isDone());
    // acknowledge all pages
    buffer.get(FIRST, 3, sizeOfPages(3)).cancel(true);
    buffer.get(SECOND, 3, sizeOfPages(3)).cancel(true);
    assertEquals(buffer.getInfo().getBuffers().stream().map(BufferInfo::getBufferedPages).collect(Collectors.toList()), Arrays.asList(1, 0));
    assertEquals(buffer.getInfo().getBuffers().stream().map(BufferInfo::getPagesSent).collect(Collectors.toList()), Arrays.asList(3L, 3L));
    // finish the buffer
    assertFalse(buffer.isFinished());
    buffer.setNoMorePages();
    assertBufferResultEquals(TYPES, getBufferResult(buffer, FIRST, 3, sizeOfPages(1), NO_WAIT), bufferResult(3, createPage(3)));
    buffer.get(FIRST, 4, 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(4L, 3L));
    assertFutureIsDone(future);
    buffer.abort(FIRST);
    buffer.abort(SECOND);
    assertQueueClosed(buffer, FIRST, 4);
    assertQueueClosed(buffer, SECOND, 3);
    assertFinished(buffer);
}