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

Example 11 with TaskSource

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

the class HttpRemoteTask method processTaskUpdate.

private synchronized void processTaskUpdate(TaskInfo newValue, List<TaskSource> sources) {
    updateTaskInfo(newValue);
    // remove acknowledged splits, which frees memory
    for (TaskSource source : sources) {
        PlanNodeId planNodeId = source.getPlanNodeId();
        boolean isTableScanSource = tableScanPlanNodeIds.contains(planNodeId);
        int removed = 0;
        long removedWeight = 0;
        for (ScheduledSplit split : source.getSplits()) {
            if (pendingSplits.remove(planNodeId, split)) {
                if (isTableScanSource) {
                    removed++;
                    removedWeight = addExact(removedWeight, split.getSplit().getSplitWeight().getRawValue());
                }
            }
        }
        if (source.isNoMoreSplits()) {
            noMoreSplits.put(planNodeId, false);
        }
        for (Lifespan lifespan : source.getNoMoreSplitsForLifespan()) {
            pendingNoMoreSplitsForLifespan.remove(planNodeId, lifespan);
        }
        if (isTableScanSource) {
            pendingSourceSplitCount -= removed;
            pendingSourceSplitsWeight -= removedWeight;
        }
    }
    // Update stats before split queue space to ensure node stats are up to date before waking up the scheduler
    updateTaskStats();
    updateSplitQueueSpace();
}

Example 12 with TaskSource

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

the class TestSystemMemoryBlocking method testTableScanSystemMemoryBlocking.

@Test
public void testTableScanSystemMemoryBlocking() {
    PlanNodeId sourceId = new PlanNodeId("source");
    final List<Type> types = ImmutableList.of(VARCHAR);
    TableScanOperator source = new TableScanOperator(driverContext.addOperatorContext(1, new PlanNodeId("test"), "values"), sourceId, (session, split, table, columns) -> new FixedPageSource(rowPagesBuilder(types).addSequencePage(10, 1).addSequencePage(10, 1).addSequencePage(10, 1).addSequencePage(10, 1).addSequencePage(10, 1).build()), new TableHandle(new ConnectorId("test"), new ConnectorTableHandle() {
    }, new ConnectorTransactionHandle() {
    }, Optional.empty()), ImmutableList.of());
    PageConsumerOperator sink = createSinkOperator(types);
    Driver driver = Driver.createDriver(driverContext, source, sink);
    assertSame(driver.getDriverContext(), driverContext);
    assertFalse(driver.isFinished());
    Split testSplit = new Split(new ConnectorId("test"), TestingTransactionHandle.create(), new TestSplit());
    driver.updateSource(new TaskSource(sourceId, ImmutableSet.of(new ScheduledSplit(0, sourceId, testSplit)), true));
    ListenableFuture<?> blocked = driver.processFor(new Duration(1, NANOSECONDS));
    // the driver shouldn't block in the first call as it will be able to move a page between source and the sink operator
    // but the operator should be blocked
    assertTrue(blocked.isDone());
    assertFalse(source.getOperatorContext().isWaitingForMemory().isDone());
    // and they should stay blocked until more memory becomes available
    for (int i = 0; i < 10; i++) {
        blocked = driver.processFor(new Duration(1, NANOSECONDS));
        assertFalse(blocked.isDone());
        assertFalse(source.getOperatorContext().isWaitingForMemory().isDone());
    }
    // free up some memory
    memoryPool.free(QUERY_ID, "test", memoryPool.getReservedBytes());
    // the operator should be unblocked
    assertTrue(source.getOperatorContext().isWaitingForMemory().isDone());
    // the driver shouldn't be blocked
    blocked = driver.processFor(new Duration(1, NANOSECONDS));
    assertTrue(blocked.isDone());
}

Example 13 with TaskSource

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

the class Driver method processNewSources.

