Search in sources :

Example 1 with NO_NODES_AVAILABLE

use of io.prestosql.spi.StandardErrorCode.NO_NODES_AVAILABLE in project hetu-core by openlookeng.

the class NodePartitioningManager method getNodePartitioningMap.

public NodePartitionMap getNodePartitioningMap(Session session, PartitioningHandle partitioningHandle, boolean isSnapshotEnabled, Integer nodeCount) {
    requireNonNull(session, "session is null");
    requireNonNull(partitioningHandle, "partitioningHandle is null");
    if (partitioningHandle.getConnectorHandle() instanceof SystemPartitioningHandle) {
        return ((SystemPartitioningHandle) partitioningHandle.getConnectorHandle()).getNodePartitionMap(session, nodeScheduler, isSnapshotEnabled, nodeCount);
    }
    CatalogName catalogName = partitioningHandle.getConnectorId().orElseThrow(() -> new IllegalArgumentException("No connector ID for partitioning handle: " + partitioningHandle));
    ConnectorNodePartitioningProvider partitioningProvider = partitioningProviders.get(catalogName);
    checkArgument(partitioningProvider != null, "No partitioning provider for connector %s", catalogName);
    ConnectorBucketNodeMap connectorBucketNodeMap = getConnectorBucketNodeMap(session, partitioningHandle);
    // safety check for crazy partitioning
    checkArgument(connectorBucketNodeMap.getBucketCount() < 1_000_000, "Too many buckets in partitioning: %s", connectorBucketNodeMap.getBucketCount());
    List<InternalNode> bucketToNode;
    if (connectorBucketNodeMap.hasFixedMapping()) {
        bucketToNode = getFixedMapping(connectorBucketNodeMap);
    } else {
        NodeSelector nodeSelector = nodeScheduler.createNodeSelector(catalogName, false, null);
        List<InternalNode> nodes;
        if (isSnapshotEnabled) {
            Integer count = nodeCount;
            if (count == null) {
                // Initial schedule: reserve some nodes
                count = calculateTaskCount(nodeSelector.selectableNodeCount());
            }
            nodes = nodeSelector.selectRandomNodes(count);
            checkCondition(nodes.size() == count, NO_NODES_AVAILABLE, "Snapshot: not enough worker nodes to resume expected number of tasks: " + count);
        } else {
            nodes = nodeSelector.allNodes();
        }
        bucketToNode = createArbitraryBucketToNode(nodes, connectorBucketNodeMap.getBucketCount());
    }
    int[] bucketToPartition = new int[connectorBucketNodeMap.getBucketCount()];
    BiMap<InternalNode, Integer> nodeToPartition = HashBiMap.create();
    int nextPartitionId = 0;
    for (int bucket = 0; bucket < bucketToNode.size(); bucket++) {
        InternalNode node = bucketToNode.get(bucket);
        Integer partitionId = nodeToPartition.get(node);
        if (partitionId == null) {
            partitionId = nextPartitionId++;
            nodeToPartition.put(node, partitionId);
        }
        bucketToPartition[bucket] = partitionId;
    }
    List<InternalNode> partitionToNode = IntStream.range(0, nodeToPartition.size()).mapToObj(partitionId -> nodeToPartition.inverse().get(partitionId)).collect(toImmutableList());
    return new NodePartitionMap(partitionToNode, bucketToPartition, getSplitToBucket(session, partitioningHandle));
}
Also used : IntStream(java.util.stream.IntStream) ConnectorSplit(io.prestosql.spi.connector.ConnectorSplit) NodeScheduler(io.prestosql.execution.scheduler.NodeScheduler) Split(io.prestosql.metadata.Split) NO_NODES_AVAILABLE(io.prestosql.spi.StandardErrorCode.NO_NODES_AVAILABLE) ArrayList(java.util.ArrayList) ConcurrentMap(java.util.concurrent.ConcurrentMap) Inject(javax.inject.Inject) SnapshotConfig.calculateTaskCount(io.prestosql.snapshot.SnapshotConfig.calculateTaskCount) Preconditions.checkArgument(com.google.common.base.Preconditions.checkArgument) Failures.checkCondition(io.prestosql.util.Failures.checkCondition) ConnectorNodePartitioningProvider(io.prestosql.spi.connector.ConnectorNodePartitioningProvider) Objects.requireNonNull(java.util.Objects.requireNonNull) Session(io.prestosql.Session) Type(io.prestosql.spi.type.Type) FixedBucketNodeMap(io.prestosql.execution.scheduler.FixedBucketNodeMap) BiMap(com.google.common.collect.BiMap) DynamicBucketNodeMap(io.prestosql.execution.scheduler.group.DynamicBucketNodeMap) InternalNode(io.prestosql.metadata.InternalNode) Collection(java.util.Collection) ImmutableList.toImmutableList(com.google.common.collect.ImmutableList.toImmutableList) CatalogName(io.prestosql.spi.connector.CatalogName) ConcurrentHashMap(java.util.concurrent.ConcurrentHashMap) ToIntFunction(java.util.function.ToIntFunction) BucketFunction(io.prestosql.spi.connector.BucketFunction) BucketNodeMap(io.prestosql.execution.scheduler.BucketNodeMap) HashBiMap(com.google.common.collect.HashBiMap) List(java.util.List) BucketPartitionFunction(io.prestosql.operator.BucketPartitionFunction) Stream(java.util.stream.Stream) ConnectorPartitionHandle(io.prestosql.spi.connector.ConnectorPartitionHandle) NodeSelector(io.prestosql.execution.scheduler.NodeSelector) ConnectorBucketNodeMap(io.prestosql.spi.connector.ConnectorBucketNodeMap) Optional(java.util.Optional) PartitionFunction(io.prestosql.operator.PartitionFunction) Collections(java.util.Collections) ConnectorNodePartitioningProvider(io.prestosql.spi.connector.ConnectorNodePartitioningProvider) ConnectorBucketNodeMap(io.prestosql.spi.connector.ConnectorBucketNodeMap) CatalogName(io.prestosql.spi.connector.CatalogName) InternalNode(io.prestosql.metadata.InternalNode) NodeSelector(io.prestosql.execution.scheduler.NodeSelector)

