use of org.apache.pinot.common.request.BrokerRequest 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);
}
}
use of org.apache.pinot.common.request.BrokerRequest in project trino by trinodb.
the class DynamicTableBuilder method buildFromPql.
public static DynamicTable buildFromPql(PinotMetadata pinotMetadata, SchemaTableName schemaTableName, PinotClient pinotClient) {
requireNonNull(pinotMetadata, "pinotMetadata is null");
requireNonNull(schemaTableName, "schemaTableName is null");
String query = schemaTableName.getTableName();
BrokerRequest request = REQUEST_COMPILER.compileToBrokerRequest(query);
PinotQuery pinotQuery = request.getPinotQuery();
QueryContext queryContext = BrokerRequestToQueryContextConverter.convert(request);
String pinotTableName = stripSuffix(request.getQuerySource().getTableName());
Optional<String> suffix = getSuffix(request.getQuerySource().getTableName());
Map<String, ColumnHandle> columnHandles = pinotMetadata.getPinotColumnHandles(pinotTableName);
List<OrderByExpression> orderBy = ImmutableList.of();
PinotTypeResolver pinotTypeResolver = new PinotTypeResolver(pinotClient, pinotTableName);
List<PinotColumnHandle> selectColumns = ImmutableList.of();
Map<String, PinotColumnNameAndTrinoType> aggregateTypes = ImmutableMap.of();
if (queryContext.getAggregationFunctions() != null) {
checkState(queryContext.getAggregationFunctions().length > 0, "Aggregation Functions is empty");
aggregateTypes = getAggregateTypes(schemaTableName, queryContext, columnHandles);
}
if (queryContext.getSelectExpressions() != null) {
checkState(!queryContext.getSelectExpressions().isEmpty(), "Pinot selections is empty");
selectColumns = getPinotColumns(schemaTableName, queryContext.getSelectExpressions(), queryContext.getAliasList(), columnHandles, pinotTypeResolver, aggregateTypes);
}
if (queryContext.getOrderByExpressions() != null) {
ImmutableList.Builder<OrderByExpression> orderByBuilder = ImmutableList.builder();
for (OrderByExpressionContext orderByExpressionContext : queryContext.getOrderByExpressions()) {
ExpressionContext expressionContext = orderByExpressionContext.getExpression();
PinotColumnHandle pinotColumnHandle = getPinotColumnHandle(schemaTableName, expressionContext, Optional.empty(), columnHandles, pinotTypeResolver, aggregateTypes);
orderByBuilder.add(new OrderByExpression(pinotColumnHandle.getExpression(), orderByExpressionContext.isAsc()));
}
orderBy = orderByBuilder.build();
}
List<PinotColumnHandle> groupByColumns = ImmutableList.of();
if (queryContext.getGroupByExpressions() != null) {
groupByColumns = getPinotColumns(schemaTableName, queryContext.getGroupByExpressions(), ImmutableList.of(), columnHandles, pinotTypeResolver, aggregateTypes);
}
Optional<String> filter = Optional.empty();
if (pinotQuery.getFilterExpression() != null) {
String formatted = formatFilter(schemaTableName, queryContext.getFilter(), columnHandles);
filter = Optional.of(formatted);
}
return new DynamicTable(pinotTableName, suffix, selectColumns, filter, groupByColumns, ImmutableList.of(), orderBy, OptionalLong.of(queryContext.getLimit()), getOffset(queryContext), query);
}
Aggregations