Examples with ServerResponse org.apache.pinot.common.proto.Server.ServerResponse used on opensource projects

Example 1 with ServerResponse

use of org.apache.pinot.common.proto.Server.ServerResponse in project presto by prestodb.

the class PinotSegmentStreamingPageSource method getNextPage.

/**
 * @return constructed page for pinot data.
 */
@Override
public Page getNextPage() {
    if (closed) {
        return null;
    }
    if (serverResponseIterator == null) {
        serverResponseIterator = queryPinot(split);
    }
    ByteBuffer byteBuffer = null;
    try {
        // So we need to check ResponseType of each ServerResponse.
        if (serverResponseIterator.hasNext()) {
            long startTimeNanos = System.nanoTime();
            ServerResponse serverResponse = serverResponseIterator.next();
            readTimeNanos += System.nanoTime() - startTimeNanos;
            final String responseType = serverResponse.getMetadataOrThrow("responseType");
            switch(responseType) {
                case CommonConstants.Query.Response.ResponseType.DATA:
                    estimatedMemoryUsageInBytes = serverResponse.getSerializedSize();
                    // Store each dataTable which will later be constructed into Pages.
                    try {
                        byteBuffer = serverResponse.getPayload().asReadOnlyByteBuffer();
                        DataTable dataTable = DataTableFactory.getDataTable(byteBuffer);
                        checkExceptions(dataTable, split, PinotSessionProperties.isMarkDataFetchExceptionsAsRetriable(session));
                        currentDataTable = new PinotSegmentPageSource.PinotDataTableWithSize(dataTable, serverResponse.getSerializedSize());
                    } catch (IOException e) {
                        throw new PinotException(PINOT_DATA_FETCH_EXCEPTION, split.getSegmentPinotQuery(), String.format("Encountered Pinot exceptions when fetching data table from Split: < %s >", split), e);
                    }
                    break;
                case CommonConstants.Query.Response.ResponseType.METADATA:
                    // The last part of the response is Metadata
                    currentDataTable = null;
                    serverResponseIterator = null;
                    close();
                    return null;
                default:
                    throw new PinotException(PINOT_UNEXPECTED_RESPONSE, split.getSegmentPinotQuery(), String.format("Encountered Pinot exceptions, unknown response type - %s", responseType));
            }
        }
        Page page = fillNextPage();
        completedPositions += currentDataTable.getDataTable().getNumberOfRows();
        return page;
    } finally {
        if (byteBuffer != null) {
            byteBuffer.clear();
        }
    }
}