use of com.facebook.presto.ScheduledSplit in project presto by prestodb.
the class LocalQueryRunner method createDrivers.
public List<Driver> createDrivers(Session session, @Language("SQL") String sql, OutputFactory outputFactory, TaskContext taskContext) {
Plan plan = createPlan(session, sql);
if (printPlan) {
System.out.println(PlanPrinter.textLogicalPlan(plan.getRoot(), plan.getTypes(), metadata, session));
}
SubPlan subplan = PlanFragmenter.createSubPlans(session, metadata, plan);
if (!subplan.getChildren().isEmpty()) {
throw new AssertionError("Expected subplan to have no children");
}
LocalExecutionPlanner executionPlanner = new LocalExecutionPlanner(metadata, sqlParser, Optional.empty(), pageSourceManager, indexManager, nodePartitioningManager, pageSinkManager, null, expressionCompiler, joinFilterFunctionCompiler, new IndexJoinLookupStats(), // make sure tests fail if compiler breaks
new CompilerConfig().setInterpreterEnabled(false), new TaskManagerConfig().setTaskConcurrency(4), spillerFactory, blockEncodingSerde, new PagesIndex.TestingFactory(), new JoinCompiler(), new LookupJoinOperators(new JoinProbeCompiler()));
// plan query
LocalExecutionPlan localExecutionPlan = executionPlanner.plan(session, subplan.getFragment().getRoot(), subplan.getFragment().getPartitioningScheme().getOutputLayout(), plan.getTypes(), 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++, tableScan.getId(), split));
}
}
sources.add(new TaskSource(tableScan.getId(), scheduledSplits.build(), true));
}
// create drivers
List<Driver> drivers = new ArrayList<>();
Map<PlanNodeId, DriverFactory> driverFactoriesBySource = new HashMap<>();
for (DriverFactory driverFactory : localExecutionPlan.getDriverFactories()) {
for (int i = 0; i < driverFactory.getDriverInstances().orElse(1); i++) {
if (driverFactory.getSourceId().isPresent()) {
checkState(driverFactoriesBySource.put(driverFactory.getSourceId().get(), driverFactory) == null);
} else {
DriverContext driverContext = taskContext.addPipelineContext(driverFactory.getPipelineId(), driverFactory.isInputDriver(), driverFactory.isOutputDriver()).addDriverContext();
Driver driver = driverFactory.createDriver(driverContext);
drivers.add(driver);
}
}
}
// add sources to the drivers
for (TaskSource source : sources) {
DriverFactory driverFactory = driverFactoriesBySource.get(source.getPlanNodeId());
checkState(driverFactory != null);
for (ScheduledSplit split : source.getSplits()) {
DriverContext driverContext = taskContext.addPipelineContext(driverFactory.getPipelineId(), driverFactory.isInputDriver(), driverFactory.isOutputDriver()).addDriverContext();
Driver driver = driverFactory.createDriver(driverContext);
driver.updateSource(new TaskSource(split.getPlanNodeId(), ImmutableSet.of(split), true));
drivers.add(driver);
}
}
for (DriverFactory driverFactory : localExecutionPlan.getDriverFactories()) {
driverFactory.close();
}
return ImmutableList.copyOf(drivers);
}
use of com.facebook.presto.ScheduledSplit in project presto by prestodb.
the class TestDriver method testBrokenOperatorAddSource.
@Test
public void testBrokenOperatorAddSource() throws Exception {
PlanNodeId sourceId = new PlanNodeId("source");
final List<Type> types = ImmutableList.of(VARCHAR, BIGINT, BIGINT);
// create a table scan operator that does not block, which will cause the driver loop to busy wait
TableScanOperator source = new NotBlockedTableScanOperator(driverContext.addOperatorContext(99, new PlanNodeId("test"), "values"), sourceId, new PageSourceProvider() {
@Override
public ConnectorPageSource createPageSource(Session session, Split split, List<ColumnHandle> columns) {
return new FixedPageSource(rowPagesBuilder(types).addSequencePage(10, 20, 30, 40).build());
}
}, types, ImmutableList.of());
BrokenOperator brokenOperator = new BrokenOperator(driverContext.addOperatorContext(0, new PlanNodeId("test"), "source"));
final Driver driver = new Driver(driverContext, source, brokenOperator);
// block thread in operator processing
Future<Boolean> driverProcessFor = executor.submit(new Callable<Boolean>() {
@Override
public Boolean call() throws Exception {
return driver.processFor(new Duration(1, TimeUnit.MILLISECONDS)).isDone();
}
});
brokenOperator.waitForLocked();
assertSame(driver.getDriverContext(), driverContext);
assertFalse(driver.isFinished());
// processFor always returns NOT_BLOCKED, because DriveLockResult was not acquired
assertTrue(driver.processFor(new Duration(1, TimeUnit.MILLISECONDS)).isDone());
assertFalse(driver.isFinished());
driver.updateSource(new TaskSource(sourceId, ImmutableSet.of(new ScheduledSplit(0, sourceId, newMockSplit())), true));
assertFalse(driver.isFinished());
// processFor always returns NOT_BLOCKED, because DriveLockResult was not acquired
assertTrue(driver.processFor(new Duration(1, TimeUnit.SECONDS)).isDone());
assertFalse(driver.isFinished());
driver.close();
assertTrue(driver.isFinished());
try {
driverProcessFor.get(1, TimeUnit.SECONDS);
fail("Expected InterruptedException");
} catch (ExecutionException e) {
checkArgument(getRootCause(e) instanceof InterruptedException, "Expected root cause exception to be an instance of InterruptedException");
}
}
use of com.facebook.presto.ScheduledSplit in project presto by prestodb.
the class TestDriver method testAddSourceFinish.
@Test
public void testAddSourceFinish() {
PlanNodeId sourceId = new PlanNodeId("source");
final List<Type> types = ImmutableList.of(VARCHAR, BIGINT, BIGINT);
TableScanOperator source = new TableScanOperator(driverContext.addOperatorContext(99, new PlanNodeId("test"), "values"), sourceId, new PageSourceProvider() {
@Override
public ConnectorPageSource createPageSource(Session session, Split split, List<ColumnHandle> columns) {
return new FixedPageSource(rowPagesBuilder(types).addSequencePage(10, 20, 30, 40).build());
}
}, types, ImmutableList.of());
PageConsumerOperator sink = createSinkOperator(source);
Driver driver = new Driver(driverContext, source, sink);
assertSame(driver.getDriverContext(), driverContext);
assertFalse(driver.isFinished());
assertFalse(driver.processFor(new Duration(1, TimeUnit.MILLISECONDS)).isDone());
assertFalse(driver.isFinished());
driver.updateSource(new TaskSource(sourceId, ImmutableSet.of(new ScheduledSplit(0, sourceId, newMockSplit())), true));
assertFalse(driver.isFinished());
assertTrue(driver.processFor(new Duration(1, TimeUnit.SECONDS)).isDone());
assertTrue(driver.isFinished());
assertTrue(sink.isFinished());
assertTrue(source.isFinished());
}
use of com.facebook.presto.ScheduledSplit in project presto by prestodb.
the class SqlTaskExecution method updateSources.
private synchronized Map<PlanNodeId, TaskSource> updateSources(List<TaskSource> sources) {
Map<PlanNodeId, TaskSource> updatedUnpartitionedSources = new HashMap<>();
// first remove any split that was already acknowledged
long currentMaxAcknowledgedSplit = this.maxAcknowledgedSplit;
sources = sources.stream().map(source -> new TaskSource(source.getPlanNodeId(), source.getSplits().stream().filter(scheduledSplit -> scheduledSplit.getSequenceId() > currentMaxAcknowledgedSplit).collect(Collectors.toSet()), source.isNoMoreSplits())).collect(toList());
// update task with new sources
for (TaskSource source : sources) {
if (partitionedDriverFactories.containsKey(source.getPlanNodeId())) {
schedulePartitionedSource(source);
} else {
scheduleUnpartitionedSource(source, updatedUnpartitionedSources);
}
}
// update maxAcknowledgedSplit
maxAcknowledgedSplit = sources.stream().flatMap(source -> source.getSplits().stream()).mapToLong(ScheduledSplit::getSequenceId).max().orElse(maxAcknowledgedSplit);
return updatedUnpartitionedSources;
}
use of com.facebook.presto.ScheduledSplit in project presto by prestodb.
the class SqlTaskExecution method schedulePartitionedSource.
private synchronized void schedulePartitionedSource(TaskSource source) {
// when the source is scheduled
if (!isSchedulingSource(source.getPlanNodeId())) {
pendingSplits.merge(source.getPlanNodeId(), source, TaskSource::update);
return;
}
DriverSplitRunnerFactory partitionedDriverFactory = partitionedDriverFactories.get(source.getPlanNodeId());
ImmutableList.Builder<DriverSplitRunner> runners = ImmutableList.builder();
for (ScheduledSplit scheduledSplit : source.getSplits()) {
// create a new driver for the split
runners.add(partitionedDriverFactory.createDriverRunner(scheduledSplit, true));
}
enqueueDrivers(false, runners.build());
if (source.isNoMoreSplits()) {
partitionedDriverFactory.setNoMoreSplits();
sourceStartOrder.remove(source.getPlanNodeId());
// schedule next source
if (!sourceStartOrder.isEmpty()) {
TaskSource nextSource = pendingSplits.get(sourceStartOrder.peek());
if (nextSource != null) {
schedulePartitionedSource(nextSource);
}
}
}
}
Aggregations