Examples with OutputBuffers com.facebook.presto.OutputBuffers used on opensource projects

Example 6 with OutputBuffers

use of com.facebook.presto.OutputBuffers in project presto by prestodb.

the class TestArbitraryOutputBuffer method testSimple.

@Test
public void testSimple() throws Exception {
    OutputBuffers outputBuffers = createInitialEmptyOutputBuffers(ARBITRARY);
    ArbitraryOutputBuffer buffer = createArbitraryBuffer(outputBuffers, sizeOfPages(10));
    // add three items
    for (int i = 0; i < 3; i++) {
        addPage(buffer, createPage(i));
    }
    outputBuffers = createInitialEmptyOutputBuffers(ARBITRARY).withBuffer(FIRST, BROADCAST_PARTITION_ID);
    // add a queue
    buffer.setOutputBuffers(outputBuffers);
    assertQueueState(buffer, 3, FIRST, 0, 0);
    // get the three elements
    assertBufferResultEquals(TYPES, getBufferResult(buffer, FIRST, 0, sizeOfPages(10), NO_WAIT), bufferResult(0, createPage(0), createPage(1), createPage(2)));
    // pages not acknowledged yet so state is the same
    assertQueueState(buffer, 0, FIRST, 3, 0);
    // acknowledge first three pages
    buffer.get(FIRST, 3, sizeOfPages(1)).cancel(true);
    // pages now acknowledged
    assertQueueState(buffer, 0, FIRST, 0, 3);
    // fill the buffer, so that it has 10 buffered pages
    for (int i = 3; i < 13; i++) {
        addPage(buffer, createPage(i));
    }
    // there is a pending read from above so one page will be assigned to the first buffer
    assertQueueState(buffer, 9, FIRST, 1, 3);
    // try to add one more page, which should block
    ListenableFuture<?> future = enqueuePage(buffer, createPage(13));
    assertFalse(future.isDone());
    assertQueueState(buffer, 10, FIRST, 1, 3);
    // remove a page
    assertBufferResultEquals(TYPES, getBufferResult(buffer, FIRST, 3, sizeOfPages(1), NO_WAIT), bufferResult(3, createPage(3)));
    // page not acknowledged yet so sent count is the same
    assertQueueState(buffer, 10, FIRST, 1, 3);
    // we should still be blocked
    assertFalse(future.isDone());
    //
    // add another buffer and verify it sees buffered pages
    outputBuffers = outputBuffers.withBuffer(SECOND, BROADCAST_PARTITION_ID);
    buffer.setOutputBuffers(outputBuffers);
    assertQueueState(buffer, 10, SECOND, 0, 0);
    assertBufferResultEquals(TYPES, getBufferResult(buffer, SECOND, 0, sizeOfPages(10), NO_WAIT), bufferResult(0, createPage(4), createPage(5), createPage(6), createPage(7), createPage(8), createPage(9), createPage(10), createPage(11), createPage(12), createPage(13)));
    // page not acknowledged yet so sent count is still zero
    assertQueueState(buffer, 0, SECOND, 10, 0);
    // acknowledge the 10 pages
    buffer.get(SECOND, 10, sizeOfPages(10)).cancel(true);
    assertQueueState(buffer, 0, SECOND, 0, 10);
    //
    // tell shared buffer there will be no more queues
    outputBuffers = outputBuffers.withNoMoreBufferIds();
    buffer.setOutputBuffers(outputBuffers);
    // buffers should see the same stats and the blocked page future from above should be done
    assertQueueState(buffer, 0, FIRST, 1, 3);
    assertQueueState(buffer, 0, SECOND, 0, 10);
    assertFutureIsDone(future);
    // add 3 more pages, buffers always show the same stats
    addPage(buffer, createPage(14));
    addPage(buffer, createPage(15));
    addPage(buffer, createPage(16));
    assertQueueState(buffer, 2, FIRST, 1, 3);
    assertQueueState(buffer, 2, SECOND, 1, 10);
    // pull one page from the second buffer
    assertBufferResultEquals(TYPES, getBufferResult(buffer, SECOND, 10, sizeOfPages(1), NO_WAIT), bufferResult(10, createPage(14)));
    assertQueueState(buffer, 2, FIRST, 1, 3);
    assertQueueState(buffer, 2, SECOND, 1, 10);
    // acknowledge the page in the first buffer and pull remaining ones
    assertBufferResultEquals(TYPES, getBufferResult(buffer, FIRST, 4, sizeOfPages(10), NO_WAIT), bufferResult(4, createPage(15), createPage(16)));
    assertQueueState(buffer, 0, FIRST, 2, 4);
    assertQueueState(buffer, 0, SECOND, 1, 10);
    //
    // finish the buffer
    assertFalse(buffer.isFinished());
    buffer.setNoMorePages();
    assertQueueState(buffer, 0, FIRST, 2, 4);
    assertQueueState(buffer, 0, SECOND, 1, 10);
    // not fully finished until all pages are consumed
    assertFalse(buffer.isFinished());
    // acknowledge the pages from the first buffer; buffer should not close automatically
    assertBufferResultEquals(TYPES, getBufferResult(buffer, FIRST, 6, sizeOfPages(10), NO_WAIT), emptyResults(TASK_INSTANCE_ID, 6, true));
    assertQueueState(buffer, 0, FIRST, 0, 6);
    assertQueueState(buffer, 0, SECOND, 1, 10);
    assertFalse(buffer.isFinished());
    // finish first queue
    buffer.abort(FIRST);
    assertQueueClosed(buffer, 0, FIRST, 6);
    assertQueueState(buffer, 0, SECOND, 1, 10);
    assertFalse(buffer.isFinished());
    // acknowledge a page from the second queue; queue should not close automatically
    assertBufferResultEquals(TYPES, getBufferResult(buffer, SECOND, 11, sizeOfPages(1), NO_WAIT), emptyResults(TASK_INSTANCE_ID, 11, true));
    assertQueueState(buffer, 0, SECOND, 0, 11);
    assertFalse(buffer.isFinished());
    // finish second queue
    buffer.abort(SECOND);
    assertQueueClosed(buffer, 0, FIRST, 6);
    assertQueueClosed(buffer, 0, SECOND, 11);
    assertFinished(buffer);
    assertBufferResultEquals(TYPES, getBufferResult(buffer, FIRST, 6, sizeOfPages(10), NO_WAIT), emptyResults(TASK_INSTANCE_ID, 6, true));
    assertBufferResultEquals(TYPES, getBufferResult(buffer, SECOND, 11, sizeOfPages(10), NO_WAIT), emptyResults(TASK_INSTANCE_ID, 11, true));
}

