Examples with TrinoThriftPageResult io.trino.plugin.thrift.api.TrinoThriftPageResult used on opensource projects

Example 1 with TrinoThriftPageResult

use of io.trino.plugin.thrift.api.TrinoThriftPageResult in project trino by trinodb.

the class ThriftIndexPageSource method getNextPage.

@Override
public Page getNextPage() {
    if (finished) {
        return null;
    }
    if (!loadAllSplits()) {
        return null;
    }
    // check if any data requests were started
    if (dataSignalFuture == null) {
        // no data requests were started, start a number of initial requests
        checkState(contexts.isEmpty() && dataRequests.isEmpty(), "some splits are already started");
        if (splits.isEmpty()) {
            // all done: no splits
            finished = true;
            return null;
        }
        for (int i = 0; i < min(lookupRequestsConcurrency, splits.size()); i++) {
            startDataFetchForNextSplit();
        }
        updateSignalAndStatusFutures();
    }
    // check if any data request is finished
    if (!dataSignalFuture.isDone()) {
        // not finished yet
        return null;
    }
    // at least one of data requests completed
    ListenableFuture<TrinoThriftPageResult> resultFuture = getAndRemoveNextCompletedRequest();
    RunningSplitContext resultContext = contexts.remove(resultFuture);
    checkState(resultContext != null, "no associated context for the request");
    TrinoThriftPageResult pageResult = getFutureValue(resultFuture);
    Page page = pageResult.toPage(outputColumnTypes);
    if (page != null) {
        long pageSize = page.getSizeInBytes();
        completedBytes += pageSize;
        stats.addIndexPageSize(pageSize);
    } else {
        stats.addIndexPageSize(0);
    }
    if (pageResult.getNextToken() != null) {
        // can get more data
        sendDataRequest(resultContext, pageResult.getNextToken());
        updateSignalAndStatusFutures();
        return page;
    }
    // are there more splits available
    if (splitIndex < splits.size()) {
        // can send data request for a new split
        startDataFetchForNextSplit();
        updateSignalAndStatusFutures();
    } else if (!dataRequests.isEmpty()) {
        // no more new splits, but some requests are still in progress, wait for them
        updateSignalAndStatusFutures();
    } else {
        // all done: no more new splits, no requests in progress
        dataSignalFuture = null;
        statusFuture = null;
        finished = true;
    }
    return page;
}

Example 2 with TrinoThriftPageResult

use of io.trino.plugin.thrift.api.TrinoThriftPageResult in project trino by trinodb.

the class ThriftIndexPageSource method sendDataRequest.

private void sendDataRequest(RunningSplitContext context, @Nullable TrinoThriftId nextToken) {
    long start = System.nanoTime();
    ListenableFuture<TrinoThriftPageResult> future = context.getClient().getRows(context.getSplit().getSplitId(), outputColumnNames, maxBytesPerResponse, new TrinoThriftNullableToken(nextToken));
    future = catchingThriftException(future);
    future.addListener(() -> readTimeNanos.addAndGet(System.nanoTime() - start), directExecutor());
    dataRequests.add(future);
    contexts.put(future, context);
}

Example 3 with TrinoThriftPageResult

use of io.trino.plugin.thrift.api.TrinoThriftPageResult in project trino by trinodb.

the class ThriftPageSource method sendDataRequestInternal.

private CompletableFuture<TrinoThriftPageResult> sendDataRequestInternal() {
    long start = System.nanoTime();
    ListenableFuture<TrinoThriftPageResult> rowsBatchFuture = client.getRows(splitId, columnNames, maxBytesPerResponse, new TrinoThriftNullableToken(nextToken));
    rowsBatchFuture = catchingThriftException(rowsBatchFuture);
    rowsBatchFuture.addListener(() -> readTimeNanos.addAndGet(System.nanoTime() - start), directExecutor());
    return toCompletableFuture(nonCancellationPropagating(rowsBatchFuture));
}

Example 4 with TrinoThriftPageResult

use of io.trino.plugin.thrift.api.TrinoThriftPageResult in project trino by trinodb.

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<TrinoThriftPageResult>> futures = IntStream.range(0, splits).mapToObj(i -> SettableFuture.<TrinoThriftPageResult>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<TrinoThriftPageResult> getRows(TrinoThriftId splitId, List<String> columns, long maxBytes, TrinoThriftNullableToken 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();
}