Examples with OperatorContext com.facebook.presto.operator.OperatorContext used on opensource projects

Example 6 with OperatorContext

use of com.facebook.presto.operator.OperatorContext in project presto by prestodb.

the class AbstractOperatorBenchmark method createTableScanOperator.

protected final OperatorFactory createTableScanOperator(int operatorId, PlanNodeId planNodeId, String tableName, String... columnNames) {
    checkArgument(session.getCatalog().isPresent(), "catalog not set");
    checkArgument(session.getSchema().isPresent(), "schema not set");
    // look up the table
    Metadata metadata = localQueryRunner.getMetadata();
    QualifiedObjectName qualifiedTableName = new QualifiedObjectName(session.getCatalog().get(), session.getSchema().get(), tableName);
    TableHandle tableHandle = metadata.getTableHandle(session, qualifiedTableName).orElse(null);
    checkArgument(tableHandle != null, "Table %s does not exist", qualifiedTableName);
    // lookup the columns
    Map<String, ColumnHandle> allColumnHandles = metadata.getColumnHandles(session, tableHandle);
    ImmutableList.Builder<ColumnHandle> columnHandlesBuilder = ImmutableList.builder();
    for (String columnName : columnNames) {
        ColumnHandle columnHandle = allColumnHandles.get(columnName);
        checkArgument(columnHandle != null, "Table %s does not have a column %s", tableName, columnName);
        columnHandlesBuilder.add(columnHandle);
    }
    List<ColumnHandle> columnHandles = columnHandlesBuilder.build();
    Split split = getLocalQuerySplit(session, tableHandle);
    return new OperatorFactory() {

        @Override
        public Operator createOperator(DriverContext driverContext) {
            OperatorContext operatorContext = driverContext.addOperatorContext(operatorId, planNodeId, "BenchmarkSource");
            ConnectorPageSource pageSource = localQueryRunner.getPageSourceManager().createPageSource(session, split, tableHandle.withDynamicFilter(TupleDomain::all), columnHandles);
            return new PageSourceOperator(pageSource, operatorContext);
        }

        @Override
        public void noMoreOperators() {
        }

        @Override
        public OperatorFactory duplicate() {
            throw new UnsupportedOperationException();
        }
    };
}

Example 7 with OperatorContext

use of com.facebook.presto.operator.OperatorContext in project presto by prestodb.

the class TaskThresholdMemoryRevokingScheduler method revokeHighMemoryTasks.

private synchronized void revokeHighMemoryTasks() {
    if (checkPending.getAndSet(false)) {
        Collection<SqlTask> sqlTasks = requireNonNull(allTasksSupplier.get());
        for (SqlTask task : sqlTasks) {
            Optional<TaskContext> taskContext = task.getTaskContext();
            if (!taskContext.isPresent()) {
                continue;
            }
            long currentTaskRevocableMemory = taskContext.get().getTaskMemoryContext().getRevocableMemory();
            if (currentTaskRevocableMemory < maxRevocableMemoryPerTask) {
                continue;
            }
            AtomicLong remainingBytesToRevokeAtomic = new AtomicLong(currentTaskRevocableMemory - maxRevocableMemoryPerTask);
            taskContext.get().accept(new VoidTraversingQueryContextVisitor<AtomicLong>() {

                @Override
                public Void visitOperatorContext(OperatorContext operatorContext, AtomicLong remainingBytesToRevoke) {
                    if (remainingBytesToRevoke.get() > 0) {
                        long revokedBytes = operatorContext.requestMemoryRevoking();
                        if (revokedBytes > 0) {
                            remainingBytesToRevoke.addAndGet(-revokedBytes);
                            log.debug("taskId=%s: requested revoking %s; remaining %s", task.getTaskId(), revokedBytes, remainingBytesToRevoke.get());
                        }
                    }
                    return null;
                }
            }, remainingBytesToRevokeAtomic);
        }
    }
}

Example 8 with OperatorContext

use of com.facebook.presto.operator.OperatorContext in project presto by prestodb.

the class MemoryRevokingScheduler method requestRevoking.

