private Split getLocalQuerySplit(TableHandle tableHandle) {
try {
List<Partition> partitions =
splitManager
.getPartitions(tableHandle, Optional.<TupleDomain<ColumnHandle>>absent())
.getPartitions();
SplitSource splitSource = splitManager.getPartitionSplits(tableHandle, partitions);
Split split = Iterables.getOnlyElement(splitSource.getNextBatch(1000));
checkState(splitSource.isFinished(), "Expected only one split for a local query");
return split;
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
throw Throwables.propagate(e);
}
}
private List<Partition> getPartitions(TableScanNode node) {
if (node.getGeneratedPartitions().isPresent()) {
return node.getGeneratedPartitions().get().getPartitions();
}
// Otherwise return all partitions
PartitionResult matchingPartitions =
splitManager.getPartitions(node.getTable(), Optional.<TupleDomain<ColumnHandle>>absent());
return matchingPartitions.getPartitions();
}
private Split getLocalQuerySplit(TableLayoutHandle handle) {
SplitSource splitSource = splitManager.getSplits(defaultSession, handle);
List<Split> splits = new ArrayList<>();
splits.addAll(getFutureValue(splitSource.getNextBatch(1000)));
while (!splitSource.isFinished()) {
splits.addAll(getFutureValue(splitSource.getNextBatch(1000)));
}
checkArgument(
splits.size() == 1,
"Expected only one split for a local query, but got %s splits",
splits.size());
return splits.get(0);
}
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 void addConnector(String catalogName, String connectorId, Connector connector) {
ConnectorMetadata connectorMetadata = connector.getMetadata();
checkState(connectorMetadata != null, "Connector %s can not provide metadata", connectorId);
ConnectorSplitManager connectorSplitManager = connector.getSplitManager();
checkState(
connectorSplitManager != null, "Connector %s does not have a split manager", connectorId);
ConnectorPageSourceProvider connectorPageSourceProvider = null;
try {
connectorPageSourceProvider = connector.getPageSourceProvider();
checkNotNull(
connectorPageSourceProvider,
"Connector %s returned a null page source provider",
connectorId);
} catch (UnsupportedOperationException ignored) {
}
if (connectorPageSourceProvider == null) {
ConnectorRecordSetProvider connectorRecordSetProvider = null;
try {
connectorRecordSetProvider = connector.getRecordSetProvider();
checkNotNull(
connectorRecordSetProvider,
"Connector %s returned a null record set provider",
connectorId);
} catch (UnsupportedOperationException ignored) {
}
checkState(
connectorRecordSetProvider != null,
"Connector %s has neither a PageSource or RecordSet provider",
connectorId);
connectorPageSourceProvider = new RecordPageSourceProvider(connectorRecordSetProvider);
}
ConnectorHandleResolver connectorHandleResolver = connector.getHandleResolver();
checkNotNull(
connectorHandleResolver, "Connector %s does not have a handle resolver", connectorId);
ConnectorPageSinkProvider connectorPageSinkProvider = null;
try {
connectorPageSinkProvider = connector.getPageSinkProvider();
checkNotNull(
connectorPageSinkProvider,
"Connector %s returned a null page sink provider",
connectorId);
} catch (UnsupportedOperationException ignored) {
}
if (connectorPageSinkProvider == null) {
ConnectorRecordSinkProvider connectorRecordSinkProvider = null;
try {
connectorRecordSinkProvider = connector.getRecordSinkProvider();
checkNotNull(
connectorRecordSinkProvider,
"Connector %s returned a null record sink provider",
connectorId);
connectorPageSinkProvider = new RecordPageSinkProvider(connectorRecordSinkProvider);
} catch (UnsupportedOperationException ignored) {
}
}
ConnectorIndexResolver indexResolver = null;
try {
indexResolver = connector.getIndexResolver();
checkNotNull(indexResolver, "Connector %s returned a null index resolver", connectorId);
} catch (UnsupportedOperationException ignored) {
}
// IMPORTANT: all the instances need to be fetched from the connector *before* we add them to
// the corresponding managers.
// Otherwise, a broken connector would leave the managers in an inconsistent state with respect
// to each other
metadataManager.addConnectorMetadata(connectorId, catalogName, connectorMetadata);
metadataManager.addInformationSchemaMetadata(
makeInformationSchemaConnectorId(connectorId),
catalogName,
new InformationSchemaMetadata(catalogName));
splitManager.addConnectorSplitManager(
makeInformationSchemaConnectorId(connectorId),
new InformationSchemaSplitManager(nodeManager));
pageSourceManager.addConnectorPageSourceProvider(
makeInformationSchemaConnectorId(connectorId),
new InformationSchemaPageSourceProvider(metadataManager, splitManager));
splitManager.addConnectorSplitManager(connectorId, connectorSplitManager);
handleResolver.addHandleResolver(connectorId, connectorHandleResolver);
pageSourceManager.addConnectorPageSourceProvider(connectorId, connectorPageSourceProvider);
if (connectorPageSinkProvider != null) {
pageSinkManager.addConnectorPageSinkProvider(connectorId, connectorPageSinkProvider);
}
if (indexResolver != null) {
indexManager.addIndexResolver(connectorId, indexResolver);
}
}
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);
}
@Override
public PlanNode rewriteTableScan(
TableScanNode node, Expression inheritedPredicate, PlanRewriter<Expression> planRewriter) {
DomainTranslator.ExtractionResult extractionResult =
DomainTranslator.fromPredicate(
inheritedPredicate, symbolAllocator.getTypes(), node.getAssignments());
Expression extractionRemainingExpression = extractionResult.getRemainingExpression();
TupleDomain tupleDomain = extractionResult.getTupleDomain();
if (node.getGeneratedPartitions().isPresent()) {
// Add back in the TupleDomain that was used to generate the previous set of Partitions if
// present
// And just for kicks, throw in the domain summary too (as that can only help prune down the
// ranges)
// The domains should never widen between each pass.
tupleDomain =
tupleDomain
.intersect(node.getGeneratedPartitions().get().getTupleDomainInput())
.intersect(node.getPartitionsDomainSummary());
}
PartitionResult matchingPartitions =
splitManager.getPartitions(node.getTable(), Optional.of(tupleDomain));
List<Partition> partitions = matchingPartitions.getPartitions();
TupleDomain undeterminedTupleDomain = matchingPartitions.getUndeterminedTupleDomain();
Expression unevaluatedDomainPredicate =
DomainTranslator.toPredicate(
undeterminedTupleDomain, ImmutableBiMap.copyOf(node.getAssignments()).inverse());
// Construct the post scan predicate. Add the unevaluated TupleDomain back first since those
// are generally cheaper to evaluate than anything we can't extract
Expression postScanPredicate =
combineConjuncts(unevaluatedDomainPredicate, extractionRemainingExpression);
// Do some early partition pruning
partitions =
ImmutableList.copyOf(
filter(
partitions, not(shouldPrunePartition(postScanPredicate, node.getAssignments()))));
GeneratedPartitions generatedPartitions = new GeneratedPartitions(tupleDomain, partitions);
PlanNode output = node;
if (!node.getGeneratedPartitions().equals(Optional.of(generatedPartitions))) {
// Only overwrite the originalConstraint if it was previously null
Expression originalConstraint =
node.getOriginalConstraint() == null
? inheritedPredicate
: node.getOriginalConstraint();
output =
new TableScanNode(
node.getId(),
node.getTable(),
node.getOutputSymbols(),
node.getAssignments(),
originalConstraint,
Optional.of(generatedPartitions));
}
if (!postScanPredicate.equals(BooleanLiteral.TRUE_LITERAL)) {
output = new FilterNode(idAllocator.getNextId(), output, postScanPredicate);
}
return output;
}