Examples with DataCenterQueryResults io.prestosql.client.DataCenterQueryResults used on opensource projects

Example 1 with DataCenterQueryResults

use of io.prestosql.client.DataCenterQueryResults in project hetu-core by openlookeng.

the class Query method getNextResult.

private synchronized DataCenterQueryResults getNextResult(long token, DataSize targetResultSize) {
    // check if the result for the token have already been created
    Optional<DataCenterQueryResults> cachedResult = getCachedDCQueryResult(token);
    if (cachedResult.isPresent()) {
        return cachedResult.get();
    }
    verify(nextToken.isPresent(), "Can not generate next result when next token is not present");
    verify(token == nextToken.getAsLong(), "Expected token to equal next token");
    URI queryHtmlUri = null;
    try {
        queryHtmlUri = new URI("http://localhost");
    } catch (URISyntaxException e) {
        log.error("get uri error: %s", e.getMessage());
    }
    // Remove as many pages as possible from the exchange until just greater than DESIRED_RESULT_BYTES
    // NOTE: it is critical that query results are created for the pages removed from the exchange
    // client while holding the lock because the query may transition to the finished state when the
    // last page is removed.  If another thread observes this state before the response is cached
    // the pages will be lost.
    List<SerializedPage> data = null;
    try {
        ImmutableList.Builder<SerializedPage> builder = new ImmutableList.Builder<>();
        long bytes = 0;
        long rows = 0;
        long targetResultBytes = targetResultSize.toBytes();
        while (bytes < targetResultBytes) {
            // at this point, origin is irrelevant, so we can safely ignore it
            SerializedPage serializedPage = exchangeClient.pollPage(null).getLeft();
            if (serializedPage == null) {
                break;
            }
            builder.add(new SerializedPage(serializedPage.getSlice().getBytes(), serializedPage.getPageCodecMarkers(), serializedPage.getPositionCount(), serializedPage.getUncompressedSizeInBytes()));
            bytes += serializedPage.getUncompressedSizeInBytes();
            rows += serializedPage.getPositionCount();
        }
        if (rows > 0) {
            // client implementations do not properly handle empty list of data
            data = builder.build();
        }
    } catch (Throwable cause) {
        queryManager.failQuery(queryId, cause);
    }
    // get the query info before returning
    // force update if query manager is closed
    QueryInfo queryInfo = queryManager.getFullQueryInfo(queryId);
    queryManager.recordHeartbeat(queryId);
    closeExchangeClientIfNecessary(queryInfo);
    // for queries with no output, return a fake result for clients that require it
    if ((queryInfo.getState() == QueryState.FINISHED) && !queryInfo.getOutputStage().isPresent()) {
        columns = ImmutableList.of(createColumn("result", BooleanType.BOOLEAN));
        PageBuilder pageBuilder = new PageBuilder(Collections.singletonList(BooleanType.BOOLEAN));
        pageBuilder.declarePosition();
        BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0);
        BooleanType.BOOLEAN.writeBoolean(blockBuilder, true);
        Page page = pageBuilder.build();
        SerializedPage serializedPage = serde.serialize(page);
        data = Collections.singletonList(new SerializedPage(serializedPage.getSlice().getBytes(), serializedPage.getPageCodecMarkers(), serializedPage.getPositionCount(), serializedPage.getUncompressedSizeInBytes()));
    } else if (queryInfo.isRunningAsync()) {
        columns = ImmutableList.of(createColumn("result", BooleanType.BOOLEAN), createColumn("runningAsync", BooleanType.BOOLEAN));
        PageBuilder pageBuilder = new PageBuilder(ImmutableList.of(BooleanType.BOOLEAN, BooleanType.BOOLEAN));
        pageBuilder.declarePosition();
        BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(0);
        BooleanType.BOOLEAN.writeBoolean(blockBuilder, true);
        BooleanType.BOOLEAN.writeBoolean(blockBuilder, true);
        Page page = pageBuilder.build();
        SerializedPage serializedPage = serde.serialize(page);
        data = Collections.singletonList(new SerializedPage(serializedPage.getSlice().getBytes(), serializedPage.getPageCodecMarkers(), serializedPage.getPositionCount(), serializedPage.getUncompressedSizeInBytes()));
    }
    // (3) Query supports and started running Async
    if ((!queryInfo.isFinalQueryInfo() && queryInfo.getState() != FAILED && !queryInfo.isRunningAsync()) || !exchangeClient.isClosed()) {
        nextToken = OptionalLong.of(token + 1);
    } else {
        nextToken = OptionalLong.empty();
    }
    URI nextResultsUri = null;
    if (nextToken.isPresent()) {
        try {
            nextResultsUri = new URI(Long.toString(nextToken.getAsLong()));
        } catch (URISyntaxException e) {
            log.error("get uri error: %s", e.getMessage());
        }
    }
    // update catalog, schema, and path
    setCatalog = queryInfo.getSetCatalog();
    setSchema = queryInfo.getSetSchema();
    setPath = queryInfo.getSetPath();
    // update setSessionProperties
    setSessionProperties = queryInfo.getSetSessionProperties();
    resetSessionProperties = queryInfo.getResetSessionProperties();
    // update setRoles
    setRoles = queryInfo.getSetRoles();
    // update preparedStatements
    addedPreparedStatements = queryInfo.getAddedPreparedStatements();
    deallocatedPreparedStatements = queryInfo.getDeallocatedPreparedStatements();
    // update startedTransactionId
    startedTransactionId = queryInfo.getStartedTransactionId();
    clearTransactionId = queryInfo.isClearTransactionId();
    // first time through, self is null
    DataCenterQueryResults queryResults = new DataCenterQueryResults(queryId.toString(), queryHtmlUri, findCancelableLeafStage(queryInfo), nextResultsUri, columns, data, toStatementStats(queryInfo), toQueryError(queryInfo), mappedCopy(queryInfo.getWarnings(), Query::toClientWarning), queryInfo.getUpdateType(), true);
    // cache the new result
    lastToken = token;
    lastDCQueryResult = queryResults;
    return queryResults;
}