Examples with TaskId io.trino.execution.TaskId used on opensource projects

Example 21 with TaskId

use of io.trino.execution.TaskId in project trino by trinodb.

the class TestDynamicFilterService method testShortCircuitOnAllTupleDomain.

@Test
public void testShortCircuitOnAllTupleDomain() {
    DynamicFilterService dynamicFilterService = createDynamicFilterService();
    DynamicFilterId filterId1 = new DynamicFilterId("df1");
    SymbolAllocator symbolAllocator = new SymbolAllocator();
    Symbol symbol1 = symbolAllocator.newSymbol("DF_SYMBOL1", INTEGER);
    Expression df1 = symbol1.toSymbolReference();
    QueryId queryId = new QueryId("query");
    StageId stageId1 = new StageId(queryId, 1);
    dynamicFilterService.registerQuery(queryId, session, ImmutableSet.of(filterId1), ImmutableSet.of(filterId1), ImmutableSet.of());
    dynamicFilterService.stageCannotScheduleMoreTasks(stageId1, 0, 2);
    DynamicFilter dynamicFilter = dynamicFilterService.createDynamicFilter(queryId, ImmutableList.of(new DynamicFilters.Descriptor(filterId1, df1)), ImmutableMap.of(symbol1, new TestingColumnHandle("probeColumnA")), symbolAllocator.getTypes());
    // dynamic filter is initially blocked
    assertTrue(dynamicFilter.getCurrentPredicate().isAll());
    assertFalse(dynamicFilter.isComplete());
    assertFalse(dynamicFilter.isBlocked().isDone());
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId1, 1, 0), ImmutableMap.of(filterId1, Domain.all(INTEGER)));
    // dynamic filter should be unblocked and completed
    assertTrue(dynamicFilter.getCurrentPredicate().isAll());
    assertTrue(dynamicFilter.isComplete());
    assertTrue(dynamicFilter.isBlocked().isDone());
}

Example 22 with TaskId

use of io.trino.execution.TaskId in project trino by trinodb.

the class TestDynamicFilterService method testDynamicFilterConsumer.

@Test
public void testDynamicFilterConsumer() {
    DynamicFilterService dynamicFilterService = createDynamicFilterService();
    DynamicFilterId filterId1 = new DynamicFilterId("df1");
    DynamicFilterId filterId2 = new DynamicFilterId("df2");
    Set<DynamicFilterId> dynamicFilters = ImmutableSet.of(filterId1, filterId2);
    QueryId queryId = new QueryId("query");
    StageId stageId = new StageId(queryId, 0);
    dynamicFilterService.registerQuery(queryId, session, dynamicFilters, dynamicFilters, ImmutableSet.of());
    dynamicFilterService.stageCannotScheduleMoreTasks(stageId, 0, 2);
    Map<DynamicFilterId, Domain> consumerCollectedFilters = new HashMap<>();
    dynamicFilterService.registerDynamicFilterConsumer(queryId, 0, dynamicFilters, domains -> domains.forEach((filter, domain) -> assertNull(consumerCollectedFilters.put(filter, domain))));
    assertTrue(consumerCollectedFilters.isEmpty());
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId, 0, 0), ImmutableMap.of(filterId1, singleValue(INTEGER, 1L)));
    assertTrue(consumerCollectedFilters.isEmpty());
    // complete only filterId1
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId, 1, 0), ImmutableMap.of(filterId1, singleValue(INTEGER, 3L), filterId2, singleValue(INTEGER, 2L)));
    assertEquals(consumerCollectedFilters, ImmutableMap.of(filterId1, multipleValues(INTEGER, ImmutableList.of(1L, 3L))));
    // register another consumer only for filterId1 after completion of filterId1
    Map<DynamicFilterId, Domain> secondConsumerCollectedFilters = new HashMap<>();
    dynamicFilterService.registerDynamicFilterConsumer(queryId, 0, ImmutableSet.of(filterId1), domains -> domains.forEach((filter, domain) -> assertNull(secondConsumerCollectedFilters.put(filter, domain))));
    assertEquals(secondConsumerCollectedFilters, ImmutableMap.of(filterId1, multipleValues(INTEGER, ImmutableList.of(1L, 3L))));
    // complete filterId2
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId, 0, 0), ImmutableMap.of(filterId2, singleValue(INTEGER, 4L)));
    assertEquals(consumerCollectedFilters, ImmutableMap.of(filterId1, multipleValues(INTEGER, ImmutableList.of(1L, 3L)), filterId2, multipleValues(INTEGER, ImmutableList.of(2L, 4L))));
    assertEquals(secondConsumerCollectedFilters, ImmutableMap.of(filterId1, multipleValues(INTEGER, ImmutableList.of(1L, 3L))));
}

