Examples with LocalExecutionPlan com.facebook.presto.sql.planner.LocalExecutionPlanner.LocalExecutionPlan used on opensource projects

Example 1 with LocalExecutionPlan

use of com.facebook.presto.sql.planner.LocalExecutionPlanner.LocalExecutionPlan 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);
}

Example 2 with LocalExecutionPlan

use of com.facebook.presto.sql.planner.LocalExecutionPlanner.LocalExecutionPlan in project presto by prestodb.

the class PrestoSparkTaskExecutorFactory method doCreate.

public <T extends PrestoSparkTaskOutput> IPrestoSparkTaskExecutor<T> doCreate(int partitionId, int attemptNumber, SerializedPrestoSparkTaskDescriptor serializedTaskDescriptor, Iterator<SerializedPrestoSparkTaskSource> serializedTaskSources, PrestoSparkTaskInputs inputs, CollectionAccumulator<SerializedTaskInfo> taskInfoCollector, CollectionAccumulator<PrestoSparkShuffleStats> shuffleStatsCollector, Class<T> outputType) {
    PrestoSparkTaskDescriptor taskDescriptor = taskDescriptorJsonCodec.fromJson(serializedTaskDescriptor.getBytes());
    ImmutableMap.Builder<String, TokenAuthenticator> extraAuthenticators = ImmutableMap.builder();
    authenticatorProviders.forEach(provider -> extraAuthenticators.putAll(provider.getTokenAuthenticators()));
    Session session = taskDescriptor.getSession().toSession(sessionPropertyManager, taskDescriptor.getExtraCredentials(), extraAuthenticators.build());
    PlanFragment fragment = taskDescriptor.getFragment();
    StageId stageId = new StageId(session.getQueryId(), fragment.getId().getId());
    // Clear the cache if the cache does not have broadcast table for current stageId.
    // We will only cache 1 HT at any time. If the stageId changes, we will drop the old cached HT
    prestoSparkBroadcastTableCacheManager.removeCachedTablesForStagesOtherThan(stageId);
    // TODO: include attemptId in taskId
    TaskId taskId = new TaskId(new StageExecutionId(stageId, 0), partitionId);
    List<TaskSource> taskSources = getTaskSources(serializedTaskSources);
    log.info("Task [%s] received %d splits.", taskId, taskSources.stream().mapToInt(taskSource -> taskSource.getSplits().size()).sum());
    OptionalLong totalSplitSize = computeAllSplitsSize(taskSources);
    if (totalSplitSize.isPresent()) {
        log.info("Total split size: %s bytes.", totalSplitSize.getAsLong());
    }
    // TODO: Remove this once we can display the plan on Spark UI.
    log.info(PlanPrinter.textPlanFragment(fragment, functionAndTypeManager, session, true));
    DataSize maxUserMemory = new DataSize(min(nodeMemoryConfig.getMaxQueryMemoryPerNode().toBytes(), getQueryMaxMemoryPerNode(session).toBytes()), BYTE);
    DataSize maxTotalMemory = new DataSize(min(nodeMemoryConfig.getMaxQueryTotalMemoryPerNode().toBytes(), getQueryMaxTotalMemoryPerNode(session).toBytes()), BYTE);
    DataSize maxBroadcastMemory = getSparkBroadcastJoinMaxMemoryOverride(session);
    if (maxBroadcastMemory == null) {
        maxBroadcastMemory = new DataSize(min(nodeMemoryConfig.getMaxQueryBroadcastMemory().toBytes(), getQueryMaxBroadcastMemory(session).toBytes()), BYTE);
    }
    MemoryPool memoryPool = new MemoryPool(new MemoryPoolId("spark-executor-memory-pool"), maxTotalMemory);
    SpillSpaceTracker spillSpaceTracker = new SpillSpaceTracker(maxQuerySpillPerNode);
    QueryContext queryContext = new QueryContext(session.getQueryId(), maxUserMemory, maxTotalMemory, maxBroadcastMemory, maxRevocableMemory, memoryPool, new TestingGcMonitor(), notificationExecutor, yieldExecutor, maxQuerySpillPerNode, spillSpaceTracker, memoryReservationSummaryJsonCodec);
    queryContext.setVerboseExceededMemoryLimitErrorsEnabled(isVerboseExceededMemoryLimitErrorsEnabled(session));
    queryContext.setHeapDumpOnExceededMemoryLimitEnabled(isHeapDumpOnExceededMemoryLimitEnabled(session));
    String heapDumpFilePath = Paths.get(getHeapDumpFileDirectory(session), format("%s_%s.hprof", session.getQueryId().getId(), stageId.getId())).toString();
    queryContext.setHeapDumpFilePath(heapDumpFilePath);
    TaskStateMachine taskStateMachine = new TaskStateMachine(taskId, notificationExecutor);
    TaskContext taskContext = queryContext.addTaskContext(taskStateMachine, session, // Plan has to be retained only if verbose memory exceeded errors are requested
    isVerboseExceededMemoryLimitErrorsEnabled(session) ? Optional.of(fragment.getRoot()) : Optional.empty(), perOperatorCpuTimerEnabled, cpuTimerEnabled, perOperatorAllocationTrackingEnabled, allocationTrackingEnabled, false);
    final double memoryRevokingThreshold = getMemoryRevokingThreshold(session);
    final double memoryRevokingTarget = getMemoryRevokingTarget(session);
    checkArgument(memoryRevokingTarget <= memoryRevokingThreshold, "memoryRevokingTarget should be less than or equal memoryRevokingThreshold, but got %s and %s respectively", memoryRevokingTarget, memoryRevokingThreshold);
    if (isSpillEnabled(session)) {
        memoryPool.addListener((pool, queryId, totalMemoryReservationBytes) -> {
            if (totalMemoryReservationBytes > queryContext.getPeakNodeTotalMemory()) {
                queryContext.setPeakNodeTotalMemory(totalMemoryReservationBytes);
            }
            if (totalMemoryReservationBytes > pool.getMaxBytes() * memoryRevokingThreshold && memoryRevokeRequestInProgress.compareAndSet(false, true)) {
                memoryRevocationExecutor.execute(() -> {
                    try {
                        AtomicLong remainingBytesToRevoke = new AtomicLong(totalMemoryReservationBytes - (long) (memoryRevokingTarget * pool.getMaxBytes()));
                        remainingBytesToRevoke.addAndGet(-MemoryRevokingSchedulerUtils.getMemoryAlreadyBeingRevoked(ImmutableList.of(taskContext), remainingBytesToRevoke.get()));
                        taskContext.accept(new VoidTraversingQueryContextVisitor<AtomicLong>() {

                            @Override
                            public Void visitOperatorContext(OperatorContext operatorContext, AtomicLong remainingBytesToRevoke) {
                                if (remainingBytesToRevoke.get() > 0) {
                                    long revokedBytes = operatorContext.requestMemoryRevoking();
                                    if (revokedBytes > 0) {
                                        memoryRevokePending.set(true);
                                        remainingBytesToRevoke.addAndGet(-revokedBytes);
                                    }
                                }
                                return null;
                            }
                        }, remainingBytesToRevoke);
                        memoryRevokeRequestInProgress.set(false);
                    } catch (Exception e) {
                        log.error(e, "Error requesting memory revoking");
                    }
                });
            }
            // Get the latest memory reservation info since it might have changed due to revoke
            long totalReservedMemory = pool.getQueryMemoryReservation(queryId) + pool.getQueryRevocableMemoryReservation(queryId);
            // If total memory usage is over maxTotalMemory and memory revoke request is not pending, fail the query with EXCEEDED_MEMORY_LIMIT error
            if (totalReservedMemory > maxTotalMemory.toBytes() && !memoryRevokeRequestInProgress.get() && !isMemoryRevokePending(taskContext)) {
                throw exceededLocalTotalMemoryLimit(maxTotalMemory, queryContext.getAdditionalFailureInfo(totalReservedMemory, 0) + format("Total reserved memory: %s, Total revocable memory: %s", succinctBytes(pool.getQueryMemoryReservation(queryId)), succinctBytes(pool.getQueryRevocableMemoryReservation(queryId))), isHeapDumpOnExceededMemoryLimitEnabled(session), Optional.ofNullable(heapDumpFilePath));
            }
        });
    }
    ImmutableMap.Builder<PlanNodeId, List<PrestoSparkShuffleInput>> shuffleInputs = ImmutableMap.builder();
    ImmutableMap.Builder<PlanNodeId, List<java.util.Iterator<PrestoSparkSerializedPage>>> pageInputs = ImmutableMap.builder();
    ImmutableMap.Builder<PlanNodeId, List<?>> broadcastInputs = ImmutableMap.builder();
    for (RemoteSourceNode remoteSource : fragment.getRemoteSourceNodes()) {
        List<PrestoSparkShuffleInput> remoteSourceRowInputs = new ArrayList<>();
        List<java.util.Iterator<PrestoSparkSerializedPage>> remoteSourcePageInputs = new ArrayList<>();
        List<List<?>> broadcastInputsList = new ArrayList<>();
        for (PlanFragmentId sourceFragmentId : remoteSource.getSourceFragmentIds()) {
            Iterator<Tuple2<MutablePartitionId, PrestoSparkMutableRow>> shuffleInput = inputs.getShuffleInputs().get(sourceFragmentId.toString());
            Broadcast<?> broadcastInput = inputs.getBroadcastInputs().get(sourceFragmentId.toString());
            List<PrestoSparkSerializedPage> inMemoryInput = inputs.getInMemoryInputs().get(sourceFragmentId.toString());
            if (shuffleInput != null) {
                checkArgument(broadcastInput == null, "single remote source is not expected to accept different kind of inputs");
                checkArgument(inMemoryInput == null, "single remote source is not expected to accept different kind of inputs");
                remoteSourceRowInputs.add(new PrestoSparkShuffleInput(sourceFragmentId.getId(), shuffleInput));
                continue;
            }
            if (broadcastInput != null) {
                checkArgument(inMemoryInput == null, "single remote source is not expected to accept different kind of inputs");
                // TODO: Enable NullifyingIterator once migrated to one task per JVM model
                // NullifyingIterator removes element from the list upon return
                // This allows GC to gradually reclaim memory
                // remoteSourcePageInputs.add(getNullifyingIterator(broadcastInput.value()));
                broadcastInputsList.add((List<?>) broadcastInput.value());
                continue;
            }
            if (inMemoryInput != null) {
                // for inmemory inputs pages can be released incrementally to save memory
                remoteSourcePageInputs.add(getNullifyingIterator(inMemoryInput));
                continue;
            }
            throw new IllegalArgumentException("Input not found for sourceFragmentId: " + sourceFragmentId);
        }
        if (!remoteSourceRowInputs.isEmpty()) {
            shuffleInputs.put(remoteSource.getId(), remoteSourceRowInputs);
        }
        if (!remoteSourcePageInputs.isEmpty()) {
            pageInputs.put(remoteSource.getId(), remoteSourcePageInputs);
        }
        if (!broadcastInputsList.isEmpty()) {
            broadcastInputs.put(remoteSource.getId(), broadcastInputsList);
        }
    }
    OutputBufferMemoryManager memoryManager = new OutputBufferMemoryManager(sinkMaxBufferSize.toBytes(), () -> queryContext.getTaskContextByTaskId(taskId).localSystemMemoryContext(), notificationExecutor);
    Optional<OutputPartitioning> preDeterminedPartition = Optional.empty();
    if (fragment.getPartitioningScheme().getPartitioning().getHandle().equals(FIXED_ARBITRARY_DISTRIBUTION)) {
        int partitionCount = getHashPartitionCount(session);
        preDeterminedPartition = Optional.of(new OutputPartitioning(new PreDeterminedPartitionFunction(partitionId % partitionCount, partitionCount), ImmutableList.of(), ImmutableList.of(), false, OptionalInt.empty()));
    }
    TempDataOperationContext tempDataOperationContext = new TempDataOperationContext(session.getSource(), session.getQueryId().getId(), session.getClientInfo(), Optional.of(session.getClientTags()), session.getIdentity());
    TempStorage tempStorage = tempStorageManager.getTempStorage(storageBasedBroadcastJoinStorage);
    Output<T> output = configureOutput(outputType, blockEncodingManager, memoryManager, getShuffleOutputTargetAverageRowSize(session), preDeterminedPartition, tempStorage, tempDataOperationContext, getStorageBasedBroadcastJoinWriteBufferSize(session));
    PrestoSparkOutputBuffer<?> outputBuffer = output.getOutputBuffer();
    LocalExecutionPlan localExecutionPlan = localExecutionPlanner.plan(taskContext, fragment.getRoot(), fragment.getPartitioningScheme(), fragment.getStageExecutionDescriptor(), fragment.getTableScanSchedulingOrder(), output.getOutputFactory(), new PrestoSparkRemoteSourceFactory(blockEncodingManager, shuffleInputs.build(), pageInputs.build(), broadcastInputs.build(), partitionId, shuffleStatsCollector, tempStorage, tempDataOperationContext, prestoSparkBroadcastTableCacheManager, stageId), taskDescriptor.getTableWriteInfo(), true);
    taskStateMachine.addStateChangeListener(state -> {
        if (state.isDone()) {
            outputBuffer.setNoMoreRows();
        }
    });
    PrestoSparkTaskExecution taskExecution = new PrestoSparkTaskExecution(taskStateMachine, taskContext, localExecutionPlan, taskExecutor, splitMonitor, notificationExecutor, memoryUpdateExecutor);
    taskExecution.start(taskSources);
    return new PrestoSparkTaskExecutor<>(taskContext, taskStateMachine, output.getOutputSupplier(), taskInfoCodec, taskInfoCollector, shuffleStatsCollector, executionExceptionFactory, output.getOutputBufferType(), outputBuffer, tempStorage, tempDataOperationContext);
}

