Examples with Assignments com.facebook.presto.spi.plan.Assignments used on opensource projects

Example 11 with Assignments

use of com.facebook.presto.spi.plan.Assignments in project presto by prestodb.

the class TestDruidQueryBase method tableScan.

protected TableScanNode tableScan(PlanBuilder planBuilder, DruidTableHandle connectorTableHandle, DruidColumnHandle... columnHandles) {
    List<VariableReferenceExpression> variables = Arrays.stream(columnHandles).map(ch -> new VariableReferenceExpression(Optional.empty(), ch.getColumnName().toLowerCase(ENGLISH), ch.getColumnType())).collect(toImmutableList());
    ImmutableMap.Builder<VariableReferenceExpression, ColumnHandle> assignments = ImmutableMap.builder();
    for (int i = 0; i < variables.size(); ++i) {
        assignments.put(variables.get(i), columnHandles[i]);
    }
    TableHandle tableHandle = new TableHandle(druidConnectorId, connectorTableHandle, TestingTransactionHandle.create(), Optional.empty());
    return new TableScanNode(Optional.empty(), planBuilder.getIdAllocator().getNextId(), tableHandle, variables, assignments.build(), TupleDomain.all(), TupleDomain.all());
}

Example 12 with Assignments

use of com.facebook.presto.spi.plan.Assignments in project presto by prestodb.

the class ScalarAggregationToJoinRewriter method rewriteScalarAggregation.

public PlanNode rewriteScalarAggregation(LateralJoinNode lateralJoinNode, AggregationNode aggregation) {
    List<VariableReferenceExpression> correlation = lateralJoinNode.getCorrelation();
    Optional<DecorrelatedNode> source = planNodeDecorrelator.decorrelateFilters(lookup.resolve(aggregation.getSource()), correlation);
    if (!source.isPresent()) {
        return lateralJoinNode;
    }
    VariableReferenceExpression nonNull = variableAllocator.newVariable("non_null", BooleanType.BOOLEAN);
    Assignments scalarAggregationSourceAssignments = Assignments.builder().putAll(identitiesAsSymbolReferences(source.get().getNode().getOutputVariables())).put(nonNull, castToRowExpression(TRUE_LITERAL)).build();
    ProjectNode scalarAggregationSourceWithNonNullableVariable = new ProjectNode(idAllocator.getNextId(), source.get().getNode(), scalarAggregationSourceAssignments);
    return rewriteScalarAggregation(lateralJoinNode, aggregation, scalarAggregationSourceWithNonNullableVariable, source.get().getCorrelatedPredicates(), nonNull);
}

Example 13 with Assignments

use of com.facebook.presto.spi.plan.Assignments in project presto by prestodb.

the class RelationPlanner method addCoercions.

private RelationPlan addCoercions(RelationPlan plan, Type[] targetColumnTypes) {
    RelationType oldRelation = plan.getDescriptor();
    List<VariableReferenceExpression> oldVisibleVariables = oldRelation.getVisibleFields().stream().map(oldRelation::indexOf).map(plan.getFieldMappings()::get).collect(toImmutableList());
    RelationType oldRelationWithVisibleFields = plan.getDescriptor().withOnlyVisibleFields();
    verify(targetColumnTypes.length == oldVisibleVariables.size());
    ImmutableList.Builder<VariableReferenceExpression> newVariables = new ImmutableList.Builder<>();
    Field[] newFields = new Field[targetColumnTypes.length];
    Assignments.Builder assignments = Assignments.builder();
    for (int i = 0; i < targetColumnTypes.length; i++) {
        VariableReferenceExpression inputVariable = oldVisibleVariables.get(i);
        Field oldField = oldRelationWithVisibleFields.getFieldByIndex(i);
        Type outputType = targetColumnTypes[i];
        if (!outputType.equals(inputVariable.getType())) {
            Expression cast = new Cast(createSymbolReference(inputVariable), outputType.getTypeSignature().toString());
            VariableReferenceExpression outputVariable = variableAllocator.newVariable(cast, outputType);
            assignments.put(outputVariable, castToRowExpression(cast));
            newVariables.add(outputVariable);
        } else {
            SymbolReference symbolReference = new SymbolReference(oldField.getNodeLocation(), inputVariable.getName());
            VariableReferenceExpression outputVariable = variableAllocator.newVariable(symbolReference, outputType);
            assignments.put(outputVariable, castToRowExpression(symbolReference));
            newVariables.add(outputVariable);
        }
        newFields[i] = new Field(oldField.getNodeLocation(), oldField.getRelationAlias(), oldField.getName(), targetColumnTypes[i], oldField.isHidden(), oldField.getOriginTable(), oldField.getOriginColumnName(), oldField.isAliased());
    }
    ProjectNode projectNode = new ProjectNode(idAllocator.getNextId(), plan.getRoot(), assignments.build());
    return new RelationPlan(projectNode, Scope.builder().withRelationType(RelationId.anonymous(), new RelationType(newFields)).build(), newVariables.build());
}

