Search in sources :

Example 6 with OutputBuffers

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);
}
Also used : OutputBuffers(io.prestosql.execution.buffer.OutputBuffers) OutputBuffers.createInitialEmptyOutputBuffers(io.prestosql.execution.buffer.OutputBuffers.createInitialEmptyOutputBuffers) OutputBufferId(io.prestosql.execution.buffer.OutputBuffers.OutputBufferId) AtomicReference(java.util.concurrent.atomic.AtomicReference) Test(org.testng.annotations.Test)

Example 7 with OutputBuffers

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());
}
Also used : SqlTask.createSqlTask(io.prestosql.execution.SqlTask.createSqlTask) OutputBuffers(io.prestosql.execution.buffer.OutputBuffers) OutputBuffers.createInitialEmptyOutputBuffers(io.prestosql.execution.buffer.OutputBuffers.createInitialEmptyOutputBuffers) BufferResult(io.prestosql.execution.buffer.BufferResult) DataSize(io.airlift.units.DataSize) Test(org.testng.annotations.Test)

Aggregations

OutputBuffers (io.prestosql.execution.buffer.OutputBuffers)7 OutputBuffers.createInitialEmptyOutputBuffers (io.prestosql.execution.buffer.OutputBuffers.createInitialEmptyOutputBuffers)6 OutputBufferId (io.prestosql.execution.buffer.OutputBuffers.OutputBufferId)3 Session (io.prestosql.Session)2 SqlQueryScheduler (io.prestosql.execution.scheduler.SqlQueryScheduler)2 SqlQueryScheduler.createSqlQueryScheduler (io.prestosql.execution.scheduler.SqlQueryScheduler.createSqlQueryScheduler)2 PrestoException (io.prestosql.spi.PrestoException)2 PartitioningHandle (io.prestosql.sql.planner.PartitioningHandle)2 Test (org.testng.annotations.Test)2 Preconditions.checkArgument (com.google.common.base.Preconditions.checkArgument)1 Preconditions.checkState (com.google.common.base.Preconditions.checkState)1 HashMultimap (com.google.common.collect.HashMultimap)1 ImmutableList (com.google.common.collect.ImmutableList)1 ImmutableList.toImmutableList (com.google.common.collect.ImmutableList.toImmutableList)1 ImmutableMap (com.google.common.collect.ImmutableMap)1 ImmutableMultimap (com.google.common.collect.ImmutableMultimap)1 ImmutableSet (com.google.common.collect.ImmutableSet)1 Multimap (com.google.common.collect.Multimap)1 Sets (com.google.common.collect.Sets)1 Sets.newConcurrentHashSet (com.google.common.collect.Sets.newConcurrentHashSet)1