Example 2 with NO_NODES_AVAILABLE

use of io.prestosql.spi.StandardErrorCode.NO_NODES_AVAILABLE in project hetu-core by openlookeng.

the class SimpleNodeSelector method computeAssignments.

@Override
public SplitPlacementResult computeAssignments(Set<Split> splits, List<RemoteTask> existingTasks, Optional<SqlStageExecution> stage) {
    Multimap<InternalNode, Split> assignment = HashMultimap.create();
    NodeMap nodeMapSlice = this.nodeMap.get().get();
    NodeAssignmentStats assignmentStats = new NodeAssignmentStats(nodeTaskMap, nodeMapSlice, existingTasks);
    ResettableRandomizedIterator<InternalNode> randomCandidates = randomizedNodes(nodeMapSlice, ImmutableSet.of());
    Set<InternalNode> blockedExactNodes = new HashSet<>();
    boolean splitWaitingForAnyNode = false;
    // splitsToBeRedistributed becomes true only when splits go through locality-based assignment
    boolean splitsToBeRedistributed = false;
    Set<Split> remainingSplits = new HashSet<>();
    // Check if the current stage has a TableScanNode which is reading the table for the 2nd time or beyond
    if (stage.isPresent() && stage.get().getStateMachine().getConsumerScanNode() != null) {
        try {
            // if node exists, get the TableScanNode and cast it as consumer
            TableScanNode consumer = stage.get().getStateMachine().getConsumerScanNode();
            // all tables part of this stage
            Map<PlanNodeId, TableInfo> tables = stage.get().getStageInfo().getTables();
            QualifiedObjectName tableName;
            for (Map.Entry<PlanNodeId, TableInfo> entry : tables.entrySet()) {
                tableName = entry.getValue().getTableName();
                if (tableSplitAssignmentInfo.getReuseTableScanMappingIdSplitAssignmentMap().containsKey(consumer.getReuseTableScanMappingId())) {
                    // compare splitkey using equals and then assign nodes accordingly.
                    HashMap<SplitKey, InternalNode> splitKeyNodeAssignment = tableSplitAssignmentInfo.getSplitKeyNodeAssignment(consumer.getReuseTableScanMappingId());
                    Set<SplitKey> splitKeySet = splitKeyNodeAssignment.keySet();
                    assignment.putAll(createConsumerScanNodeAssignment(tableName, splits, splitKeySet, splitKeyNodeAssignment));
                    for (Map.Entry<InternalNode, Split> nodeAssignmentEntry : assignment.entries()) {
                        InternalNode node = nodeAssignmentEntry.getKey();
                        assignmentStats.addAssignedSplit(node);
                    }
                }
            }
            log.debug("Consumer:: Assignment size is " + assignment.size() + " ,Assignment is " + assignment + " ,Assignment Stats is " + assignmentStats);
        } catch (NotImplementedException e) {
            log.error("Not a Hive Split! Other Connector Splits not supported currently. Error: " + e);
            throw new UnsupportedOperationException("Not a Hive Split! Other Connector Splits not supported currently. Error: " + e);
        }
    } else {
        // optimizedLocalScheduling enables prioritized assignment of splits to local nodes when splits contain locality information
        if (optimizedLocalScheduling) {
            // should not hit for consumer case
            for (Split split : splits) {
                if (split.isRemotelyAccessible() && !split.getAddresses().isEmpty()) {
                    List<InternalNode> candidateNodes = selectExactNodes(nodeMapSlice, split.getAddresses(), includeCoordinator);
                    Optional<InternalNode> chosenNode = candidateNodes.stream().filter(ownerNode -> assignmentStats.getTotalSplitCount(ownerNode) < maxSplitsPerNode).min(comparingInt(assignmentStats::getTotalSplitCount));
                    if (chosenNode.isPresent()) {
                        assignment.put(chosenNode.get(), split);
                        // check later
                        assignmentStats.addAssignedSplit(chosenNode.get());
                        splitsToBeRedistributed = true;
                        continue;
                    }
                }
                remainingSplits.add(split);
            }
        } else {
            remainingSplits = splits;
        }
        for (Split split : remainingSplits) {
            randomCandidates.reset();
            List<InternalNode> candidateNodes;
            if (!split.isRemotelyAccessible()) {
                candidateNodes = selectExactNodes(nodeMapSlice, split.getAddresses(), includeCoordinator);
            } else {
                candidateNodes = selectNodes(minCandidates, randomCandidates);
            }
            if (candidateNodes.isEmpty()) {
                log.debug("No nodes available to schedule %s. Available nodes %s", split, nodeMapSlice.getNodesByHost().keys());
                throw new PrestoException(NO_NODES_AVAILABLE, "No nodes available to run query");
            }
            InternalNode chosenNode = null;
            int min = Integer.MAX_VALUE;
            for (InternalNode node : candidateNodes) {
                int totalSplitCount = assignmentStats.getTotalSplitCount(node);
                if (totalSplitCount < min && totalSplitCount < maxSplitsPerNode) {
                    chosenNode = node;
                    min = totalSplitCount;
                }
            }
            if (chosenNode == null) {
                // min is guaranteed to be MAX_VALUE at this line
                for (InternalNode node : candidateNodes) {
                    int totalSplitCount = assignmentStats.getQueuedSplitCountForStage(node);
                    if (totalSplitCount < min && totalSplitCount < maxPendingSplitsPerTask) {
                        chosenNode = node;
                        min = totalSplitCount;
                    }
                }
            }
            if (chosenNode != null) {
                assignment.put(chosenNode, split);
                assignmentStats.addAssignedSplit(chosenNode);
            } else {
                if (split.isRemotelyAccessible()) {
                    splitWaitingForAnyNode = true;
                } else // Exact node set won't matter, if a split is waiting for any node
                if (!splitWaitingForAnyNode) {
                    blockedExactNodes.addAll(candidateNodes);
                }
            }
        }
    }
    ListenableFuture<?> blocked;
    if (splitWaitingForAnyNode) {
        blocked = toWhenHasSplitQueueSpaceFuture(existingTasks, calculateLowWatermark(maxPendingSplitsPerTask));
    } else {
        blocked = toWhenHasSplitQueueSpaceFuture(blockedExactNodes, existingTasks, calculateLowWatermark(maxPendingSplitsPerTask));
    }
    if (!stage.isPresent() || stage.get().getStateMachine().getConsumerScanNode() == null) {
        if (splitsToBeRedistributed) {
            // skip for consumer
            equateDistribution(assignment, assignmentStats, nodeMapSlice);
        }
    }
    // Check if the current stage has a TableScanNode which is reading the table for the 1st time
    if (stage.isPresent() && stage.get().getStateMachine().getProducerScanNode() != null) {
        // if node exists, get the TableScanNode and annotate it as producer
        saveProducerScanNodeAssignment(stage, assignment, assignmentStats);
    }
    // Check if its CTE node and its feeder
    if (stage.isPresent() && stage.get().getFragment().getFeederCTEId().isPresent()) {
        updateFeederNodeAndSplitCount(stage.get(), assignment);
    }
    return new SplitPlacementResult(blocked, assignment);
}
Also used : NodeScheduler.randomizedNodes(io.prestosql.execution.scheduler.NodeScheduler.randomizedNodes) ListenableFuture(com.google.common.util.concurrent.ListenableFuture) Logger(io.airlift.log.Logger) Supplier(com.google.common.base.Supplier) HashMap(java.util.HashMap) Split(io.prestosql.metadata.Split) NO_NODES_AVAILABLE(io.prestosql.spi.StandardErrorCode.NO_NODES_AVAILABLE) Multimap(com.google.common.collect.Multimap) AtomicReference(java.util.concurrent.atomic.AtomicReference) QualifiedObjectName(io.prestosql.spi.connector.QualifiedObjectName) InetAddress(java.net.InetAddress) HashSet(java.util.HashSet) HashMultimap(com.google.common.collect.HashMultimap) ImmutableList(com.google.common.collect.ImmutableList) Map(java.util.Map) Objects.requireNonNull(java.util.Objects.requireNonNull) Suppliers(com.google.common.base.Suppliers) NodeTaskMap(io.prestosql.execution.NodeTaskMap) InternalNodeManager(io.prestosql.metadata.InternalNodeManager) PlanNodeId(io.prestosql.spi.plan.PlanNodeId) TableInfo(io.prestosql.execution.TableInfo) PrestoException(io.prestosql.spi.PrestoException) Comparator.comparingInt(java.util.Comparator.comparingInt) ImmutableSet(com.google.common.collect.ImmutableSet) NotImplementedException(sun.reflect.generics.reflectiveObjects.NotImplementedException) HostAddress(io.prestosql.spi.HostAddress) Iterator(java.util.Iterator) NodeScheduler.calculateLowWatermark(io.prestosql.execution.scheduler.NodeScheduler.calculateLowWatermark) IndexedPriorityQueue(io.prestosql.execution.resourcegroups.IndexedPriorityQueue) InternalNode(io.prestosql.metadata.InternalNode) Collection(java.util.Collection) TableScanNode(io.prestosql.spi.plan.TableScanNode) Set(java.util.Set) NodeScheduler.toWhenHasSplitQueueSpaceFuture(io.prestosql.execution.scheduler.NodeScheduler.toWhenHasSplitQueueSpaceFuture) UnknownHostException(java.net.UnknownHostException) Collectors(java.util.stream.Collectors) SetMultimap(com.google.common.collect.SetMultimap) NodeScheduler.selectNodes(io.prestosql.execution.scheduler.NodeScheduler.selectNodes) List(java.util.List) SplitKey(io.prestosql.execution.SplitKey) GENERIC_INTERNAL_ERROR(io.prestosql.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR) Optional(java.util.Optional) NodeScheduler.selectDistributionNodes(io.prestosql.execution.scheduler.NodeScheduler.selectDistributionNodes) VisibleForTesting(com.google.common.annotations.VisibleForTesting) SqlStageExecution(io.prestosql.execution.SqlStageExecution) NodeScheduler.selectExactNodes(io.prestosql.execution.scheduler.NodeScheduler.selectExactNodes) RemoteTask(io.prestosql.execution.RemoteTask) SplitKey(io.prestosql.execution.SplitKey) NotImplementedException(sun.reflect.generics.reflectiveObjects.NotImplementedException) PrestoException(io.prestosql.spi.PrestoException) PlanNodeId(io.prestosql.spi.plan.PlanNodeId) TableInfo(io.prestosql.execution.TableInfo) HashSet(java.util.HashSet) QualifiedObjectName(io.prestosql.spi.connector.QualifiedObjectName) TableScanNode(io.prestosql.spi.plan.TableScanNode) InternalNode(io.prestosql.metadata.InternalNode) Split(io.prestosql.metadata.Split) HashMap(java.util.HashMap) Map(java.util.Map) NodeTaskMap(io.prestosql.execution.NodeTaskMap)

