private Function<Symbol, Optional<Symbol>> translateGroupIdSymbols(GroupIdNode node) {
List<Symbol> commonGroupingColumns = node.getCommonGroupingColumns();
return symbol -> {
if (node.getIdentityMappings().containsKey(symbol)) {
return Optional.of(node.getIdentityMappings().get(symbol));
}
if (commonGroupingColumns.contains(symbol)) {
return Optional.of(symbol);
}
return Optional.empty();
};
}
@Override
public PlanNode visitGroupId(GroupIdNode node, RewriteContext<Set<Symbol>> context) {
checkState(
node.getDistinctGroupingColumns()
.stream()
.allMatch(column -> context.get().contains(column)));
PlanNode source = context.rewrite(node.getSource(), ImmutableSet.copyOf(context.get()));
List<Symbol> requiredSymbols =
context
.get()
.stream()
.filter(symbol -> !symbol.equals(node.getGroupIdSymbol()))
.collect(toImmutableList());
return new GroupIdNode(
node.getId(), source, requiredSymbols, node.getGroupingSets(), node.getGroupIdSymbol());
}