Examples with AsyncQueryResponse org.apache.pinot.core.transport.AsyncQueryResponse used on opensource projects

Example 1 with AsyncQueryResponse

use of org.apache.pinot.core.transport.AsyncQueryResponse in project trino by trinodb.

the class PinotQueryClient method queryPinotServerForDataTable.

public Map<ServerInstance, DataTable> queryPinotServerForDataTable(String query, String serverHost, List<String> segments, long connectionTimeoutInMillis, int pinotRetryCount) {
    // TODO: separate into offline and realtime methods
    BrokerRequest brokerRequest;
    try {
        brokerRequest = REQUEST_COMPILER.compileToBrokerRequest(query);
    } catch (SqlCompilationException e) {
        throw new PinotException(PINOT_INVALID_PQL_GENERATED, Optional.of(query), format("Parsing error with on %s, Error = %s", serverHost, e.getMessage()), e);
    }
    ServerInstance serverInstance = pinotHostMapper.getServerInstance(serverHost);
    Map<ServerInstance, List<String>> routingTable = new HashMap<>();
    routingTable.put(serverInstance, new ArrayList<>(segments));
    String tableName = brokerRequest.getQuerySource().getTableName();
    String rawTableName = TableNameBuilder.extractRawTableName(tableName);
    Map<ServerInstance, List<String>> offlineRoutingTable = TableNameBuilder.isOfflineTableResource(tableName) ? routingTable : null;
    Map<ServerInstance, List<String>> realtimeRoutingTable = TableNameBuilder.isRealtimeTableResource(tableName) ? routingTable : null;
    BrokerRequest offlineBrokerRequest = TableNameBuilder.isOfflineTableResource(tableName) ? brokerRequest : null;
    BrokerRequest realtimeBrokerRequest = TableNameBuilder.isRealtimeTableResource(tableName) ? brokerRequest : null;
    AsyncQueryResponse asyncQueryResponse = doWithRetries(pinotRetryCount, requestId -> queryRouter.submitQuery(requestId, rawTableName, offlineBrokerRequest, offlineRoutingTable, realtimeBrokerRequest, realtimeRoutingTable, connectionTimeoutInMillis));
    try {
        Map<ServerRoutingInstance, ServerResponse> response = asyncQueryResponse.getResponse();
        Map<ServerInstance, DataTable> dataTableMap = new HashMap<>();
        for (Map.Entry<ServerRoutingInstance, ServerResponse> entry : response.entrySet()) {
            ServerResponse serverResponse = entry.getValue();
            DataTable dataTable = serverResponse.getDataTable();
            dataTableMap.put(toServerInstance(entry.getKey()), dataTable);
        }
        return dataTableMap;
    } catch (InterruptedException e) {
        throw new PinotException(PINOT_EXCEPTION, Optional.of(query), "Pinot query execution was interrupted", e);
    }
}