use of io.trino.sql.analyzer.FieldId in project trino by trinodb.
the class QueryPlanner method planGroupingSets.
private GroupingSetsPlan planGroupingSets(PlanBuilder subPlan, QuerySpecification node, GroupingSetAnalysis groupingSetAnalysis) {
Map<Symbol, Symbol> groupingSetMappings = new LinkedHashMap<>();
// Compute a set of artificial columns that will contain the values of the original columns
// filtered by whether the column is included in the grouping set
// This will become the basis for the scope for any column references
Symbol[] fields = new Symbol[subPlan.getTranslations().getFieldSymbols().size()];
for (FieldId field : groupingSetAnalysis.getAllFields()) {
Symbol input = subPlan.getTranslations().getFieldSymbols().get(field.getFieldIndex());
Symbol output = symbolAllocator.newSymbol(input, "gid");
fields[field.getFieldIndex()] = output;
groupingSetMappings.put(output, input);
}
Map<ScopeAware<Expression>, Symbol> complexExpressions = new HashMap<>();
for (Expression expression : groupingSetAnalysis.getComplexExpressions()) {
if (!complexExpressions.containsKey(scopeAwareKey(expression, analysis, subPlan.getScope()))) {
Symbol input = subPlan.translate(expression);
Symbol output = symbolAllocator.newSymbol(expression, analysis.getType(expression), "gid");
complexExpressions.put(scopeAwareKey(expression, analysis, subPlan.getScope()), output);
groupingSetMappings.put(output, input);
}
}
// For the purpose of "distinct", we need to canonicalize column references that may have varying
// syntactic forms (e.g., "t.a" vs "a"). Thus we need to enumerate grouping sets based on the underlying
// fieldId associated with each column reference expression.
// The catch is that simple group-by expressions can be arbitrary expressions (this is a departure from the SQL specification).
// But, they don't affect the number of grouping sets or the behavior of "distinct" . We can compute all the candidate
// grouping sets in terms of fieldId, dedup as appropriate and then cross-join them with the complex expressions.
// This tracks the grouping sets before complex expressions are considered.
// It's also used to compute the descriptors needed to implement grouping()
List<Set<FieldId>> columnOnlyGroupingSets = enumerateGroupingSets(groupingSetAnalysis);
if (node.getGroupBy().isPresent() && node.getGroupBy().get().isDistinct()) {
columnOnlyGroupingSets = columnOnlyGroupingSets.stream().distinct().collect(toImmutableList());
}
// translate from FieldIds to Symbols
List<List<Symbol>> sets = columnOnlyGroupingSets.stream().map(set -> set.stream().map(FieldId::getFieldIndex).map(index -> fields[index]).collect(toImmutableList())).collect(toImmutableList());
// combine (cartesian product) with complex expressions
List<List<Symbol>> groupingSets = sets.stream().map(set -> ImmutableList.<Symbol>builder().addAll(set).addAll(complexExpressions.values()).build()).collect(toImmutableList());
// Generate GroupIdNode (multiple grouping sets) or ProjectNode (single grouping set)
PlanNode groupId;
Optional<Symbol> groupIdSymbol = Optional.empty();
if (groupingSets.size() > 1) {
groupIdSymbol = Optional.of(symbolAllocator.newSymbol("groupId", BIGINT));
groupId = new GroupIdNode(idAllocator.getNextId(), subPlan.getRoot(), groupingSets, groupingSetMappings, subPlan.getRoot().getOutputSymbols(), groupIdSymbol.get());
} else {
Assignments.Builder assignments = Assignments.builder();
assignments.putIdentities(subPlan.getRoot().getOutputSymbols());
groupingSetMappings.forEach((key, value) -> assignments.put(key, value.toSymbolReference()));
groupId = new ProjectNode(idAllocator.getNextId(), subPlan.getRoot(), assignments.build());
}
subPlan = new PlanBuilder(subPlan.getTranslations().withNewMappings(complexExpressions, Arrays.asList(fields)), groupId);
return new GroupingSetsPlan(subPlan, columnOnlyGroupingSets, groupingSets, groupIdSymbol);
}
use of io.trino.sql.analyzer.FieldId in project trino by trinodb.
the class GroupingOperationRewriter method rewriteGroupingOperation.
public static Expression rewriteGroupingOperation(GroupingOperation expression, List<Set<Integer>> groupingSets, Map<NodeRef<Expression>, ResolvedField> columnReferenceFields, Optional<Symbol> groupIdSymbol) {
requireNonNull(groupIdSymbol, "groupIdSymbol is null");
// See SQL:2011:4.16.2 and SQL:2011:6.9.10.
if (groupingSets.size() == 1) {
return new LongLiteral("0");
} else {
checkState(groupIdSymbol.isPresent(), "groupId symbol is missing");
RelationId relationId = columnReferenceFields.get(NodeRef.of(expression.getGroupingColumns().get(0))).getFieldId().getRelationId();
List<Integer> columns = expression.getGroupingColumns().stream().map(NodeRef::of).peek(groupingColumn -> checkState(columnReferenceFields.containsKey(groupingColumn), "the grouping column is not in the columnReferencesField map")).map(columnReferenceFields::get).map(ResolvedField::getFieldId).map(fieldId -> translateFieldToInteger(fieldId, relationId)).collect(toImmutableList());
List<Expression> groupingResults = groupingSets.stream().map(groupingSet -> String.valueOf(calculateGrouping(groupingSet, columns))).map(LongLiteral::new).collect(toImmutableList());
// It is necessary to add a 1 to the groupId because the underlying array is indexed starting at 1
return new SubscriptExpression(new ArrayConstructor(groupingResults), new ArithmeticBinaryExpression(ADD, groupIdSymbol.get().toSymbolReference(), new GenericLiteral("BIGINT", "1")));
}
}
Aggregations