@Override
protected List<Driver> createDrivers(TaskContext taskContext) {
if (lookupSourceSupplier == null) {
OperatorFactory ordersTableScan =
createTableScanOperator(0, "orders", "orderkey", "totalprice");
HashBuilderOperatorFactory hashBuilder =
new HashBuilderOperatorFactory(1, ordersTableScan.getTypes(), Ints.asList(0), 1_500_000);
DriverContext driverContext = taskContext.addPipelineContext(false, false).addDriverContext();
Driver driver =
new DriverFactory(false, false, ordersTableScan, hashBuilder).createDriver(driverContext);
while (!driver.isFinished()) {
driver.process();
}
lookupSourceSupplier = hashBuilder.getLookupSourceSupplier();
}
OperatorFactory lineItemTableScan =
createTableScanOperator(0, "lineitem", "orderkey", "quantity");
OperatorFactory joinOperator =
LookupJoinOperators.innerJoin(
1, lookupSourceSupplier, lineItemTableScan.getTypes(), Ints.asList(0));
NullOutputOperatorFactory output = new NullOutputOperatorFactory(2, joinOperator.getTypes());
DriverFactory driverFactory =
new DriverFactory(true, true, lineItemTableScan, joinOperator, output);
DriverContext driverContext = taskContext.addPipelineContext(true, true).addDriverContext();
Driver driver = driverFactory.createDriver(driverContext);
return ImmutableList.of(driver);
}
@Test(dataProvider = "hashEnabledValues")
public void testProbeOuterJoin(
boolean parallelBuild, boolean probeHashEnabled, boolean buildHashEnabled) throws Exception {
TaskContext taskContext = createTaskContext();
// build
List<Type> buildTypes = ImmutableList.<Type>of(VARCHAR, BIGINT, BIGINT);
RowPagesBuilder buildPages =
rowPagesBuilder(buildHashEnabled, Ints.asList(0), ImmutableList.of(VARCHAR, BIGINT, BIGINT))
.addSequencePage(10, 20, 30, 40);
LookupSourceSupplier lookupSourceSupplier =
buildHash(parallelBuild, taskContext, Ints.asList(0), buildPages);
// probe
List<Type> probeTypes = ImmutableList.<Type>of(VARCHAR, BIGINT, BIGINT);
RowPagesBuilder probePages = rowPagesBuilder(probeHashEnabled, Ints.asList(0), probeTypes);
List<Page> probeInput = probePages.addSequencePage(15, 20, 1020, 2020).build();
OperatorFactory joinOperatorFactory =
LookupJoinOperators.probeOuterJoin(
0,
new PlanNodeId("test"),
lookupSourceSupplier,
probePages.getTypes(),
Ints.asList(0),
probePages.getHashChannel());
Operator joinOperator =
joinOperatorFactory.createOperator(
taskContext.addPipelineContext(true, true).addDriverContext());
// expected
// expected
MaterializedResult expected =
MaterializedResult.resultBuilder(taskContext.getSession(), concat(probeTypes, buildTypes))
.row("20", 1020, 2020, "20", 30, 40)
.row("21", 1021, 2021, "21", 31, 41)
.row("22", 1022, 2022, "22", 32, 42)
.row("23", 1023, 2023, "23", 33, 43)
.row("24", 1024, 2024, "24", 34, 44)
.row("25", 1025, 2025, "25", 35, 45)
.row("26", 1026, 2026, "26", 36, 46)
.row("27", 1027, 2027, "27", 37, 47)
.row("28", 1028, 2028, "28", 38, 48)
.row("29", 1029, 2029, "29", 39, 49)
.row("30", 1030, 2030, null, null, null)
.row("31", 1031, 2031, null, null, null)
.row("32", 1032, 2032, null, null, null)
.row("33", 1033, 2033, null, null, null)
.row("34", 1034, 2034, null, null, null)
.build();
assertOperatorEquals(
joinOperator, probeInput, expected, true, getHashChannels(probePages, buildPages));
}
public synchronized boolean reserveMemory(long bytes) {
boolean result = taskContext.reserveMemory(bytes);
if (result) {
memoryReservation.getAndAdd(bytes);
}
return result;
}
@Override
protected List<Driver> createDrivers(TaskContext taskContext) {
DriverFactory driverFactory = createDriverFactory();
DriverContext driverContext = taskContext.addPipelineContext(true, true).addDriverContext();
Driver driver = driverFactory.createDriver(driverContext);
return ImmutableList.of(driver);
}
@Test(dataProvider = "hashEnabledValues")
public void testOuterJoinWithNullOnBothSides(
boolean parallelBuild, boolean probeHashEnabled, boolean buildHashEnabled) throws Exception {
TaskContext taskContext = createTaskContext();
// build
RowPagesBuilder buildPages =
rowPagesBuilder(buildHashEnabled, Ints.asList(0), ImmutableList.of(VARCHAR))
.row("a")
.row((String) null)
.row((String) null)
.row("a")
.row("b");
LookupSourceSupplier lookupSourceSupplier =
buildHash(parallelBuild, taskContext, Ints.asList(0), buildPages);
// probe
List<Type> probeTypes = ImmutableList.<Type>of(VARCHAR);
RowPagesBuilder probePages = rowPagesBuilder(probeHashEnabled, Ints.asList(0), probeTypes);
List<Page> probeInput = probePages.row("a").row("b").row((String) null).row("c").build();
OperatorFactory joinOperatorFactory =
LookupJoinOperators.probeOuterJoin(
0,
new PlanNodeId("test"),
lookupSourceSupplier,
probePages.getTypes(),
Ints.asList(0),
probePages.getHashChannel());
Operator joinOperator =
joinOperatorFactory.createOperator(
taskContext.addPipelineContext(true, true).addDriverContext());
// expected
MaterializedResult expected =
MaterializedResult.resultBuilder(
taskContext.getSession(), concat(probeTypes, buildPages.getTypes()))
.row("a", "a")
.row("a", "a")
.row("b", "b")
.row(null, null)
.row("c", null)
.build();
assertOperatorEquals(
joinOperator, probeInput, expected, true, getHashChannels(probePages, buildPages));
}
@Override
public TaskInfo getTaskInfo(boolean full) {
try (SetThreadName setThreadName = new SetThreadName("Task-%s", taskId)) {
checkTaskCompletion();
TaskState state = taskStateMachine.getState();
List<FailureInfo> failures = ImmutableList.of();
if (state == TaskState.FAILED) {
failures = toFailures(taskStateMachine.getFailureCauses());
}
return new TaskInfo(
taskStateMachine.getTaskId(),
nextTaskInfoVersion.getAndIncrement(),
state,
location,
lastHeartbeat.get(),
sharedBuffer.getInfo(),
getNoMoreSplits(),
taskContext.getTaskStats(),
failures,
taskContext.getOutputItems());
}
}
public DataSize getOperatorPreAllocatedMemory() {
return taskContext.getOperatorPreAllocatedMemory();
}
public DataSize getMaxMemorySize() {
return taskContext.getMaxMemorySize();
}
public boolean isDone() {
return taskContext.isDone();
}
public void failed(Throwable cause) {
taskContext.failed(cause);
}
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));
}
}
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);
}
private static LookupSourceSupplier buildHash(
boolean parallelBuild,
TaskContext taskContext,
List<Integer> hashChannels,
RowPagesBuilder buildPages) {
if (parallelBuild) {
ParallelHashBuilder parallelHashBuilder =
new ParallelHashBuilder(
buildPages.getTypes(),
hashChannels,
buildPages.getHashChannel(),
100,
PARTITION_COUNT);
// collect input data
DriverContext collectDriverContext =
taskContext.addPipelineContext(true, true).addDriverContext();
ValuesOperatorFactory valuesOperatorFactory =
new ValuesOperatorFactory(
0, new PlanNodeId("test"), buildPages.getTypes(), buildPages.build());
OperatorFactory collectOperatorFactory =
parallelHashBuilder.getCollectOperatorFactory(1, new PlanNodeId("test"));
Driver driver =
new Driver(
collectDriverContext,
valuesOperatorFactory.createOperator(collectDriverContext),
collectOperatorFactory.createOperator(collectDriverContext));
while (!driver.isFinished()) {
driver.process();
}
// build hash tables
PipelineContext buildPipeline = taskContext.addPipelineContext(true, true);
OperatorFactory buildOperatorFactory =
parallelHashBuilder.getBuildOperatorFactory(new PlanNodeId("test"));
for (int i = 0; i < PARTITION_COUNT; i++) {
DriverContext buildDriverContext = buildPipeline.addDriverContext();
Driver buildDriver =
new Driver(buildDriverContext, buildOperatorFactory.createOperator(buildDriverContext));
while (!buildDriver.isFinished()) {
buildDriver.process();
}
}
return parallelHashBuilder.getLookupSourceSupplier();
} else {
DriverContext driverContext = taskContext.addPipelineContext(true, true).addDriverContext();
ValuesOperatorFactory valuesOperatorFactory =
new ValuesOperatorFactory(
0, new PlanNodeId("test"), buildPages.getTypes(), buildPages.build());
HashBuilderOperatorFactory hashBuilderOperatorFactory =
new HashBuilderOperatorFactory(
1,
new PlanNodeId("test"),
buildPages.getTypes(),
hashChannels,
buildPages.getHashChannel(),
100);
Driver driver =
new Driver(
driverContext,
valuesOperatorFactory.createOperator(driverContext),
hashBuilderOperatorFactory.createOperator(driverContext));
while (!driver.isFinished()) {
driver.process();
}
return hashBuilderOperatorFactory.getLookupSourceSupplier();
}
}
public boolean isCpuTimerEnabled() {
return taskContext.isCpuTimerEnabled();
}
public Session getSession() {
return taskContext.getSession();
}
public TaskId getTaskId() {
return taskContext.getTaskId();
}
public void start() {
taskContext.start();
}
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);
}
/*
select orderkey, quantity, totalprice
from lineitem join orders using (orderkey)
*/
@Override
protected List<Driver> createDrivers(TaskContext taskContext) {
ImmutableList.Builder<OperatorFactory> driversBuilder = ImmutableList.builder();
driversBuilder.add(ordersTableScan);
OperatorFactory source = ordersTableScan;
Optional<Integer> hashChannel = Optional.empty();
if (hashEnabled) {
source =
createHashProjectOperator(1, new PlanNodeId("test"), ImmutableList.of(BIGINT, DOUBLE));
driversBuilder.add(source);
hashChannel = Optional.of(2);
}
// hash build
HashBuilderOperatorFactory hashBuilder =
new HashBuilderOperatorFactory(
2,
new PlanNodeId("test"),
source.getTypes(),
ImmutableMap.of(),
Ints.asList(0),
hashChannel,
false,
Optional.empty(),
1_500_000,
1);
driversBuilder.add(hashBuilder);
DriverFactory hashBuildDriverFactory =
new DriverFactory(true, false, driversBuilder.build(), OptionalInt.empty());
Driver hashBuildDriver =
hashBuildDriverFactory.createDriver(
taskContext.addPipelineContext(true, false).addDriverContext());
hashBuildDriverFactory.close();
// join
ImmutableList.Builder<OperatorFactory> joinDriversBuilder = ImmutableList.builder();
joinDriversBuilder.add(lineItemTableScan);
source = lineItemTableScan;
hashChannel = Optional.empty();
if (hashEnabled) {
source =
createHashProjectOperator(1, new PlanNodeId("test"), ImmutableList.of(BIGINT, BIGINT));
joinDriversBuilder.add(source);
hashChannel = Optional.of(2);
}
OperatorFactory joinOperator =
LookupJoinOperators.innerJoin(
2,
new PlanNodeId("test"),
hashBuilder.getLookupSourceFactory(),
source.getTypes(),
Ints.asList(0),
hashChannel,
false);
joinDriversBuilder.add(joinOperator);
joinDriversBuilder.add(
new NullOutputOperatorFactory(3, new PlanNodeId("test"), joinOperator.getTypes()));
DriverFactory joinDriverFactory =
new DriverFactory(true, true, joinDriversBuilder.build(), OptionalInt.empty());
Driver joinDriver =
joinDriverFactory.createDriver(
taskContext.addPipelineContext(true, true).addDriverContext());
joinDriverFactory.close();
return ImmutableList.of(hashBuildDriver, joinDriver);
}