Examples with QueryId io.trino.spi.QueryId used on opensource projects

Example 21 with QueryId

use of io.trino.spi.QueryId in project trino by trinodb.

the class TestDynamicFilterService method testDynamicFilterConsumerCallbackCount.

@Test
public void testDynamicFilterConsumerCallbackCount() {
    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<>();
    AtomicInteger callbackCount = new AtomicInteger();
    dynamicFilterService.registerDynamicFilterConsumer(queryId, 0, dynamicFilters, domains -> {
        callbackCount.getAndIncrement();
        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), filterId2, singleValue(INTEGER, 2L)));
    assertTrue(consumerCollectedFilters.isEmpty());
    // complete both filterId1 and filterId2
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId, 1, 0), ImmutableMap.of(filterId1, singleValue(INTEGER, 3L), filterId2, singleValue(INTEGER, 4L)));
    assertEquals(consumerCollectedFilters, ImmutableMap.of(filterId1, multipleValues(INTEGER, ImmutableList.of(1L, 3L)), filterId2, multipleValues(INTEGER, ImmutableList.of(2L, 4L))));
    // both filters should be received in single callback
    assertEquals(callbackCount.get(), 1);
    // register another consumer after both filters have been collected
    Map<DynamicFilterId, Domain> secondConsumerCollectedFilters = new HashMap<>();
    AtomicInteger secondCallbackCount = new AtomicInteger();
    dynamicFilterService.registerDynamicFilterConsumer(queryId, 0, dynamicFilters, domains -> {
        secondCallbackCount.getAndIncrement();
        domains.forEach((filter, domain) -> assertNull(secondConsumerCollectedFilters.put(filter, domain)));
    });
    assertEquals(secondConsumerCollectedFilters, ImmutableMap.of(filterId1, multipleValues(INTEGER, ImmutableList.of(1L, 3L)), filterId2, multipleValues(INTEGER, ImmutableList.of(2L, 4L))));
    // both filters should be received by second consumer in single callback
    assertEquals(secondCallbackCount.get(), 1);
    // first consumer should not receive callback again since it already got the completed filter
    assertEquals(callbackCount.get(), 1);
}

Example 22 with QueryId

use of io.trino.spi.QueryId in project trino by trinodb.

the class TestDynamicFilterService method testMultipleAttempts.

@Test
public void testMultipleAttempts() {
    DynamicFilterService dynamicFilterService = createDynamicFilterService();
    DynamicFilterId filterId = new DynamicFilterId("df");
    QueryId queryId = new QueryId("query");
    StageId stageId = new StageId(queryId, 0);
    dynamicFilterService.registerQuery(queryId, session, ImmutableSet.of(filterId), ImmutableSet.of(filterId), ImmutableSet.of());
    dynamicFilterService.stageCannotScheduleMoreTasks(stageId, 0, 3);
    assertFalse(dynamicFilterService.getSummary(queryId, filterId).isPresent());
    // Collect DF from 2 tasks
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId, 0, 0), ImmutableMap.of(filterId, singleValue(INTEGER, 1L)));
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId, 1, 0), ImmutableMap.of(filterId, singleValue(INTEGER, 2L)));
    assertFalse(dynamicFilterService.getSummary(queryId, filterId).isPresent());
    // Register query retry
    dynamicFilterService.registerQueryRetry(queryId, 1);
    dynamicFilterService.stageCannotScheduleMoreTasks(stageId, 1, 3);
    // Ignore update from previous attempt
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId, 2, 0), ImmutableMap.of(filterId, singleValue(INTEGER, 3L)));
    assertFalse(dynamicFilterService.getSummary(queryId, filterId).isPresent());
    // Collect DF from 3 tasks in new attempt
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId, 0, 1), ImmutableMap.of(filterId, singleValue(INTEGER, 4L)));
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId, 1, 1), ImmutableMap.of(filterId, singleValue(INTEGER, 5L)));
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId, 2, 1), ImmutableMap.of(filterId, singleValue(INTEGER, 6L)));
    assertEquals(dynamicFilterService.getSummary(queryId, filterId), Optional.of(multipleValues(INTEGER, ImmutableList.of(4L, 5L, 6L))));
    DynamicFiltersStats stats = dynamicFilterService.getDynamicFilteringStats(queryId, session);
    assertEquals(stats.getDynamicFiltersCompleted(), 1);
    assertEquals(stats.getLazyDynamicFilters(), 1);
    assertEquals(stats.getReplicatedDynamicFilters(), 0);
    assertEquals(stats.getDynamicFilterDomainStats(), ImmutableList.of(new DynamicFilterDomainStats(filterId, getSimplifiedDomainString(4L, 6L, 3, INTEGER))));
}

