use of com.facebook.presto.operator.TaskContext in project presto by prestodb.
the class QueryContext method addTaskContext.
public TaskContext addTaskContext(TaskStateMachine taskStateMachine, Session session, Optional<PlanNode> taskPlan, boolean perOperatorCpuTimerEnabled, boolean cpuTimerEnabled, boolean perOperatorAllocationTrackingEnabled, boolean allocationTrackingEnabled, boolean legacyLifespanCompletionCondition) {
TaskContext taskContext = TaskContext.createTaskContext(this, taskStateMachine, gcMonitor, notificationExecutor, yieldExecutor, session, queryMemoryContext.newMemoryTrackingContext(), taskPlan, perOperatorCpuTimerEnabled, cpuTimerEnabled, perOperatorAllocationTrackingEnabled, allocationTrackingEnabled, legacyLifespanCompletionCondition);
taskContexts.put(taskStateMachine.getTaskId(), taskContext);
return taskContext;
}
use of com.facebook.presto.operator.TaskContext in project presto by prestodb.
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), new DataSize(1, GIGABYTE), new DataSize(2, GIGABYTE), memoryPool, new TestingGcMonitor(), localQueryRunner.getExecutor(), localQueryRunner.getScheduler(), new DataSize(4, GIGABYTE), spillSpaceTracker, listJsonCodec(TaskMemoryReservationSummary.class));
TaskContext taskContext = queryContext.addTaskContext(new TaskStateMachine(new TaskId("query", 0, 0, 0), localQueryRunner.getExecutor()), localQueryRunner.getDefaultSession(), Optional.empty(), false, false, false, false, false);
// 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();
}
use of com.facebook.presto.operator.TaskContext in project presto by prestodb.
the class TestSpatialJoinOperator method testDistributedSpatialJoin.
@Test
public void testDistributedSpatialJoin() {
TaskContext taskContext = createTaskContext();
DriverContext driverContext = taskContext.addPipelineContext(0, true, true, true).addDriverContext();
RowPagesBuilder buildPages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR, INTEGER));
addGeometryPartitionRows(buildPages, POLYGON_A, "A");
buildPages.row(null, "null", null);
buildPages.pageBreak();
addGeometryPartitionRows(buildPages, POLYGON_B, "B");
addGeometryPartitionRows(buildPages, POLYGON_C, "C");
addGeometryPartitionRows(buildPages, POLYGON_D, "D");
RowPagesBuilder probePages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR, INTEGER));
addGeometryPartitionRows(probePages, POINT_X, "x");
probePages.row(null, "null", null);
addGeometryPartitionRows(probePages, POINT_Y, "y");
probePages.pageBreak();
addGeometryPartitionRows(probePages, POINT_Z, "z");
addGeometryPartitionRows(probePages, POINT_W, "w");
addGeometryPartitionRows(probePages, POINT_V, "v");
addGeometryPartitionRows(probePages, MULTIPOINT_U, "u");
probePages.pageBreak();
addGeometryPartitionRows(probePages, MULTIPOINT_T, "t");
addGeometryPartitionRows(probePages, POINT_S, "s");
addGeometryPartitionRows(probePages, MULTIPOINT_R, "r");
addGeometryPartitionRows(probePages, POINT_Q, "q");
MaterializedResult expected = resultBuilder(taskContext.getSession(), ImmutableList.of(VARCHAR, VARCHAR)).row("x", "A").row("y", "A").row("y", "B").row("z", "B").row("v", "C").row("u", "C").row("t", "C").row("s", "D").row("r", "C").row("r", "D").build();
PagesSpatialIndexFactory pagesSpatialIndexFactory = buildIndex(driverContext, (build, probe, r) -> build.intersects(probe), Optional.empty(), Optional.of(2), Optional.of(KDB_TREE_JSON), Optional.empty(), buildPages);
OperatorFactory joinOperatorFactory = new SpatialJoinOperatorFactory(2, new PlanNodeId("test"), INNER, probePages.getTypes(), Ints.asList(1), 0, Optional.of(2), pagesSpatialIndexFactory);
assertOperatorEqualsIgnoreOrder(joinOperatorFactory, driverContext, probePages.build(), expected);
}
use of com.facebook.presto.operator.TaskContext in project presto by prestodb.
the class TestSpatialJoinOperator method testDistributedSpatialSelfJoin.
@Test
public void testDistributedSpatialSelfJoin() {
TaskContext taskContext = createTaskContext();
DriverContext driverContext = taskContext.addPipelineContext(0, true, true, true).addDriverContext();
RowPagesBuilder pages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR, INTEGER));
addGeometryPartitionRows(pages, POLYGON_A, "A");
pages.row(null, "null", null);
pages.pageBreak();
addGeometryPartitionRows(pages, POLYGON_B, "B");
MaterializedResult expected = resultBuilder(taskContext.getSession(), ImmutableList.of(VARCHAR, VARCHAR)).row("A", "A").row("A", "B").row("B", "A").row("B", "B").build();
PagesSpatialIndexFactory pagesSpatialIndexFactory = buildIndex(driverContext, (build, probe, r) -> build.intersects(probe), Optional.empty(), Optional.of(2), Optional.of(KDB_TREE_JSON), Optional.empty(), pages);
OperatorFactory joinOperatorFactory = new SpatialJoinOperatorFactory(2, new PlanNodeId("test"), INNER, pages.getTypes(), Ints.asList(1), 0, Optional.of(2), pagesSpatialIndexFactory);
assertOperatorEqualsIgnoreOrder(joinOperatorFactory, driverContext, pages.build(), expected);
}
use of com.facebook.presto.operator.TaskContext in project presto by prestodb.
the class TestSpatialJoinOperator method testYield.
@Test
public void testYield() {
// create a filter function that yields for every probe match
// verify we will yield #match times totally
TaskContext taskContext = createTaskContext();
DriverContext driverContext = taskContext.addPipelineContext(0, true, true, false).addDriverContext();
// force a yield for every match
AtomicInteger filterFunctionCalls = new AtomicInteger();
InternalJoinFilterFunction filterFunction = new TestInternalJoinFilterFunction(((leftPosition, leftPage, rightPosition, rightPage) -> {
filterFunctionCalls.incrementAndGet();
driverContext.getYieldSignal().forceYieldForTesting();
return true;
}));
RowPagesBuilder buildPages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR)).row(POLYGON_A, "A").pageBreak().row(POLYGON_B, "B");
PagesSpatialIndexFactory pagesSpatialIndexFactory = buildIndex(driverContext, (build, probe, r) -> build.contains(probe), Optional.empty(), Optional.of(filterFunction), buildPages);
// 10 points in polygon A (x0...x9)
// 10 points in polygons A and B (y0...y9)
// 10 points in polygon B (z0...z9)
// 40 total matches
RowPagesBuilder probePages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR));
for (int i = 0; i < 10; i++) {
probePages.row(stPoint(1 + 0.1 * i, 1 + 0.1 * i), "x" + i);
}
for (int i = 0; i < 10; i++) {
probePages.row(stPoint(4.5 + 0.01 * i, 4.5 + 0.01 * i), "y" + i);
}
for (int i = 0; i < 10; i++) {
probePages.row(stPoint(6 + 0.1 * i, 6 + 0.1 * i), "z" + i);
}
List<Page> probeInput = probePages.build();
OperatorFactory joinOperatorFactory = new SpatialJoinOperatorFactory(2, new PlanNodeId("test"), INNER, probePages.getTypes(), Ints.asList(1), 0, Optional.empty(), pagesSpatialIndexFactory);
Operator operator = joinOperatorFactory.createOperator(driverContext);
assertTrue(operator.needsInput());
operator.addInput(probeInput.get(0));
operator.finish();
// we will yield 40 times due to filterFunction
for (int i = 0; i < 40; i++) {
driverContext.getYieldSignal().setWithDelay(5 * SECONDS.toNanos(1), driverContext.getYieldExecutor());
assertNull(operator.getOutput());
assertEquals(filterFunctionCalls.get(), i + 1, "Expected join to stop processing (yield) after calling filter function once");
driverContext.getYieldSignal().reset();
}
// delayed yield is not going to prevent operator from producing a page now (yield won't be forced because filter function won't be called anymore)
driverContext.getYieldSignal().setWithDelay(5 * SECONDS.toNanos(1), driverContext.getYieldExecutor());
Page output = operator.getOutput();
assertNotNull(output);
// make sure we have 40 matches
assertEquals(output.getPositionCount(), 40);
}
Aggregations