public PartitionedOutputBufferManager(PartitioningHandle partitioningHandle, int partitionCount, Consumer<OutputBuffers> outputBufferTarget) { checkArgument(partitionCount >= 1, "partitionCount must be at least 1"); ImmutableMap.Builder<OutputBufferId, Integer> partitions = ImmutableMap.builder(); for (int partition = 0; partition < partitionCount; partition++) { partitions.put(new OutputBufferId(partition), partition); } OutputBuffers outputBuffers = createInitialEmptyOutputBuffers(requireNonNull(partitioningHandle, "partitioningHandle is null")) .withBuffers(partitions.build()) .withNoMoreBufferIds(); outputBufferTarget.accept(outputBuffers); this.outputBuffers = outputBuffers.getBuffers(); }
outputBuffers = outputBuffers.withBuffer(newBuffer, BROADCAST_PARTITION_ID); outputBuffers = outputBuffers.withNoMoreBufferIds();
@Test public void testAddQueueAfterCreation() { PartitionedOutputBuffer buffer = createPartitionedBuffer( createInitialEmptyOutputBuffers(PARTITIONED) .withBuffer(FIRST, 0) .withNoMoreBufferIds(), sizeOfPages(10)); assertFalse(buffer.isFinished()); try { buffer.setOutputBuffers(createInitialEmptyOutputBuffers(PARTITIONED) .withBuffer(FIRST, 0) .withBuffer(SECOND, 0) .withNoMoreBufferIds()); fail("Expected IllegalStateException from addQueue after noMoreQueues has been called"); } catch (IllegalArgumentException ignored) { } }
@Test public void testAddQueueAfterNoMoreQueues() throws Exception { SharedBuffer sharedBuffer = new SharedBuffer(TASK_ID, TASK_INSTANCE_ID, stateNotificationExecutor, sizeOfPages(10)); assertFalse(sharedBuffer.isFinished()); // tell buffer no more queues will be added sharedBuffer.setOutputBuffers(INITIAL_EMPTY_OUTPUT_BUFFERS.withNoMoreBufferIds()); assertFalse(sharedBuffer.isFinished()); // set no more queues a second time to assure that we don't get an exception or such sharedBuffer.setOutputBuffers(INITIAL_EMPTY_OUTPUT_BUFFERS.withNoMoreBufferIds()); assertFalse(sharedBuffer.isFinished()); // set no more queues a third time to assure that we don't get an exception or such sharedBuffer.setOutputBuffers(INITIAL_EMPTY_OUTPUT_BUFFERS.withNoMoreBufferIds()); assertFalse(sharedBuffer.isFinished()); try { OutputBuffers outputBuffers = INITIAL_EMPTY_OUTPUT_BUFFERS .withBuffer(FOO, 0) .withNoMoreBufferIds(); sharedBuffer.setOutputBuffers(outputBuffers); fail("Expected IllegalStateException from addQueue after noMoreQueues has been called"); } catch (IllegalStateException ignored) { } }
@Test public void testAddQueueAfterNoMoreQueues() { BroadcastOutputBuffer buffer = createBroadcastBuffer(createInitialEmptyOutputBuffers(BROADCAST), sizeOfPages(10)); assertFalse(buffer.isFinished()); // tell buffer no more queues will be added buffer.setOutputBuffers(createInitialEmptyOutputBuffers(BROADCAST).withNoMoreBufferIds()); assertTrue(buffer.isFinished()); // set no more queues a second time to assure that we don't get an exception or such buffer.setOutputBuffers(createInitialEmptyOutputBuffers(BROADCAST).withNoMoreBufferIds()); assertTrue(buffer.isFinished()); // set no more queues a third time to assure that we don't get an exception or such buffer.setOutputBuffers(createInitialEmptyOutputBuffers(BROADCAST).withNoMoreBufferIds()); assertTrue(buffer.isFinished()); }
@Test public void testEmptyQuery() throws Exception { SqlTask sqlTask = createInitialTask(); TaskInfo taskInfo = sqlTask.updateTask(TEST_SESSION, Optional.of(PLAN_FRAGMENT), ImmutableList.<TaskSource>of(), INITIAL_EMPTY_OUTPUT_BUFFERS); assertEquals(taskInfo.getState(), TaskState.RUNNING); taskInfo = sqlTask.getTaskInfo(); assertEquals(taskInfo.getState(), TaskState.RUNNING); taskInfo = sqlTask.updateTask(TEST_SESSION, Optional.of(PLAN_FRAGMENT), ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.<ScheduledSplit>of(), true)), INITIAL_EMPTY_OUTPUT_BUFFERS.withNoMoreBufferIds()); assertEquals(taskInfo.getState(), TaskState.FINISHED); taskInfo = sqlTask.getTaskInfo(); assertEquals(taskInfo.getState(), TaskState.FINISHED); }
private PartitionedOutputBuffer newTestingOutputBuffer(ScheduledExecutorService taskNotificationExecutor) { return new PartitionedOutputBuffer( "task-id", new StateMachine<>("bufferState", taskNotificationExecutor, OPEN, TERMINAL_BUFFER_STATES), createInitialEmptyOutputBuffers(PARTITIONED) .withBuffer(OUTPUT_BUFFER_ID, 0) .withNoMoreBufferIds(), new DataSize(1, MEGABYTE), () -> new SimpleLocalMemoryContext(newSimpleAggregatedMemoryContext(), "test"), taskNotificationExecutor); }
@Test public void testAddQueueAfterCreation() { BroadcastOutputBuffer buffer = createBroadcastBuffer( createInitialEmptyOutputBuffers(BROADCAST) .withBuffer(FIRST, BROADCAST_PARTITION_ID) .withNoMoreBufferIds(), sizeOfPages(10)); assertFalse(buffer.isFinished()); try { buffer.setOutputBuffers(createInitialEmptyOutputBuffers(BROADCAST) .withBuffer(FIRST, BROADCAST_PARTITION_ID) .withBuffer(SECOND, BROADCAST_PARTITION_ID) .withNoMoreBufferIds()); fail("Expected IllegalStateException from addQueue after noMoreQueues has been called"); } catch (IllegalArgumentException ignored) { } }
@Test public void testAddQueueAfterDestroy() throws Exception { SharedBuffer sharedBuffer = new SharedBuffer(TASK_ID, TASK_INSTANCE_ID, stateNotificationExecutor, sizeOfPages(10)); assertFalse(sharedBuffer.isFinished()); // destroy buffer sharedBuffer.destroy(); assertFinished(sharedBuffer); // set no more queues to assure that we don't get an exception or such sharedBuffer.setOutputBuffers(INITIAL_EMPTY_OUTPUT_BUFFERS.withNoMoreBufferIds()); assertFinished(sharedBuffer); // set no more queues a second time to assure that we don't get an exception or such sharedBuffer.setOutputBuffers(INITIAL_EMPTY_OUTPUT_BUFFERS.withNoMoreBufferIds()); assertFinished(sharedBuffer); // add queue calls after finish should be ignored sharedBuffer.setOutputBuffers(INITIAL_EMPTY_OUTPUT_BUFFERS.withBuffer(FOO, 0).withNoMoreBufferIds()); }
@Test public void testEmptyQuery() { SqlTask sqlTask = createInitialTask(); TaskInfo taskInfo = sqlTask.updateTask(TEST_SESSION, Optional.of(PLAN_FRAGMENT), ImmutableList.of(), createInitialEmptyOutputBuffers(PARTITIONED) .withNoMoreBufferIds(), OptionalInt.empty()); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING); taskInfo = sqlTask.getTaskInfo(); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING); taskInfo = sqlTask.updateTask(TEST_SESSION, Optional.of(PLAN_FRAGMENT), ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(), true)), createInitialEmptyOutputBuffers(PARTITIONED) .withNoMoreBufferIds(), OptionalInt.empty()); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.FINISHED); taskInfo = sqlTask.getTaskInfo(); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.FINISHED); }
private PartitionedOutputOperator createPartitionedOutputOperator() { PartitionFunction partitionFunction = new LocalPartitionGenerator(new InterpretedHashGenerator(ImmutableList.of(BIGINT), new int[] {0}), PARTITION_COUNT); PagesSerdeFactory serdeFactory = new PagesSerdeFactory(new BlockEncodingManager(new TypeRegistry()), false); OutputBuffers buffers = createInitialEmptyOutputBuffers(PARTITIONED); for (int partition = 0; partition < PARTITION_COUNT; partition++) { buffers = buffers.withBuffer(new OutputBuffers.OutputBufferId(partition), partition); } PartitionedOutputBuffer buffer = createPartitionedBuffer( buffers.withNoMoreBufferIds(), new DataSize(Long.MAX_VALUE, BYTE)); // don't let output buffer block PartitionedOutputFactory operatorFactory = new PartitionedOutputFactory( partitionFunction, ImmutableList.of(0), ImmutableList.of(Optional.empty()), false, OptionalInt.empty(), buffer, new DataSize(1, GIGABYTE)); return (PartitionedOutputOperator) operatorFactory .createOutputOperator(0, new PlanNodeId("plan-node-0"), TYPES, Function.identity(), serdeFactory) .createOperator(createDriverContext()); }
outputBuffers = outputBuffers.withBuffer(newBuffer, newBuffer.getId()); outputBuffers = outputBuffers.withNoMoreBufferIds();
@Test public void testAddQueueAfterCreation() { ArbitraryOutputBuffer buffer = createArbitraryBuffer( createInitialEmptyOutputBuffers(ARBITRARY) .withBuffer(FIRST, BROADCAST_PARTITION_ID) .withNoMoreBufferIds(), sizeOfPages(10)); assertFalse(buffer.isFinished()); try { buffer.setOutputBuffers(createInitialEmptyOutputBuffers(ARBITRARY) .withBuffer(FIRST, BROADCAST_PARTITION_ID) .withBuffer(SECOND, BROADCAST_PARTITION_ID) .withNoMoreBufferIds()); fail("Expected IllegalStateException from addQueue after noMoreQueues has been called"); } catch (IllegalArgumentException ignored) { } }
@Test public void testAbort() throws Exception { SqlTask sqlTask = createInitialTask(); TaskInfo taskInfo = sqlTask.updateTask(TEST_SESSION, Optional.of(PLAN_FRAGMENT), ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(SPLIT), true)), INITIAL_EMPTY_OUTPUT_BUFFERS.withBuffer(OUT, 0).withNoMoreBufferIds()); assertEquals(taskInfo.getState(), TaskState.RUNNING); taskInfo = sqlTask.getTaskInfo(); assertEquals(taskInfo.getState(), TaskState.RUNNING); sqlTask.abortTaskResults(OUT); taskInfo = sqlTask.getTaskInfo(taskInfo.getState()).get(1, SECONDS); assertEquals(taskInfo.getState(), TaskState.FINISHED); taskInfo = sqlTask.getTaskInfo(); assertEquals(taskInfo.getState(), TaskState.FINISHED); }
@Test public void testEmptyQuery() { try (SqlTaskManager sqlTaskManager = createSqlTaskManager(new TaskManagerConfig())) { TaskId taskId = TASK_ID; TaskInfo taskInfo = createTask(sqlTaskManager, taskId, createInitialEmptyOutputBuffers(PARTITIONED).withNoMoreBufferIds()); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING); taskInfo = createTask(sqlTaskManager, taskId, ImmutableSet.of(), createInitialEmptyOutputBuffers(PARTITIONED).withNoMoreBufferIds()); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.FINISHED); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getTaskStatus().getState(), TaskState.FINISHED); } }
private SqlStageExecution createSqlStageExecution(StageExecutionPlan tableScanPlan, NodeTaskMap nodeTaskMap) { StageId stageId = new StageId(new QueryId("query"), 0); SqlStageExecution stage = SqlStageExecution.createSqlStageExecution(stageId, locationFactory.createStageLocation(stageId), tableScanPlan.getFragment(), new MockRemoteTaskFactory(queryExecutor, scheduledExecutor), TEST_SESSION, true, nodeTaskMap, queryExecutor, new NoOpFailureDetector(), new SplitSchedulerStats()); stage.setOutputBuffers(createInitialEmptyOutputBuffers(PARTITIONED) .withBuffer(OUT, 0) .withNoMoreBufferIds()); return stage; }
private SqlStageExecution createSqlStageExecution(StageExecutionPlan tableScanPlan, NodeTaskMap nodeTaskMap) { StageId stageId = new StageId(new QueryId("query"), "stage"); SqlStageExecution stage = new SqlStageExecution(stageId, locationFactory.createStageLocation(stageId), tableScanPlan.getFragment(), new MockRemoteTaskFactory(executor), TEST_SESSION, nodeTaskMap, executor); stage.setOutputBuffers(INITIAL_EMPTY_OUTPUT_BUFFERS .withBuffer(OUT, 0) .withNoMoreBufferIds()); return stage; }
@Test public void testInvalidConstructorArg() { try { createBroadcastBuffer(createInitialEmptyOutputBuffers(BROADCAST).withBuffer(FIRST, BROADCAST_PARTITION_ID).withNoMoreBufferIds(), new DataSize(0, BYTE)); fail("Expected IllegalStateException"); } catch (IllegalArgumentException ignored) { } try { createBroadcastBuffer(createInitialEmptyOutputBuffers(BROADCAST), new DataSize(0, BYTE)); fail("Expected IllegalStateException"); } catch (IllegalArgumentException ignored) { } }
@Test public void testAbortResults() throws Exception { try (SqlTaskManager sqlTaskManager = createSqlTaskManager(new TaskManagerConfig())) { TaskId taskId = TASK_ID; TaskInfo taskInfo = sqlTaskManager.updateTask(TEST_SESSION, taskId, Optional.of(PLAN_FRAGMENT), ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(SPLIT), true)), INITIAL_EMPTY_OUTPUT_BUFFERS.withBuffer(OUT, 0).withNoMoreBufferIds()); assertEquals(taskInfo.getState(), TaskState.RUNNING); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getState(), TaskState.RUNNING); sqlTaskManager.abortTaskResults(taskId, OUT); taskInfo = sqlTaskManager.getTaskInfo(taskId, taskInfo.getState()).get(1, TimeUnit.SECONDS); assertEquals(taskInfo.getState(), TaskState.FINISHED); taskInfo = sqlTaskManager.getTaskInfo(taskId); assertEquals(taskInfo.getState(), TaskState.FINISHED); } }
@Test(expectedExceptions = IllegalStateException.class, expectedExceptionsMessageRegExp = ".*does not contain.*\\[0]") public void testSetFinalBuffersWihtoutDeclaringUsedBuffer() { BroadcastOutputBuffer buffer = createBroadcastBuffer(createInitialEmptyOutputBuffers(BROADCAST), sizeOfPages(10)); assertFalse(buffer.isFinished()); // get a page from a buffer that doesn't exist yet ListenableFuture<BufferResult> future = buffer.get(FIRST, 0L, sizeOfPages(1)); assertFalse(future.isDone()); // add a page and set no more pages addPage(buffer, createPage(33)); buffer.setNoMorePages(); // read the page assertTrue(future.isDone()); assertBufferResultEquals(TYPES, getFuture(future, NO_WAIT), bufferResult(0, createPage(33))); // acknowledge the page and verify we are finished assertBufferResultEquals(TYPES, getBufferResult(buffer, FIRST, 1, sizeOfPages(10), NO_WAIT), emptyResults(TASK_INSTANCE_ID, 1, true)); buffer.abort(FIRST); // set final buffers to a set that does not contain the buffer, which will fail buffer.setOutputBuffers(createInitialEmptyOutputBuffers(BROADCAST).withNoMoreBufferIds()); }