private void requestRevoking(MemoryPoolId memoryPoolId, ArrayList<SqlTask> sqlTasks, long remainingBytesToRevoke) {
    VoidTraversingQueryContextVisitor<AtomicLong> visitor = new VoidTraversingQueryContextVisitor<AtomicLong>() {

        @Override
        public Void visitPipelineContext(PipelineContext pipelineContext, AtomicLong remainingBytesToRevoke) {
            if (remainingBytesToRevoke.get() <= 0) {
                // exit immediately if no work needs to be done
                return null;
            }
            return super.visitPipelineContext(pipelineContext, remainingBytesToRevoke);
        }

        @Override
        public Void visitOperatorContext(OperatorContext operatorContext, AtomicLong remainingBytesToRevoke) {
            if (remainingBytesToRevoke.get() > 0) {
                long revokedBytes = operatorContext.requestMemoryRevoking();
                if (revokedBytes > 0) {
                    remainingBytesToRevoke.addAndGet(-revokedBytes);
                    log.debug("memoryPool=%s: requested revoking %s; remaining %s", memoryPoolId, revokedBytes, remainingBytesToRevoke.get());
                }
            }
            return null;
        }
    };
    // Sort the tasks into their traversal order
    log.debug("Ordering by %s", spillingStrategy);
    sortTasksToTraversalOrder(sqlTasks, spillingStrategy);
    AtomicLong remainingBytesToRevokeAtomic = new AtomicLong(remainingBytesToRevoke);
    for (SqlTask task : sqlTasks) {
        Optional<TaskContext> taskContext = task.getTaskContext();
        if (taskContext.isPresent()) {
            taskContext.get().accept(visitor, remainingBytesToRevokeAtomic);
            if (remainingBytesToRevokeAtomic.get() <= 0) {
                // No further revoking required
                return;
            }
        }
    }
}

Example 9 with OperatorContext

use of com.facebook.presto.operator.OperatorContext in project presto by prestodb.

the class MemoryRevokingScheduler method revokeQueryMemory.

private void revokeQueryMemory(QueryContext queryContext, long maxTotalMemory) {
    QueryId queryId = queryContext.getQueryId();
    MemoryPool memoryPool = queryContext.getMemoryPool();
    // get a fresh value for queryTotalMemory in case it's changed (e.g. by a previous revocation request)
    long queryTotalMemory = getTotalQueryMemoryReservation(queryId, memoryPool);
    // order tasks by decreasing revocableMemory so that we don't spill more tasks than needed
    SortedMap<Long, TaskContext> queryTaskContextsMap = new TreeMap<>(Comparator.reverseOrder());
    queryContext.getAllTaskContexts().forEach(taskContext -> queryTaskContextsMap.put(taskContext.getTaskMemoryContext().getRevocableMemory(), taskContext));
    AtomicLong remainingBytesToRevoke = new AtomicLong(queryTotalMemory - maxTotalMemory);
    Collection<TaskContext> queryTaskContexts = queryTaskContextsMap.values();
    remainingBytesToRevoke.addAndGet(-MemoryRevokingSchedulerUtils.getMemoryAlreadyBeingRevoked(queryTaskContexts, remainingBytesToRevoke.get()));
    for (TaskContext taskContext : queryTaskContexts) {
        if (remainingBytesToRevoke.get() <= 0) {
            break;
        }
        taskContext.accept(new VoidTraversingQueryContextVisitor<AtomicLong>() {

            @Override
            public Void visitOperatorContext(OperatorContext operatorContext, AtomicLong remainingBytesToRevoke) {
                if (remainingBytesToRevoke.get() > 0) {
                    long revokedBytes = operatorContext.requestMemoryRevoking();
                    if (revokedBytes > 0) {
                        remainingBytesToRevoke.addAndGet(-revokedBytes);
                        log.debug("taskId=%s: requested revoking %s; remaining %s", taskContext.getTaskId(), revokedBytes, remainingBytesToRevoke);
                    }
                }
                return null;
            }
        }, remainingBytesToRevoke);
    }
}

Example 10 with OperatorContext

use of com.facebook.presto.operator.OperatorContext in project presto by prestodb.

the class MemoryRevokingSchedulerUtils method getMemoryAlreadyBeingRevoked.

public static long getMemoryAlreadyBeingRevoked(Collection<TaskContext> taskContexts, long targetRevokingLimit) {
    TraversingQueryContextVisitor<Void, Long> visitor = new TraversingQueryContextVisitor<Void, Long>() {

        @Override
        public Long visitOperatorContext(OperatorContext operatorContext, Void context) {
            if (operatorContext.isMemoryRevokingRequested()) {
                return operatorContext.getReservedRevocableBytes();
            }
            return 0L;
        }

        @Override
        public Long mergeResults(List<Long> childrenResults) {
            return childrenResults.stream().mapToLong(i -> i).sum();
        }
    };
    long currentRevoking = 0;
    for (TaskContext taskContext : taskContexts) {
        currentRevoking += taskContext.accept(visitor, null);
        if (currentRevoking > targetRevokingLimit) {
            // Return early, target value exceeded and revoking will not occur
            return currentRevoking;
        }
    }
    return currentRevoking;
}