Example 3 with LocalExecutionPlan

use of com.facebook.presto.sql.planner.LocalExecutionPlanner.LocalExecutionPlan in project presto by prestodb.

the class TestLocalExecutionPlanner method getLocalExecutionPlan.

private LocalExecutionPlan getLocalExecutionPlan(Session session) {
    SubPlan subPlan = runner.inTransaction(session, transactionSession -> {
        Plan plan = runner.createPlan(transactionSession, "SELECT avg(totalprice) FROM orders", OPTIMIZED_AND_VALIDATED, false, WarningCollector.NOOP);
        return runner.createSubPlans(transactionSession, plan, false);
    });
    // Expect only one child sub plan doing partial aggregation.
    assertEquals(subPlan.getChildren().size(), 1);
    PlanFragment leafFragment = subPlan.getChildren().get(0).getFragment();
    return createTestingPlanner().plan(createTaskContext(EXECUTOR, SCHEDULED_EXECUTOR, session), leafFragment.getRoot(), leafFragment.getPartitioningScheme(), leafFragment.getStageExecutionDescriptor(), leafFragment.getTableScanSchedulingOrder(), new TestingOutputBuffer(), new TestingRemoteSourceFactory(), new TableWriteInfo(Optional.empty(), Optional.empty(), Optional.empty()));
}

