use of com.facebook.presto.pinot.PinotColumnHandle in project presto by prestodb.
the class PinotQueryGeneratorContext method toSqlQuery.
/**
* Convert the current context to a Pinot SQL
*/
public PinotQueryGenerator.GeneratedPinotQuery toSqlQuery(PinotConfig pinotConfig, ConnectorSession session) {
int nonAggregateShortQueryLimit = PinotSessionProperties.getNonAggregateLimitForBrokerQueries(session);
boolean isQueryShort = (hasAggregation() || hasGroupBy()) || limit.orElse(Integer.MAX_VALUE) < nonAggregateShortQueryLimit;
boolean forBroker = !PinotSessionProperties.isForbidBrokerQueries(session) && isQueryShort;
String groupByExpressions = groupByColumns.stream().map(x -> selections.get(x).getDefinition()).collect(Collectors.joining(", "));
String selectExpressions = outputs.stream().filter(o -> !groupByColumns.contains(o)).map(o -> updateSelection(selections.get(o).getDefinition(), session)).collect(Collectors.joining(", "));
String expressions = (groupByExpressions.isEmpty()) ? selectExpressions : (selectExpressions.isEmpty()) ? groupByExpressions : groupByExpressions + ", " + selectExpressions;
String tableName = from.orElseThrow(() -> new PinotException(PINOT_QUERY_GENERATOR_FAILURE, Optional.empty(), "Table name not encountered yet"));
// Rules for limit:
// - If its a selection query:
// + given limit or configured limit
// - Else if has group by:
// + default limit or configured top limit
// - Aggregation only query limit is ignored.
// - Fail if limit is invalid
int queryLimit = -1;
if (!hasAggregation() && !hasGroupBy()) {
if (!limit.isPresent() && forBroker) {
throw new PinotException(PINOT_QUERY_GENERATOR_FAILURE, Optional.empty(), "Broker non aggregate queries have to have a limit");
} else {
queryLimit = limit.orElse(PinotSessionProperties.getLimitLargerForSegment(session));
}
} else if (hasGroupBy()) {
if (limit.isPresent()) {
queryLimit = limit.getAsInt();
} else {
queryLimit = PinotSessionProperties.getTopNLarge(session);
}
}
String limitClause = "";
if (queryLimit > 0) {
limitClause = " LIMIT " + queryLimit;
}
String query = generatePinotQueryHelper(forBroker, expressions, tableName, limitClause);
LinkedHashMap<VariableReferenceExpression, PinotColumnHandle> assignments = getAssignments(true);
List<Integer> indices = getIndicesMappingFromPinotSchemaToPrestoSchema(query, assignments);
return new PinotQueryGenerator.GeneratedPinotQuery(tableName, query, PinotQueryGenerator.PinotQueryFormat.SQL, indices, groupByColumns.size(), filter.isPresent(), isQueryShort);
}
use of com.facebook.presto.pinot.PinotColumnHandle in project presto by prestodb.
the class PinotQueryGeneratorContext method toPqlQuery.
/**
* Convert the current context to a PQL
*/
public PinotQueryGenerator.GeneratedPinotQuery toPqlQuery(PinotConfig pinotConfig, ConnectorSession session) {
int nonAggregateShortQueryLimit = PinotSessionProperties.getNonAggregateLimitForBrokerQueries(session);
boolean isQueryShort = hasAggregation() || limit.orElse(Integer.MAX_VALUE) < nonAggregateShortQueryLimit;
boolean forBroker = !PinotSessionProperties.isForbidBrokerQueries(session) && isQueryShort;
if (!pinotConfig.isAllowMultipleAggregations() && aggregations > 1 && hasGroupBy()) {
throw new PinotException(PINOT_QUERY_GENERATOR_FAILURE, Optional.empty(), "Multiple aggregates in the presence of group by is forbidden");
}
if (hasLimit() && aggregations > 1 && hasGroupBy()) {
throw new PinotException(PINOT_QUERY_GENERATOR_FAILURE, Optional.empty(), "Multiple aggregates in the presence of group by and limit is forbidden");
}
String expressions = outputs.stream().filter(// remove the group by columns from the query as Pinot barfs if the group by column is an expression
o -> !groupByColumns.contains(o)).map(o -> updateSelection(selections.get(o).getDefinition(), session)).collect(Collectors.joining(", "));
if (expressions.isEmpty()) {
throw new PinotException(PINOT_QUERY_GENERATOR_FAILURE, Optional.empty(), "Empty PQL expressions: " + toString());
}
String tableName = from.orElseThrow(() -> new PinotException(PINOT_QUERY_GENERATOR_FAILURE, Optional.empty(), "Table name not encountered yet"));
// Rules for limit:
// - If its a selection query:
// + given limit or configured limit
// - Else if has group by:
// + ensure that only one aggregation
// + default limit or configured top limit
// - Fail if limit is invalid
String limitKeyWord = "";
int queryLimit = -1;
if (!hasAggregation()) {
if (!limit.isPresent() && forBroker) {
throw new PinotException(PINOT_QUERY_GENERATOR_FAILURE, Optional.empty(), "Broker non aggregate queries have to have a limit");
} else {
queryLimit = limit.orElse(PinotSessionProperties.getLimitLargerForSegment(session));
}
limitKeyWord = "LIMIT";
} else if (hasGroupBy()) {
limitKeyWord = "TOP";
if (limit.isPresent()) {
if (aggregations > 1) {
throw new PinotException(PINOT_QUERY_GENERATOR_FAILURE, Optional.of(generatePinotQueryHelper(forBroker, expressions, tableName, "")), "Pinot has weird semantics with group by and multiple aggregation functions and limits");
} else {
queryLimit = limit.getAsInt();
}
} else {
queryLimit = PinotSessionProperties.getTopNLarge(session);
}
}
String limitClause = "";
if (!limitKeyWord.isEmpty()) {
limitClause = " " + limitKeyWord + " " + queryLimit;
}
String query = generatePinotQueryHelper(forBroker, expressions, tableName, limitClause);
LinkedHashMap<VariableReferenceExpression, PinotColumnHandle> assignments = getAssignments(false);
List<Integer> indices = getIndicesMappingFromPinotSchemaToPrestoSchema(query, assignments);
return new PinotQueryGenerator.GeneratedPinotQuery(tableName, query, PinotQueryGenerator.PinotQueryFormat.PQL, indices, groupByColumns.size(), filter.isPresent(), isQueryShort);
}
use of com.facebook.presto.pinot.PinotColumnHandle in project presto by prestodb.
the class PinotQueryGeneratorContext method getIndicesMappingFromPinotSchemaToPrestoSchema.
private List<Integer> getIndicesMappingFromPinotSchemaToPrestoSchema(String query, Map<VariableReferenceExpression, PinotColumnHandle> assignments) {
LinkedHashMap<VariableReferenceExpression, Selection> expressionsInPinotOrder = new LinkedHashMap<>();
for (VariableReferenceExpression groupByColumn : groupByColumns) {
Selection groupByColumnDefinition = selections.get(groupByColumn);
if (groupByColumnDefinition == null) {
throw new IllegalStateException(format("Group By column (%s) definition not found in input selections: %s", groupByColumn, Joiner.on(",").withKeyValueSeparator(":").join(selections)));
}
expressionsInPinotOrder.put(groupByColumn, groupByColumnDefinition);
}
for (VariableReferenceExpression outputColumn : outputs) {
Selection outputColumnDefinition = selections.get(outputColumn);
if (outputColumnDefinition == null) {
throw new IllegalStateException(format("Output column (%s) definition not found in input selections: %s", outputColumn, Joiner.on(",").withKeyValueSeparator(":").join(selections)));
}
expressionsInPinotOrder.put(outputColumn, outputColumnDefinition);
}
if (useSqlSyntax) {
checkSupported(assignments.size() <= expressionsInPinotOrder.keySet().stream().filter(key -> !hiddenColumnSet.contains(key)).count(), "Expected returned expressions %s is a superset of selections %s", Joiner.on(",").withKeyValueSeparator(":").join(expressionsInPinotOrder), Joiner.on(",").withKeyValueSeparator("=").join(assignments));
} else {
checkSupported(assignments.size() == expressionsInPinotOrder.keySet().stream().filter(key -> !hiddenColumnSet.contains(key)).count(), "Expected returned expressions %s to match selections %s", Joiner.on(",").withKeyValueSeparator(":").join(expressionsInPinotOrder), Joiner.on(",").withKeyValueSeparator("=").join(assignments));
}
Map<VariableReferenceExpression, Integer> assignmentToIndex = new HashMap<>();
Iterator<Map.Entry<VariableReferenceExpression, PinotColumnHandle>> assignmentsIterator = assignments.entrySet().iterator();
for (int i = 0; i < assignments.size(); i++) {
VariableReferenceExpression key = assignmentsIterator.next().getKey();
Integer previous = assignmentToIndex.put(key, i);
if (previous != null) {
throw new PinotException(PINOT_UNSUPPORTED_EXPRESSION, Optional.of(query), format("Expected Pinot column handle %s to occur only once, but we have: %s", key, Joiner.on(",").withKeyValueSeparator("=").join(assignments)));
}
}
ImmutableList.Builder<Integer> outputIndices = ImmutableList.builder();
for (Map.Entry<VariableReferenceExpression, Selection> expression : expressionsInPinotOrder.entrySet()) {
Integer index;
if (hiddenColumnSet.contains(expression.getKey())) {
// negative output index means to skip this value returned by pinot at query time
index = -1;
} else {
index = assignmentToIndex.get(expression.getKey());
}
if (index == null) {
if (useSqlSyntax) {
// negative output index means to skip this value returned by pinot at query time
index = -1;
} else {
throw new PinotException(PINOT_UNSUPPORTED_EXPRESSION, Optional.of(query), format("Expected to find a Pinot column handle for the expression %s, but we have %s", expression, Joiner.on(",").withKeyValueSeparator(":").join(assignmentToIndex)));
}
}
outputIndices.add(index);
}
return outputIndices.build();
}
Aggregations