Search in sources :

Example 1 with StageStateMachine

use of io.prestosql.execution.StageStateMachine in project hetu-core by openlookeng.

the class TestDynamicFilterUtil method registerDf.

public static void registerDf(String filterId, Session session, JoinNode.DistributionType joinType, DynamicFilterService dynamicFilterService) {
    JoinNode node = mock(JoinNode.class);
    HashMap<String, Symbol> dfs = new HashMap<>();
    List<JoinNode.EquiJoinClause> criteria = new ArrayList<JoinNode.EquiJoinClause>();
    Symbol right = new Symbol("rightCol");
    Symbol left = new Symbol("leftCol");
    JoinNode.EquiJoinClause clause = new JoinNode.EquiJoinClause(left, right);
    criteria.add(clause);
    dfs.put(filterId, right);
    when(node.getCriteria()).thenReturn(criteria);
    when(node.getDynamicFilters()).thenReturn(dfs);
    when(node.getDistributionType()).thenReturn(Optional.of(joinType));
    RemoteSourceNode leftNode = mock(RemoteSourceNode.class);
    when(node.getLeft()).thenReturn(leftNode);
    HashSet<TaskId> tasks = new HashSet<>();
    tasks.add(new TaskId("task1.0"));
    tasks.add(new TaskId("task1.1"));
    StageStateMachine stateMachine = mock(StageStateMachine.class);
    when(stateMachine.getSession()).thenReturn(session);
    InternalNode worker = mock(InternalNode.class);
    InternalNode worker2 = mock(InternalNode.class);
    HashSet<InternalNode> workers = new HashSet<>();
    when(worker.getNodeIdentifier()).thenReturn("w1");
    when(worker2.getNodeIdentifier()).thenReturn("w2");
    workers.add(worker);
    workers.add(worker2);
    dynamicFilterService.registerTasks(node, tasks, workers, stateMachine);
}
Also used : TaskId(io.prestosql.execution.TaskId) HashMap(java.util.HashMap) JoinNode(io.prestosql.spi.plan.JoinNode) Symbol(io.prestosql.spi.plan.Symbol) ArrayList(java.util.ArrayList) Matchers.anyString(org.mockito.Matchers.anyString) RemoteSourceNode(io.prestosql.sql.planner.plan.RemoteSourceNode) StageStateMachine(io.prestosql.execution.StageStateMachine) InternalNode(io.prestosql.metadata.InternalNode) HashSet(java.util.HashSet)

Example 2 with StageStateMachine