Example 14 with Assignments

use of com.facebook.presto.spi.plan.Assignments in project presto by prestodb.

the class RewriteAggregationIfToFilter method apply.

@Override
public Result apply(AggregationNode aggregationNode, Captures captures, Context context) {
    ProjectNode sourceProject = captures.get(CHILD);
    Set<Aggregation> aggregationsToRewrite = aggregationNode.getAggregations().values().stream().filter(aggregation -> shouldRewriteAggregation(aggregation, sourceProject)).collect(toImmutableSet());
    if (aggregationsToRewrite.isEmpty()) {
        return Result.empty();
    }
    context.getSession().getRuntimeStats().addMetricValue(REWRITE_AGGREGATION_IF_TO_FILTER_APPLIED, 1);
    // Get the corresponding assignments in the input project.
    // The aggregationReferences only has the aggregations to rewrite, thus the sourceAssignments only has IF/CAST(IF) expressions with NULL false results.
    // Multiple aggregations may reference the same input. We use a map to dedup them based on the VariableReferenceExpression, so that we only do the rewrite once per input
    // IF expression.
    // The order of sourceAssignments determines the order of generating the new variables for the IF conditions and results. We use a sorted map to get a deterministic
    // order based on the name of the VariableReferenceExpressions.
    Map<VariableReferenceExpression, RowExpression> sourceAssignments = aggregationsToRewrite.stream().map(aggregation -> (VariableReferenceExpression) aggregation.getArguments().get(0)).collect(toImmutableSortedMap(VariableReferenceExpression::compareTo, identity(), variable -> sourceProject.getAssignments().get(variable), (left, right) -> left));
    Assignments.Builder newAssignments = Assignments.builder();
    newAssignments.putAll(sourceProject.getAssignments());
    // Map from the aggregation reference to the IF condition reference which will be put in the mask.
    Map<VariableReferenceExpression, VariableReferenceExpression> aggregationReferenceToConditionReference = new HashMap<>();
    // Map from the aggregation reference to the IF result reference. This only contains the aggregates where the IF can be safely unwrapped.
    // E.g., SUM(IF(CARDINALITY(array) > 0, array[1])) will not be included in this map as array[1] can return errors if we unwrap the IF.
    Map<VariableReferenceExpression, VariableReferenceExpression> aggregationReferenceToIfResultReference = new HashMap<>();
    AggregationIfToFilterRewriteStrategy rewriteStrategy = getAggregationIfToFilterRewriteStrategy(context.getSession());
    for (Map.Entry<VariableReferenceExpression, RowExpression> entry : sourceAssignments.entrySet()) {
        VariableReferenceExpression outputVariable = entry.getKey();
        RowExpression rowExpression = entry.getValue();
        SpecialFormExpression ifExpression = (SpecialFormExpression) ((rowExpression instanceof CallExpression) ? ((CallExpression) rowExpression).getArguments().get(0) : rowExpression);
        RowExpression condition = ifExpression.getArguments().get(0);
        VariableReferenceExpression conditionReference = context.getVariableAllocator().newVariable(condition);
        newAssignments.put(conditionReference, condition);
        aggregationReferenceToConditionReference.put(outputVariable, conditionReference);
        if (canUnwrapIf(ifExpression, rewriteStrategy)) {
            RowExpression trueResult = ifExpression.getArguments().get(1);
            if (rowExpression instanceof CallExpression) {
                // Wrap the result with CAST().
                trueResult = new CallExpression(((CallExpression) rowExpression).getDisplayName(), ((CallExpression) rowExpression).getFunctionHandle(), rowExpression.getType(), ImmutableList.of(trueResult));
            }
            VariableReferenceExpression ifResultReference = context.getVariableAllocator().newVariable(trueResult);
            newAssignments.put(ifResultReference, trueResult);
            aggregationReferenceToIfResultReference.put(outputVariable, ifResultReference);
        }
    }
    // Build new aggregations.
    ImmutableMap.Builder<VariableReferenceExpression, Aggregation> aggregations = ImmutableMap.builder();
    // Stores the masks used to build the filter predicates. Use set to dedup the predicates.
    ImmutableSortedSet.Builder<VariableReferenceExpression> masks = ImmutableSortedSet.naturalOrder();
    for (Map.Entry<VariableReferenceExpression, Aggregation> entry : aggregationNode.getAggregations().entrySet()) {
        VariableReferenceExpression output = entry.getKey();
        Aggregation aggregation = entry.getValue();
        if (!aggregationsToRewrite.contains(aggregation)) {
            aggregations.put(output, aggregation);
            continue;
        }
        VariableReferenceExpression aggregationReference = (VariableReferenceExpression) aggregation.getArguments().get(0);
        CallExpression callExpression = aggregation.getCall();
        VariableReferenceExpression ifResultReference = aggregationReferenceToIfResultReference.get(aggregationReference);
        if (ifResultReference != null) {
            callExpression = new CallExpression(callExpression.getSourceLocation(), callExpression.getDisplayName(), callExpression.getFunctionHandle(), callExpression.getType(), ImmutableList.of(ifResultReference));
        }
        VariableReferenceExpression mask = aggregationReferenceToConditionReference.get(aggregationReference);
        aggregations.put(output, new Aggregation(callExpression, Optional.empty(), aggregation.getOrderBy(), aggregation.isDistinct(), Optional.of(aggregationReferenceToConditionReference.get(aggregationReference))));
        masks.add(mask);
    }
    RowExpression predicate = TRUE_CONSTANT;
    if (!aggregationNode.hasNonEmptyGroupingSet() && aggregationsToRewrite.size() == aggregationNode.getAggregations().size()) {
        // All aggregations are rewritten by this rule. We can add a filter with all the masks to make the query more efficient.
        predicate = or(masks.build());
    }
    return Result.ofPlanNode(new AggregationNode(aggregationNode.getSourceLocation(), context.getIdAllocator().getNextId(), new FilterNode(aggregationNode.getSourceLocation(), context.getIdAllocator().getNextId(), new ProjectNode(context.getIdAllocator().getNextId(), sourceProject.getSource(), newAssignments.build()), predicate), aggregations.build(), aggregationNode.getGroupingSets(), aggregationNode.getPreGroupedVariables(), aggregationNode.getStep(), aggregationNode.getHashVariable(), aggregationNode.getGroupIdVariable()));
}

