use of io.trino.sql.planner.Symbol in project trino by trinodb.
the class AggregationDecorrelation method rewriteWithMasks.
public static Map<Symbol, Aggregation> rewriteWithMasks(Map<Symbol, Aggregation> aggregations, Map<Symbol, Symbol> masks) {
ImmutableMap.Builder<Symbol, Aggregation> rewritten = ImmutableMap.builder();
for (Map.Entry<Symbol, Aggregation> entry : aggregations.entrySet()) {
Symbol symbol = entry.getKey();
Aggregation aggregation = entry.getValue();
rewritten.put(symbol, new Aggregation(aggregation.getResolvedFunction(), aggregation.getArguments(), aggregation.isDistinct(), aggregation.getFilter(), aggregation.getOrderingScheme(), Optional.of(masks.get(symbol))));
}
return rewritten.buildOrThrow();
}
use of io.trino.sql.planner.Symbol in project trino by trinodb.
the class MultipleDistinctAggregationToMarkDistinct method apply.
@Override
public Result apply(AggregationNode parent, Captures captures, Context context) {
if (!SystemSessionProperties.useMarkDistinct(context.getSession())) {
return Result.empty();
}
// the distinct marker for the given set of input columns
Map<Set<Symbol>, Symbol> markers = new HashMap<>();
Map<Symbol, Aggregation> newAggregations = new HashMap<>();
PlanNode subPlan = parent.getSource();
for (Map.Entry<Symbol, Aggregation> entry : parent.getAggregations().entrySet()) {
Aggregation aggregation = entry.getValue();
if (aggregation.isDistinct() && aggregation.getFilter().isEmpty() && aggregation.getMask().isEmpty()) {
Set<Symbol> inputs = aggregation.getArguments().stream().map(Symbol::from).collect(toSet());
Symbol marker = markers.get(inputs);
if (marker == null) {
marker = context.getSymbolAllocator().newSymbol(Iterables.getLast(inputs).getName(), BOOLEAN, "distinct");
markers.put(inputs, marker);
ImmutableSet.Builder<Symbol> distinctSymbols = ImmutableSet.<Symbol>builder().addAll(parent.getGroupingKeys()).addAll(inputs);
parent.getGroupIdSymbol().ifPresent(distinctSymbols::add);
subPlan = new MarkDistinctNode(context.getIdAllocator().getNextId(), subPlan, marker, ImmutableList.copyOf(distinctSymbols.build()), Optional.empty());
}
// remove the distinct flag and set the distinct marker
newAggregations.put(entry.getKey(), new Aggregation(aggregation.getResolvedFunction(), aggregation.getArguments(), false, aggregation.getFilter(), aggregation.getOrderingScheme(), Optional.of(marker)));
} else {
newAggregations.put(entry.getKey(), aggregation);
}
}
return Result.ofPlanNode(new AggregationNode(parent.getId(), subPlan, newAggregations, parent.getGroupingSets(), ImmutableList.of(), parent.getStep(), parent.getHashSymbol(), parent.getGroupIdSymbol()));
}
use of io.trino.sql.planner.Symbol in project trino by trinodb.
the class PruneApplyColumns method pushDownProjectOff.
@Override
protected Optional<PlanNode> pushDownProjectOff(Context context, ApplyNode applyNode, Set<Symbol> referencedOutputs) {
// remove unused apply node
if (intersection(applyNode.getSubqueryAssignments().getSymbols(), referencedOutputs).isEmpty()) {
return Optional.of(applyNode.getInput());
}
// extract referenced assignments
ImmutableSet.Builder<Symbol> requiredAssignmentsSymbols = ImmutableSet.builder();
Assignments.Builder newSubqueryAssignments = Assignments.builder();
for (Map.Entry<Symbol, Expression> entry : applyNode.getSubqueryAssignments().entrySet()) {
if (referencedOutputs.contains(entry.getKey())) {
requiredAssignmentsSymbols.addAll(extractUnique(entry.getValue()));
newSubqueryAssignments.put(entry);
}
}
// prune subquery symbols
Optional<PlanNode> newSubquery = restrictOutputs(context.getIdAllocator(), applyNode.getSubquery(), requiredAssignmentsSymbols.build());
// prune input symbols
Set<Symbol> requiredInputSymbols = ImmutableSet.<Symbol>builder().addAll(referencedOutputs).addAll(applyNode.getCorrelation()).addAll(requiredAssignmentsSymbols.build()).build();
Optional<PlanNode> newInput = restrictOutputs(context.getIdAllocator(), applyNode.getInput(), requiredInputSymbols);
boolean pruned = newSubquery.isPresent() || newInput.isPresent() || newSubqueryAssignments.build().size() < applyNode.getSubqueryAssignments().size();
if (pruned) {
return Optional.of(new ApplyNode(applyNode.getId(), newInput.orElse(applyNode.getInput()), newSubquery.orElse(applyNode.getSubquery()), newSubqueryAssignments.build(), applyNode.getCorrelation(), applyNode.getOriginSubquery()));
}
return Optional.empty();
}
use of io.trino.sql.planner.Symbol in project trino by trinodb.
the class PushLimitThroughProject method apply.
@Override
public Result apply(LimitNode parent, Captures captures, Context context) {
ProjectNode projectNode = captures.get(CHILD);
// Do not push down if the projection is made up of symbol references and exclusive dereferences. This prevents
// undoing of PushDownDereferencesThroughLimit. We still push limit in the case of overlapping dereferences since
// it enables PushDownDereferencesThroughLimit rule to push optimal dereferences.
Set<Expression> projections = ImmutableSet.copyOf(projectNode.getAssignments().getExpressions());
if (!extractRowSubscripts(projections, false, context.getSession(), typeAnalyzer, context.getSymbolAllocator().getTypes()).isEmpty() && exclusiveDereferences(projections, context.getSession(), typeAnalyzer, context.getSymbolAllocator().getTypes())) {
return Result.empty();
}
// for a LimitNode without ties and pre-sorted inputs, simply reorder the nodes
if (!parent.isWithTies() && !parent.requiresPreSortedInputs()) {
return Result.ofPlanNode(transpose(parent, projectNode));
}
// for a LimitNode with ties, the tiesResolvingScheme must be rewritten in terms of symbols before projection
SymbolMapper.Builder symbolMapper = SymbolMapper.builder();
Set<Symbol> symbolsForRewrite = ImmutableSet.<Symbol>builder().addAll(parent.getPreSortedInputs()).addAll(parent.getTiesResolvingScheme().map(OrderingScheme::getOrderBy).orElse(ImmutableList.of())).build();
for (Symbol symbol : symbolsForRewrite) {
Expression expression = projectNode.getAssignments().get(symbol);
// if a symbol results from some computation, the translation fails
if (!(expression instanceof SymbolReference)) {
return Result.empty();
}
symbolMapper.put(symbol, Symbol.from(expression));
}
LimitNode mappedLimitNode = symbolMapper.build().map(parent, projectNode.getSource());
return Result.ofPlanNode(projectNode.replaceChildren(ImmutableList.of(mappedLimitNode)));
}
use of io.trino.sql.planner.Symbol in project trino by trinodb.
the class PushPartialAggregationThroughExchange method split.
private PlanNode split(AggregationNode node, Context context) {
// otherwise, add a partial and final with an exchange in between
Map<Symbol, AggregationNode.Aggregation> intermediateAggregation = new HashMap<>();
Map<Symbol, AggregationNode.Aggregation> finalAggregation = new HashMap<>();
for (Map.Entry<Symbol, AggregationNode.Aggregation> entry : node.getAggregations().entrySet()) {
AggregationNode.Aggregation originalAggregation = entry.getValue();
ResolvedFunction resolvedFunction = originalAggregation.getResolvedFunction();
AggregationFunctionMetadata functionMetadata = plannerContext.getMetadata().getAggregationFunctionMetadata(resolvedFunction);
List<Type> intermediateTypes = functionMetadata.getIntermediateTypes().stream().map(plannerContext.getTypeManager()::getType).collect(toImmutableList());
Type intermediateType = intermediateTypes.size() == 1 ? intermediateTypes.get(0) : RowType.anonymous(intermediateTypes);
Symbol intermediateSymbol = context.getSymbolAllocator().newSymbol(resolvedFunction.getSignature().getName(), intermediateType);
checkState(originalAggregation.getOrderingScheme().isEmpty(), "Aggregate with ORDER BY does not support partial aggregation");
intermediateAggregation.put(intermediateSymbol, new AggregationNode.Aggregation(resolvedFunction, originalAggregation.getArguments(), originalAggregation.isDistinct(), originalAggregation.getFilter(), originalAggregation.getOrderingScheme(), originalAggregation.getMask()));
// rewrite final aggregation in terms of intermediate function
finalAggregation.put(entry.getKey(), new AggregationNode.Aggregation(resolvedFunction, ImmutableList.<Expression>builder().add(intermediateSymbol.toSymbolReference()).addAll(originalAggregation.getArguments().stream().filter(LambdaExpression.class::isInstance).collect(toImmutableList())).build(), false, Optional.empty(), Optional.empty(), Optional.empty()));
}
PlanNode partial = new AggregationNode(context.getIdAllocator().getNextId(), node.getSource(), intermediateAggregation, node.getGroupingSets(), // through the exchange may or may not preserve these properties. Hence, it is safest to drop preGroupedSymbols here.
ImmutableList.of(), PARTIAL, node.getHashSymbol(), node.getGroupIdSymbol());
return new AggregationNode(node.getId(), partial, finalAggregation, node.getGroupingSets(), // through the exchange may or may not preserve these properties. Hence, it is safest to drop preGroupedSymbols here.
ImmutableList.of(), FINAL, node.getHashSymbol(), node.getGroupIdSymbol());
}
Aggregations