@GuardedBy("exclusiveLock")
private void processNewSources() {
    checkLockHeld("Lock must be held to call processNewSources");
    // only update if the driver is still alive
    if (state.get() != State.ALIVE) {
        return;
    }
    TaskSource sourceUpdate = pendingTaskSourceUpdates.getAndSet(null);
    if (sourceUpdate == null) {
        return;
    }
    // merge the current source and the specified source update
    TaskSource newSource = currentTaskSource.update(sourceUpdate);
    // if the update contains no new data, just return
    if (newSource == currentTaskSource) {
        return;
    }
    // determine new splits to add
    Set<ScheduledSplit> newSplits = Sets.difference(newSource.getSplits(), currentTaskSource.getSplits());
    // add new splits
    SourceOperator sourceOperator = this.sourceOperator.orElseThrow(VerifyException::new);
    for (ScheduledSplit newSplit : newSplits) {
        Split split = newSplit.getSplit();
        if (fragmentResultCacheContext.get().isPresent() && !(split.getConnectorSplit() instanceof RemoteSplit)) {
            checkState(!this.cachedResult.get().isPresent());
            this.fragmentResultCacheContext.set(this.fragmentResultCacheContext.get().map(context -> context.updateRuntimeInformation(split.getConnectorSplit())));
            Optional<Iterator<Page>> pages = fragmentResultCacheContext.get().get().getFragmentResultCacheManager().get(fragmentResultCacheContext.get().get().getHashedCanonicalPlanFragment(), split);
            sourceOperator.getOperatorContext().getRuntimeStats().addMetricValue(pages.isPresent() ? RuntimeMetricName.FRAGMENT_RESULT_CACHE_HIT : RuntimeMetricName.FRAGMENT_RESULT_CACHE_MISS, 1);
            this.cachedResult.set(pages);
            this.split.set(split);
        }
        Supplier<Optional<UpdatablePageSource>> pageSource = sourceOperator.addSplit(split);
        deleteOperator.ifPresent(deleteOperator -> deleteOperator.setPageSource(pageSource));
    }
    // set no more splits
    if (newSource.isNoMoreSplits()) {
        sourceOperator.noMoreSplits();
    }
    currentTaskSource = newSource;
}

Example 14 with TaskSource

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

the class IndexLoader method streamIndexDataForSingleKey.

public IndexedData streamIndexDataForSingleKey(UpdateRequest updateRequest) {
    Page indexKeyTuple = updateRequest.getPage().getRegion(0, 1);
    PageBuffer pageBuffer = new PageBuffer(100);
    DriverFactory driverFactory = indexBuildDriverFactoryProvider.createStreaming(pageBuffer, indexKeyTuple);
    Driver driver = driverFactory.createDriver(pipelineContext.addDriverContext());
    PageRecordSet pageRecordSet = new PageRecordSet(keyTypes, indexKeyTuple);
    PlanNodeId planNodeId = driverFactory.getSourceId().get();
    ScheduledSplit split = new ScheduledSplit(0, planNodeId, new Split(INDEX_CONNECTOR_ID, new ConnectorTransactionHandle() {
    }, new IndexSplit(pageRecordSet)));
    driver.updateSource(new TaskSource(planNodeId, ImmutableSet.of(split), true));
    return new StreamingIndexedData(outputTypes, keyTypes, indexKeyTuple, pageBuffer, driver);
}

Example 15 with TaskSource

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

the class PrestoSparkTaskExecution method scheduleDriversForSplitLifeCycle.

private synchronized void scheduleDriversForSplitLifeCycle(List<TaskSource> sources) {
    checkArgument(sources.stream().allMatch(TaskSource::isNoMoreSplits), "All task sources are expected to be final");
    ListMultimap<PlanNodeId, ScheduledSplit> splits = ArrayListMultimap.create();
    for (TaskSource taskSource : sources) {
        splits.putAll(taskSource.getPlanNodeId(), taskSource.getSplits());
    }
    for (PlanNodeId planNodeId : schedulingOrder) {
        DriverSplitRunnerFactory driverSplitRunnerFactory = driverRunnerFactoriesWithSplitLifeCycle.get(planNodeId);
        List<ScheduledSplit> planNodeSplits = splits.get(planNodeId);
        scheduleTableScanSource(driverSplitRunnerFactory, planNodeSplits);
    }
}