Examples with DynamicFilterId io.trino.sql.planner.plan.DynamicFilterId used on opensource projects

Example 51 with DynamicFilterId

use of io.trino.sql.planner.plan.DynamicFilterId in project trino by trinodb.

the class TestDynamicFilterService method testDynamicFilterSummaryCompletion.

@Test
public void testDynamicFilterSummaryCompletion() {
    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());
    // assert initial dynamic filtering stats
    DynamicFiltersStats stats = dynamicFilterService.getDynamicFilteringStats(queryId, session);
    assertEquals(stats.getTotalDynamicFilters(), 1);
    assertEquals(stats.getDynamicFiltersCompleted(), 0);
    assertEquals(stats.getLazyDynamicFilters(), 1);
    assertEquals(stats.getReplicatedDynamicFilters(), 0);
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId, 0, 0), ImmutableMap.of(filterId, singleValue(INTEGER, 1L)));
    assertFalse(dynamicFilterService.getSummary(queryId, filterId).isPresent());
    stats = dynamicFilterService.getDynamicFilteringStats(queryId, session);
    assertEquals(stats.getDynamicFiltersCompleted(), 0);
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId, 1, 0), ImmutableMap.of(filterId, singleValue(INTEGER, 2L)));
    assertFalse(dynamicFilterService.getSummary(queryId, filterId).isPresent());
    stats = dynamicFilterService.getDynamicFilteringStats(queryId, session);
    assertEquals(stats.getDynamicFiltersCompleted(), 0);
    dynamicFilterService.addTaskDynamicFilters(new TaskId(stageId, 2, 0), ImmutableMap.of(filterId, singleValue(INTEGER, 3L)));
    Optional<Domain> summary = dynamicFilterService.getSummary(queryId, filterId);
    assertTrue(summary.isPresent());
    assertEquals(summary.get(), multipleValues(INTEGER, ImmutableList.of(1L, 2L, 3L)));
    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(1L, 3L, 3, INTEGER))));
}

Example 52 with DynamicFilterId

use of io.trino.sql.planner.plan.DynamicFilterId in project trino by trinodb.

the class DynamicFilterService method collectDynamicFilters.

private void collectDynamicFilters(StageId stageId, Optional<Set<DynamicFilterId>> selectedFilters) {
    DynamicFilterContext context = dynamicFilterContexts.get(stageId.getQueryId());
    if (context == null) {
        // query has been removed
        return;
    }
    OptionalInt stageNumberOfTasks = context.getNumberOfTasks(stageId);
    Map<DynamicFilterId, List<Domain>> newDynamicFilters = context.getTaskDynamicFilters(stageId, selectedFilters).entrySet().stream().filter(stageDomains -> {
        if (stageDomains.getValue().stream().anyMatch(Domain::isAll)) {
            // if one of the domains is all, we don't need to get dynamic filters from all tasks
            return true;
        }
        if (!stageDomains.getValue().isEmpty() && context.getReplicatedDynamicFilters().contains(stageDomains.getKey())) {
            // for replicated dynamic filters it's enough to get dynamic filter from a single task
            checkState(stageDomains.getValue().size() == 1, "Replicated dynamic filter should be collected from single task");
            return true;
        }
        // check if all tasks of a dynamic filter source have reported dynamic filter summary
        return stageNumberOfTasks.isPresent() && stageDomains.getValue().size() == stageNumberOfTasks.getAsInt();
    }).collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
    context.addDynamicFilters(newDynamicFilters);
}

Example 53 with DynamicFilterId

use of io.trino.sql.planner.plan.DynamicFilterId 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);
    }
}

Example 54 with DynamicFilterId

use of io.trino.sql.planner.plan.DynamicFilterId in project trino by trinodb.

the class DynamicFilterService method getSourceStageInnerLazyDynamicFilters.

@VisibleForTesting
static Set<DynamicFilterId> getSourceStageInnerLazyDynamicFilters(PlanFragment plan) {
    if (!plan.getPartitioning().equals(SOURCE_DISTRIBUTION)) {
        // dynamic filtering collecting task can be added.
        return ImmutableSet.of();
    }
    PlanNode planNode = plan.getRoot();
    Set<DynamicFilterId> innerStageDynamicFilters = intersection(getProducedDynamicFilters(planNode), getConsumedDynamicFilters(planNode));
    Set<DynamicFilterId> replicatedDynamicFilters = getReplicatedDynamicFilters(planNode);
    return ImmutableSet.copyOf(intersection(innerStageDynamicFilters, replicatedDynamicFilters));
}

Example 55 with DynamicFilterId

use of io.trino.sql.planner.plan.DynamicFilterId in project trino by trinodb.

the class JoinMatcher method matchDynamicFilters.

private boolean matchDynamicFilters(JoinNode joinNode, SymbolAliases symbolAliases) {
    if (expectedDynamicFilter.isEmpty()) {
        return true;
    }
    Set<DynamicFilterId> dynamicFilterIds = joinNode.getDynamicFilters().keySet();
    List<DynamicFilters.Descriptor> descriptors = searchFrom(joinNode.getLeft()).where(FilterNode.class::isInstance).findAll().stream().flatMap(filterNode -> extractExpressions(filterNode).stream()).flatMap(expression -> extractDynamicFilters(expression).getDynamicConjuncts().stream()).filter(descriptor -> dynamicFilterIds.contains(descriptor.getId())).collect(toImmutableList());
    Map<DynamicFilterId, Symbol> idToBuildSymbolMap = joinNode.getDynamicFilters();
    Set<Expression> actual = new HashSet<>();
    for (DynamicFilters.Descriptor descriptor : descriptors) {
        Expression probe = descriptor.getInput();
        Symbol build = idToBuildSymbolMap.get(descriptor.getId());
        if (build == null) {
            return false;
        }
        Expression expression;
        if (descriptor.isNullAllowed()) {
            expression = new NotExpression(new ComparisonExpression(IS_DISTINCT_FROM, probe, build.toSymbolReference()));
        } else {
            expression = new ComparisonExpression(descriptor.getOperator(), probe, build.toSymbolReference());
        }
        actual.add(expression);
    }
    Set<Expression> expected = expectedDynamicFilter.get().stream().map(pattern -> pattern.getExpression(symbolAliases)).collect(toImmutableSet());
    return expected.equals(actual);
}