public List<Driver> createDrivers( @Language("SQL") String sql, OutputFactory outputFactory, TaskContext taskContext) { Statement statement = SqlParser.createStatement(sql); if (printPlan) { assertFormattedSql(statement); } PlanNodeIdAllocator idAllocator = new PlanNodeIdAllocator(); FeaturesConfig featuresConfig = new FeaturesConfig().setExperimentalSyntaxEnabled(true); PlanOptimizersFactory planOptimizersFactory = new PlanOptimizersFactory(metadata, splitManager, indexManager, featuresConfig); QueryExplainer queryExplainer = new QueryExplainer( session, planOptimizersFactory.get(), metadata, featuresConfig.isExperimentalSyntaxEnabled()); Analyzer analyzer = new Analyzer( session, metadata, Optional.of(queryExplainer), featuresConfig.isExperimentalSyntaxEnabled()); Analysis analysis = analyzer.analyze(statement); Plan plan = new LogicalPlanner(session, planOptimizersFactory.get(), idAllocator, metadata) .plan(analysis); if (printPlan) { System.out.println(PlanPrinter.textLogicalPlan(plan.getRoot(), plan.getTypes(), metadata)); } SubPlan subplan = new DistributedLogicalPlanner(session, metadata, idAllocator).createSubPlans(plan, true); assertTrue(subplan.getChildren().isEmpty(), "Expected subplan to have no children"); LocalExecutionPlanner executionPlanner = new LocalExecutionPlanner( new NodeInfo(new NodeConfig().setEnvironment("test").setNodeId("test-node")), metadata, dataStreamProvider, indexManager, storageManager, recordSinkManager, null, compiler); // plan query LocalExecutionPlan localExecutionPlan = executionPlanner.plan( session, subplan.getFragment().getRoot(), plan.getTypes(), outputFactory); // generate sources List<TaskSource> sources = new ArrayList<>(); long sequenceId = 0; for (PlanNode sourceNode : subplan.getFragment().getSources()) { if (sourceNode instanceof ValuesNode) { continue; } TableScanNode tableScan = (TableScanNode) sourceNode; SplitSource splitSource = splitManager.getPartitionSplits(tableScan.getTable(), getPartitions(tableScan)); ImmutableSet.Builder<ScheduledSplit> scheduledSplits = ImmutableSet.builder(); while (!splitSource.isFinished()) { try { for (Split split : splitSource.getNextBatch(1000)) { scheduledSplits.add(new ScheduledSplit(sequenceId++, split)); } } catch (InterruptedException e) { Thread.currentThread().interrupt(); throw Throwables.propagate(e); } } sources.add(new TaskSource(tableScan.getId(), scheduledSplits.build(), true)); } // create drivers List<Driver> drivers = new ArrayList<>(); Map<PlanNodeId, Driver> driversBySource = new HashMap<>(); for (DriverFactory driverFactory : localExecutionPlan.getDriverFactories()) { DriverContext driverContext = taskContext .addPipelineContext(driverFactory.isInputDriver(), driverFactory.isOutputDriver()) .addDriverContext(); Driver driver = driverFactory.createDriver(driverContext); drivers.add(driver); for (PlanNodeId sourceId : driver.getSourceIds()) { driversBySource.put(sourceId, driver); } driverFactory.close(); } // add sources to the drivers for (TaskSource source : sources) { for (Driver driver : driversBySource.values()) { driver.updateSource(source); } } return ImmutableList.copyOf(drivers); }
private SqlTaskExecution( TaskStateMachine taskStateMachine, TaskContext taskContext, SharedBuffer sharedBuffer, PlanFragment fragment, LocalExecutionPlanner planner, TaskExecutor taskExecutor, QueryMonitor queryMonitor, Executor notificationExecutor) { this.taskStateMachine = checkNotNull(taskStateMachine, "taskStateMachine is null"); this.taskId = taskStateMachine.getTaskId(); this.taskContext = checkNotNull(taskContext, "taskContext is null"); this.sharedBuffer = checkNotNull(sharedBuffer, "sharedBuffer is null"); this.taskExecutor = checkNotNull(taskExecutor, "driverExecutor is null"); this.notificationExecutor = checkNotNull(notificationExecutor, "notificationExecutor is null"); this.queryMonitor = checkNotNull(queryMonitor, "queryMonitor is null"); try (SetThreadName ignored = new SetThreadName("Task-%s", taskId)) { List<DriverFactory> driverFactories; try { OutputFactory outputOperatorFactory; if (fragment.getOutputPartitioning() == OutputPartitioning.NONE) { outputOperatorFactory = new TaskOutputFactory(sharedBuffer); } else if (fragment.getOutputPartitioning() == OutputPartitioning.HASH) { outputOperatorFactory = new PartitionedOutputFactory(sharedBuffer); } else { throw new PrestoException( NOT_SUPPORTED, format("OutputPartitioning %s is not supported", fragment.getOutputPartitioning())); } LocalExecutionPlan localExecutionPlan = planner.plan( taskContext.getSession(), fragment.getRoot(), fragment.getOutputLayout(), fragment.getSymbols(), fragment.getDistribution(), outputOperatorFactory); driverFactories = localExecutionPlan.getDriverFactories(); } catch (Throwable e) { // planning failed taskStateMachine.failed(e); throw Throwables.propagate(e); } // index driver factories DriverSplitRunnerFactory partitionedDriverFactory = null; ImmutableList.Builder<DriverSplitRunnerFactory> unpartitionedDriverFactories = ImmutableList.builder(); for (DriverFactory driverFactory : driverFactories) { if (driverFactory.getSourceIds().contains(fragment.getPartitionedSource())) { checkState( partitionedDriverFactory == null, "multiple partitioned sources are not supported"); partitionedDriverFactory = new DriverSplitRunnerFactory(driverFactory); } else { unpartitionedDriverFactories.add(new DriverSplitRunnerFactory(driverFactory)); } } this.unpartitionedDriverFactories = unpartitionedDriverFactories.build(); if (fragment.getDistribution() == PlanDistribution.SOURCE) { checkArgument( partitionedDriverFactory != null, "Fragment is partitioned, but no partitioned driver found"); } this.partitionedSourceId = fragment.getPartitionedSource(); this.partitionedDriverFactory = partitionedDriverFactory; // don't register the task if it is already completed (most likely failed during planning // above) if (!taskStateMachine.getState().isDone()) { taskHandle = taskExecutor.addTask(taskId); taskStateMachine.addStateChangeListener( new RemoveTaskHandleWhenDone(taskExecutor, taskHandle)); } else { taskHandle = null; } sharedBuffer.addStateChangeListener( new CheckTaskCompletionOnBufferFinish(SqlTaskExecution.this)); } }
private SqlTaskExecution( Session session, TaskId taskId, URI location, PlanFragment fragment, OutputBuffers outputBuffers, LocalExecutionPlanner planner, DataSize maxBufferSize, TaskExecutor taskExecutor, DataSize maxTaskMemoryUsage, DataSize operatorPreAllocatedMemory, QueryMonitor queryMonitor, Executor notificationExecutor, boolean cpuTimerEnabled) { try (SetThreadName setThreadName = new SetThreadName("Task-%s", taskId)) { this.taskId = checkNotNull(taskId, "taskId is null"); this.location = checkNotNull(location, "location is null"); this.taskExecutor = checkNotNull(taskExecutor, "driverExecutor is null"); this.notificationExecutor = checkNotNull(notificationExecutor, "notificationExecutor is null"); this.taskStateMachine = new TaskStateMachine(taskId, notificationExecutor); taskStateMachine.addStateChangeListener( new StateChangeListener<TaskState>() { @Override public void stateChanged(TaskState taskState) { if (taskState.isDone()) { SqlTaskExecution.this.taskExecutor.removeTask(taskHandle); // make sure buffers are cleaned up if (taskState != TaskState.FAILED) { // don't close buffers for a failed query // closed buffers signal to upstream tasks that everything finished cleanly sharedBuffer.destroy(); } } } }); this.taskContext = new TaskContext( taskStateMachine, notificationExecutor, session, checkNotNull(maxTaskMemoryUsage, "maxTaskMemoryUsage is null"), checkNotNull(operatorPreAllocatedMemory, "operatorPreAllocatedMemory is null"), cpuTimerEnabled); this.sharedBuffer = new SharedBuffer( taskId, notificationExecutor, checkNotNull(maxBufferSize, "maxBufferSize is null"), outputBuffers); sharedBuffer.addStateChangeListener( new StateChangeListener<QueueState>() { @Override public void stateChanged(QueueState taskState) { if (taskState == QueueState.FINISHED) { checkTaskCompletion(); } } }); this.queryMonitor = checkNotNull(queryMonitor, "queryMonitor is null"); taskHandle = taskExecutor.addTask(taskId); LocalExecutionPlan localExecutionPlan = planner.plan( session, fragment.getRoot(), fragment.getSymbols(), new TaskOutputFactory(sharedBuffer)); List<DriverFactory> driverFactories = localExecutionPlan.getDriverFactories(); // index driver factories DriverSplitRunnerFactory partitionedDriverFactory = null; ImmutableList.Builder<DriverSplitRunnerFactory> unpartitionedDriverFactories = ImmutableList.builder(); for (DriverFactory driverFactory : driverFactories) { if (driverFactory.getSourceIds().contains(fragment.getPartitionedSource())) { checkState( partitionedDriverFactory == null, "multiple partitioned sources are not supported"); partitionedDriverFactory = new DriverSplitRunnerFactory(driverFactory); } else { unpartitionedDriverFactories.add(new DriverSplitRunnerFactory(driverFactory)); } } this.unpartitionedDriverFactories = unpartitionedDriverFactories.build(); if (fragment.getDistribution() == PlanDistribution.SOURCE) { checkArgument( partitionedDriverFactory != null, "Fragment is partitioned, but no partitioned driver found"); } this.partitionedSourceId = fragment.getPartitionedSource(); this.partitionedDriverFactory = partitionedDriverFactory; } }
public List<Driver> createDrivers( Session session, @Language("SQL") String sql, OutputFactory outputFactory, TaskContext taskContext) { Statement statement = sqlParser.createStatement(sql); assertFormattedSql(sqlParser, statement); PlanNodeIdAllocator idAllocator = new PlanNodeIdAllocator(); FeaturesConfig featuresConfig = new FeaturesConfig() .setExperimentalSyntaxEnabled(true) .setDistributedIndexJoinsEnabled(false) .setOptimizeHashGeneration(true); PlanOptimizersFactory planOptimizersFactory = new PlanOptimizersFactory(metadata, sqlParser, indexManager, featuresConfig, true); QueryExplainer queryExplainer = new QueryExplainer( planOptimizersFactory.get(), metadata, accessControl, sqlParser, dataDefinitionTask, featuresConfig.isExperimentalSyntaxEnabled()); Analyzer analyzer = new Analyzer( session, metadata, sqlParser, accessControl, Optional.of(queryExplainer), featuresConfig.isExperimentalSyntaxEnabled()); Analysis analysis = analyzer.analyze(statement); Plan plan = new LogicalPlanner(session, planOptimizersFactory.get(), idAllocator, metadata) .plan(analysis); if (printPlan) { System.out.println( PlanPrinter.textLogicalPlan(plan.getRoot(), plan.getTypes(), metadata, session)); } SubPlan subplan = new PlanFragmenter().createSubPlans(plan); if (!subplan.getChildren().isEmpty()) { throw new AssertionError("Expected subplan to have no children"); } LocalExecutionPlanner executionPlanner = new LocalExecutionPlanner( metadata, sqlParser, pageSourceManager, indexManager, pageSinkManager, null, compiler, new IndexJoinLookupStats(), new CompilerConfig() .setInterpreterEnabled(false), // make sure tests fail if compiler breaks new TaskManagerConfig().setTaskDefaultConcurrency(4)); // plan query LocalExecutionPlan localExecutionPlan = executionPlanner.plan( session, subplan.getFragment().getRoot(), subplan.getFragment().getOutputLayout(), plan.getTypes(), subplan.getFragment().getDistribution(), outputFactory); // generate sources List<TaskSource> sources = new ArrayList<>(); long sequenceId = 0; for (TableScanNode tableScan : findTableScanNodes(subplan.getFragment().getRoot())) { TableLayoutHandle layout = tableScan.getLayout().get(); SplitSource splitSource = splitManager.getSplits(session, layout); ImmutableSet.Builder<ScheduledSplit> scheduledSplits = ImmutableSet.builder(); while (!splitSource.isFinished()) { for (Split split : getFutureValue(splitSource.getNextBatch(1000))) { scheduledSplits.add(new ScheduledSplit(sequenceId++, split)); } } sources.add(new TaskSource(tableScan.getId(), scheduledSplits.build(), true)); } // create drivers List<Driver> drivers = new ArrayList<>(); Map<PlanNodeId, Driver> driversBySource = new HashMap<>(); for (DriverFactory driverFactory : localExecutionPlan.getDriverFactories()) { for (int i = 0; i < driverFactory.getDriverInstances(); i++) { DriverContext driverContext = taskContext .addPipelineContext(driverFactory.isInputDriver(), driverFactory.isOutputDriver()) .addDriverContext(); Driver driver = driverFactory.createDriver(driverContext); drivers.add(driver); for (PlanNodeId sourceId : driver.getSourceIds()) { driversBySource.put(sourceId, driver); } } driverFactory.close(); } // add sources to the drivers for (TaskSource source : sources) { for (Driver driver : driversBySource.values()) { driver.updateSource(source); } } return ImmutableList.copyOf(drivers); }