use of io.prestosql.execution.StageStateMachine 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(",")));
        }
    }
}
Also used : ScheduledFuture(java.util.concurrent.ScheduledFuture) DynamicFilter(io.prestosql.spi.dynamicfilter.DynamicFilter) Inject(com.google.inject.Inject) StateStore(io.prestosql.spi.statestore.StateStore) DynamicFilterUtils.findFilterNodeInStage(io.prestosql.utils.DynamicFilterUtils.findFilterNodeInStage) CallExpression(io.prestosql.spi.relation.CallExpression) PreDestroy(javax.annotation.PreDestroy) Preconditions.checkArgument(com.google.common.base.Preconditions.checkArgument) FilterNode(io.prestosql.spi.plan.FilterNode) ByteArrayInputStream(java.io.ByteArrayInputStream) Map(java.util.Map) SystemSessionProperties.getDynamicFilteringDataType(io.prestosql.SystemSessionProperties.getDynamicFilteringDataType) PrestoException(io.prestosql.spi.PrestoException) ImmutableSet(com.google.common.collect.ImmutableSet) ImmutableMap(com.google.common.collect.ImmutableMap) StateMap(io.prestosql.spi.statestore.StateMap) Predicate(java.util.function.Predicate) Collection(java.util.Collection) ConcurrentHashMap(java.util.concurrent.ConcurrentHashMap) Set(java.util.Set) DynamicFilters(io.prestosql.sql.DynamicFilters) PlanNode(io.prestosql.spi.plan.PlanNode) StateSet(io.prestosql.spi.statestore.StateSet) CopyOnWriteArraySet(java.util.concurrent.CopyOnWriteArraySet) Threads.threadsNamed(io.airlift.concurrent.Threads.threadsNamed) Collectors(java.util.stream.Collectors) DynamicFilterUtils.getDynamicFilterDataType(io.prestosql.utils.DynamicFilterUtils.getDynamicFilterDataType) Executors(java.util.concurrent.Executors) Preconditions.checkState(com.google.common.base.Preconditions.checkState) Type(io.prestosql.spi.dynamicfilter.DynamicFilter.Type) List(java.util.List) PostConstruct(javax.annotation.PostConstruct) Entry(java.util.Map.Entry) GENERIC_INTERNAL_ERROR(io.prestosql.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR) Optional(java.util.Optional) LOCAL(io.prestosql.spi.dynamicfilter.DynamicFilter.Type.LOCAL) TaskId(io.prestosql.execution.TaskId) ByteArrayOutputStream(java.io.ByteArrayOutputStream) Logger(io.airlift.log.Logger) HASHSET(io.prestosql.spi.dynamicfilter.DynamicFilter.DataType.HASHSET) HashMap(java.util.HashMap) Supplier(java.util.function.Supplier) SemiJoinNode(io.prestosql.sql.planner.plan.SemiJoinNode) ArrayList(java.util.ArrayList) HashSet(java.util.HashSet) BloomFilter(io.prestosql.spi.util.BloomFilter) DataType(io.prestosql.spi.dynamicfilter.DynamicFilter.DataType) Objects.requireNonNull(java.util.Objects.requireNonNull) Session(io.prestosql.Session) BloomFilterDynamicFilter(io.prestosql.spi.dynamicfilter.BloomFilterDynamicFilter) ScheduledExecutorService(java.util.concurrent.ScheduledExecutorService) QueryId(io.prestosql.spi.QueryId) JoinNode(io.prestosql.spi.plan.JoinNode) Symbol(io.prestosql.spi.plan.Symbol) StageStateMachine(io.prestosql.execution.StageStateMachine) StateCollection(io.prestosql.spi.statestore.StateCollection) InternalNode(io.prestosql.metadata.InternalNode) IOException(java.io.IOException) BLOOM_FILTER(io.prestosql.spi.dynamicfilter.DynamicFilter.DataType.BLOOM_FILTER) VariableReferenceExpression(io.prestosql.spi.relation.VariableReferenceExpression) TimeUnit(java.util.concurrent.TimeUnit) GLOBAL(io.prestosql.spi.dynamicfilter.DynamicFilter.Type.GLOBAL) DynamicFilterUtils(io.prestosql.utils.DynamicFilterUtils) DynamicFilterUtils.createKey(io.prestosql.utils.DynamicFilterUtils.createKey) ColumnHandle(io.prestosql.spi.connector.ColumnHandle) StateStoreProvider(io.prestosql.statestore.StateStoreProvider) RowExpression(io.prestosql.spi.relation.RowExpression) MAP(io.prestosql.spi.statestore.StateCollection.Type.MAP) DynamicFilterFactory(io.prestosql.spi.dynamicfilter.DynamicFilterFactory) Collections(java.util.Collections) SET(io.prestosql.spi.statestore.StateCollection.Type.SET) TaskId(io.prestosql.execution.TaskId) Symbol(io.prestosql.spi.plan.Symbol) SemiJoinNode(io.prestosql.sql.planner.plan.SemiJoinNode) JoinNode(io.prestosql.spi.plan.JoinNode) StateStore(io.prestosql.spi.statestore.StateStore) SemiJoinNode(io.prestosql.sql.planner.plan.SemiJoinNode) Map(java.util.Map) ImmutableMap(com.google.common.collect.ImmutableMap) StateMap(io.prestosql.spi.statestore.StateMap) ConcurrentHashMap(java.util.concurrent.ConcurrentHashMap) HashMap(java.util.HashMap)

Aggregations

StageStateMachine (io.prestosql.execution.StageStateMachine)2 TaskId (io.prestosql.execution.TaskId)2 InternalNode (io.prestosql.metadata.InternalNode)2 JoinNode (io.prestosql.spi.plan.JoinNode)2 Symbol (io.prestosql.spi.plan.Symbol)2 ArrayList (java.util.ArrayList)2 HashMap (java.util.HashMap)2 HashSet (java.util.HashSet)2 Preconditions.checkArgument (com.google.common.base.Preconditions.checkArgument)1 Preconditions.checkState (com.google.common.base.Preconditions.checkState)1 ImmutableMap (com.google.common.collect.ImmutableMap)1 ImmutableSet (com.google.common.collect.ImmutableSet)1 Inject (com.google.inject.Inject)1 Threads.threadsNamed (io.airlift.concurrent.Threads.threadsNamed)1 Logger (io.airlift.log.Logger)1 Session (io.prestosql.Session)1 SystemSessionProperties.getDynamicFilteringDataType (io.prestosql.SystemSessionProperties.getDynamicFilteringDataType)1 PrestoException (io.prestosql.spi.PrestoException)1 QueryId (io.prestosql.spi.QueryId)1 GENERIC_INTERNAL_ERROR (io.prestosql.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR)1