SqlStageExecution stage = createSqlStageExecution( stageId, locationFactory.createStageLocation(stageId), plan.getFragment(), remoteTaskFactory,
taskId -> createSqlTask( taskId, locationFactory.createLocalTaskLocation(taskId), nodeInfo.getNodeId(), queryContexts.getUnchecked(taskId.getQueryId()),
private synchronized void updateNodes(MemoryPoolAssignmentsRequest assignments) { ImmutableSet.Builder<Node> builder = ImmutableSet.builder(); Set<Node> aliveNodes = builder .addAll(nodeManager.getNodes(ACTIVE)) .addAll(nodeManager.getNodes(SHUTTING_DOWN)) .build(); ImmutableSet<String> aliveNodeIds = aliveNodes.stream() .map(Node::getNodeIdentifier) .collect(toImmutableSet()); // Remove nodes that don't exist anymore // Make a copy to materialize the set difference Set<String> deadNodes = ImmutableSet.copyOf(difference(nodes.keySet(), aliveNodeIds)); nodes.keySet().removeAll(deadNodes); // Add new nodes for (Node node : aliveNodes) { if (!nodes.containsKey(node.getNodeIdentifier())) { nodes.put(node.getNodeIdentifier(), new RemoteNodeMemory(node, httpClient, memoryInfoCodec, assignmentsRequestJsonCodec, locationFactory.createMemoryInfoLocation(node))); } } // If work isn't scheduled on the coordinator (the current node) there is no point // in polling or updating (when moving queries to the reserved pool) its memory pools if (!isWorkScheduledOnCoordinator) { nodes.remove(nodeManager.getCurrentNode().getNodeIdentifier()); } // Schedule refresh for (RemoteNodeMemory node : nodes.values()) { node.asyncRefresh(assignments); } }
private <T extends Statement> DataDefinitionExecution<T> createDataDefinitionExecution( String query, Session session, ResourceGroupId resourceGroup, T statement, List<Expression> parameters, WarningCollector warningCollector) { @SuppressWarnings("unchecked") DataDefinitionTask<T> task = (DataDefinitionTask<T>) tasks.get(statement.getClass()); checkArgument(task != null, "no task for statement: %s", statement.getClass().getSimpleName()); QueryStateMachine stateMachine = QueryStateMachine.begin( query, session, locationFactory.createQueryLocation(session.getQueryId()), resourceGroup, task.isTransactionControl(), transactionManager, accessControl, executor, metadata, warningCollector); stateMachine.setUpdateType(task.getName()); return new DataDefinitionExecution<>(task, statement, transactionManager, metadata, accessControl, stateMachine, parameters); } }
taskId, node.getNodeIdentifier(), locationFactory.createTaskLocation(node, taskId), fragment, initialSplits,
query, session, locationFactory.createQueryLocation(session.getQueryId()), resourceGroup, preparedQuery,
@Override public RemoteTask createRemoteTask(Session session, TaskId taskId, Node node, int partition, PlanFragment fragment, Multimap<PlanNodeId, Split> initialSplits, OutputBuffers outputBuffers, PartitionedSplitCountTracker partitionedSplitCountTracker) { return new HttpRemoteTask(session, taskId, node.getNodeIdentifier(), partition, locationFactory.createTaskLocation(node, taskId), fragment, initialSplits, outputBuffers, httpClient, executor, errorScheduledExecutor, minErrorDuration, taskInfoRefreshMaxWait, taskInfoCodec, taskUpdateRequestCodec, partitionedSplitCountTracker ); } }
private SqlStageExecution createSqlStageExecution(StageExecutionPlan tableScanPlan, NodeTaskMap nodeTaskMap) { StageId stageId = new StageId(new QueryId("query"), 0); SqlStageExecution stage = SqlStageExecution.createSqlStageExecution(stageId, locationFactory.createStageLocation(stageId), tableScanPlan.getFragment(), new MockRemoteTaskFactory(queryExecutor, scheduledExecutor), TEST_SESSION, true, nodeTaskMap, queryExecutor, new NoOpFailureDetector(), new SplitSchedulerStats()); stage.setOutputBuffers(createInitialEmptyOutputBuffers(PARTITIONED) .withBuffer(OUT, 0) .withNoMoreBufferIds()); return stage; }
session, query, locationFactory.createQueryLocation(queryId), Optional.ofNullable(selectionContext).map(SelectionContext::getResourceGroupId), queryExecutor,
private void updateNodes(MemoryPoolAssignmentsRequest assignments) { ImmutableSet.Builder builder = new ImmutableSet.Builder(); Set<Node> aliveNodes = builder .addAll(nodeManager.getNodes(ACTIVE)) .addAll(nodeManager.getNodes(SHUTTING_DOWN)) .build(); ImmutableSet<String> aliveNodeIds = aliveNodes.stream() .map(Node::getNodeIdentifier) .collect(toImmutableSet()); // Remove nodes that don't exist anymore // Make a copy to materialize the set difference Set<String> deadNodes = ImmutableSet.copyOf(difference(nodes.keySet(), aliveNodeIds)); nodes.keySet().removeAll(deadNodes); // Add new nodes for (Node node : aliveNodes) { if (!nodes.containsKey(node.getNodeIdentifier())) { nodes.put(node.getNodeIdentifier(), new RemoteNodeMemory(httpClient, memoryInfoCodec, assignmentsRequestJsonCodec, locationFactory.createMemoryInfoLocation(node))); } } // Schedule refresh for (RemoteNodeMemory node : nodes.values()) { node.asyncRefresh(assignments); } }
@Override public SqlTask load(TaskId taskId) throws Exception { return new SqlTask( taskId, locationFactory.createLocalTaskLocation(taskId), queryContexts.getUnchecked(taskId.getQueryId()), sqlTaskExecutionFactory, taskNotificationExecutor, sqlTask -> { finishedTaskStats.merge(sqlTask.getIoStats()); return null; }, maxBufferSize ); } });
SqlStageExecution stage = new SqlStageExecution( stageId, locationFactory.createStageLocation(stageId), plan.getFragment(), remoteTaskFactory,
@Override public SqlQueryExecution createQueryExecution(QueryId queryId, String query, Session session, Statement statement) { String executionPolicyName = SystemSessionProperties.getExecutionPolicy(session); ExecutionPolicy executionPolicy = executionPolicies.get(executionPolicyName); checkArgument(executionPolicy != null, "No execution policy %s", executionPolicy); return new SqlQueryExecution( queryId, query, session, locationFactory.createQueryLocation(queryId), statement, transactionManager, metadata, accessControl, sqlParser, splitManager, nodePartitioningManager, nodeScheduler, planOptimizers, remoteTaskFactory, locationFactory, scheduleSplitBatchSize, experimentalSyntaxEnabled, executor, nodeTaskMap, queryExplainer, executionPolicy); } }
private SqlStageExecution createSqlStageExecution(StageExecutionPlan tableScanPlan, NodeTaskMap nodeTaskMap) { StageId stageId = new StageId(new QueryId("query"), "stage"); SqlStageExecution stage = new SqlStageExecution(stageId, locationFactory.createStageLocation(stageId), tableScanPlan.getFragment(), new MockRemoteTaskFactory(executor), TEST_SESSION, nodeTaskMap, executor); stage.setOutputBuffers(INITIAL_EMPTY_OUTPUT_BUFFERS .withBuffer(OUT, 0) .withNoMoreBufferIds()); return stage; }
@Override public DataDefinitionExecution<?> createQueryExecution( QueryId queryId, String query, Session session, Statement statement) { URI self = locationFactory.createQueryLocation(queryId); DataDefinitionTask<Statement> task = getTask(statement); checkArgument(task != null, "no task for statement: %s", statement.getClass().getSimpleName()); QueryStateMachine stateMachine = QueryStateMachine.begin(queryId, query, session, self, task.isTransactionControl(), transactionManager, executor); stateMachine.setUpdateType(task.getName()); return new DataDefinitionExecution<>(task, statement, transactionManager, metadata, accessControl, stateMachine); }
URI self = locationFactory.createQueryLocation(queryId); QueryExecution execution = new FailedQueryExecution(queryId, query, session, self, transactionManager, queryExecutor, e);