PlanNode newRoot = fragment.getRoot(); if (!fragment.getPartitioning().isSingleNode()) { PartitioningHandleReassigner partitioningHandleReassigner = new PartitioningHandleReassigner(fragment.getPartitioning(), metadata, session); newRoot = SimplePlanRewriter.rewriteWith(partitioningHandleReassigner, newRoot); PartitioningScheme outputPartitioningScheme = fragment.getPartitioningScheme(); Partitioning newOutputPartitioning = outputPartitioningScheme.getPartitioning(); if (outputPartitioningScheme.getPartitioning().getHandle().getConnectorId().isPresent()) { PlanFragment newFragment = new PlanFragment( fragment.getId(), newRoot, fragment.getSymbols(), fragment.getPartitioning(), fragment.getPartitionedSources(), new PartitioningScheme( newOutputPartitioning, outputPartitioningScheme.isReplicateNullsAndAny(), outputPartitioningScheme.getBucketToPartition()), fragment.getStageExecutionDescriptor(), fragment.getStatsAndCosts()); childrenBuilder.add(reassignPartitioningHandleIfNecessaryHelper(session, child, fragment.getPartitioning()));
private synchronized List<TaskSource> getSources() { return Stream.concat(planFragment.getPartitionedSourceNodes().stream(), planFragment.getRemoteSourceNodes().stream()) .filter(Objects::nonNull) .map(PlanNode::getId) .map(this::getSource) .filter(Objects::nonNull) .collect(toImmutableList()); }
.collect(toImmutableList()); this.partitionedSourceNodes = findSources(root, partitionedSources); findRemoteSourceNodes(root, remoteSourceNodes); this.remoteSourceNodes = remoteSourceNodes.build();
private static void printFragmentNodes(StringBuilder output, PlanFragment fragment, PlanNodeIdGenerator idGenerator) { String clusterId = "cluster_" + fragment.getId(); output.append("subgraph ") .append(clusterId) .append(" {") .append('\n'); output.append(format("label = \"%s\"", fragment.getPartitioning())) .append('\n'); PlanNode plan = fragment.getRoot(); plan.accept(new NodePrinter(output, idGenerator), null); output.append("}") .append('\n'); }
fragment.getId(), fragment.getPartitioning())); PartitioningScheme partitioningScheme = fragment.getPartitioningScheme(); builder.append(indentString(1)) .append(format("Output layout: [%s]\n", formatHash(partitioningScheme.getHashColumn()))); builder.append(indentString(1)).append(format("Grouped Execution: %s\n", fragment.getStageExecutionDescriptor().isStageGroupedExecution())); .flatMap(f -> f.getSymbols().entrySet().stream()) .distinct() .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue))); builder.append(textLogicalPlan(fragment.getRoot(), typeProvider, Optional.of(fragment.getStageExecutionDescriptor()), functionRegistry, fragment.getStatsAndCosts(), session, planNodeStats, 1, verbose)) .append("\n");
fragment.getId(), fragment.getPartitioning())); PartitioningScheme partitioningScheme = fragment.getPartitioningScheme(); builder.append(indentString(1)) .append(format("Output layout: [%s]\n", formatHash(partitioningScheme.getHashColumn()))); builder.append(indentString(1)).append(format("Grouped Execution: %s\n", fragment.getStageExecutionStrategy().isAnyScanGroupedExecution())); .flatMap(f -> f.getSymbols().entrySet().stream()) .distinct() .collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue))); builder.append(textLogicalPlan(fragment.getRoot(), typeProvider, Optional.of(fragment.getStageExecutionStrategy()), functionRegistry, fragment.getStatsAndCosts(), session, planNodeStats, 1, verbose)) .append("\n");
localExecutionPlan = planner.plan( taskContext, fragment.getRoot(), TypeProvider.copyOf(fragment.getSymbols()), fragment.getPartitioningScheme(), fragment.getStageExecutionStrategy(), fragment.getPartitionedSources(), outputBuffer);
localExecutionPlan = planner.plan( taskContext, fragment.getRoot(), TypeProvider.copyOf(fragment.getSymbols()), fragment.getPartitioningScheme(), fragment.getStageExecutionDescriptor(), fragment.getPartitionedSources(), outputBuffer);
public PhasedExecutionSchedule(Collection<SqlStageExecution> stages) { List<Set<PlanFragmentId>> phases = extractPhases(stages.stream().map(SqlStageExecution::getFragment).collect(toImmutableList())); Map<PlanFragmentId, SqlStageExecution> stagesByFragmentId = stages.stream().collect(toImmutableMap(stage -> stage.getFragment().getId(), identity())); // create a mutable list of mutable sets of stages, so we can remove completed stages schedulePhases = new ArrayList<>(); for (Set<PlanFragmentId> phase : phases) { schedulePhases.add(phase.stream() .map(stagesByFragmentId::get) .collect(Collectors.toCollection(HashSet::new))); } }
public PlanFragment withGroupedExecution(List<PlanNodeId> capableTableScanNodes) { return new PlanFragment(id, root, symbols, partitioning, partitionedSources, partitioningScheme, groupedExecution(capableTableScanNodes), statsAndCosts); }
PartitioningHandle partitioningHandle = plan.getFragment().getPartitioning(); if (partitioningHandle.equals(SOURCE_DISTRIBUTION)) { checkArgument(!plan.getFragment().getStageExecutionStrategy().isAnyScanGroupedExecution()); stageSchedulers.put(stageId, newSourcePartitionedSchedulerAsStageScheduler(stage, planNodeId, splitSource, placementPolicy, splitBatchSize)); bucketToPartition = Optional.of(new int[1]); List<PlanNodeId> schedulingOrder = plan.getFragment().getPartitionedSources(); ConnectorId connectorId = partitioningHandle.getConnectorId().orElseThrow(IllegalStateException::new); List<ConnectorPartitionHandle> connectorPartitionHandles; boolean groupedExecutionForStage = plan.getFragment().getStageExecutionStrategy().isAnyScanGroupedExecution(); if (groupedExecutionForStage) { connectorPartitionHandles = nodePartitioningManager.listPartitionHandles(session, partitioningHandle); NodePartitionMap nodePartitionMap = partitioningCache.apply(plan.getFragment().getPartitioning()); if (groupedExecutionForStage) { checkState(connectorPartitionHandles.size() == nodePartitionMap.getBucketToPartition().length); stage, splitSources, plan.getFragment().getStageExecutionStrategy(), schedulingOrder, stageNodeList, NodePartitionMap nodePartitionMap = partitioningCache.apply(plan.getFragment().getPartitioning()); List<Node> partitionToNode = nodePartitionMap.getPartitionToNode(); stageLinkages.put(stageId, new StageLinkage(plan.getFragment().getId(), parent, childStages));
PartitioningHandle partitioningHandle = plan.getFragment().getPartitioning(); if (partitioningHandle.equals(SOURCE_DISTRIBUTION)) { checkArgument(!plan.getFragment().getStageExecutionDescriptor().isStageGroupedExecution()); stageSchedulers.put(stageId, newSourcePartitionedSchedulerAsStageScheduler(stage, planNodeId, splitSource, placementPolicy, splitBatchSize)); bucketToPartition = Optional.of(new int[1]); List<PlanNodeId> schedulingOrder = plan.getFragment().getPartitionedSources(); ConnectorId connectorId = partitioningHandle.getConnectorId().orElseThrow(IllegalStateException::new); List<ConnectorPartitionHandle> connectorPartitionHandles; boolean groupedExecutionForStage = plan.getFragment().getStageExecutionDescriptor().isStageGroupedExecution(); if (groupedExecutionForStage) { connectorPartitionHandles = nodePartitioningManager.listPartitionHandles(session, partitioningHandle); NodePartitionMap nodePartitionMap = partitioningCache.apply(plan.getFragment().getPartitioning()); if (groupedExecutionForStage) { checkState(connectorPartitionHandles.size() == nodePartitionMap.getBucketToPartition().length); stage, splitSources, plan.getFragment().getStageExecutionDescriptor(), schedulingOrder, stageNodeList, NodePartitionMap nodePartitionMap = partitioningCache.apply(plan.getFragment().getPartitioning()); List<Node> partitionToNode = nodePartitionMap.getPartitionToNode(); stageLinkages.put(stageId, new StageLinkage(plan.getFragment().getId(), parent, childStages));
private static boolean hasSourceDistributedStage(Set<SqlStageExecution> phase) { return phase.stream().anyMatch(stage -> !stage.getFragment().getPartitionedSources().isEmpty()); }
private static PlanFragment createUnionPlanFragment(String name, PlanFragment... fragments) { PlanNode planNode = new UnionNode( new PlanNodeId(name + "_id"), Stream.of(fragments) .map(fragment -> new RemoteSourceNode(new PlanNodeId(fragment.getId().toString()), fragment.getId(), fragment.getPartitioningScheme().getOutputLayout(), Optional.empty(), REPARTITION)) .collect(toImmutableList()), ImmutableListMultimap.of(), ImmutableList.of()); return createFragment(planNode); }
private static void reconstructStatsAndCosts( StageInfo stage, ImmutableMap.Builder<PlanNodeId, PlanNodeStatsEstimate> planNodeStats, ImmutableMap.Builder<PlanNodeId, PlanNodeCostEstimate> planNodeCosts) { PlanFragment planFragment = stage.getPlan(); if (planFragment != null) { planNodeStats.putAll(planFragment.getStatsAndCosts().getStats()); planNodeCosts.putAll(planFragment.getStatsAndCosts().getCosts()); } for (StageInfo subStage : stage.getSubStages()) { reconstructStatsAndCosts(subStage, planNodeStats, planNodeCosts); } }
stateMachine.setColumns(outputStageExecutionPlan.getFieldNames(), outputStageExecutionPlan.getFragment().getTypes()); PartitioningHandle partitioningHandle = plan.getRoot().getFragment().getPartitioningScheme().getPartitioning().getHandle(); OutputBuffers rootOutputBuffers = createInitialEmptyOutputBuffers(partitioningHandle) .withBuffer(OUTPUT_BUFFER_ID, BROADCAST_PARTITION_ID)
public StageLinkage(PlanFragmentId fragmentId, ExchangeLocationsConsumer parent, Set<SqlStageExecution> children) { this.currentStageFragmentId = fragmentId; this.parent = parent; this.childOutputBufferManagers = children.stream() .map(childStage -> { PartitioningHandle partitioningHandle = childStage.getFragment().getPartitioningScheme().getPartitioning().getHandle(); if (partitioningHandle.equals(FIXED_BROADCAST_DISTRIBUTION)) { return new BroadcastOutputBufferManager(childStage::setOutputBuffers); } else if (partitioningHandle.equals(SCALED_WRITER_DISTRIBUTION)) { return new ScaledOutputBufferManager(childStage::setOutputBuffers); } else { int partitionCount = Ints.max(childStage.getFragment().getPartitioningScheme().getBucketToPartition().get()) + 1; return new PartitionedOutputBufferManager(partitioningHandle, partitionCount, childStage::setOutputBuffers); } }) .collect(toImmutableSet()); this.childStageIds = children.stream() .map(SqlStageExecution::getStageId) .collect(toImmutableSet()); }
private SqlStageExecution(StageStateMachine stateMachine, RemoteTaskFactory remoteTaskFactory, NodeTaskMap nodeTaskMap, boolean summarizeTaskInfo, Executor executor, FailureDetector failureDetector) { this.stateMachine = stateMachine; this.remoteTaskFactory = requireNonNull(remoteTaskFactory, "remoteTaskFactory is null"); this.nodeTaskMap = requireNonNull(nodeTaskMap, "nodeTaskMap is null"); this.summarizeTaskInfo = summarizeTaskInfo; this.executor = requireNonNull(executor, "executor is null"); this.failureDetector = requireNonNull(failureDetector, "failureDetector is null"); ImmutableMap.Builder<PlanFragmentId, RemoteSourceNode> fragmentToExchangeSource = ImmutableMap.builder(); for (RemoteSourceNode remoteSourceNode : stateMachine.getFragment().getRemoteSourceNodes()) { for (PlanFragmentId planFragmentId : remoteSourceNode.getSourceFragmentIds()) { fragmentToExchangeSource.put(planFragmentId, remoteSourceNode); } } this.exchangeSources = fragmentToExchangeSource.build(); }
if (fragment.getPartitionedSourceNodes().stream().anyMatch(TableScanNode.class::isInstance)) { rawInputDataSize += taskStats.getRawInputDataSize().toBytes(); rawInputPositions += taskStats.getRawInputPositions();