Examples with DynamicFilters io.trino.sql.DynamicFilters used on opensource projects

Example 1 with DynamicFilters

use of io.trino.sql.DynamicFilters 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 2 with DynamicFilters

use of io.trino.sql.DynamicFilters in project trino by trinodb.

the class DynamicFilterService method registerQuery.

public void registerQuery(SqlQueryExecution sqlQueryExecution, SubPlan fragmentedPlan) {
    PlanNode queryPlan = sqlQueryExecution.getQueryPlan().getRoot();
    Set<DynamicFilterId> dynamicFilters = getProducedDynamicFilters(queryPlan);
    Set<DynamicFilterId> replicatedDynamicFilters = getReplicatedDynamicFilters(queryPlan);
    Set<DynamicFilterId> lazyDynamicFilters = fragmentedPlan.getAllFragments().stream().flatMap(plan -> getLazyDynamicFilters(plan).stream()).collect(toImmutableSet());
    // register query only if it contains dynamic filters
    if (!dynamicFilters.isEmpty()) {
        registerQuery(sqlQueryExecution.getQueryId(), sqlQueryExecution.getSession(), dynamicFilters, lazyDynamicFilters, replicatedDynamicFilters);
    }
}