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);
}
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(",")));
}
}
}
Aggregations