Search in sources :

Example 1 with PrestoThriftId

use of io.prestosql.plugin.thrift.api.PrestoThriftId in project hetu-core by openlookeng.

the class ThriftTpchService method getRowsInternal.

private static PrestoThriftPageResult getRowsInternal(ConnectorPageSource pageSource, String tableName, List<String> columnNames, @Nullable PrestoThriftId nextToken) {
    // very inefficient implementation as it needs to re-generate all previous results to get the next page
    int skipPages = nextToken != null ? Ints.fromByteArray(nextToken.getId()) : 0;
    skipPages(pageSource, skipPages);
    Page page = null;
    while (!pageSource.isFinished() && page == null) {
        page = pageSource.getNextPage();
        skipPages++;
    }
    PrestoThriftId newNextToken = pageSource.isFinished() ? null : new PrestoThriftId(Ints.toByteArray(skipPages));
    return toThriftPage(page, types(tableName, columnNames), newNextToken);
}
Also used : PrestoThriftId(io.prestosql.plugin.thrift.api.PrestoThriftId) Page(io.prestosql.spi.Page)

Example 2 with PrestoThriftId

use of io.prestosql.plugin.thrift.api.PrestoThriftId in project hetu-core by openlookeng.

the class TestThriftIndexPageSource method testGetNextPageTwoConcurrentRequests.

