Examples with MemoryPool io.prestosql.memory.MemoryPool used on opensource projects

Example 1 with MemoryPool

use of io.prestosql.memory.MemoryPool in project hetu-core by openlookeng.

the class MemoryLocalQueryRunner method execute.

public List<Page> execute(@Language("SQL") String query) {
    MemoryPool memoryPool = new MemoryPool(new MemoryPoolId("test"), new DataSize(2, GIGABYTE));
    SpillSpaceTracker spillSpaceTracker = new SpillSpaceTracker(new DataSize(1, GIGABYTE));
    QueryContext queryContext = new QueryContext(new QueryId("test"), new DataSize(1, GIGABYTE), new DataSize(2, GIGABYTE), memoryPool, new TestingGcMonitor(), localQueryRunner.getExecutor(), localQueryRunner.getScheduler(), new DataSize(4, GIGABYTE), spillSpaceTracker, NOOP_SNAPSHOT_UTILS);
    TaskContext taskContext = queryContext.addTaskContext(new TaskStateMachine(new TaskId("query", 0, 0), localQueryRunner.getExecutor()), localQueryRunner.getDefaultSession(), false, false, OptionalInt.empty(), Optional.empty(), TESTING_SERDE_FACTORY);
    // Use NullOutputFactory to avoid coping out results to avoid affecting benchmark results
    ImmutableList.Builder<Page> output = ImmutableList.builder();
    List<Driver> drivers = localQueryRunner.createDrivers(query, new PageConsumerOperator.PageConsumerOutputFactory(types -> output::add), taskContext);
    boolean done = false;
    while (!done) {
        boolean processed = false;
        for (Driver driver : drivers) {
            if (!driver.isFinished()) {
                driver.process();
                processed = true;
            }
        }
        done = !processed;
    }
    return output.build();
}

Example 2 with MemoryPool

use of io.prestosql.memory.MemoryPool in project hetu-core by openlookeng.

the class TestSqlTask method createInitialTask.

private SqlTask createInitialTask() {
    TaskId taskId = new TaskId("query", 0, nextTaskId.incrementAndGet());
    String instanceId = "0-query_test_instance_id";
    URI location = URI.create("fake://task/" + taskId);
    QueryContext queryContext = new QueryContext(new QueryId("query"), new DataSize(1, MEGABYTE), new DataSize(2, MEGABYTE), new MemoryPool(new MemoryPoolId("test"), new DataSize(1, GIGABYTE)), new TestingGcMonitor(), taskNotificationExecutor, driverYieldExecutor, new DataSize(1, MEGABYTE), new SpillSpaceTracker(new DataSize(1, GIGABYTE)), NOOP_SNAPSHOT_UTILS);
    queryContext.addTaskContext(new TaskStateMachine(taskId, taskNotificationExecutor), testSessionBuilder().build(), false, false, OptionalInt.empty(), Optional.empty(), TESTING_SERDE_FACTORY);
    return createSqlTask(taskId, instanceId, location, "fake", queryContext, sqlTaskExecutionFactory, taskNotificationExecutor, Functions.identity(), new DataSize(32, MEGABYTE), new CounterStat(), createTestMetadataManager());
}

Example 3 with MemoryPool

use of io.prestosql.memory.MemoryPool in project hetu-core by openlookeng.

the class TestMemoryRevokingScheduler method setUp.

@BeforeMethod
public void setUp() {
    memoryPool = new MemoryPool(GENERAL_POOL, new DataSize(10, BYTE));
    TaskExecutor taskExecutor = new TaskExecutor(8, 16, 3, 4, Ticker.systemTicker());
    taskExecutor.start();
    // Must be single threaded
    executor = newScheduledThreadPool(1, threadsNamed("task-notification-%s"));
    scheduledExecutor = newScheduledThreadPool(2, threadsNamed("task-notification-%s"));
    LocalExecutionPlanner planner = createTestingPlanner();
    sqlTaskExecutionFactory = new SqlTaskExecutionFactory(executor, taskExecutor, planner, createTestSplitMonitor(), new TaskManagerConfig(), createTestMetadataManager());
    allOperatorContexts = null;
}

Example 4 with MemoryPool

use of io.prestosql.memory.MemoryPool in project hetu-core by openlookeng.

the class BufferTestUtils method newTestingTaskContext.

static TaskContext newTestingTaskContext() {
    ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();
    QueryContext queryContext = new QueryContext(new QueryId("query_id"), new DataSize(1, MEGABYTE), new DataSize(2, MEGABYTE), new MemoryPool(new MemoryPoolId("test"), new DataSize(1, GIGABYTE)), new TestingGcMonitor(), executor, executor, new DataSize(1, MEGABYTE), new SpillSpaceTracker(new DataSize(1, GIGABYTE)), NOOP_SNAPSHOT_UTILS);
    TaskId taskId = TaskId.valueOf("query_id.1.2");
    TaskContext taskContext = queryContext.addTaskContext(new TaskStateMachine(taskId, executor), TEST_SNAPSHOT_SESSION, false, false, OptionalInt.empty(), Optional.empty(), TESTING_SERDE_FACTORY);
    taskContext.getSnapshotManager().setTotalComponents(100);
    return taskContext;
}

Example 5 with MemoryPool

use of io.prestosql.memory.MemoryPool in project hetu-core by openlookeng.

the class SqlTaskManager method updateMemoryPoolAssignments.

@Override
public synchronized void updateMemoryPoolAssignments(MemoryPoolAssignmentsRequest assignments) {
    if (coordinatorId != null && coordinatorId.equals(assignments.getCoordinatorId()) && assignments.getVersion() <= currentMemoryPoolAssignmentVersion) {
        return;
    }
    currentMemoryPoolAssignmentVersion = assignments.getVersion();
    if (coordinatorId != null && !coordinatorId.equals(assignments.getCoordinatorId())) {
    // Disable the warning for multi-coordinator feature
    }
    coordinatorId = assignments.getCoordinatorId();
    for (MemoryPoolAssignment assignment : assignments.getAssignments()) {
        if (assignment.getPoolId().equals(GENERAL_POOL)) {
            queryContexts.getUnchecked(assignment.getQueryId()).setMemoryPool(localMemoryManager.getGeneralPool());
        } else if (assignment.getPoolId().equals(RESERVED_POOL)) {
            MemoryPool reservedPool = localMemoryManager.getReservedPool().orElseThrow(() -> new IllegalArgumentException(format("Cannot move %s to the reserved pool as the reserved pool is not enabled", assignment.getQueryId())));
            queryContexts.getUnchecked(assignment.getQueryId()).setMemoryPool(reservedPool);
        } else {
            throw new IllegalArgumentException(format("Cannot move %s to %s as the target memory pool id is invalid", assignment.getQueryId(), assignment.getPoolId()));
        }
    }
}