Examples with SET io.prestosql.spi.statestore.StateCollection.Type.SET used on opensource projects

Example 1 with SET

use of io.prestosql.spi.statestore.StateCollection.Type.SET in project hetu-core by openlookeng.

the class LocalDynamicFilter method create.

public static Optional<LocalDynamicFilter> create(JoinNode planNode, int partitionCount, Session session, TaskId taskId, StateStoreProvider stateStoreProvider) {
    Set<String> joinDynamicFilters = planNode.getDynamicFilters().keySet();
    // Mapping from probe-side dynamic filters' IDs to their matching probe symbols.
    Multimap<String, Symbol> localProbeSymbols = MultimapBuilder.treeKeys().arrayListValues().build();
    PlanNode buildNode;
    DynamicFilter.Type localType = DynamicFilter.Type.LOCAL;
    List<FilterNode> filterNodes = findFilterNodeInStage(planNode);
    if (filterNodes.isEmpty()) {
        buildNode = planNode.getRight();
        mapProbeSymbolsFromCriteria(planNode.getDynamicFilters(), localProbeSymbols, planNode.getCriteria());
        localType = DynamicFilter.Type.GLOBAL;
    } else {
        buildNode = planNode.getRight();
        for (FilterNode filterNode : filterNodes) {
            mapProbeSymbols(filterNode.getPredicate(), joinDynamicFilters, localProbeSymbols);
        }
    }
    final List<Symbol> buildSideSymbols = buildNode.getOutputSymbols();
    Map<String, Integer> localBuildChannels = planNode.getDynamicFilters().entrySet().stream().filter(entry -> localProbeSymbols.containsKey(entry.getKey())).collect(toMap(// Dynamic filter ID
    entry -> entry.getKey(), // Build-side channel index
    entry -> {
        Symbol buildSymbol = entry.getValue();
        int buildChannelIndex = buildSideSymbols.indexOf(buildSymbol);
        verify(buildChannelIndex >= 0);
        return buildChannelIndex;
    }));
    if (localBuildChannels.isEmpty()) {
        return Optional.empty();
    }
    return Optional.of(new LocalDynamicFilter(localProbeSymbols, localBuildChannels, partitionCount, localType, session, taskId, stateStoreProvider));
}

Example 2 with SET

use of io.prestosql.spi.statestore.StateCollection.Type.SET in project hetu-core by openlookeng.

the class DynamicFilterService method registerTasksHelper.

private void registerTasksHelper(PlanNode node, Symbol buildSymbol, Map<String, Symbol> dynamicFiltersMap, Set<TaskId> taskIds, Set<InternalNode> workers, StageStateMachine stateMachine) {
    final StateStore stateStore = stateStoreProvider.getStateStore();
    String queryId = stateMachine.getSession().getQueryId().toString();
    for (Map.Entry<String, Symbol> entry : dynamicFiltersMap.entrySet()) {
        Symbol buildSymbolToCheck = buildSymbol != null ? buildSymbol : node.getOutputSymbols().contains(entry.getValue()) ? entry.getValue() : null;
        if (buildSymbolToCheck != null && entry.getValue().getName().equals(buildSymbol.getName())) {
            String filterId = entry.getKey();
            stateStore.createStateCollection(createKey(DynamicFilterUtils.TASKSPREFIX, filterId, queryId), SET);
            stateStore.createStateCollection(createKey(DynamicFilterUtils.PARTIALPREFIX, filterId, queryId), SET);
            dynamicFilters.putIfAbsent(queryId, new ConcurrentHashMap<>());
            Map<String, DynamicFilterRegistryInfo> filters = dynamicFilters.get(queryId);
            if (node instanceof JoinNode) {
                filters.put(filterId, extractDynamicFilterRegistryInfo((JoinNode) node, stateMachine.getSession(), filterId));
            } else if (node instanceof SemiJoinNode) {
                filters.put(filterId, extractDynamicFilterRegistryInfo((SemiJoinNode) node, stateMachine.getSession()));
            }
            dynamicFiltersToTask.putIfAbsent(filterId + "-" + queryId, new CopyOnWriteArraySet<>());
            CopyOnWriteArraySet<TaskId> taskSet = dynamicFiltersToTask.get(filterId + "-" + queryId);
            taskSet.addAll(taskIds);
            log.debug("registerTasks source " + filterId + " filters:" + filters + ", workers: " + workers.stream().map(x -> x.getNodeIdentifier()).collect(Collectors.joining(",")) + ", taskIds: " + taskIds.stream().map(TaskId::toString).collect(Collectors.joining(",")));
        }
    }
}