Example 23 with QueryId

use of io.trino.spi.QueryId in project trino by trinodb.

the class TestDynamicFilterService method testMultipleColumnMapping.

@Test
public void testMultipleColumnMapping() {
    DynamicFilterService dynamicFilterService = createDynamicFilterService();
    DynamicFilterId filterId1 = new DynamicFilterId("df1");
    SymbolAllocator symbolAllocator = new SymbolAllocator();
    Symbol symbol1 = symbolAllocator.newSymbol("DF_SYMBOL1", INTEGER);
    Symbol symbol2 = symbolAllocator.newSymbol("DF_SYMBOL2", INTEGER);
    Expression df1 = symbol1.toSymbolReference();
    Expression df2 = symbol2.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, 1);
    TestingColumnHandle column1 = new TestingColumnHandle("probeColumnA");
    TestingColumnHandle column2 = new TestingColumnHandle("probeColumnB");
    DynamicFilter dynamicFilter = dynamicFilterService.createDynamicFilter(queryId, ImmutableList.of(new DynamicFilters.Descriptor(filterId1, df1), new DynamicFilters.Descriptor(filterId1, df2)), ImmutableMap.of(symbol1, column1, symbol2, column2), symbolAllocator.getTypes());
    assertEquals(dynamicFilter.getColumnsCovered(), Set.of(column1, column2), "columns covered");
    Domain domain = singleValue(INTEGER, 1L);
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId1, 0, 0), ImmutableMap.of(filterId1, domain));
    assertEquals(dynamicFilter.getCurrentPredicate(), TupleDomain.withColumnDomains(ImmutableMap.of(column1, domain, column2, domain)));
}

Example 24 with QueryId

use of io.trino.spi.QueryId in project trino by trinodb.

the class TestDynamicFilterService method testDynamicFilter.

