use of io.prestosql.execution.buffer.OutputBuffers in project hetu-core by openlookeng.
the class TestBroadcastOutputBufferManager method test.
@Test
public void test() {
AtomicReference<OutputBuffers> outputBufferTarget = new AtomicReference<>();
BroadcastOutputBufferManager hashOutputBufferManager = new BroadcastOutputBufferManager(outputBufferTarget::set);
assertEquals(outputBufferTarget.get(), createInitialEmptyOutputBuffers(BROADCAST));
hashOutputBufferManager.addOutputBuffers(ImmutableList.of(new OutputBufferId(0)), false);
OutputBuffers expectedOutputBuffers = createInitialEmptyOutputBuffers(BROADCAST).withBuffer(new OutputBufferId(0), BROADCAST_PARTITION_ID);
assertEquals(outputBufferTarget.get(), expectedOutputBuffers);
hashOutputBufferManager.addOutputBuffers(ImmutableList.of(new OutputBufferId(1), new OutputBufferId(2)), false);
expectedOutputBuffers = expectedOutputBuffers.withBuffer(new OutputBufferId(1), BROADCAST_PARTITION_ID);
expectedOutputBuffers = expectedOutputBuffers.withBuffer(new OutputBufferId(2), BROADCAST_PARTITION_ID);
assertEquals(outputBufferTarget.get(), expectedOutputBuffers);
// set no more buffers
hashOutputBufferManager.addOutputBuffers(ImmutableList.of(new OutputBufferId(3)), true);
expectedOutputBuffers = expectedOutputBuffers.withBuffer(new OutputBufferId(3), BROADCAST_PARTITION_ID);
expectedOutputBuffers = expectedOutputBuffers.withNoMoreBufferIds();
assertEquals(outputBufferTarget.get(), expectedOutputBuffers);
// try to add another buffer, which should not result in an error
// and output buffers should not change
hashOutputBufferManager.addOutputBuffers(ImmutableList.of(new OutputBufferId(5)), false);
assertEquals(outputBufferTarget.get(), expectedOutputBuffers);
// try to set no more buffers again, which should not result in an error
// and output buffers should not change
hashOutputBufferManager.addOutputBuffers(ImmutableList.of(new OutputBufferId(6)), true);
assertEquals(outputBufferTarget.get(), expectedOutputBuffers);
}
use of io.prestosql.execution.buffer.OutputBuffers in project hetu-core by openlookeng.
the class TestSqlTask method testBufferCloseOnFinish.
@Test
public void testBufferCloseOnFinish() throws Exception {
SqlTask sqlTask = createInitialTask();
OutputBuffers outputBuffers = createInitialEmptyOutputBuffers(PARTITIONED).withBuffer(OUT, 0).withNoMoreBufferIds();
updateTask(sqlTask, EMPTY_SOURCES, outputBuffers);
ListenableFuture<BufferResult> bufferResult = sqlTask.getTaskResults(OUT, 0, new DataSize(1, MEGABYTE));
assertFalse(bufferResult.isDone());
// close the sources (no splits will ever be added)
updateTask(sqlTask, ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(), true)), outputBuffers);
// finish the task by calling abort on it
sqlTask.abortTaskResults(OUT);
// buffer will be closed by cancel event (wait for event to fire)
bufferResult.get(1, SECONDS);
// verify the buffer is closed
bufferResult = sqlTask.getTaskResults(OUT, 0, new DataSize(1, MEGABYTE));
assertTrue(bufferResult.isDone());
assertTrue(bufferResult.get().isBufferComplete());
}
Aggregations