Example 15 with Assignments

use of com.facebook.presto.spi.plan.Assignments in project presto by prestodb.

the class InlineProjections method apply.

@Override
public Result apply(ProjectNode parent, Captures captures, Context context) {
    ProjectNode child = captures.get(CHILD);
    // Do not inline remote projections, or if parent and child has different locality
    if (parent.getLocality().equals(REMOTE) || child.getLocality().equals(REMOTE) || !parent.getLocality().equals(child.getLocality())) {
        return Result.empty();
    }
    Sets.SetView<VariableReferenceExpression> targets = extractInliningTargets(parent, child, context);
    if (targets.isEmpty()) {
        return Result.empty();
    }
    // inline the expressions
    Assignments assignments = child.getAssignments().filter(targets::contains);
    Map<VariableReferenceExpression, RowExpression> parentAssignments = parent.getAssignments().entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, entry -> inlineReferences(entry.getValue(), assignments, context.getVariableAllocator().getTypes())));
    // Synthesize identity assignments for the inputs of expressions that were inlined
    // to place in the child projection.
    // If all assignments end up becoming identity assignments, they'll get pruned by
    // other rules
    Set<VariableReferenceExpression> inputs = child.getAssignments().entrySet().stream().filter(entry -> targets.contains(entry.getKey())).map(Map.Entry::getValue).flatMap(expression -> extractInputs(expression, context.getVariableAllocator().getTypes()).stream()).collect(toSet());
    Builder childAssignments = Assignments.builder();
    for (Map.Entry<VariableReferenceExpression, RowExpression> assignment : child.getAssignments().entrySet()) {
        if (!targets.contains(assignment.getKey())) {
            childAssignments.put(assignment);
        }
    }
    boolean allTranslated = child.getAssignments().entrySet().stream().map(Map.Entry::getValue).noneMatch(OriginalExpressionUtils::isExpression);
    for (VariableReferenceExpression input : inputs) {
        if (allTranslated) {
            childAssignments.put(input, input);
        } else {
            childAssignments.put(identityAsSymbolReference(input));
        }
    }
    return Result.ofPlanNode(new ProjectNode(parent.getSourceLocation(), parent.getId(), new ProjectNode(parent.getSourceLocation(), child.getId(), child.getSource(), childAssignments.build(), child.getLocality()), Assignments.copyOf(parentAssignments), parent.getLocality()));
}