@Test
public void testDynamicFilter() {
    DynamicFilterService dynamicFilterService = createDynamicFilterService();
    DynamicFilterId filterId1 = new DynamicFilterId("df1");
    DynamicFilterId filterId2 = new DynamicFilterId("df2");
    DynamicFilterId filterId3 = new DynamicFilterId("df3");
    SymbolAllocator symbolAllocator = new SymbolAllocator();
    Symbol symbol1 = symbolAllocator.newSymbol("DF_SYMBOL1", INTEGER);
    Symbol symbol2 = symbolAllocator.newSymbol("DF_SYMBOL2", INTEGER);
    Symbol symbol3 = symbolAllocator.newSymbol("DF_SYMBOL3", INTEGER);
    Expression df1 = symbol1.toSymbolReference();
    Expression df2 = symbol2.toSymbolReference();
    Expression df3 = symbol3.toSymbolReference();
    QueryId queryId = new QueryId("query");
    StageId stageId1 = new StageId(queryId, 1);
    StageId stageId2 = new StageId(queryId, 2);
    StageId stageId3 = new StageId(queryId, 3);
    dynamicFilterService.registerQuery(queryId, session, ImmutableSet.of(filterId1, filterId2, filterId3), ImmutableSet.of(filterId1, filterId2, filterId3), ImmutableSet.of());
    dynamicFilterService.stageCannotScheduleMoreTasks(stageId1, 0, 2);
    dynamicFilterService.stageCannotScheduleMoreTasks(stageId2, 0, 2);
    dynamicFilterService.stageCannotScheduleMoreTasks(stageId3, 0, 2);
    DynamicFilter dynamicFilter = dynamicFilterService.createDynamicFilter(queryId, ImmutableList.of(new DynamicFilters.Descriptor(filterId1, df1), new DynamicFilters.Descriptor(filterId2, df2), new DynamicFilters.Descriptor(filterId3, df3)), ImmutableMap.of(symbol1, new TestingColumnHandle("probeColumnA"), symbol2, new TestingColumnHandle("probeColumnA"), symbol3, new TestingColumnHandle("probeColumnB")), symbolAllocator.getTypes());
    assertEquals(dynamicFilter.getColumnsCovered(), Set.of(new TestingColumnHandle("probeColumnA"), new TestingColumnHandle("probeColumnB")), "columns covered");
    assertTrue(dynamicFilter.getCurrentPredicate().isAll());
    assertFalse(dynamicFilter.isComplete());
    assertTrue(dynamicFilter.isAwaitable());
    // assert initial dynamic filtering stats
    DynamicFiltersStats stats = dynamicFilterService.getDynamicFilteringStats(queryId, session);
    assertEquals(stats.getTotalDynamicFilters(), 3);
    assertEquals(stats.getDynamicFiltersCompleted(), 0);
    assertEquals(stats.getLazyDynamicFilters(), 3);
    assertEquals(stats.getReplicatedDynamicFilters(), 0);
    // dynamic filter should be blocked waiting for tuple domain to be provided
    CompletableFuture<?> blockedFuture = dynamicFilter.isBlocked();
    assertFalse(blockedFuture.isDone());
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId1, 0, 0), ImmutableMap.of(filterId1, singleValue(INTEGER, 1L)));
    // tuple domain from two tasks are needed for dynamic filter to be narrowed down
    assertTrue(dynamicFilter.getCurrentPredicate().isAll());
    assertFalse(dynamicFilter.isComplete());
    assertTrue(dynamicFilter.isAwaitable());
    assertFalse(blockedFuture.isDone());
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId1, 1, 0), ImmutableMap.of(filterId1, singleValue(INTEGER, 2L)));
    // dynamic filter (id1) has been collected as tuple domains from two tasks have been provided
    assertEquals(dynamicFilter.getCurrentPredicate(), TupleDomain.withColumnDomains(ImmutableMap.of(new TestingColumnHandle("probeColumnA"), multipleValues(INTEGER, ImmutableList.of(1L, 2L)))));
    assertTrue(blockedFuture.isDone());
    assertFalse(blockedFuture.isCompletedExceptionally());
    stats = dynamicFilterService.getDynamicFilteringStats(queryId, session);
    assertEquals(stats.getDynamicFiltersCompleted(), 1);
    // there are still more dynamic filters to be collected
    assertFalse(dynamicFilter.isComplete());
    assertTrue(dynamicFilter.isAwaitable());
    blockedFuture = dynamicFilter.isBlocked();
    assertFalse(blockedFuture.isDone());
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId2, 0, 0), ImmutableMap.of(filterId2, singleValue(INTEGER, 2L)));
    // tuple domain from two tasks (stage 2) are needed for dynamic filter to be narrowed down
    assertEquals(dynamicFilter.getCurrentPredicate(), TupleDomain.withColumnDomains(ImmutableMap.of(new TestingColumnHandle("probeColumnA"), multipleValues(INTEGER, ImmutableList.of(1L, 2L)))));
    assertFalse(dynamicFilter.isComplete());
    assertTrue(dynamicFilter.isAwaitable());
    assertFalse(blockedFuture.isDone());
    stats = dynamicFilterService.getDynamicFilteringStats(queryId, session);
    assertEquals(stats.getDynamicFiltersCompleted(), 1);
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId2, 1, 0), ImmutableMap.of(filterId2, singleValue(INTEGER, 3L)));
    // dynamic filter (id2) has been collected as tuple domains from two tasks have been provided
    assertEquals(dynamicFilter.getCurrentPredicate(), TupleDomain.withColumnDomains(ImmutableMap.of(new TestingColumnHandle("probeColumnA"), singleValue(INTEGER, 2L))));
    assertTrue(blockedFuture.isDone());
    assertFalse(blockedFuture.isCompletedExceptionally());
    stats = dynamicFilterService.getDynamicFilteringStats(queryId, session);
    assertEquals(stats.getDynamicFiltersCompleted(), 2);
    // there are still more dynamic filters to be collected for columns A and B
    assertFalse(dynamicFilter.isComplete());
    assertTrue(dynamicFilter.isAwaitable());
    blockedFuture = dynamicFilter.isBlocked();
    assertFalse(blockedFuture.isDone());
    // make sure dynamic filter on just column A is now completed
    DynamicFilter dynamicFilterColumnA = dynamicFilterService.createDynamicFilter(queryId, ImmutableList.of(new DynamicFilters.Descriptor(filterId1, df1), new DynamicFilters.Descriptor(filterId2, df2)), ImmutableMap.of(symbol1, new TestingColumnHandle("probeColumnA"), symbol2, new TestingColumnHandle("probeColumnA")), symbolAllocator.getTypes());
    assertEquals(dynamicFilterColumnA.getColumnsCovered(), Set.of(new TestingColumnHandle("probeColumnA")), "columns covered");
    assertTrue(dynamicFilterColumnA.isComplete());
    assertFalse(dynamicFilterColumnA.isAwaitable());
    assertTrue(dynamicFilterColumnA.isBlocked().isDone());
    assertEquals(dynamicFilterColumnA.getCurrentPredicate(), TupleDomain.withColumnDomains(ImmutableMap.of(new TestingColumnHandle("probeColumnA"), singleValue(INTEGER, 2L))));
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId3, 0, 0), ImmutableMap.of(filterId3, none(INTEGER)));
    // tuple domain from two tasks (stage 3) are needed for dynamic filter to be narrowed down
    assertEquals(dynamicFilter.getCurrentPredicate(), TupleDomain.withColumnDomains(ImmutableMap.of(new TestingColumnHandle("probeColumnA"), singleValue(INTEGER, 2L))));
    assertFalse(dynamicFilter.isComplete());
    assertTrue(dynamicFilter.isAwaitable());
    assertFalse(blockedFuture.isDone());
    stats = dynamicFilterService.getDynamicFilteringStats(queryId, session);
    assertEquals(stats.getDynamicFiltersCompleted(), 2);
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId3, 1, 0), ImmutableMap.of(filterId3, none(INTEGER)));
    // "none" dynamic filter (id3) has been collected for column B as tuple domains from two tasks have been provided
    assertEquals(dynamicFilter.getCurrentPredicate(), TupleDomain.none());
    assertTrue(blockedFuture.isDone());
    assertFalse(blockedFuture.isCompletedExceptionally());
    stats = dynamicFilterService.getDynamicFilteringStats(queryId, session);
    assertEquals(stats.getDynamicFiltersCompleted(), 3);
    assertEquals(stats.getLazyDynamicFilters(), 3);
    assertEquals(stats.getReplicatedDynamicFilters(), 0);
    assertEquals(ImmutableSet.copyOf(stats.getDynamicFilterDomainStats()), ImmutableSet.of(new DynamicFilterDomainStats(filterId1, getSimplifiedDomainString(1L, 2L, 2, INTEGER)), new DynamicFilterDomainStats(filterId2, getSimplifiedDomainString(2L, 3L, 2, INTEGER)), new DynamicFilterDomainStats(filterId3, Domain.none(INTEGER).toString(session.toConnectorSession()))));
    // all dynamic filters have been collected, no need for more requests
    assertTrue(dynamicFilter.isComplete());
    assertFalse(dynamicFilter.isAwaitable());
    assertTrue(dynamicFilter.isBlocked().isDone());
}