Example 3 with NO_NODES_AVAILABLE

use of io.prestosql.spi.StandardErrorCode.NO_NODES_AVAILABLE in project hetu-core by openlookeng.

the class SqlQueryScheduler method createStages.

private List<SqlStageExecution> createStages(ExchangeLocationsConsumer parent, AtomicInteger nextStageId, LocationFactory locationFactory, StageExecutionPlan plan, NodeScheduler nodeScheduler, RemoteTaskFactory remoteTaskFactory, Session session, int splitBatchSize, BiFunction<PartitioningHandle, Integer, NodePartitionMap> partitioningCache, NodePartitioningManager nodePartitioningManager, ExecutorService queryExecutor, ScheduledExecutorService schedulerExecutor, FailureDetector failureDetector, NodeTaskMap nodeTaskMap, ImmutableMap.Builder<StageId, StageScheduler> stageSchedulers, ImmutableMap.Builder<StageId, StageLinkage> stageLinkages, boolean isSnapshotEnabled, QuerySnapshotManager snapshotManager, Map<StageId, Integer> stageTaskCounts) {
    ImmutableList.Builder<SqlStageExecution> localStages = ImmutableList.builder();
    StageId stageId = new StageId(queryStateMachine.getQueryId(), nextStageId.getAndIncrement());
    SqlStageExecution stageExecution = createSqlStageExecution(stageId, locationFactory.createStageLocation(stageId), plan.getFragment(), plan.getTables(), remoteTaskFactory, session, summarizeTaskInfo, nodeTaskMap, queryExecutor, failureDetector, schedulerStats, dynamicFilterService, snapshotManager);
    localStages.add(stageExecution);
    Optional<int[]> bucketToPartition;
    PartitioningHandle partitioningHandle = plan.getFragment().getPartitioning();
    boolean keepConsumerOnFeederNodes = !plan.getFragment().getFeederCTEId().isPresent() && plan.getFragment().getFeederCTEParentId().isPresent();
    if (partitioningHandle.equals(SOURCE_DISTRIBUTION)) {
        // nodes are selected dynamically based on the constraints of the splits and the system load
        Entry<PlanNodeId, SplitSource> entry = Iterables.getOnlyElement(plan.getSplitSources().entrySet());
        PlanNodeId planNodeId = entry.getKey();
        SplitSource splitSource = entry.getValue();
        CatalogName catalogName = splitSource.getCatalogName();
        if (isInternalSystemConnector(catalogName)) {
            catalogName = null;
        }
        NodeSelector nodeSelector = nodeScheduler.createNodeSelector(catalogName, keepConsumerOnFeederNodes, feederScheduledNodes);
        if (isSnapshotEnabled) {
            // When snapshot is enabled, then no task can be added after the query started running,
            // otherwise assumptions about how many "input channels" may be broken.
            nodeSelector.lockDownNodes();
        }
        SplitPlacementPolicy placementPolicy = new DynamicSplitPlacementPolicy(nodeSelector, stageExecution::getAllTasks);
        checkArgument(!plan.getFragment().getStageExecutionDescriptor().isStageGroupedExecution());
        stageSchedulers.put(stageId, newSourcePartitionedSchedulerAsStageScheduler(stageExecution, planNodeId, splitSource, placementPolicy, splitBatchSize, session, heuristicIndexerManager));
        bucketToPartition = Optional.of(new int[1]);
    } else if (partitioningHandle.equals(SCALED_WRITER_DISTRIBUTION)) {
        bucketToPartition = Optional.of(new int[1]);
    } else {
        Map<PlanNodeId, SplitSource> splitSources = plan.getSplitSources();
        if (!splitSources.isEmpty()) {
            // contains local source
            List<PlanNodeId> schedulingOrder = plan.getFragment().getPartitionedSources();
            CatalogName catalogName = partitioningHandle.getConnectorId().orElseThrow(IllegalStateException::new);
            List<ConnectorPartitionHandle> connectorPartitionHandles;
            boolean groupedExecutionForStage = plan.getFragment().getStageExecutionDescriptor().isStageGroupedExecution();
            if (groupedExecutionForStage) {
                connectorPartitionHandles = nodePartitioningManager.listPartitionHandles(session, partitioningHandle);
                checkState(!ImmutableList.of(NOT_PARTITIONED).equals(connectorPartitionHandles));
            } else {
                connectorPartitionHandles = ImmutableList.of(NOT_PARTITIONED);
            }
            BucketNodeMap bucketNodeMap;
            List<InternalNode> stageNodeList;
            if (plan.getFragment().getRemoteSourceNodes().stream().allMatch(node -> node.getExchangeType() == REPLICATE)) {
                // no remote source
                boolean dynamicLifespanSchedule = plan.getFragment().getStageExecutionDescriptor().isDynamicLifespanSchedule();
                if (isSnapshotEnabled) {
                    NodeSelector nodeSelector = nodeScheduler.createNodeSelector(catalogName, keepConsumerOnFeederNodes, feederScheduledNodes);
                    int nodeCount;
                    if (stageTaskCounts != null) {
                        // Resuming: need to create same number of tasks as old stage.
                        nodeCount = stageTaskCounts.get(stageId);
                    } else {
                        // Scheduling: reserve some nodes for resuming
                        nodeCount = calculateTaskCount(nodeSelector.selectableNodeCount());
                    }
                    stageNodeList = new ArrayList<>(nodeSelector.selectRandomNodes(nodeCount));
                    checkCondition(stageNodeList.size() == nodeCount, NO_NODES_AVAILABLE, "Snapshot: not enough worker nodes to resume expected number of tasks: " + nodeCount);
                    // Make sure bucketNodeMap uses the same node list
                    bucketNodeMap = nodePartitioningManager.getBucketNodeMap(session, partitioningHandle, dynamicLifespanSchedule, stageNodeList);
                } else {
                    bucketNodeMap = nodePartitioningManager.getBucketNodeMap(session, partitioningHandle, dynamicLifespanSchedule);
                    stageNodeList = new ArrayList<>(nodeScheduler.createNodeSelector(catalogName, keepConsumerOnFeederNodes, feederScheduledNodes).allNodes());
                }
                // verify execution is consistent with planner's decision on dynamic lifespan schedule
                verify(bucketNodeMap.isDynamic() == dynamicLifespanSchedule);
                Collections.shuffle(stageNodeList);
                bucketToPartition = Optional.empty();
            } else {
                // cannot use dynamic lifespan schedule
                verify(!plan.getFragment().getStageExecutionDescriptor().isDynamicLifespanSchedule());
                // remote source requires nodePartitionMap
                NodePartitionMap nodePartitionMap = partitioningCache.apply(plan.getFragment().getPartitioning(), stageTaskCounts == null ? null : stageTaskCounts.get(stageId));
                if (groupedExecutionForStage) {
                    checkState(connectorPartitionHandles.size() == nodePartitionMap.getBucketToPartition().length);
                }
                stageNodeList = nodePartitionMap.getPartitionToNode();
                bucketNodeMap = nodePartitionMap.asBucketNodeMap();
                bucketToPartition = Optional.of(nodePartitionMap.getBucketToPartition());
            }
            stageSchedulers.put(stageId, new FixedSourcePartitionedScheduler(stageExecution, splitSources, plan.getFragment().getStageExecutionDescriptor(), schedulingOrder, stageNodeList, bucketNodeMap, splitBatchSize, getConcurrentLifespansPerNode(session), nodeScheduler.createNodeSelector(catalogName, keepConsumerOnFeederNodes, feederScheduledNodes), connectorPartitionHandles, session, heuristicIndexerManager));
        } else {
            // all sources are remote
            NodePartitionMap nodePartitionMap = partitioningCache.apply(plan.getFragment().getPartitioning(), stageTaskCounts == null ? null : stageTaskCounts.get(stageId));
            List<InternalNode> partitionToNode = nodePartitionMap.getPartitionToNode();
            // todo this should asynchronously wait a standard timeout period before failing
            checkCondition(!partitionToNode.isEmpty(), NO_NODES_AVAILABLE, "No worker nodes available");
            stageSchedulers.put(stageId, new FixedCountScheduler(stageExecution, partitionToNode));
            bucketToPartition = Optional.of(nodePartitionMap.getBucketToPartition());
        }
    }
    ImmutableSet.Builder<SqlStageExecution> childStagesBuilder = ImmutableSet.builder();
    for (StageExecutionPlan subStagePlan : plan.getSubStages()) {
        if (visitedPlanFrags.contains(subStagePlan.getFragment().getId())) {
            continue;
        }
        visitedPlanFrags.add(subStagePlan.getFragment().getId());
        List<SqlStageExecution> subTree = createStages(stageExecution::addExchangeLocations, nextStageId, locationFactory, subStagePlan.withBucketToPartition(bucketToPartition), nodeScheduler, remoteTaskFactory, session, splitBatchSize, partitioningCache, nodePartitioningManager, queryExecutor, schedulerExecutor, failureDetector, nodeTaskMap, stageSchedulers, stageLinkages, isSnapshotEnabled, snapshotManager, stageTaskCounts);
        localStages.addAll(subTree);
        SqlStageExecution childStage = subTree.get(0);
        childStagesBuilder.add(childStage);
        Optional<RemoteSourceNode> parentNode = plan.getFragment().getRemoteSourceNodes().stream().filter(x -> x.getSourceFragmentIds().contains(childStage.getFragment().getId())).findAny();
        checkArgument(parentNode.isPresent(), "Couldn't find parent of a CTE node");
        childStage.setParentId(parentNode.get().getId());
    }
    Set<SqlStageExecution> childStages = childStagesBuilder.build();
    stageExecution.addStateChangeListener(newState -> {
        if (newState.isDone() && newState != StageState.RESCHEDULING) {
            // Snapshot: For "rescheduling", tasks are already cancelled (for resume)
            childStages.forEach(SqlStageExecution::cancel);
        }
    });
    stageLinkages.put(stageId, new StageLinkage(plan.getFragment().getId(), parent, childStages));
    if (partitioningHandle.equals(SCALED_WRITER_DISTRIBUTION)) {
        Supplier<Collection<TaskStatus>> sourceTasksProvider = () -> childStages.stream().map(SqlStageExecution::getAllTasks).flatMap(Collection::stream).map(RemoteTask::getTaskStatus).collect(toList());
        Supplier<Collection<TaskStatus>> writerTasksProvider = () -> stageExecution.getAllTasks().stream().map(RemoteTask::getTaskStatus).collect(toList());
        ScaledWriterScheduler scheduler = new ScaledWriterScheduler(stageExecution, sourceTasksProvider, writerTasksProvider, nodeScheduler.createNodeSelector(null, keepConsumerOnFeederNodes, feederScheduledNodes), schedulerExecutor, getWriterMinSize(session), isSnapshotEnabled, stageTaskCounts != null ? stageTaskCounts.get(stageId) : null);
        whenAllStages(childStages, StageState::isDone).addListener(scheduler::finish, directExecutor());
        stageSchedulers.put(stageId, scheduler);
    }
    return localStages.build();
}
Also used : CANCELED(io.prestosql.execution.StageState.CANCELED) SCHEDULED(io.prestosql.execution.StageState.SCHEDULED) PlanFragmentId(io.prestosql.sql.planner.plan.PlanFragmentId) NO_NODES_AVAILABLE(io.prestosql.spi.StandardErrorCode.NO_NODES_AVAILABLE) FIXED_BROADCAST_DISTRIBUTION(io.prestosql.sql.planner.SystemPartitioningHandle.FIXED_BROADCAST_DISTRIBUTION) StageExecutionPlan(io.prestosql.sql.planner.StageExecutionPlan) Map(java.util.Map) SystemSessionProperties.getWriterMinSize(io.prestosql.SystemSessionProperties.getWriterMinSize) HeuristicIndexerManager(io.prestosql.heuristicindex.HeuristicIndexerManager) PlanNodeId(io.prestosql.spi.plan.PlanNodeId) TaskStatus(io.prestosql.execution.TaskStatus) ImmutableList.toImmutableList(com.google.common.collect.ImmutableList.toImmutableList) Set(java.util.Set) LocationFactory(io.prestosql.execution.LocationFactory) MILLISECONDS(java.util.concurrent.TimeUnit.MILLISECONDS) MoreExecutors.directExecutor(com.google.common.util.concurrent.MoreExecutors.directExecutor) ImmutableMap.toImmutableMap(com.google.common.collect.ImmutableMap.toImmutableMap) StageState(io.prestosql.execution.StageState) ConnectorPartitionHandle(io.prestosql.spi.connector.ConnectorPartitionHandle) GENERIC_INTERNAL_ERROR(io.prestosql.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR) SetThreadName(io.airlift.concurrent.SetThreadName) Iterables(com.google.common.collect.Iterables) RESUMABLE_FAILURE(io.prestosql.execution.StageState.RESUMABLE_FAILURE) Supplier(java.util.function.Supplier) SCALED_WRITER_DISTRIBUTION(io.prestosql.sql.planner.SystemPartitioningHandle.SCALED_WRITER_DISTRIBUTION) QueryStateMachine(io.prestosql.execution.QueryStateMachine) ArrayList(java.util.ArrayList) CatalogName.isInternalSystemConnector(io.prestosql.spi.connector.CatalogName.isInternalSystemConnector) Session(io.prestosql.Session) ScheduledExecutorService(java.util.concurrent.ScheduledExecutorService) ImmutableSet.toImmutableSet(com.google.common.collect.ImmutableSet.toImmutableSet) NodeTaskMap(io.prestosql.execution.NodeTaskMap) SplitSource(io.prestosql.split.SplitSource) FINISHED(io.prestosql.execution.StageState.FINISHED) StageId(io.prestosql.execution.StageId) InternalNode(io.prestosql.metadata.InternalNode) Sets.newConcurrentHashSet(com.google.common.collect.Sets.newConcurrentHashSet) QuerySnapshotManager(io.prestosql.snapshot.QuerySnapshotManager) ResourceGroupInfo(io.prestosql.server.ResourceGroupInfo) PartitioningHandle(io.prestosql.sql.planner.PartitioningHandle) QueryState(io.prestosql.execution.QueryState) NodePartitionMap(io.prestosql.sql.planner.NodePartitionMap) SqlStageExecution.createSqlStageExecution(io.prestosql.execution.SqlStageExecution.createSqlStageExecution) SqlStageExecution(io.prestosql.execution.SqlStageExecution) ABORTED(io.prestosql.execution.StageState.ABORTED) FailureDetector(io.prestosql.failuredetector.FailureDetector) RemoteTask(io.prestosql.execution.RemoteTask) FAILED(io.prestosql.execution.StageState.FAILED) SystemSessionProperties(io.prestosql.SystemSessionProperties) BiFunction(java.util.function.BiFunction) SettableFuture(com.google.common.util.concurrent.SettableFuture) SourcePartitionedScheduler.newSourcePartitionedSchedulerAsStageScheduler(io.prestosql.execution.scheduler.SourcePartitionedScheduler.newSourcePartitionedSchedulerAsStageScheduler) Duration(io.airlift.units.Duration) Preconditions.checkArgument(com.google.common.base.Preconditions.checkArgument) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) URI(java.net.URI) Collectors.toSet(java.util.stream.Collectors.toSet) PrestoException(io.prestosql.spi.PrestoException) ImmutableSet(com.google.common.collect.ImmutableSet) ImmutableMap(com.google.common.collect.ImmutableMap) Predicate(java.util.function.Predicate) Collection(java.util.Collection) CatalogName(io.prestosql.spi.connector.CatalogName) ConcurrentHashMap(java.util.concurrent.ConcurrentHashMap) RemoteTaskFactory(io.prestosql.execution.RemoteTaskFactory) UUID(java.util.UUID) String.format(java.lang.String.format) Preconditions.checkState(com.google.common.base.Preconditions.checkState) List(java.util.List) SOURCE_DISTRIBUTION(io.prestosql.sql.planner.SystemPartitioningHandle.SOURCE_DISTRIBUTION) StageInfo(io.prestosql.execution.StageInfo) Entry(java.util.Map.Entry) HttpUriBuilder.uriBuilderFrom(io.airlift.http.client.HttpUriBuilder.uriBuilderFrom) Function.identity(java.util.function.Function.identity) Optional(java.util.Optional) MoreFutures.whenAnyComplete(io.airlift.concurrent.MoreFutures.whenAnyComplete) BasicStageStats(io.prestosql.execution.BasicStageStats) NodePartitioningManager(io.prestosql.sql.planner.NodePartitioningManager) NOT_PARTITIONED(io.prestosql.spi.connector.NotPartitionedPartitionHandle.NOT_PARTITIONED) RUNNING(io.prestosql.execution.StageState.RUNNING) TaskId(io.prestosql.execution.TaskId) ListenableFuture(com.google.common.util.concurrent.ListenableFuture) Logger(io.airlift.log.Logger) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) RemoteSourceNode(io.prestosql.sql.planner.plan.RemoteSourceNode) HashMap(java.util.HashMap) OutputBuffers(io.prestosql.execution.buffer.OutputBuffers) TaskLocation(io.prestosql.operator.TaskLocation) HashSet(java.util.HashSet) SnapshotConfig.calculateTaskCount(io.prestosql.snapshot.SnapshotConfig.calculateTaskCount) ImmutableList(com.google.common.collect.ImmutableList) SystemSessionProperties.isReuseTableScanEnabled(io.prestosql.SystemSessionProperties.isReuseTableScanEnabled) OutputBufferId(io.prestosql.execution.buffer.OutputBuffers.OutputBufferId) Verify.verify(com.google.common.base.Verify.verify) Failures.checkCondition(io.prestosql.util.Failures.checkCondition) Objects.requireNonNull(java.util.Objects.requireNonNull) TimeStat(io.airlift.stats.TimeStat) REPLICATE(io.prestosql.sql.planner.plan.ExchangeNode.Type.REPLICATE) ExecutorService(java.util.concurrent.ExecutorService) Ints(com.google.common.primitives.Ints) DynamicFilterService(io.prestosql.dynamicfilter.DynamicFilterService) MoreFutures.tryGetFutureValue(io.airlift.concurrent.MoreFutures.tryGetFutureValue) SystemSessionProperties.getConcurrentLifespansPerNode(io.prestosql.SystemSessionProperties.getConcurrentLifespansPerNode) Collectors.toList(java.util.stream.Collectors.toList) Collections(java.util.Collections) SECONDS(java.util.concurrent.TimeUnit.SECONDS) BasicStageStats.aggregateBasicStageStats(io.prestosql.execution.BasicStageStats.aggregateBasicStageStats) ImmutableList.toImmutableList(com.google.common.collect.ImmutableList.toImmutableList) ImmutableList(com.google.common.collect.ImmutableList) StageExecutionPlan(io.prestosql.sql.planner.StageExecutionPlan) StageId(io.prestosql.execution.StageId) ArrayList(java.util.ArrayList) SqlStageExecution.createSqlStageExecution(io.prestosql.execution.SqlStageExecution.createSqlStageExecution) SqlStageExecution(io.prestosql.execution.SqlStageExecution) PlanNodeId(io.prestosql.spi.plan.PlanNodeId) RemoteSourceNode(io.prestosql.sql.planner.plan.RemoteSourceNode) ImmutableSet.toImmutableSet(com.google.common.collect.ImmutableSet.toImmutableSet) ImmutableSet(com.google.common.collect.ImmutableSet) ImmutableList.toImmutableList(com.google.common.collect.ImmutableList.toImmutableList) ArrayList(java.util.ArrayList) List(java.util.List) ImmutableList(com.google.common.collect.ImmutableList) Collectors.toList(java.util.stream.Collectors.toList) NodePartitionMap(io.prestosql.sql.planner.NodePartitionMap) Collection(java.util.Collection) CatalogName(io.prestosql.spi.connector.CatalogName) PartitioningHandle(io.prestosql.sql.planner.PartitioningHandle) SplitSource(io.prestosql.split.SplitSource) Map(java.util.Map) ImmutableMap.toImmutableMap(com.google.common.collect.ImmutableMap.toImmutableMap) NodeTaskMap(io.prestosql.execution.NodeTaskMap) NodePartitionMap(io.prestosql.sql.planner.NodePartitionMap) ImmutableMap(com.google.common.collect.ImmutableMap) ConcurrentHashMap(java.util.concurrent.ConcurrentHashMap) HashMap(java.util.HashMap)

