use of io.prestosql.sql.planner.plan.RemoteSourceNode 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.sql.planner.plan.RemoteSourceNode in project hetu-core by openlookeng.
the class DistributedExecutionPlanner method collectSources.
private List<MarkerSplitSource> collectSources(Map<PlanFragmentId, Object> leftmostSources, Object source) {
if (source instanceof ValuesNode) {
// TODO-cp-I2X9J6: should we worry about dependencies about Values operators, when it's the "left" of a join?
return ImmutableList.of();
}
if (source instanceof RemoteSourceNode) {
List<PlanFragmentId> fragments = ((RemoteSourceNode) source).getSourceFragmentIds();
if (fragments.size() == 1) {
return collectSources(leftmostSources, leftmostSources.get(fragments.get(0)));
}
List<MarkerSplitSource> sources = new ArrayList<>();
for (PlanFragmentId id : fragments) {
sources.addAll(collectSources(leftmostSources, leftmostSources.get(id)));
}
// Adding all these sources as "union dependencies" for each other, to make sure they produce the same set of markers.
for (MarkerSplitSource unionSource : sources) {
unionSource.addUnionSources(sources);
}
return sources;
}
// Must be a split source
return ImmutableList.of((MarkerSplitSource) source);
}
use of io.prestosql.sql.planner.plan.RemoteSourceNode in project hetu-core by openlookeng.
the class AllAtOnceExecutionSchedule method getPreferredScheduleOrder.
@VisibleForTesting
static List<PlanFragmentId> getPreferredScheduleOrder(Collection<PlanFragment> fragments) {
// determine output fragment
Set<PlanFragmentId> remoteSources = fragments.stream().map(PlanFragment::getRemoteSourceNodes).flatMap(Collection::stream).map(RemoteSourceNode::getSourceFragmentIds).flatMap(Collection::stream).collect(toImmutableSet());
Set<PlanFragment> rootFragments = fragments.stream().filter(fragment -> !remoteSources.contains(fragment.getId())).collect(toImmutableSet());
checkArgument(rootFragments.size() == 1, "Expected one root fragment, but found: " + rootFragments);
Visitor visitor = new Visitor(fragments);
visitor.processFragment(getOnlyElement(rootFragments).getId());
return visitor.getSchedulerOrder();
}
use of io.prestosql.sql.planner.plan.RemoteSourceNode in project hetu-core by openlookeng.
the class TestSqlStageExecution method createExchangePlanFragment.
private static PlanFragment createExchangePlanFragment() {
PlanNode planNode = new RemoteSourceNode(new PlanNodeId("exchange"), ImmutableList.of(new PlanFragmentId("source")), ImmutableList.of(new Symbol("column")), Optional.empty(), REPARTITION);
ImmutableMap.Builder<Symbol, Type> types = ImmutableMap.builder();
for (Symbol symbol : planNode.getOutputSymbols()) {
types.put(symbol, VARCHAR);
}
return new PlanFragment(new PlanFragmentId("exchange_fragment_id"), planNode, types.build(), SOURCE_DISTRIBUTION, ImmutableList.of(planNode.getId()), new PartitioningScheme(Partitioning.create(SINGLE_DISTRIBUTION, ImmutableList.of()), planNode.getOutputSymbols()), ungroupedExecution(), StatsAndCosts.empty(), Optional.empty(), Optional.empty(), Optional.empty());
}
use of io.prestosql.sql.planner.plan.RemoteSourceNode in project hetu-core by openlookeng.
the class SqlStageExecution method addExchangeLocations.
public synchronized void addExchangeLocations(PlanFragmentId fragmentId, Set<RemoteTask> sourceTasks, boolean noMoreExchangeLocations) {
requireNonNull(fragmentId, "fragmentId is null");
requireNonNull(sourceTasks, "sourceTasks is null");
RemoteSourceNode remoteSource = exchangeSources.get(fragmentId);
checkArgument(remoteSource != null, "Unknown remote source %s. Known sources are %s", fragmentId, exchangeSources.keySet());
this.sourceTasks.putAll(remoteSource.getId(), sourceTasks);
for (RemoteTask task : getAllTasks()) {
ImmutableMultimap.Builder<PlanNodeId, Split> newSplits = ImmutableMultimap.builder();
for (RemoteTask sourceTask : sourceTasks) {
newSplits.put(remoteSource.getId(), newConnectSplit(task.getTaskId(), sourceTask));
}
task.addSplits(newSplits.build());
}
if (noMoreExchangeLocations) {
completeSourceFragments.add(fragmentId);
// is the source now complete?
if (completeSourceFragments.containsAll(remoteSource.getSourceFragmentIds())) {
completeSources.add(remoteSource.getId());
for (RemoteTask task : getAllTasks()) {
task.noMoreSplits(remoteSource.getId());
}
}
}
}
Aggregations