Example 25 with QueryId

use of io.trino.spi.QueryId in project trino by trinodb.

the class TestDynamicFilterService method testDynamicFilterCoercion.

@Test
public void testDynamicFilterCoercion() {
    DynamicFilterService dynamicFilterService = createDynamicFilterService();
    DynamicFilterId filterId1 = new DynamicFilterId("df1");
    SymbolAllocator symbolAllocator = new SymbolAllocator();
    Symbol symbol1 = symbolAllocator.newSymbol("DF_SYMBOL1", INTEGER);
    Expression df1 = new Cast(symbol1.toSymbolReference(), toSqlType(BIGINT));
    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, 1);
    DynamicFilter dynamicFilter = dynamicFilterService.createDynamicFilter(queryId, ImmutableList.of(new DynamicFilters.Descriptor(filterId1, df1)), ImmutableMap.of(symbol1, new TestingColumnHandle("probeColumnA")), symbolAllocator.getTypes());
    assertEquals(dynamicFilter.getColumnsCovered(), Set.of(new TestingColumnHandle("probeColumnA")), "columns covered");
    assertFalse(dynamicFilter.isComplete());
    assertTrue(dynamicFilter.getCurrentPredicate().isAll());
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId1, 0, 0), ImmutableMap.of(filterId1, multipleValues(BIGINT, ImmutableList.of(1L, 2L, 3L))));
    assertTrue(dynamicFilter.isComplete());
    assertEquals(dynamicFilter.getCurrentPredicate(), TupleDomain.withColumnDomains(ImmutableMap.of(new TestingColumnHandle("probeColumnA"), multipleValues(INTEGER, ImmutableList.of(1L, 2L, 3L)))));
}