use of com.facebook.presto.testing.PageConsumerOperator.PageConsumerOutputFactory in project presto by prestodb.
the class LocalQueryRunner method executeInternal.
private MaterializedResult executeInternal(Session session, @Language("SQL") String sql) {
lock.readLock().lock();
try (Closer closer = Closer.create()) {
AtomicReference<MaterializedResult.Builder> builder = new AtomicReference<>();
PageConsumerOutputFactory outputFactory = new PageConsumerOutputFactory(types -> {
builder.compareAndSet(null, MaterializedResult.resultBuilder(session, types));
return builder.get()::page;
});
TaskContext taskContext = createTaskContext(executor, session);
List<Driver> drivers = createDrivers(session, sql, outputFactory, taskContext);
drivers.forEach(closer::register);
boolean done = false;
while (!done) {
boolean processed = false;
for (Driver driver : drivers) {
if (!driver.isFinished()) {
driver.process();
processed = true;
}
}
done = !processed;
}
verify(builder.get() != null, "Output operator was not created");
return builder.get().build();
} catch (IOException e) {
throw Throwables.propagate(e);
} finally {
lock.readLock().unlock();
}
}
use of com.facebook.presto.testing.PageConsumerOperator.PageConsumerOutputFactory in project presto by prestodb.
the class TestMemoryPools method testBlocking.
@Test
public void testBlocking() throws Exception {
Session session = testSessionBuilder().setCatalog("tpch").setSchema("tiny").setSystemProperty("task_default_concurrency", "1").build();
LocalQueryRunner localQueryRunner = queryRunnerWithInitialTransaction(session);
// add tpch
localQueryRunner.createCatalog("tpch", new TpchConnectorFactory(1), ImmutableMap.of());
// reserve all the memory in the pool
MemoryPool pool = new MemoryPool(new MemoryPoolId("test"), new DataSize(10, MEGABYTE));
QueryId fakeQueryId = new QueryId("fake");
assertTrue(pool.tryReserve(fakeQueryId, TEN_MEGABYTES));
MemoryPool systemPool = new MemoryPool(new MemoryPoolId("testSystem"), new DataSize(10, MEGABYTE));
QueryContext queryContext = new QueryContext(new QueryId("query"), new DataSize(10, MEGABYTE), pool, systemPool, localQueryRunner.getExecutor());
// discard all output
OutputFactory outputFactory = new PageConsumerOutputFactory(types -> (page -> {
}));
TaskContext taskContext = createTaskContext(queryContext, localQueryRunner.getExecutor(), session);
List<Driver> drivers = localQueryRunner.createDrivers("SELECT COUNT(*) FROM orders JOIN lineitem USING (orderkey)", outputFactory, taskContext);
// run driver, until it blocks
while (!isWaitingForMemory(drivers)) {
for (Driver driver : drivers) {
driver.process();
}
}
// driver should be blocked waiting for memory
for (Driver driver : drivers) {
assertFalse(driver.isFinished());
}
assertTrue(pool.getFreeBytes() <= 0);
pool.free(fakeQueryId, TEN_MEGABYTES);
do {
assertFalse(isWaitingForMemory(drivers));
boolean progress = false;
for (Driver driver : drivers) {
ListenableFuture<?> blocked = driver.process();
progress = progress | blocked.isDone();
}
// query should not block
assertTrue(progress);
} while (!drivers.stream().allMatch(Driver::isFinished));
}
use of com.facebook.presto.testing.PageConsumerOperator.PageConsumerOutputFactory in project presto by prestodb.
the class TestMemoryPools method setupConsumeRevocableMemory.
private RevocableMemoryOperator setupConsumeRevocableMemory(DataSize reservedPerPage, long numberOfPages) {
AtomicReference<RevocableMemoryOperator> createOperator = new AtomicReference<>();
setUp(() -> {
DriverContext driverContext = taskContext.addPipelineContext(0, false, false, false).addDriverContext();
OperatorContext revokableOperatorContext = driverContext.addOperatorContext(Integer.MAX_VALUE, new PlanNodeId("revokable_operator"), TableScanOperator.class.getSimpleName());
OutputFactory outputFactory = new PageConsumerOutputFactory(types -> (page -> {
}));
Operator outputOperator = outputFactory.createOutputOperator(2, new PlanNodeId("output"), ImmutableList.of(), Function.identity(), Optional.empty(), new TestingPagesSerdeFactory()).createOperator(driverContext);
RevocableMemoryOperator revocableMemoryOperator = new RevocableMemoryOperator(revokableOperatorContext, reservedPerPage, numberOfPages);
createOperator.set(revocableMemoryOperator);
Driver driver = Driver.createDriver(driverContext, revocableMemoryOperator, outputOperator);
return ImmutableList.of(driver);
});
return createOperator.get();
}
use of com.facebook.presto.testing.PageConsumerOperator.PageConsumerOutputFactory in project presto by prestodb.
the class LocalQueryRunner method executeInternal.
private MaterializedResultWithPlan executeInternal(Session session, @Language("SQL") String sql, WarningCollector warningCollector) {
lock.readLock().lock();
try (Closer closer = Closer.create()) {
AtomicReference<MaterializedResult.Builder> builder = new AtomicReference<>();
PageConsumerOutputFactory outputFactory = new PageConsumerOutputFactory(types -> {
builder.compareAndSet(null, MaterializedResult.resultBuilder(session, types));
return builder.get()::page;
});
Plan plan = createPlan(session, sql, warningCollector);
TaskContext taskContext = TestingTaskContext.builder(notificationExecutor, yieldExecutor, session).setMaxSpillSize(nodeSpillConfig.getMaxSpillPerNode()).setQueryMaxSpillSize(nodeSpillConfig.getQueryMaxSpillPerNode()).setQueryMaxTotalMemory(getQueryMaxTotalMemoryPerNode(session)).setTaskPlan(plan.getRoot()).build();
taskContext.getQueryContext().setVerboseExceededMemoryLimitErrorsEnabled(isVerboseExceededMemoryLimitErrorsEnabled(session));
taskContext.getQueryContext().setHeapDumpOnExceededMemoryLimitEnabled(isHeapDumpOnExceededMemoryLimitEnabled(session));
String heapDumpFilePath = Paths.get(getHeapDumpFileDirectory(session), format("%s_%s.hprof", session.getQueryId().getId(), taskContext.getTaskId().getStageExecutionId().getStageId().getId())).toString();
taskContext.getQueryContext().setHeapDumpFilePath(heapDumpFilePath);
List<Driver> drivers = createDrivers(session, plan, outputFactory, taskContext);
drivers.forEach(closer::register);
boolean done = false;
while (!done) {
boolean processed = false;
for (Driver driver : drivers) {
if (alwaysRevokeMemory) {
driver.getDriverContext().getOperatorContexts().stream().filter(operatorContext -> operatorContext.getOperatorStats().getRevocableMemoryReservation().getValue() > 0).forEach(OperatorContext::requestMemoryRevoking);
}
if (!driver.isFinished()) {
driver.process();
processed = true;
}
}
done = !processed;
}
verify(builder.get() != null, "Output operator was not created");
return new MaterializedResultWithPlan(builder.get().build(), plan);
} catch (IOException e) {
throw new UncheckedIOException(e);
} finally {
lock.readLock().unlock();
}
}
Aggregations