Example 23 with TaskId

use of io.trino.execution.TaskId in project trino by trinodb.

the class TestHttpRemoteTask method testOutboundDynamicFilters.

@Test(timeOut = 30_000)
public void testOutboundDynamicFilters() throws Exception {
    DynamicFilterId filterId1 = new DynamicFilterId("df1");
    DynamicFilterId filterId2 = new DynamicFilterId("df2");
    SymbolAllocator symbolAllocator = new SymbolAllocator();
    Symbol symbol1 = symbolAllocator.newSymbol("DF_SYMBOL1", BIGINT);
    Symbol symbol2 = symbolAllocator.newSymbol("DF_SYMBOL2", BIGINT);
    SymbolReference df1 = symbol1.toSymbolReference();
    SymbolReference df2 = symbol2.toSymbolReference();
    ColumnHandle handle1 = new TestingColumnHandle("column1");
    ColumnHandle handle2 = new TestingColumnHandle("column2");
    QueryId queryId = new QueryId("test");
    TestingTaskResource testingTaskResource = new TestingTaskResource(new AtomicLong(System.nanoTime()), FailureScenario.NO_FAILURE);
    DynamicFilterService dynamicFilterService = new DynamicFilterService(PLANNER_CONTEXT.getMetadata(), PLANNER_CONTEXT.getFunctionManager(), new TypeOperators(), newDirectExecutorService());
    dynamicFilterService.registerQuery(queryId, TEST_SESSION, ImmutableSet.of(filterId1, filterId2), ImmutableSet.of(filterId1, filterId2), ImmutableSet.of());
    dynamicFilterService.stageCannotScheduleMoreTasks(new StageId(queryId, 1), 0, 1);
    DynamicFilter dynamicFilter = dynamicFilterService.createDynamicFilter(queryId, ImmutableList.of(new DynamicFilters.Descriptor(filterId1, df1), new DynamicFilters.Descriptor(filterId2, df2)), ImmutableMap.of(symbol1, handle1, symbol2, handle2), symbolAllocator.getTypes());
    // make sure initial dynamic filter is collected
    CompletableFuture<?> future = dynamicFilter.isBlocked();
    dynamicFilterService.addTaskDynamicFilters(new TaskId(new StageId(queryId.getId(), 1), 1, 0), ImmutableMap.of(filterId1, Domain.singleValue(BIGINT, 1L)));
    future.get();
    assertEquals(dynamicFilter.getCurrentPredicate(), TupleDomain.withColumnDomains(ImmutableMap.of(handle1, Domain.singleValue(BIGINT, 1L))));
    // Create remote task after dynamic filter is created to simulate new nodes joining
    HttpRemoteTaskFactory httpRemoteTaskFactory = createHttpRemoteTaskFactory(testingTaskResource, dynamicFilterService);
    RemoteTask remoteTask = createRemoteTask(httpRemoteTaskFactory, ImmutableSet.of(filterId1, filterId2));
    testingTaskResource.setInitialTaskInfo(remoteTask.getTaskInfo());
    remoteTask.start();
    assertEventually(new Duration(10, SECONDS), () -> assertEquals(testingTaskResource.getDynamicFiltersSentCounter(), 1L));
    assertEquals(testingTaskResource.getCreateOrUpdateCounter(), 1L);
    // schedule a couple of splits to trigger task updates
    addSplit(remoteTask, testingTaskResource, 1);
    addSplit(remoteTask, testingTaskResource, 2);
    // make sure dynamic filter was sent in task updates only once
    assertEquals(testingTaskResource.getDynamicFiltersSentCounter(), 1L);
    assertEquals(testingTaskResource.getCreateOrUpdateCounter(), 3L);
    assertEquals(testingTaskResource.getLatestDynamicFilterFromCoordinator(), ImmutableMap.of(filterId1, Domain.singleValue(BIGINT, 1L)));
    future = dynamicFilter.isBlocked();
    dynamicFilterService.addTaskDynamicFilters(new TaskId(new StageId(queryId.getId(), 1), 1, 0), ImmutableMap.of(filterId2, Domain.singleValue(BIGINT, 2L)));
    future.get();
    assertEquals(dynamicFilter.getCurrentPredicate(), TupleDomain.withColumnDomains(ImmutableMap.of(handle1, Domain.singleValue(BIGINT, 1L), handle2, Domain.singleValue(BIGINT, 2L))));
    // dynamic filter should be sent even though there were no further splits scheduled
    assertEventually(new Duration(10, SECONDS), () -> assertEquals(testingTaskResource.getDynamicFiltersSentCounter(), 2L));
    assertEquals(testingTaskResource.getCreateOrUpdateCounter(), 4L);
    // previously sent dynamic filter should not be repeated
    assertEquals(testingTaskResource.getLatestDynamicFilterFromCoordinator(), ImmutableMap.of(filterId2, Domain.singleValue(BIGINT, 2L)));
    httpRemoteTaskFactory.stop();
    dynamicFilterService.stop();
}

