use of com.facebook.presto.spi.relation.VariableReferenceExpression in project presto by prestodb.
the class PushPartialAggregationThroughExchange method apply.
@Override
public Result apply(AggregationNode aggregationNode, Captures captures, Context context) {
ExchangeNode exchangeNode = captures.get(EXCHANGE_NODE);
boolean decomposable = isDecomposable(aggregationNode, functionAndTypeManager);
if (aggregationNode.getStep().equals(SINGLE) && aggregationNode.hasEmptyGroupingSet() && aggregationNode.hasNonEmptyGroupingSet() && exchangeNode.getType() == REPARTITION) {
// single-step aggregation w/ empty grouping sets in a partitioned stage, so we need a partial that will produce
// the default intermediates for the empty grouping set that will be routed to the appropriate final aggregation.
// TODO: technically, AddExchanges generates a broken plan that this rule "fixes"
checkState(decomposable, "Distributed aggregation with empty grouping set requires partial but functions are not decomposable");
return Result.ofPlanNode(split(aggregationNode, context));
}
PartialAggregationStrategy partialAggregationStrategy = getPartialAggregationStrategy(context.getSession());
if (!decomposable || partialAggregationStrategy == NEVER || partialAggregationStrategy == AUTOMATIC && partialAggregationNotUseful(aggregationNode, exchangeNode, context) && aggregationNode.getGroupingKeys().size() == 1) {
return Result.empty();
}
// the cardinality of the stream (i.e., gather or repartition)
if ((exchangeNode.getType() != GATHER && exchangeNode.getType() != REPARTITION) || exchangeNode.getPartitioningScheme().isReplicateNullsAndAny()) {
return Result.empty();
}
if (exchangeNode.getType() == REPARTITION) {
// if partitioning columns are not a subset of grouping keys,
// we can't push this through
List<VariableReferenceExpression> partitioningColumns = exchangeNode.getPartitioningScheme().getPartitioning().getArguments().stream().filter(VariableReferenceExpression.class::isInstance).map(VariableReferenceExpression.class::cast).collect(Collectors.toList());
if (!aggregationNode.getGroupingKeys().containsAll(partitioningColumns)) {
return Result.empty();
}
}
// currently, we only support plans that don't use pre-computed hash functions
if (aggregationNode.getHashVariable().isPresent() || exchangeNode.getPartitioningScheme().getHashColumn().isPresent()) {
return Result.empty();
}
switch(aggregationNode.getStep()) {
case SINGLE:
// Split it into a FINAL on top of a PARTIAL and
return Result.ofPlanNode(split(aggregationNode, context));
case PARTIAL:
// Push it underneath each branch of the exchange
return Result.ofPlanNode(pushPartial(aggregationNode, exchangeNode, context));
default:
return Result.empty();
}
}
use of com.facebook.presto.spi.relation.VariableReferenceExpression in project presto by prestodb.
the class PushPartialAggregationThroughExchange method pushPartial.
private PlanNode pushPartial(AggregationNode aggregation, ExchangeNode exchange, Context context) {
List<PlanNode> partials = new ArrayList<>();
for (int i = 0; i < exchange.getSources().size(); i++) {
PlanNode source = exchange.getSources().get(i);
SymbolMapper.Builder mappingsBuilder = SymbolMapper.builder();
for (int outputIndex = 0; outputIndex < exchange.getOutputVariables().size(); outputIndex++) {
VariableReferenceExpression output = exchange.getOutputVariables().get(outputIndex);
VariableReferenceExpression input = exchange.getInputs().get(i).get(outputIndex);
if (!output.equals(input)) {
mappingsBuilder.put(output, input);
}
}
SymbolMapper symbolMapper = mappingsBuilder.build();
AggregationNode mappedPartial = symbolMapper.map(aggregation, source, context.getIdAllocator());
Assignments.Builder assignments = Assignments.builder();
for (VariableReferenceExpression output : aggregation.getOutputVariables()) {
VariableReferenceExpression input = symbolMapper.map(output);
assignments.put(output, input);
}
partials.add(new ProjectNode(exchange.getSourceLocation(), context.getIdAllocator().getNextId(), mappedPartial, assignments.build(), LOCAL));
}
for (PlanNode node : partials) {
verify(aggregation.getOutputVariables().equals(node.getOutputVariables()));
}
// Since this exchange source is now guaranteed to have the same symbols as the inputs to the the partial
// aggregation, we don't need to rewrite symbols in the partitioning function
List<VariableReferenceExpression> aggregationOutputs = aggregation.getOutputVariables();
PartitioningScheme partitioning = new PartitioningScheme(exchange.getPartitioningScheme().getPartitioning(), aggregationOutputs, exchange.getPartitioningScheme().getHashColumn(), exchange.getPartitioningScheme().isReplicateNullsAndAny(), exchange.getPartitioningScheme().getBucketToPartition());
return new ExchangeNode(aggregation.getSourceLocation(), context.getIdAllocator().getNextId(), exchange.getType(), exchange.getScope(), partitioning, partials, ImmutableList.copyOf(Collections.nCopies(partials.size(), aggregationOutputs)), exchange.isEnsureSourceOrdering(), Optional.empty());
}
use of com.facebook.presto.spi.relation.VariableReferenceExpression in project presto by prestodb.
the class PushProjectionThroughExchange method apply.
@Override
public Result apply(ProjectNode project, Captures captures, Context context) {
ExchangeNode exchange = captures.get(CHILD);
Set<VariableReferenceExpression> partitioningColumns = exchange.getPartitioningScheme().getPartitioning().getVariableReferences();
ImmutableList.Builder<PlanNode> newSourceBuilder = ImmutableList.builder();
ImmutableList.Builder<List<VariableReferenceExpression>> inputsBuilder = ImmutableList.builder();
for (int i = 0; i < exchange.getSources().size(); i++) {
Map<VariableReferenceExpression, VariableReferenceExpression> outputToInputMap = extractExchangeOutputToInput(exchange, i);
Assignments.Builder projections = Assignments.builder();
ImmutableList.Builder<VariableReferenceExpression> inputs = ImmutableList.builder();
// Need to retain the partition keys for the exchange
partitioningColumns.stream().map(outputToInputMap::get).forEach(variable -> {
projections.put(variable, variable);
inputs.add(variable);
});
if (exchange.getPartitioningScheme().getHashColumn().isPresent()) {
// Need to retain the hash symbol for the exchange
VariableReferenceExpression hashVariable = exchange.getPartitioningScheme().getHashColumn().get();
projections.put(hashVariable, hashVariable);
inputs.add(hashVariable);
}
if (exchange.getOrderingScheme().isPresent()) {
// need to retain ordering columns for the exchange
exchange.getOrderingScheme().get().getOrderByVariables().stream().filter(variable -> !partitioningColumns.contains(variable)).map(outputToInputMap::get).forEach(variable -> {
projections.put(variable, variable);
inputs.add(variable);
});
}
for (Map.Entry<VariableReferenceExpression, RowExpression> projection : project.getAssignments().entrySet()) {
RowExpression translatedExpression = RowExpressionVariableInliner.inlineVariables(outputToInputMap, projection.getValue());
VariableReferenceExpression variable = context.getVariableAllocator().newVariable(translatedExpression);
projections.put(variable, translatedExpression);
inputs.add(variable);
}
newSourceBuilder.add(new ProjectNode(project.getSourceLocation(), context.getIdAllocator().getNextId(), exchange.getSources().get(i), projections.build(), project.getLocality()));
inputsBuilder.add(inputs.build());
}
// Construct the output symbols in the same order as the sources
ImmutableList.Builder<VariableReferenceExpression> outputBuilder = ImmutableList.builder();
partitioningColumns.forEach(outputBuilder::add);
exchange.getPartitioningScheme().getHashColumn().ifPresent(outputBuilder::add);
if (exchange.getOrderingScheme().isPresent()) {
exchange.getOrderingScheme().get().getOrderByVariables().stream().filter(variable -> !partitioningColumns.contains(variable)).forEach(outputBuilder::add);
}
for (Map.Entry<VariableReferenceExpression, RowExpression> projection : project.getAssignments().entrySet()) {
outputBuilder.add(projection.getKey());
}
// outputBuilder contains all partition and hash symbols so simply swap the output layout
PartitioningScheme partitioningScheme = new PartitioningScheme(exchange.getPartitioningScheme().getPartitioning(), outputBuilder.build(), exchange.getPartitioningScheme().getHashColumn(), exchange.getPartitioningScheme().isReplicateNullsAndAny(), exchange.getPartitioningScheme().getBucketToPartition());
PlanNode result = new ExchangeNode(exchange.getSourceLocation(), exchange.getId(), exchange.getType(), exchange.getScope(), partitioningScheme, newSourceBuilder.build(), inputsBuilder.build(), exchange.isEnsureSourceOrdering(), exchange.getOrderingScheme());
// we need to strip unnecessary symbols (hash, partitioning columns).
return Result.ofPlanNode(restrictOutputs(context.getIdAllocator(), result, ImmutableSet.copyOf(project.getOutputVariables()), true).orElse(result));
}
use of com.facebook.presto.spi.relation.VariableReferenceExpression in project presto by prestodb.
the class SingleDistinctAggregationToGroupBy method apply.
@Override
public Result apply(AggregationNode aggregation, Captures captures, Context context) {
List<Set<RowExpression>> argumentSets = extractArgumentSets(aggregation).collect(Collectors.toList());
Set<VariableReferenceExpression> variables = Iterables.getOnlyElement(argumentSets).stream().map(OriginalExpressionUtils::castToExpression).map(context.getVariableAllocator()::toVariableReference).collect(Collectors.toSet());
return Result.ofPlanNode(new AggregationNode(aggregation.getSourceLocation(), aggregation.getId(), new AggregationNode(aggregation.getSourceLocation(), context.getIdAllocator().getNextId(), aggregation.getSource(), ImmutableMap.of(), singleGroupingSet(ImmutableList.<VariableReferenceExpression>builder().addAll(aggregation.getGroupingKeys()).addAll(variables).build()), ImmutableList.of(), SINGLE, Optional.empty(), Optional.empty()), // remove DISTINCT flag from function calls
aggregation.getAggregations().entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e -> removeDistinct(e.getValue()))), aggregation.getGroupingSets(), emptyList(), aggregation.getStep(), aggregation.getHashVariable(), aggregation.getGroupIdVariable()));
}
use of com.facebook.presto.spi.relation.VariableReferenceExpression in project presto by prestodb.
the class TransformCorrelatedScalarSubquery method apply.
@Override
public Result apply(LateralJoinNode lateralJoinNode, Captures captures, Context context) {
PlanNode subquery = context.getLookup().resolve(lateralJoinNode.getSubquery());
if (!searchFrom(subquery, context.getLookup()).where(EnforceSingleRowNode.class::isInstance).recurseOnlyWhen(ProjectNode.class::isInstance).matches()) {
return Result.empty();
}
PlanNode rewrittenSubquery = searchFrom(subquery, context.getLookup()).where(EnforceSingleRowNode.class::isInstance).recurseOnlyWhen(ProjectNode.class::isInstance).removeFirst();
if (isAtMostScalar(rewrittenSubquery, context.getLookup())) {
return Result.ofPlanNode(new LateralJoinNode(lateralJoinNode.getSourceLocation(), context.getIdAllocator().getNextId(), lateralJoinNode.getInput(), rewrittenSubquery, lateralJoinNode.getCorrelation(), lateralJoinNode.getType(), lateralJoinNode.getOriginSubqueryError()));
}
VariableReferenceExpression unique = context.getVariableAllocator().newVariable("unique", BIGINT);
LateralJoinNode rewrittenLateralJoinNode = new LateralJoinNode(lateralJoinNode.getSourceLocation(), context.getIdAllocator().getNextId(), new AssignUniqueId(lateralJoinNode.getSourceLocation(), context.getIdAllocator().getNextId(), lateralJoinNode.getInput(), unique), rewrittenSubquery, lateralJoinNode.getCorrelation(), lateralJoinNode.getType(), lateralJoinNode.getOriginSubqueryError());
VariableReferenceExpression isDistinct = context.getVariableAllocator().newVariable("is_distinct", BooleanType.BOOLEAN);
MarkDistinctNode markDistinctNode = new MarkDistinctNode(rewrittenLateralJoinNode.getSourceLocation(), context.getIdAllocator().getNextId(), rewrittenLateralJoinNode, isDistinct, rewrittenLateralJoinNode.getInput().getOutputVariables(), Optional.empty());
FilterNode filterNode = new FilterNode(markDistinctNode.getSourceLocation(), context.getIdAllocator().getNextId(), markDistinctNode, castToRowExpression(new SimpleCaseExpression(createSymbolReference(isDistinct), ImmutableList.of(new WhenClause(TRUE_LITERAL, TRUE_LITERAL)), Optional.of(new Cast(new FunctionCall(QualifiedName.of("fail"), ImmutableList.of(new LongLiteral(Integer.toString(SUBQUERY_MULTIPLE_ROWS.toErrorCode().getCode())), new StringLiteral("Scalar sub-query has returned multiple rows"))), BOOLEAN)))));
return Result.ofPlanNode(new ProjectNode(context.getIdAllocator().getNextId(), filterNode, identityAssignmentsAsSymbolReferences(lateralJoinNode.getOutputVariables())));
}
Aggregations