Example 7 with OutputBuffers

use of com.facebook.presto.OutputBuffers in project presto by prestodb.

the class TestBroadcastOutputBuffer method testAddQueueAfterNoMoreQueues.

@Test
public void testAddQueueAfterNoMoreQueues() throws Exception {
    BroadcastOutputBuffer buffer = createBroadcastBuffer(createInitialEmptyOutputBuffers(BROADCAST), sizeOfPages(10));
    assertFalse(buffer.isFinished());
    // tell buffer no more queues will be added
    buffer.setOutputBuffers(createInitialEmptyOutputBuffers(BROADCAST).withNoMoreBufferIds());
    assertFalse(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());
    assertFalse(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());
    assertFalse(buffer.isFinished());
    try {
        OutputBuffers outputBuffers = createInitialEmptyOutputBuffers(BROADCAST).withBuffer(FIRST, BROADCAST_PARTITION_ID).withNoMoreBufferIds();
        buffer.setOutputBuffers(outputBuffers);
        fail("Expected IllegalStateException from addQueue after noMoreQueues has been called");
    } catch (IllegalArgumentException ignored) {
    }
}

Example 8 with OutputBuffers

use of com.facebook.presto.OutputBuffers in project presto by prestodb.

the class TestBroadcastOutputBuffer method testSimple.

