Examples with TaskSource com.facebook.presto.execution.TaskSource used on opensource projects

Example 6 with TaskSource

use of com.facebook.presto.execution.TaskSource in project presto by prestodb.

the class HttpRemoteTask method getSource.

private synchronized TaskSource getSource(PlanNodeId planNodeId) {
    Set<ScheduledSplit> splits = pendingSplits.get(planNodeId);
    boolean pendingNoMoreSplits = Boolean.TRUE.equals(this.noMoreSplits.get(planNodeId));
    boolean noMoreSplits = this.noMoreSplits.containsKey(planNodeId);
    Set<Lifespan> noMoreSplitsForLifespan = pendingNoMoreSplitsForLifespan.get(planNodeId);
    TaskSource element = null;
    if (!splits.isEmpty() || !noMoreSplitsForLifespan.isEmpty() || pendingNoMoreSplits) {
        element = new TaskSource(planNodeId, splits, noMoreSplitsForLifespan, noMoreSplits);
    }
    return element;
}

Example 7 with TaskSource

use of com.facebook.presto.execution.TaskSource in project presto by prestodb.

the class LocalQueryRunner method createDrivers.

private List<Driver> createDrivers(Session session, Plan plan, OutputFactory outputFactory, TaskContext taskContext) {
    if (printPlan) {
        System.out.println(PlanPrinter.textLogicalPlan(plan.getRoot(), plan.getTypes(), metadata.getFunctionAndTypeManager(), plan.getStatsAndCosts(), session, 0, false));
    }
    SubPlan subplan = createSubPlans(session, plan, true);
    if (!subplan.getChildren().isEmpty()) {
        throw new AssertionError("Expected subplan to have no children");
    }
    LocalExecutionPlanner executionPlanner = new LocalExecutionPlanner(metadata, Optional.empty(), pageSourceManager, indexManager, partitioningProviderManager, nodePartitioningManager, pageSinkManager, distributedMetadataManager, expressionCompiler, pageFunctionCompiler, joinFilterFunctionCompiler, new IndexJoinLookupStats(), new TaskManagerConfig().setTaskConcurrency(4), new MemoryManagerConfig(), spillerFactory, singleStreamSpillerFactory, partitioningSpillerFactory, blockEncodingManager, new PagesIndex.TestingFactory(false), joinCompiler, new LookupJoinOperators(), new OrderingCompiler(), jsonCodec(TableCommitContext.class), new RowExpressionDeterminismEvaluator(metadata), new NoOpFragmentResultCacheManager(), new ObjectMapper(), standaloneSpillerFactory);
    // plan query
    StageExecutionDescriptor stageExecutionDescriptor = subplan.getFragment().getStageExecutionDescriptor();
    StreamingPlanSection streamingPlanSection = extractStreamingSections(subplan);
    checkState(streamingPlanSection.getChildren().isEmpty(), "expected no materialized exchanges");
    StreamingSubPlan streamingSubPlan = streamingPlanSection.getPlan();
    LocalExecutionPlan localExecutionPlan = executionPlanner.plan(taskContext, stageExecutionDescriptor, subplan.getFragment().getRoot(), subplan.getFragment().getPartitioningScheme(), subplan.getFragment().getTableScanSchedulingOrder(), outputFactory, Optional.empty(), new UnsupportedRemoteSourceFactory(), createTableWriteInfo(streamingSubPlan, metadata, session), false);
    // generate sources
    List<TaskSource> sources = new ArrayList<>();
    long sequenceId = 0;
    for (TableScanNode tableScan : findTableScanNodes(subplan.getFragment().getRoot())) {
        SplitSource splitSource = splitManager.getSplits(session, tableScan.getTable(), getSplitSchedulingStrategy(stageExecutionDescriptor, tableScan.getId()), WarningCollector.NOOP);
        ImmutableSet.Builder<ScheduledSplit> scheduledSplits = ImmutableSet.builder();
        while (!splitSource.isFinished()) {
            for (Split split : getNextBatch(splitSource)) {
                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(), false).addDriverContext();
                Driver driver = driverFactory.createDriver(driverContext);
                drivers.add(driver);
            }
        }
    }
    // add sources to the drivers
    Set<PlanNodeId> tableScanPlanNodeIds = ImmutableSet.copyOf(subplan.getFragment().getTableScanSchedulingOrder());
    for (TaskSource source : sources) {
        DriverFactory driverFactory = driverFactoriesBySource.get(source.getPlanNodeId());
        checkState(driverFactory != null);
        boolean partitioned = tableScanPlanNodeIds.contains(driverFactory.getSourceId().get());
        for (ScheduledSplit split : source.getSplits()) {
            DriverContext driverContext = taskContext.addPipelineContext(driverFactory.getPipelineId(), driverFactory.isInputDriver(), driverFactory.isOutputDriver(), partitioned).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.noMoreDrivers();
    }
    return ImmutableList.copyOf(drivers);
}

Example 8 with TaskSource

use of com.facebook.presto.execution.TaskSource in project presto by prestodb.

the class TestDriver method testBrokenOperatorAddSource.

private void testBrokenOperatorAddSource(DriverContext driverContext) 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, (session, split, table, columns) -> new FixedPageSource(rowPagesBuilder(types).addSequencePage(10, 20, 30, 40).build()), TESTING_TABLE_HANDLE, ImmutableList.of());
    BrokenOperator brokenOperator = new BrokenOperator(driverContext.addOperatorContext(0, new PlanNodeId("test"), "source"));
    final Driver driver = Driver.createDriver(driverContext, source, brokenOperator);
    // block thread in operator processing
    Future<Boolean> driverProcessFor = executor.submit(() -> 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) {
        assertDriverInterrupted(e.getCause());
    }
}

Example 9 with TaskSource

use of com.facebook.presto.execution.TaskSource in project presto by prestodb.

the class TestDriver method testAddSourceFinish.

private void testAddSourceFinish(DriverContext driverContext) {
    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, (session, split, table, columns) -> new FixedPageSource(rowPagesBuilder(types).addSequencePage(10, 20, 30, 40).build()), TESTING_TABLE_HANDLE, ImmutableList.of());
    PageConsumerOperator sink = createSinkOperator(driverContext, types);
    Driver driver = Driver.createDriver(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());
}

Example 10 with TaskSource

use of com.facebook.presto.execution.TaskSource in project presto by prestodb.

the class TestDriver method processSourceDriver.

private void processSourceDriver(DriverContext driverContext) {
    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, (session, split, table, columns) -> new FixedPageSource(rowPagesBuilder(types).addSequencePage(10, 20, 30, 40).build()), TESTING_TABLE_HANDLE, ImmutableList.of());
    PageConsumerOperator sink = createSinkOperator(driverContext, types);
    Driver driver = Driver.createDriver(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());
}