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); }
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); }