private static void printSubPlan(
SubPlan plan,
Map<PlanFragmentId, PlanFragment> fragmentsById,
PlanNodeIdGenerator idGenerator,
StringBuilder output) {
PlanFragment fragment = plan.getFragment();
printFragmentNodes(output, fragment, idGenerator);
fragment.getRoot().accept(new EdgePrinter(output, fragmentsById, idGenerator), null);
for (SubPlan child : plan.getChildren()) {
printSubPlan(child, fragmentsById, idGenerator, output);
}
}
public static String printDistributed(SubPlan plan) {
List<PlanFragment> fragments = plan.getAllFragments();
Map<PlanFragmentId, PlanFragment> fragmentsById =
Maps.uniqueIndex(fragments, PlanFragment.idGetter());
PlanNodeIdGenerator idGenerator = new PlanNodeIdGenerator();
StringBuilder output = new StringBuilder();
output.append("digraph distributed_plan {\n");
printSubPlan(plan, fragmentsById, idGenerator, output);
output.append("}\n");
return output.toString();
}
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);
}