@Test
public void testGetNextPageTwoConcurrentRequests() throws Exception {
    final int splits = 3;
    final int lookupRequestsConcurrency = 2;
    final int rowsPerSplit = 1;
    List<SettableFuture<PrestoThriftPageResult>> futures = IntStream.range(0, splits).mapToObj(i -> SettableFuture.<PrestoThriftPageResult>create()).collect(toImmutableList());
    List<CountDownLatch> signals = IntStream.range(0, splits).mapToObj(i -> new CountDownLatch(1)).collect(toImmutableList());
    TestingThriftService client = new TestingThriftService(rowsPerSplit, false, false) {

        @Override
        public ListenableFuture<PrestoThriftPageResult> getRows(PrestoThriftId splitId, List<String> columns, long maxBytes, PrestoThriftNullableToken nextToken) {
            int key = Ints.fromByteArray(splitId.getId());
            signals.get(key).countDown();
            return futures.get(key);
        }
    };
    ThriftConnectorStats stats = new ThriftConnectorStats();
    long pageSizeReceived = 0;
    ThriftIndexPageSource pageSource = new ThriftIndexPageSource((context, headers) -> client, ImmutableMap.of(), stats, new ThriftIndexHandle(new SchemaTableName("default", "table1"), TupleDomain.all()), ImmutableList.of(column("a", INTEGER)), ImmutableList.of(column("b", INTEGER)), new InMemoryRecordSet(ImmutableList.of(INTEGER), generateKeys(0, splits)), MAX_BYTES_PER_RESPONSE, lookupRequestsConcurrency);
    assertNull(pageSource.getNextPage());
    assertEquals((long) stats.getIndexPageSize().getAllTime().getTotal(), 0);
    signals.get(0).await(1, SECONDS);
    signals.get(1).await(1, SECONDS);
    signals.get(2).await(1, SECONDS);
    assertEquals(signals.get(0).getCount(), 0, "first request wasn't sent");
    assertEquals(signals.get(1).getCount(), 0, "second request wasn't sent");
    assertEquals(signals.get(2).getCount(), 1, "third request shouldn't be sent");
    // at this point first two requests were sent
    assertFalse(pageSource.isFinished());
    assertNull(pageSource.getNextPage());
    assertEquals((long) stats.getIndexPageSize().getAllTime().getTotal(), 0);
    // completing the second request
    futures.get(1).set(pageResult(20, null));
    Page page = pageSource.getNextPage();
    pageSizeReceived += page.getSizeInBytes();
    assertEquals((long) stats.getIndexPageSize().getAllTime().getTotal(), pageSizeReceived);
    assertNotNull(page);
    assertEquals(page.getPositionCount(), 1);
    assertEquals(page.getBlock(0).getInt(0, 0), 20);
    // not complete yet
    assertFalse(pageSource.isFinished());
    // once one of the requests completes the next one should be sent
    signals.get(2).await(1, SECONDS);
    assertEquals(signals.get(2).getCount(), 0, "third request wasn't sent");
    // completing the first request
    futures.get(0).set(pageResult(10, null));
    page = pageSource.getNextPage();
    assertNotNull(page);
    pageSizeReceived += page.getSizeInBytes();
    assertEquals((long) stats.getIndexPageSize().getAllTime().getTotal(), pageSizeReceived);
    assertEquals(page.getPositionCount(), 1);
    assertEquals(page.getBlock(0).getInt(0, 0), 10);
    // still not complete
    assertFalse(pageSource.isFinished());
    // completing the third request
    futures.get(2).set(pageResult(30, null));
    page = pageSource.getNextPage();
    assertNotNull(page);
    pageSizeReceived += page.getSizeInBytes();
    assertEquals((long) stats.getIndexPageSize().getAllTime().getTotal(), pageSizeReceived);
    assertEquals(page.getPositionCount(), 1);
    assertEquals(page.getBlock(0).getInt(0, 0), 30);
    // finished now
    assertTrue(pageSource.isFinished());
    // after completion
    assertNull(pageSource.getNextPage());
    pageSource.close();
}
Also used : SettableFuture(com.google.common.util.concurrent.SettableFuture) IntStream(java.util.stream.IntStream) Collections.shuffle(java.util.Collections.shuffle) ListenableFuture(com.google.common.util.concurrent.ListenableFuture) Assert.assertNull(org.testng.Assert.assertNull) PrestoThriftNullableTableMetadata(io.prestosql.plugin.thrift.api.PrestoThriftNullableTableMetadata) Assert.assertEquals(org.testng.Assert.assertEquals) Test(org.testng.annotations.Test) CompletableFuture(java.util.concurrent.CompletableFuture) INTEGER(io.prestosql.spi.type.IntegerType.INTEGER) SettableFuture(com.google.common.util.concurrent.SettableFuture) PrestoThriftNullableSchemaName(io.prestosql.plugin.thrift.api.PrestoThriftNullableSchemaName) PrestoThriftSchemaTableName(io.prestosql.plugin.thrift.api.PrestoThriftSchemaTableName) ArrayList(java.util.ArrayList) PrestoThriftPageResult(io.prestosql.plugin.thrift.api.PrestoThriftPageResult) InMemoryRecordSet(io.prestosql.spi.connector.InMemoryRecordSet) SchemaTableName(io.prestosql.spi.connector.SchemaTableName) ImmutableList(com.google.common.collect.ImmutableList) Type(io.prestosql.spi.type.Type) PrestoThriftService(io.prestosql.plugin.thrift.api.PrestoThriftService) Assert.assertFalse(org.testng.Assert.assertFalse) Block(io.prestosql.spi.block.Block) PrestoThriftBlock.integerData(io.prestosql.plugin.thrift.api.PrestoThriftBlock.integerData) PrestoThriftTupleDomain(io.prestosql.plugin.thrift.api.PrestoThriftTupleDomain) PrestoThriftServiceException(io.prestosql.plugin.thrift.api.PrestoThriftServiceException) Futures.immediateFuture(com.google.common.util.concurrent.Futures.immediateFuture) PrestoThriftInteger(io.prestosql.plugin.thrift.api.datatypes.PrestoThriftInteger) ImmutableMap(com.google.common.collect.ImmutableMap) PrestoThriftId(io.prestosql.plugin.thrift.api.PrestoThriftId) TupleDomain(io.prestosql.spi.predicate.TupleDomain) ImmutableList.toImmutableList(com.google.common.collect.ImmutableList.toImmutableList) Page(io.prestosql.spi.Page) PrestoThriftSplitBatch(io.prestosql.plugin.thrift.api.PrestoThriftSplitBatch) Assert.assertNotNull(org.testng.Assert.assertNotNull) Ints(com.google.common.primitives.Ints) PrestoThriftSplit(io.prestosql.plugin.thrift.api.PrestoThriftSplit) PrestoThriftNullableToken(io.prestosql.plugin.thrift.api.PrestoThriftNullableToken) PrestoThriftNullableColumnSet(io.prestosql.plugin.thrift.api.PrestoThriftNullableColumnSet) CountDownLatch(java.util.concurrent.CountDownLatch) List(java.util.List) Assert.assertTrue(org.testng.Assert.assertTrue) Collections(java.util.Collections) SECONDS(java.util.concurrent.TimeUnit.SECONDS) PrestoThriftId(io.prestosql.plugin.thrift.api.PrestoThriftId) PrestoThriftNullableToken(io.prestosql.plugin.thrift.api.PrestoThriftNullableToken) PrestoThriftPageResult(io.prestosql.plugin.thrift.api.PrestoThriftPageResult) Page(io.prestosql.spi.Page) CountDownLatch(java.util.concurrent.CountDownLatch) PrestoThriftSchemaTableName(io.prestosql.plugin.thrift.api.PrestoThriftSchemaTableName) SchemaTableName(io.prestosql.spi.connector.SchemaTableName) InMemoryRecordSet(io.prestosql.spi.connector.InMemoryRecordSet) ArrayList(java.util.ArrayList) ImmutableList(com.google.common.collect.ImmutableList) ImmutableList.toImmutableList(com.google.common.collect.ImmutableList.toImmutableList) List(java.util.List) Test(org.testng.annotations.Test)

Example 3 with PrestoThriftId

use of io.prestosql.plugin.thrift.api.PrestoThriftId in project hetu-core by openlookeng.

the class ThriftIndexedTpchService method getIndexSplitsSync.

