use of io.prestosql.operator.TaskContext in project hetu-core by openlookeng.
the class TestSpatialJoinOperator method testEmptyBuildLeftJoin.
@Test
public void testEmptyBuildLeftJoin() {
TaskContext taskContext = createTaskContext();
RowPagesBuilder buildPages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR));
RowPagesBuilder probePages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR)).row(POINT_X, "x").row(null, "null").row(POINT_Y, "y").pageBreak().row(POINT_Z, "z").pageBreak().row(POINT_W, "w");
MaterializedResult expected = resultBuilder(taskContext.getSession(), ImmutableList.of(VARCHAR, VARCHAR)).row("x", null).row("null", null).row("y", null).row("z", null).row("w", null).build();
assertSpatialJoin(taskContext, LEFT, buildPages, probePages, expected);
}
use of io.prestosql.operator.TaskContext in project hetu-core by openlookeng.
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)).row(POLYGON_A, "A", 1).row(POLYGON_A, "A", 2).row(null, "null", null).pageBreak().row(POLYGON_B, "B", 0).row(POLYGON_B, "B", 2);
RowPagesBuilder probePages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR, INTEGER)).row(POINT_X, "x", 2).row(null, "null", null).row(POINT_Y, "y", 2).pageBreak().row(POINT_Z, "z", 0);
MaterializedResult expected = resultBuilder(taskContext.getSession(), ImmutableList.of(VARCHAR, VARCHAR)).row("x", "A").row("y", "A").row("y", "B").row("z", "B").build();
PagesSpatialIndexFactory pagesSpatialIndexFactory = buildIndex(driverContext, (build, probe, r) -> build.contains(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);
assertOperatorEquals(joinOperatorFactory, driverContext, probePages.build(), expected);
}
use of io.prestosql.operator.TaskContext in project hetu-core by openlookeng.
the class TestSqlTaskExecution method testMarker.
@Test(dataProvider = "executionStrategies", timeOut = 20_000)
public void testMarker(PipelineExecutionStrategy executionStrategy) throws Exception {
// This test is a copy of testSimple(). Only difference is some Marker task sources are added to the input,
// and to verify if they can be converted to pages, and these pages can be consumed correctly.
ScheduledExecutorService taskNotificationExecutor = newScheduledThreadPool(10, threadsNamed("task-notification-%s"));
ScheduledExecutorService driverYieldExecutor = newScheduledThreadPool(2, threadsNamed("driver-yield-%s"));
TaskExecutor taskExecutor = new TaskExecutor(5, 10, 3, 4, Ticker.systemTicker());
taskExecutor.start();
try {
TaskStateMachine taskStateMachine = new TaskStateMachine(TaskId.valueOf("query.1.1"), taskNotificationExecutor);
PartitionedOutputBuffer outputBuffer = newTestingOutputBuffer(taskNotificationExecutor);
OutputBufferConsumer outputBufferConsumer = new OutputBufferConsumer(outputBuffer, OUTPUT_BUFFER_ID);
//
// test initialization: simple task with 1 pipeline
//
// pipeline 0 ... pipeline id
// partitioned ... partitioned/unpartitioned pipeline
// grouped ... execution strategy (in grouped test)
// ungrouped ... execution strategy (in ungrouped test)
//
// TaskOutput
// |
// Scan
//
// See #testComplex for all the bahaviors that are tested. Not all of them apply here.
TestingScanOperatorFactory testingScanOperatorFactory = new TestingScanOperatorFactory(0, TABLE_SCAN_NODE_ID, ImmutableList.of(VARCHAR));
TaskOutputOperatorFactory taskOutputOperatorFactory = new TaskOutputOperatorFactory(1, TABLE_SCAN_NODE_ID, outputBuffer, Function.identity());
LocalExecutionPlan localExecutionPlan = new LocalExecutionPlan(ImmutableList.of(new DriverFactory(0, true, true, ImmutableList.of(testingScanOperatorFactory, taskOutputOperatorFactory), OptionalInt.empty(), executionStrategy)), ImmutableList.of(TABLE_SCAN_NODE_ID), executionStrategy == GROUPED_EXECUTION ? StageExecutionDescriptor.fixedLifespanScheduleGroupedExecution(ImmutableList.of(TABLE_SCAN_NODE_ID)) : StageExecutionDescriptor.ungroupedExecution(), Optional.empty());
TaskContext taskContext = newTestingTaskContext(taskNotificationExecutor, driverYieldExecutor, taskStateMachine, true);
SqlTaskExecution sqlTaskExecution = SqlTaskExecution.createSqlTaskExecution(taskStateMachine, taskContext, outputBuffer, ImmutableList.of(), localExecutionPlan, taskExecutor, taskNotificationExecutor, createTestSplitMonitor());
//
// test body
assertEquals(taskStateMachine.getState(), TaskState.RUNNING);
switch(executionStrategy) {
case UNGROUPED_EXECUTION:
// add source for pipeline
sqlTaskExecution.addSources(ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(newScheduledSplit(0, TABLE_SCAN_NODE_ID, Lifespan.taskWide(), 100000, 123)), false)));
// assert that partial task result is produced
outputBufferConsumer.consume(123, ASSERT_WAIT_TIMEOUT);
// add marker source
sqlTaskExecution.addSources(ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(newMarkerSplit(1, TABLE_SCAN_NODE_ID, Lifespan.taskWide(), 1)), false)));
// assert that partial task result is produced
outputBufferConsumer.consume(1, ASSERT_WAIT_TIMEOUT);
// pause operator execution to make sure that
// * operatorFactory will be closed even though operator can't execute
// * completedDriverGroups will NOT include the newly scheduled driver group while pause is in place
testingScanOperatorFactory.getPauser().pause();
// add source for pipeline, mark as no more splits
sqlTaskExecution.addSources(ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(newScheduledSplit(2, TABLE_SCAN_NODE_ID, Lifespan.taskWide(), 200000, 300), newScheduledSplit(3, TABLE_SCAN_NODE_ID, Lifespan.taskWide(), 300000, 200)), false)));
// add marker source
sqlTaskExecution.addSources(ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(newMarkerSplit(4, TABLE_SCAN_NODE_ID, Lifespan.taskWide(), 2)), true)));
// resume operator execution: needed before "isOverallNoMoreOperators" can become true
testingScanOperatorFactory.getPauser().resume();
// assert that pipeline will have no more drivers
waitUntilEquals(testingScanOperatorFactory::isOverallNoMoreOperators, true, ASSERT_WAIT_TIMEOUT);
// assert that no DriverGroup is fully completed
assertEquals(taskContext.getCompletedDriverGroups(), ImmutableSet.of());
// assert that task result is produced
outputBufferConsumer.consume(300 + 200 + 1, ASSERT_WAIT_TIMEOUT);
outputBufferConsumer.assertBufferComplete(ASSERT_WAIT_TIMEOUT);
break;
case GROUPED_EXECUTION:
// add source for pipeline (driver group [1, 5]), mark driver group [1] as noMoreSplits
sqlTaskExecution.addSources(ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(newScheduledSplit(0, TABLE_SCAN_NODE_ID, Lifespan.driverGroup(1), 0, 1), newScheduledSplit(1, TABLE_SCAN_NODE_ID, Lifespan.driverGroup(5), 100000, 10)), false)));
// add marker source
sqlTaskExecution.addSources(ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(newMarkerSplit(2, TABLE_SCAN_NODE_ID, Lifespan.driverGroup(1), 1)), ImmutableSet.of(Lifespan.driverGroup(1)), false)));
// assert that pipeline will have no more drivers for driver group [1]
waitUntilEquals(testingScanOperatorFactory::getDriverGroupsWithNoMoreOperators, ImmutableSet.of(Lifespan.driverGroup(1)), ASSERT_WAIT_TIMEOUT);
// assert that partial result is produced for both driver groups
outputBufferConsumer.consume(1 + 10 + 1, ASSERT_WAIT_TIMEOUT);
// assert that driver group [1] is fully completed
waitUntilEquals(taskContext::getCompletedDriverGroups, ImmutableSet.of(Lifespan.driverGroup(1)), ASSERT_WAIT_TIMEOUT);
// pause operator execution to make sure that
// * operatorFactory will be closed even though operator can't execute
// * completedDriverGroups will NOT include the newly scheduled driver group while pause is in place
testingScanOperatorFactory.getPauser().pause();
// add source for pipeline (driver group [5]), mark driver group [5] as noMoreSplits
sqlTaskExecution.addSources(ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(newScheduledSplit(3, TABLE_SCAN_NODE_ID, Lifespan.driverGroup(5), 200000, 300)), false)));
// add marker source
sqlTaskExecution.addSources(ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(newMarkerSplit(4, TABLE_SCAN_NODE_ID, Lifespan.driverGroup(5), 2)), ImmutableSet.of(Lifespan.driverGroup(5)), false)));
// resume operator execution
testingScanOperatorFactory.getPauser().resume();
// assert that pipeline will have no more drivers for driver group [1, 5]
waitUntilEquals(testingScanOperatorFactory::getDriverGroupsWithNoMoreOperators, ImmutableSet.of(Lifespan.driverGroup(1), Lifespan.driverGroup(5)), ASSERT_WAIT_TIMEOUT);
// assert that partial result is produced
outputBufferConsumer.consume(300 + 1, ASSERT_WAIT_TIMEOUT);
// assert that driver group [1, 5] is fully completed
waitUntilEquals(taskContext::getCompletedDriverGroups, ImmutableSet.of(Lifespan.driverGroup(1), Lifespan.driverGroup(5)), ASSERT_WAIT_TIMEOUT);
// pause operator execution to make sure that
testingScanOperatorFactory.getPauser().pause();
// add source for pipeline (driver group [7]), mark pipeline as noMoreSplits without explicitly marking driver group 7
sqlTaskExecution.addSources(ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(newScheduledSplit(5, TABLE_SCAN_NODE_ID, Lifespan.driverGroup(7), 300000, 45), newScheduledSplit(6, TABLE_SCAN_NODE_ID, Lifespan.driverGroup(7), 400000, 54)), false)));
// add marker source
sqlTaskExecution.addSources(ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(newMarkerSplit(7, TABLE_SCAN_NODE_ID, Lifespan.driverGroup(7), 3)), true)));
// resume operator execution
testingScanOperatorFactory.getPauser().resume();
// assert that pipeline will have no more drivers for driver group [1, 5, 7]
waitUntilEquals(testingScanOperatorFactory::getDriverGroupsWithNoMoreOperators, ImmutableSet.of(Lifespan.driverGroup(1), Lifespan.driverGroup(5), Lifespan.driverGroup(7)), ASSERT_WAIT_TIMEOUT);
// assert that pipeline will have no more drivers
waitUntilEquals(testingScanOperatorFactory::isOverallNoMoreOperators, true, ASSERT_WAIT_TIMEOUT);
// assert that result is produced
outputBufferConsumer.consume(45 + 54 + 1, ASSERT_WAIT_TIMEOUT);
outputBufferConsumer.assertBufferComplete(ASSERT_WAIT_TIMEOUT);
// assert that driver group [1, 5, 7] is fully completed
waitUntilEquals(taskContext::getCompletedDriverGroups, ImmutableSet.of(Lifespan.driverGroup(1), Lifespan.driverGroup(5), Lifespan.driverGroup(7)), ASSERT_WAIT_TIMEOUT);
break;
default:
throw new UnsupportedOperationException();
}
// complete the task by calling abort on it
outputBufferConsumer.abort();
TaskState taskState = taskStateMachine.getStateChange(TaskState.RUNNING).get(10, SECONDS);
assertEquals(taskState, TaskState.FINISHED);
} finally {
taskExecutor.stop();
taskNotificationExecutor.shutdownNow();
driverYieldExecutor.shutdown();
}
}
use of io.prestosql.operator.TaskContext in project hetu-core by openlookeng.
the class TestMemoryRevokingScheduler method testScheduleMemoryRevoking.
@Test
public void testScheduleMemoryRevoking() throws Exception {
SqlTask sqlTask1 = newSqlTask();
SqlTask sqlTask2 = newSqlTask();
TaskContext taskContext1 = sqlTask1.getQueryContext().addTaskContext(new TaskStateMachine(new TaskId("q1", 1, 1), executor), session, false, false, OptionalInt.empty(), Optional.empty(), TESTING_SERDE_FACTORY);
PipelineContext pipelineContext11 = taskContext1.addPipelineContext(0, false, false, false);
DriverContext driverContext111 = pipelineContext11.addDriverContext();
OperatorContext operatorContext1 = driverContext111.addOperatorContext(1, new PlanNodeId("na"), "na");
OperatorContext operatorContext2 = driverContext111.addOperatorContext(2, new PlanNodeId("na"), "na");
DriverContext driverContext112 = pipelineContext11.addDriverContext();
OperatorContext operatorContext3 = driverContext112.addOperatorContext(3, new PlanNodeId("na"), "na");
TaskContext taskContext2 = sqlTask2.getQueryContext().addTaskContext(new TaskStateMachine(new TaskId("q2", 1, 1), executor), session, false, false, OptionalInt.empty(), Optional.empty(), TESTING_SERDE_FACTORY);
PipelineContext pipelineContext21 = taskContext2.addPipelineContext(1, false, false, false);
DriverContext driverContext211 = pipelineContext21.addDriverContext();
OperatorContext operatorContext4 = driverContext211.addOperatorContext(4, new PlanNodeId("na"), "na");
OperatorContext operatorContext5 = driverContext211.addOperatorContext(5, new PlanNodeId("na"), "na");
Collection<SqlTask> tasks = ImmutableList.of(sqlTask1, sqlTask2);
MemoryRevokingScheduler scheduler = new MemoryRevokingScheduler(singletonList(memoryPool), () -> tasks, executor, 1.0, 1.0, false, new DataSize(512, MEGABYTE).toBytes());
allOperatorContexts = ImmutableSet.of(operatorContext1, operatorContext2, operatorContext3, operatorContext4, operatorContext5);
assertMemoryRevokingNotRequested();
requestMemoryRevoking(scheduler);
assertEquals(10, memoryPool.getFreeBytes());
assertMemoryRevokingNotRequested();
LocalMemoryContext revocableMemory1 = operatorContext1.localRevocableMemoryContext();
LocalMemoryContext revocableMemory3 = operatorContext3.localRevocableMemoryContext();
LocalMemoryContext revocableMemory4 = operatorContext4.localRevocableMemoryContext();
LocalMemoryContext revocableMemory5 = operatorContext5.localRevocableMemoryContext();
revocableMemory1.setBytes(3);
revocableMemory3.setBytes(6);
assertEquals(1, memoryPool.getFreeBytes());
requestMemoryRevoking(scheduler);
// we are still good - no revoking needed
assertMemoryRevokingNotRequested();
revocableMemory4.setBytes(7);
assertEquals(-6, memoryPool.getFreeBytes());
requestMemoryRevoking(scheduler);
// we need to revoke 3 and 6
assertMemoryRevokingRequestedFor(operatorContext1, operatorContext3);
// yet another revoking request should not change anything
requestMemoryRevoking(scheduler);
assertMemoryRevokingRequestedFor(operatorContext1, operatorContext3);
// lets revoke some bytes
revocableMemory1.setBytes(0);
operatorContext1.resetMemoryRevokingRequested();
requestMemoryRevoking(scheduler);
assertMemoryRevokingRequestedFor(operatorContext3);
assertEquals(-3, memoryPool.getFreeBytes());
// and allocate some more
revocableMemory5.setBytes(3);
assertEquals(-6, memoryPool.getFreeBytes());
requestMemoryRevoking(scheduler);
// we are still good with just OC3 in process of revoking
assertMemoryRevokingRequestedFor(operatorContext3);
// and allocate some more
revocableMemory5.setBytes(4);
assertEquals(-7, memoryPool.getFreeBytes());
requestMemoryRevoking(scheduler);
// no we have to trigger revoking for OC4
assertMemoryRevokingRequestedFor(operatorContext3, operatorContext4);
}
use of io.prestosql.operator.TaskContext in project hetu-core by openlookeng.
the class TestMemoryRevokingScheduler method createContexts.
private OperatorContext createContexts(SqlTask sqlTask) {
TaskContext taskContext = sqlTask.getQueryContext().addTaskContext(new TaskStateMachine(new TaskId("q", 1, 1), executor), session, false, false, OptionalInt.empty(), Optional.empty(), TESTING_SERDE_FACTORY);
PipelineContext pipelineContext = taskContext.addPipelineContext(0, false, false, false);
DriverContext driverContext = pipelineContext.addDriverContext();
OperatorContext operatorContext = driverContext.addOperatorContext(1, new PlanNodeId("na"), "na");
return operatorContext;
}
Aggregations