Aggregations

Preconditions.checkArgument (com.google.common.base.Preconditions.checkArgument)2 ImmutableList (com.google.common.collect.ImmutableList)2 ImmutableList.toImmutableList (com.google.common.collect.ImmutableList.toImmutableList)2 ImmutableSet (com.google.common.collect.ImmutableSet)2 ListenableFuture (com.google.common.util.concurrent.ListenableFuture)2 Logger (io.airlift.log.Logger)2 Session (io.prestosql.Session)2 NodeTaskMap (io.prestosql.execution.NodeTaskMap)2 RemoteTask (io.prestosql.execution.RemoteTask)2 SqlStageExecution (io.prestosql.execution.SqlStageExecution)2 InternalNode (io.prestosql.metadata.InternalNode)2 Split (io.prestosql.metadata.Split)2 NO_NODES_AVAILABLE (io.prestosql.spi.StandardErrorCode.NO_NODES_AVAILABLE)2 Collection (java.util.Collection)2 List (java.util.List)2 Objects.requireNonNull (java.util.Objects.requireNonNull)2 Optional (java.util.Optional)2 VisibleForTesting (com.google.common.annotations.VisibleForTesting)1 Preconditions.checkState (com.google.common.base.Preconditions.checkState)1 Supplier (com.google.common.base.Supplier)1