@Override
protected PrestoThriftSplitBatch getIndexSplitsSync(PrestoThriftSchemaTableName schemaTableName, List<String> indexColumnNames, PrestoThriftPageResult keys, int maxSplitCount, PrestoThriftNullableToken nextToken) throws PrestoThriftServiceException {
    checkArgument(NUMBER_OF_INDEX_SPLITS <= maxSplitCount, "maxSplitCount for lookup splits is too low");
    checkArgument(nextToken.getToken() == null, "no continuation is supported for lookup splits");
    int totalKeys = keys.getRowCount();
    int partSize = totalKeys / NUMBER_OF_INDEX_SPLITS;
    List<PrestoThriftSplit> splits = new ArrayList<>(NUMBER_OF_INDEX_SPLITS);
    for (int splitIndex = 0; splitIndex < NUMBER_OF_INDEX_SPLITS; splitIndex++) {
        int begin = partSize * splitIndex;
        int end = partSize * (splitIndex + 1);
        if (splitIndex + 1 == NUMBER_OF_INDEX_SPLITS) {
            // add remainder to the last split
            end = totalKeys;
        }
        if (begin == end) {
            // split is empty, skip it
            continue;
        }
        SplitInfo splitInfo = indexSplit(schemaTableName.getSchemaName(), schemaTableName.getTableName(), indexColumnNames, thriftPageToList(keys, begin, end));
        splits.add(new PrestoThriftSplit(new PrestoThriftId(SPLIT_INFO_CODEC.toJsonBytes(splitInfo)), ImmutableList.of()));
    }
    return new PrestoThriftSplitBatch(splits, null);
}
Also used : PrestoThriftSplitBatch(io.prestosql.plugin.thrift.api.PrestoThriftSplitBatch) PrestoThriftId(io.prestosql.plugin.thrift.api.PrestoThriftId) ArrayList(java.util.ArrayList) PrestoThriftSplit(io.prestosql.plugin.thrift.api.PrestoThriftSplit)

Example 4 with PrestoThriftId

use of io.prestosql.plugin.thrift.api.PrestoThriftId in project hetu-core by openlookeng.

the class ThriftTpchService method getSplitsSync.

private static PrestoThriftSplitBatch getSplitsSync(PrestoThriftSchemaTableName schemaTableName, int maxSplitCount, PrestoThriftNullableToken nextToken) {
    int totalParts = DEFAULT_NUMBER_OF_SPLITS;
    // last sent part
    int partNumber = nextToken.getToken() == null ? 0 : Ints.fromByteArray(nextToken.getToken().getId());
    int numberOfSplits = min(maxSplitCount, totalParts - partNumber);
    List<PrestoThriftSplit> splits = new ArrayList<>(numberOfSplits);
    for (int i = 0; i < numberOfSplits; i++) {
        SplitInfo splitInfo = normalSplit(schemaTableName.getSchemaName(), schemaTableName.getTableName(), partNumber + 1, totalParts);
        splits.add(new PrestoThriftSplit(new PrestoThriftId(SPLIT_INFO_CODEC.toJsonBytes(splitInfo)), ImmutableList.of()));
        partNumber++;
    }
    PrestoThriftId newNextToken = partNumber < totalParts ? new PrestoThriftId(Ints.toByteArray(partNumber)) : null;
    return new PrestoThriftSplitBatch(splits, newNextToken);
}
Also used : PrestoThriftSplitBatch(io.prestosql.plugin.thrift.api.PrestoThriftSplitBatch) PrestoThriftId(io.prestosql.plugin.thrift.api.PrestoThriftId) ArrayList(java.util.ArrayList) PrestoThriftSplit(io.prestosql.plugin.thrift.api.PrestoThriftSplit)

Aggregations

PrestoThriftId (io.prestosql.plugin.thrift.api.PrestoThriftId)4 PrestoThriftSplit (io.prestosql.plugin.thrift.api.PrestoThriftSplit)3 PrestoThriftSplitBatch (io.prestosql.plugin.thrift.api.PrestoThriftSplitBatch)3 ArrayList (java.util.ArrayList)3 Page (io.prestosql.spi.Page)2 ImmutableList (com.google.common.collect.ImmutableList)1 ImmutableList.toImmutableList (com.google.common.collect.ImmutableList.toImmutableList)1 ImmutableMap (com.google.common.collect.ImmutableMap)1 Ints (com.google.common.primitives.Ints)1 Futures.immediateFuture (com.google.common.util.concurrent.Futures.immediateFuture)1 ListenableFuture (com.google.common.util.concurrent.ListenableFuture)1 SettableFuture (com.google.common.util.concurrent.SettableFuture)1 PrestoThriftBlock.integerData (io.prestosql.plugin.thrift.api.PrestoThriftBlock.integerData)1 PrestoThriftNullableColumnSet (io.prestosql.plugin.thrift.api.PrestoThriftNullableColumnSet)1 PrestoThriftNullableSchemaName (io.prestosql.plugin.thrift.api.PrestoThriftNullableSchemaName)1 PrestoThriftNullableTableMetadata (io.prestosql.plugin.thrift.api.PrestoThriftNullableTableMetadata)1 PrestoThriftNullableToken (io.prestosql.plugin.thrift.api.PrestoThriftNullableToken)1 PrestoThriftPageResult (io.prestosql.plugin.thrift.api.PrestoThriftPageResult)1 PrestoThriftSchemaTableName (io.prestosql.plugin.thrift.api.PrestoThriftSchemaTableName)1 PrestoThriftService (io.prestosql.plugin.thrift.api.PrestoThriftService)1