Example 24 with TaskId

use of io.trino.execution.TaskId in project trino by trinodb.

the class TestFaultTolerantStageScheduler method testHappyPath.

@Test
public void testHappyPath() throws Exception {
    TestingRemoteTaskFactory remoteTaskFactory = new TestingRemoteTaskFactory();
    TestingTaskSourceFactory taskSourceFactory = createTaskSourceFactory(5, 2);
    TestingNodeSupplier nodeSupplier = TestingNodeSupplier.create(ImmutableMap.of(NODE_1, ImmutableList.of(CATALOG), NODE_2, ImmutableList.of(CATALOG), NODE_3, ImmutableList.of(CATALOG)));
    setupNodeAllocatorService(nodeSupplier);
    TestingExchange sinkExchange = new TestingExchange(false);
    TestingExchange sourceExchange1 = new TestingExchange(false);
    TestingExchange sourceExchange2 = new TestingExchange(false);
    try (NodeAllocator nodeAllocator = nodeAllocatorService.getNodeAllocator(SESSION, 1)) {
        FaultTolerantStageScheduler scheduler = createFaultTolerantTaskScheduler(remoteTaskFactory, taskSourceFactory, nodeAllocator, TaskLifecycleListener.NO_OP, Optional.of(sinkExchange), ImmutableMap.of(SOURCE_FRAGMENT_ID_1, sourceExchange1, SOURCE_FRAGMENT_ID_2, sourceExchange2), 2);
        ListenableFuture<Void> blocked = scheduler.isBlocked();
        assertUnblocked(blocked);
        scheduler.schedule();
        blocked = scheduler.isBlocked();
        // blocked on first source exchange
        assertBlocked(blocked);
        sourceExchange1.setSourceHandles(ImmutableList.of(new TestingExchangeSourceHandle(0, 1)));
        // still blocked on the second source exchange
        assertBlocked(blocked);
        assertFalse(scheduler.isBlocked().isDone());
        sourceExchange2.setSourceHandles(ImmutableList.of(new TestingExchangeSourceHandle(0, 1)));
        // now unblocked
        assertUnblocked(blocked);
        assertUnblocked(scheduler.isBlocked());
        scheduler.schedule();
        blocked = scheduler.isBlocked();
        // blocked on node allocation
        assertBlocked(blocked);
        // not all tasks have been enumerated yet
        assertFalse(sinkExchange.isNoMoreSinks());
        Map<TaskId, TestingRemoteTask> tasks = remoteTaskFactory.getTasks();
        // one task per node
        assertThat(tasks).hasSize(3);
        assertThat(tasks).containsKey(getTaskId(0, 0));
        assertThat(tasks).containsKey(getTaskId(1, 0));
        assertThat(tasks).containsKey(getTaskId(2, 0));
        TestingRemoteTask task = tasks.get(getTaskId(0, 0));
        // fail task for partition 0
        task.fail(new RuntimeException("some failure"));
        assertUnblocked(blocked);
        assertUnblocked(scheduler.isBlocked());
        // schedule more tasks
        scheduler.schedule();
        tasks = remoteTaskFactory.getTasks();
        assertThat(tasks).hasSize(4);
        assertThat(tasks).containsKey(getTaskId(3, 0));
        blocked = scheduler.isBlocked();
        // blocked on task scheduling
        assertBlocked(blocked);
        // finish some task
        assertThat(tasks).containsKey(getTaskId(1, 0));
        tasks.get(getTaskId(1, 0)).finish();
        assertUnblocked(blocked);
        assertUnblocked(scheduler.isBlocked());
        assertThat(sinkExchange.getFinishedSinkHandles()).contains(new TestingExchangeSinkHandle(1));
        // this will schedule failed task
        scheduler.schedule();
        blocked = scheduler.isBlocked();
        // blocked on task scheduling
        assertBlocked(blocked);
        tasks = remoteTaskFactory.getTasks();
        assertThat(tasks).hasSize(5);
        assertThat(tasks).containsKey(getTaskId(0, 1));
        // finish some task
        tasks = remoteTaskFactory.getTasks();
        assertThat(tasks).containsKey(getTaskId(3, 0));
        tasks.get(getTaskId(3, 0)).finish();
        assertThat(sinkExchange.getFinishedSinkHandles()).contains(new TestingExchangeSinkHandle(1), new TestingExchangeSinkHandle(3));
        assertUnblocked(blocked);
        // schedule the last task
        scheduler.schedule();
        tasks = remoteTaskFactory.getTasks();
        assertThat(tasks).hasSize(6);
        assertThat(tasks).containsKey(getTaskId(4, 0));
        // not finished yet, will be finished when all tasks succeed
        assertFalse(scheduler.isFinished());
        blocked = scheduler.isBlocked();
        // blocked on task scheduling
        assertBlocked(blocked);
        tasks = remoteTaskFactory.getTasks();
        assertThat(tasks).containsKey(getTaskId(4, 0));
        // finish remaining tasks
        tasks.get(getTaskId(0, 1)).finish();
        tasks.get(getTaskId(2, 0)).finish();
        tasks.get(getTaskId(4, 0)).finish();
        // now it's not blocked and finished
        assertUnblocked(blocked);
        assertUnblocked(scheduler.isBlocked());
        assertThat(sinkExchange.getFinishedSinkHandles()).contains(new TestingExchangeSinkHandle(0), new TestingExchangeSinkHandle(1), new TestingExchangeSinkHandle(2), new TestingExchangeSinkHandle(3), new TestingExchangeSinkHandle(4));
        assertTrue(scheduler.isFinished());
    }
}

Example 25 with TaskId

use of io.trino.execution.TaskId in project trino by trinodb.

the class TestFixedCountScheduler method testSingleNode.

@Test
public void testSingleNode() {
    FixedCountScheduler nodeScheduler = new FixedCountScheduler((node, partition) -> Optional.of(taskFactory.createTableScanTask(new TaskId(new StageId("test", 1), 1, 0), node, ImmutableList.of(), new PartitionedSplitCountTracker(delta -> {
    }))), generateRandomNodes(1));
    ScheduleResult result = nodeScheduler.schedule();
    assertTrue(result.isFinished());
    assertTrue(result.getBlocked().isDone());
    assertEquals(result.getNewTasks().size(), 1);
    assertTrue(result.getNewTasks().iterator().next().getNodeId().equals("other 0"));
}