Example 4 with LocalExecutionPlan

use of com.facebook.presto.sql.planner.LocalExecutionPlanner.LocalExecutionPlan in project presto by prestodb.

the class TestSqlTaskExecution method testSimple.

@Test(dataProvider = "executionStrategies", timeOut = 20_000)
public void testSimple(PipelineExecutionStrategy executionStrategy) throws Exception {
    ScheduledExecutorService taskNotificationExecutor = newScheduledThreadPool(10, threadsNamed("task-notification-%s"));
    ScheduledExecutorService driverYieldExecutor = newScheduledThreadPool(2, threadsNamed("driver-yield-%s"));
    TaskExecutor taskExecutor = new TaskExecutor(5, 10, 3, 4, TASK_FAIR, Ticker.systemTicker());
    taskExecutor.start();
    try {
        TaskStateMachine taskStateMachine = new TaskStateMachine(TASK_ID, taskNotificationExecutor);
        PartitionedOutputBuffer outputBuffer = newTestingOutputBuffer(taskNotificationExecutor);
        OutputBufferConsumer outputBufferConsumer = new OutputBufferConsumer(outputBuffer, OUTPUT_BUFFER_ID);
        // 
        // test initialization: simple task with 1 pipeline
        // 
        // pipeline 0  ... pipeline id
        // partitioned ... partitioned/unpartitioned pipeline
        // grouped   ... execution strategy (in grouped test)
        // ungrouped  ... execution strategy (in ungrouped test)
        // 
        // TaskOutput
        // |
        // Scan
        // 
        // See #testComplex for all the bahaviors that are tested. Not all of them apply here.
        TestingScanOperatorFactory testingScanOperatorFactory = new TestingScanOperatorFactory(0, TABLE_SCAN_NODE_ID, ImmutableList.of(VARCHAR));
        TaskOutputOperatorFactory taskOutputOperatorFactory = new TaskOutputOperatorFactory(1, TABLE_SCAN_NODE_ID, outputBuffer, Function.identity(), new PagesSerdeFactory(new BlockEncodingManager(), false));
        LocalExecutionPlan localExecutionPlan = new LocalExecutionPlan(ImmutableList.of(new DriverFactory(0, true, true, ImmutableList.of(testingScanOperatorFactory, taskOutputOperatorFactory), OptionalInt.empty(), executionStrategy, Optional.empty())), ImmutableList.of(TABLE_SCAN_NODE_ID), executionStrategy == GROUPED_EXECUTION ? StageExecutionDescriptor.fixedLifespanScheduleGroupedExecution(ImmutableList.of(TABLE_SCAN_NODE_ID), 8) : StageExecutionDescriptor.ungroupedExecution());
        TaskContext taskContext = newTestingTaskContext(taskNotificationExecutor, driverYieldExecutor, taskStateMachine);
        SqlTaskExecution sqlTaskExecution = SqlTaskExecution.createSqlTaskExecution(taskStateMachine, taskContext, outputBuffer, ImmutableList.of(), localExecutionPlan, taskExecutor, taskNotificationExecutor, createTestSplitMonitor());
        // 
        // test body
        assertEquals(taskStateMachine.getState(), TaskState.RUNNING);
        switch(executionStrategy) {
            case UNGROUPED_EXECUTION:
                // add source for pipeline
                sqlTaskExecution.addSources(ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(newScheduledSplit(0, TABLE_SCAN_NODE_ID, Lifespan.taskWide(), 100000, 123)), false)));
                // assert that partial task result is produced
                outputBufferConsumer.consume(123, ASSERT_WAIT_TIMEOUT);
                // pause operator execution to make sure that
                // * operatorFactory will be closed even though operator can't execute
                // * completedDriverGroups will NOT include the newly scheduled driver group while pause is in place
                testingScanOperatorFactory.getPauser().pause();
                // add source for pipeline, mark as no more splits
                sqlTaskExecution.addSources(ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(newScheduledSplit(1, TABLE_SCAN_NODE_ID, Lifespan.taskWide(), 200000, 300), newScheduledSplit(2, TABLE_SCAN_NODE_ID, Lifespan.taskWide(), 300000, 200)), true)));
                // assert that pipeline will have no more drivers
                waitUntilEquals(testingScanOperatorFactory::isOverallNoMoreOperators, true, ASSERT_WAIT_TIMEOUT);
                // assert that no DriverGroup is fully completed
                assertEquals(taskContext.getCompletedDriverGroups(), ImmutableSet.of());
                // resume operator execution
                testingScanOperatorFactory.getPauser().resume();
                // assert that task result is produced
                outputBufferConsumer.consume(300 + 200, ASSERT_WAIT_TIMEOUT);
                outputBufferConsumer.assertBufferComplete(ASSERT_WAIT_TIMEOUT);
                break;
            case GROUPED_EXECUTION:
                // add source for pipeline (driver group [1, 5]), mark driver group [1] as noMoreSplits
                sqlTaskExecution.addSources(ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(newScheduledSplit(0, TABLE_SCAN_NODE_ID, Lifespan.driverGroup(1), 0, 1), newScheduledSplit(1, TABLE_SCAN_NODE_ID, Lifespan.driverGroup(5), 100000, 10)), ImmutableSet.of(Lifespan.driverGroup(1)), false)));
                // assert that pipeline will have no more drivers for driver group [1]
                waitUntilEquals(testingScanOperatorFactory::getDriverGroupsWithNoMoreOperators, ImmutableSet.of(Lifespan.driverGroup(1)), ASSERT_WAIT_TIMEOUT);
                // assert that partial result is produced for both driver groups
                outputBufferConsumer.consume(1 + 10, ASSERT_WAIT_TIMEOUT);
                // assert that driver group [1] is fully completed
                waitUntilEquals(taskContext::getCompletedDriverGroups, ImmutableSet.of(Lifespan.driverGroup(1)), ASSERT_WAIT_TIMEOUT);
                // pause operator execution to make sure that
                // * operatorFactory will be closed even though operator can't execute
                // * completedDriverGroups will NOT include the newly scheduled driver group while pause is in place
                testingScanOperatorFactory.getPauser().pause();
                // add source for pipeline (driver group [5]), mark driver group [5] as noMoreSplits
                sqlTaskExecution.addSources(ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(newScheduledSplit(2, TABLE_SCAN_NODE_ID, Lifespan.driverGroup(5), 200000, 300)), ImmutableSet.of(Lifespan.driverGroup(5)), false)));
                // assert that pipeline will have no more drivers for driver group [1, 5]
                waitUntilEquals(testingScanOperatorFactory::getDriverGroupsWithNoMoreOperators, ImmutableSet.of(Lifespan.driverGroup(1), Lifespan.driverGroup(5)), ASSERT_WAIT_TIMEOUT);
                // assert that driver group [5] is NOT YET fully completed
                assertEquals(taskContext.getCompletedDriverGroups(), ImmutableSet.of(Lifespan.driverGroup(1)));
                // resume operator execution
                testingScanOperatorFactory.getPauser().resume();
                // assert that partial result is produced
                outputBufferConsumer.consume(300, ASSERT_WAIT_TIMEOUT);
                // assert that driver group [1, 5] is fully completed
                waitUntilEquals(taskContext::getCompletedDriverGroups, ImmutableSet.of(Lifespan.driverGroup(1), Lifespan.driverGroup(5)), ASSERT_WAIT_TIMEOUT);
                // pause operator execution to make sure that
                // * operatorFactory will be closed even though operator can't execute
                // * completedDriverGroups will NOT include the newly scheduled driver group while pause is in place
                testingScanOperatorFactory.getPauser().pause();
                // add source for pipeline (driver group [7]), mark pipeline as noMoreSplits without explicitly marking driver group 7
                sqlTaskExecution.addSources(ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(newScheduledSplit(3, TABLE_SCAN_NODE_ID, Lifespan.driverGroup(7), 300000, 45), newScheduledSplit(4, TABLE_SCAN_NODE_ID, Lifespan.driverGroup(7), 400000, 54)), ImmutableSet.of(), true)));
                // assert that pipeline will have no more drivers for driver group [1, 5, 7]
                waitUntilEquals(testingScanOperatorFactory::getDriverGroupsWithNoMoreOperators, ImmutableSet.of(Lifespan.driverGroup(1), Lifespan.driverGroup(5), Lifespan.driverGroup(7)), ASSERT_WAIT_TIMEOUT);
                // assert that pipeline will have no more drivers
                waitUntilEquals(testingScanOperatorFactory::isOverallNoMoreOperators, true, ASSERT_WAIT_TIMEOUT);
                // assert that driver group [1, 5] is fully completed
                assertEquals(taskContext.getCompletedDriverGroups(), ImmutableSet.of(Lifespan.driverGroup(1), Lifespan.driverGroup(5)));
                // resume operator execution
                testingScanOperatorFactory.getPauser().resume();
                // assert driver group [7] is not completed before output buffer is consumed
                MILLISECONDS.sleep(1000);
                assertEquals(taskContext.getCompletedDriverGroups(), ImmutableSet.of(Lifespan.driverGroup(1), Lifespan.driverGroup(5)));
                // assert that result is produced
                outputBufferConsumer.consume(45 + 54, ASSERT_WAIT_TIMEOUT);
                outputBufferConsumer.assertBufferComplete(ASSERT_WAIT_TIMEOUT);
                // assert that driver group [1, 5, 7] is fully completed
                waitUntilEquals(taskContext::getCompletedDriverGroups, ImmutableSet.of(Lifespan.driverGroup(1), Lifespan.driverGroup(5), Lifespan.driverGroup(7)), ASSERT_WAIT_TIMEOUT);
                break;
            default:
                throw new UnsupportedOperationException();
        }
        // complete the task by calling abort on it
        outputBufferConsumer.abort();
        TaskState taskState = taskStateMachine.getStateChange(TaskState.RUNNING).get(10, SECONDS);
        assertEquals(taskState, TaskState.FINISHED);
    } finally {
        taskExecutor.stop();
        taskNotificationExecutor.shutdownNow();
        driverYieldExecutor.shutdown();
    }
}

Example 5 with LocalExecutionPlan

use of com.facebook.presto.sql.planner.LocalExecutionPlanner.LocalExecutionPlan 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);
}