@Test
public void testSimple() throws Exception {
    OutputBuffers outputBuffers = createInitialEmptyOutputBuffers(BROADCAST);
    BroadcastOutputBuffer buffer = createBroadcastBuffer(outputBuffers, sizeOfPages(10));
    // add three items
    for (int i = 0; i < 3; i++) {
        addPage(buffer, createPage(i));
    }
    outputBuffers = createInitialEmptyOutputBuffers(BROADCAST).withBuffer(FIRST, BROADCAST_PARTITION_ID);
    // add a queue
    buffer.setOutputBuffers(outputBuffers);
    assertQueueState(buffer, FIRST, 3, 0);
    // get the three elements
    assertBufferResultEquals(TYPES, getBufferResult(buffer, FIRST, 0, sizeOfPages(10), NO_WAIT), bufferResult(0, createPage(0), createPage(1), createPage(2)));
    // pages not acknowledged yet so state is the same
    assertQueueState(buffer, FIRST, 3, 0);
    // acknowledge first three pages
    buffer.get(FIRST, 3, sizeOfPages(10)).cancel(true);
    // pages now acknowledged
    assertQueueState(buffer, FIRST, 0, 3);
    // fill the buffer (we already added 3 pages)
    for (int i = 3; i < 10; i++) {
        addPage(buffer, createPage(i));
    }
    assertQueueState(buffer, FIRST, 7, 3);
    // try to add one more page, which should block
    ListenableFuture<?> future = enqueuePage(buffer, createPage(10));
    assertFalse(future.isDone());
    assertQueueState(buffer, FIRST, 8, 3);
    // remove a page
    assertBufferResultEquals(TYPES, getBufferResult(buffer, FIRST, 3, sizeOfPages(1), NO_WAIT), bufferResult(3, createPage(3)));
    // page not acknowledged yet so sent count is the same
    assertQueueState(buffer, FIRST, 8, 3);
    // we should still be blocked
    assertFalse(future.isDone());
    //
    // add another buffer and verify it sees all pages
    outputBuffers = outputBuffers.withBuffer(SECOND, BROADCAST_PARTITION_ID);
    buffer.setOutputBuffers(outputBuffers);
    assertQueueState(buffer, SECOND, 11, 0);
    assertBufferResultEquals(TYPES, getBufferResult(buffer, SECOND, 0, sizeOfPages(10), NO_WAIT), bufferResult(0, createPage(0), createPage(1), createPage(2), createPage(3), createPage(4), createPage(5), createPage(6), createPage(7), createPage(8), createPage(9)));
    // page not acknowledged yet so sent count is still zero
    assertQueueState(buffer, SECOND, 11, 0);
    // acknowledge the 10 pages
    buffer.get(SECOND, 10, sizeOfPages(10)).cancel(true);
    assertQueueState(buffer, SECOND, 1, 10);
    //
    // tell shared buffer there will be no more queues
    outputBuffers = outputBuffers.withNoMoreBufferIds();
    buffer.setOutputBuffers(outputBuffers);
    // queues consumed the first three pages, so they should be dropped now and the blocked page future from above should be done
    assertQueueState(buffer, FIRST, 8, 3);
    assertQueueState(buffer, SECOND, 1, 10);
    assertFutureIsDone(future);
    // we should be able to add 3 more pages (the third will be queued)
    // although the first queue fetched the 4th page, the page has not been acknowledged yet
    addPage(buffer, createPage(11));
    addPage(buffer, createPage(12));
    future = enqueuePage(buffer, createPage(13));
    assertFalse(future.isDone());
    assertQueueState(buffer, FIRST, 11, 3);
    assertQueueState(buffer, SECOND, 4, 10);
    // acknowledge the receipt of the 3rd page and try to remove the 4th page from the first queue
    assertBufferResultEquals(TYPES, getBufferResult(buffer, FIRST, 4, sizeOfPages(1), NO_WAIT), bufferResult(4, createPage(4)));
    // the blocked page future above should be done
    assertFutureIsDone(future);
    assertQueueState(buffer, FIRST, 10, 4);
    assertQueueState(buffer, SECOND, 4, 10);
    //
    // finish the buffer
    assertFalse(buffer.isFinished());
    buffer.setNoMorePages();
    assertQueueState(buffer, FIRST, 10, 4);
    assertQueueState(buffer, SECOND, 4, 10);
    // not fully finished until all pages are consumed
    assertFalse(buffer.isFinished());
    // remove a page, not finished
    assertBufferResultEquals(TYPES, getBufferResult(buffer, FIRST, 5, sizeOfPages(1), NO_WAIT), bufferResult(5, createPage(5)));
    assertQueueState(buffer, FIRST, 9, 5);
    assertQueueState(buffer, SECOND, 4, 10);
    assertFalse(buffer.isFinished());
    // remove all remaining pages from first queue, should not be finished
    BufferResult x = getBufferResult(buffer, FIRST, 6, sizeOfPages(10), NO_WAIT);
    assertBufferResultEquals(TYPES, x, bufferResult(6, createPage(6), createPage(7), createPage(8), createPage(9), createPage(10), createPage(11), createPage(12), createPage(13)));
    assertQueueState(buffer, FIRST, 8, 6);
    assertBufferResultEquals(TYPES, getBufferResult(buffer, FIRST, 14, sizeOfPages(10), NO_WAIT), emptyResults(TASK_INSTANCE_ID, 14, true));
    // finish first queue
    buffer.abort(FIRST);
    assertQueueClosed(buffer, FIRST, 14);
    assertQueueState(buffer, SECOND, 4, 10);
    assertFalse(buffer.isFinished());
    // remove all remaining pages from second queue, should be finished
    assertBufferResultEquals(TYPES, getBufferResult(buffer, SECOND, 10, sizeOfPages(10), NO_WAIT), bufferResult(10, createPage(10), createPage(11), createPage(12), createPage(13)));
    assertQueueState(buffer, SECOND, 4, 10);
    assertBufferResultEquals(TYPES, getBufferResult(buffer, SECOND, 14, sizeOfPages(10), NO_WAIT), emptyResults(TASK_INSTANCE_ID, 14, true));
    buffer.abort(SECOND);
    assertQueueClosed(buffer, FIRST, 14);
    assertQueueClosed(buffer, SECOND, 14);
    assertFinished(buffer);
    assertBufferResultEquals(TYPES, getBufferResult(buffer, FIRST, 14, sizeOfPages(10), NO_WAIT), emptyResults(TASK_INSTANCE_ID, 14, true));
    assertBufferResultEquals(TYPES, getBufferResult(buffer, SECOND, 14, sizeOfPages(10), NO_WAIT), emptyResults(TASK_INSTANCE_ID, 14, true));
}

Example 9 with OutputBuffers

use of com.facebook.presto.OutputBuffers in project presto by prestodb.

the class TestBroadcastOutputBufferManager method test.

@Test
public void test() throws Exception {
    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);
}

Example 10 with OutputBuffers

use of com.facebook.presto.OutputBuffers in project presto by prestodb.

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());
}