Examples with Assignments io.trino.sql.planner.plan.Assignments used on opensource projects

Example 41 with Assignments

use of io.trino.sql.planner.plan.Assignments in project trino by trinodb.

the class PushProjectionIntoTableScan method apply.

@Override
public Result apply(ProjectNode project, Captures captures, Context context) {
    TableScanNode tableScan = captures.get(TABLE_SCAN);
    // Extract translatable components from projection expressions. Prepare a mapping from these internal
    // expression nodes to corresponding ConnectorExpression translations.
    Map<NodeRef<Expression>, ConnectorExpression> partialTranslations = project.getAssignments().getMap().entrySet().stream().flatMap(expression -> extractPartialTranslations(expression.getValue(), context.getSession(), typeAnalyzer, context.getSymbolAllocator().getTypes(), plannerContext).entrySet().stream()).collect(toImmutableMap(Map.Entry::getKey, Map.Entry::getValue, (first, ignore) -> first));
    List<NodeRef<Expression>> nodesForPartialProjections = ImmutableList.copyOf(partialTranslations.keySet());
    List<ConnectorExpression> connectorPartialProjections = ImmutableList.copyOf(partialTranslations.values());
    Map<String, Symbol> inputVariableMappings = tableScan.getAssignments().keySet().stream().collect(toImmutableMap(Symbol::getName, identity()));
    Map<String, ColumnHandle> assignments = inputVariableMappings.entrySet().stream().collect(toImmutableMap(Entry::getKey, entry -> tableScan.getAssignments().get(entry.getValue())));
    Optional<ProjectionApplicationResult<TableHandle>> result = plannerContext.getMetadata().applyProjection(context.getSession(), tableScan.getTable(), connectorPartialProjections, assignments);
    if (result.isEmpty()) {
        return Result.empty();
    }
    List<ConnectorExpression> newConnectorPartialProjections = result.get().getProjections();
    checkState(newConnectorPartialProjections.size() == connectorPartialProjections.size(), "Mismatch between input and output projections from the connector: expected %s but got %s", connectorPartialProjections.size(), newConnectorPartialProjections.size());
    List<Symbol> newScanOutputs = new ArrayList<>();
    Map<Symbol, ColumnHandle> newScanAssignments = new HashMap<>();
    Map<String, Symbol> variableMappings = new HashMap<>();
    for (Assignment assignment : result.get().getAssignments()) {
        Symbol symbol = context.getSymbolAllocator().newSymbol(assignment.getVariable(), assignment.getType());
        newScanOutputs.add(symbol);
        newScanAssignments.put(symbol, assignment.getColumn());
        variableMappings.put(assignment.getVariable(), symbol);
    }
    // Translate partial connector projections back to new partial projections
    List<Expression> newPartialProjections = newConnectorPartialProjections.stream().map(expression -> ConnectorExpressionTranslator.translate(context.getSession(), expression, plannerContext, variableMappings, new LiteralEncoder(plannerContext))).collect(toImmutableList());
    // Map internal node references to new partial projections
    ImmutableMap.Builder<NodeRef<Expression>, Expression> nodesToNewPartialProjectionsBuilder = ImmutableMap.builder();
    for (int i = 0; i < nodesForPartialProjections.size(); i++) {
        nodesToNewPartialProjectionsBuilder.put(nodesForPartialProjections.get(i), newPartialProjections.get(i));
    }
    Map<NodeRef<Expression>, Expression> nodesToNewPartialProjections = nodesToNewPartialProjectionsBuilder.buildOrThrow();
    // Stitch partial translations to form new complete projections
    Assignments.Builder newProjectionAssignments = Assignments.builder();
    project.getAssignments().entrySet().forEach(entry -> {
        newProjectionAssignments.put(entry.getKey(), replaceExpression(entry.getValue(), nodesToNewPartialProjections));
    });
    Optional<PlanNodeStatsEstimate> newStatistics = tableScan.getStatistics().map(statistics -> {
        PlanNodeStatsEstimate.Builder builder = PlanNodeStatsEstimate.builder();
        builder.setOutputRowCount(statistics.getOutputRowCount());
        for (int i = 0; i < connectorPartialProjections.size(); i++) {
            ConnectorExpression inputConnectorExpression = connectorPartialProjections.get(i);
            ConnectorExpression resultConnectorExpression = newConnectorPartialProjections.get(i);
            if (!(resultConnectorExpression instanceof Variable)) {
                continue;
            }
            String resultVariableName = ((Variable) resultConnectorExpression).getName();
            Expression inputExpression = ConnectorExpressionTranslator.translate(context.getSession(), inputConnectorExpression, plannerContext, inputVariableMappings, new LiteralEncoder(plannerContext));
            SymbolStatsEstimate symbolStatistics = scalarStatsCalculator.calculate(inputExpression, statistics, context.getSession(), context.getSymbolAllocator().getTypes());
            builder.addSymbolStatistics(variableMappings.get(resultVariableName), symbolStatistics);
        }
        return builder.build();
    });
    verifyTablePartitioning(context, tableScan, result.get().getHandle());
    return Result.ofPlanNode(new ProjectNode(context.getIdAllocator().getNextId(), new TableScanNode(tableScan.getId(), result.get().getHandle(), newScanOutputs, newScanAssignments, TupleDomain.all(), newStatistics, tableScan.isUpdateTarget(), tableScan.getUseConnectorNodePartitioning()), newProjectionAssignments.build()));
}

Example 42 with Assignments

use of io.trino.sql.planner.plan.Assignments in project trino by trinodb.

the class PushDownDereferenceThroughProject method apply.

@Override
public Result apply(ProjectNode node, Captures captures, Context context) {
    ProjectNode child = captures.get(CHILD);
    // Extract dereferences from project node assignments for pushdown
    Set<SubscriptExpression> dereferences = extractRowSubscripts(node.getAssignments().getExpressions(), false, context.getSession(), typeAnalyzer, context.getSymbolAllocator().getTypes());
    // Exclude dereferences on symbols being synthesized within child
    dereferences = dereferences.stream().filter(expression -> child.getSource().getOutputSymbols().contains(getBase(expression))).collect(toImmutableSet());
    if (dereferences.isEmpty()) {
        return Result.empty();
    }
    // Create new symbols for dereference expressions
    Assignments dereferenceAssignments = Assignments.of(dereferences, context.getSession(), context.getSymbolAllocator(), typeAnalyzer);
    // Rewrite project node assignments using new symbols for dereference expressions
    Map<Expression, SymbolReference> mappings = HashBiMap.create(dereferenceAssignments.getMap()).inverse().entrySet().stream().collect(toImmutableMap(Map.Entry::getKey, entry -> entry.getValue().toSymbolReference()));
    Assignments assignments = node.getAssignments().rewrite(expression -> replaceExpression(expression, mappings));
    return Result.ofPlanNode(new ProjectNode(context.getIdAllocator().getNextId(), new ProjectNode(context.getIdAllocator().getNextId(), child.getSource(), Assignments.builder().putAll(child.getAssignments()).putAll(dereferenceAssignments).build()), assignments));
}

Example 43 with Assignments

use of io.trino.sql.planner.plan.Assignments in project trino by trinodb.

the class PushDownDereferenceThroughSemiJoin method apply.

@Override
public Result apply(ProjectNode projectNode, Captures captures, Context context) {
    SemiJoinNode semiJoinNode = captures.get(CHILD);
    // Extract dereferences from project node assignments for pushdown
    Set<SubscriptExpression> dereferences = extractRowSubscripts(projectNode.getAssignments().getExpressions(), false, context.getSession(), typeAnalyzer, context.getSymbolAllocator().getTypes());
    // All dereferences can be assumed on the symbols coming from source, since filteringSource output is not propagated,
    // and semiJoinOutput is of type boolean. We exclude pushdown of dereferences on sourceJoinSymbol.
    dereferences = dereferences.stream().filter(expression -> !getBase(expression).equals(semiJoinNode.getSourceJoinSymbol())).collect(toImmutableSet());
    if (dereferences.isEmpty()) {
        return Result.empty();
    }
    // Create new symbols for dereference expressions
    Assignments dereferenceAssignments = Assignments.of(dereferences, context.getSession(), context.getSymbolAllocator(), typeAnalyzer);
    // Rewrite project node assignments using new symbols for dereference expressions
    Map<Expression, SymbolReference> mappings = HashBiMap.create(dereferenceAssignments.getMap()).inverse().entrySet().stream().collect(toImmutableMap(Map.Entry::getKey, entry -> entry.getValue().toSymbolReference()));
    Assignments assignments = projectNode.getAssignments().rewrite(expression -> replaceExpression(expression, mappings));
    PlanNode newSource = new ProjectNode(context.getIdAllocator().getNextId(), semiJoinNode.getSource(), Assignments.builder().putIdentities(semiJoinNode.getSource().getOutputSymbols()).putAll(dereferenceAssignments).build());
    PlanNode newSemiJoin = semiJoinNode.replaceChildren(ImmutableList.of(newSource, semiJoinNode.getFilteringSource()));
    return Result.ofPlanNode(new ProjectNode(context.getIdAllocator().getNextId(), newSemiJoin, assignments));
}

Example 44 with Assignments

use of io.trino.sql.planner.plan.Assignments in project trino by trinodb.

the class PushDownDereferenceThroughUnnest method apply.

@Override
public Result apply(ProjectNode projectNode, Captures captures, Context context) {
    UnnestNode unnestNode = captures.get(CHILD);
    // Extract dereferences from project node's assignments and unnest node's filter
    ImmutableList.Builder<Expression> expressionsBuilder = ImmutableList.builder();
    expressionsBuilder.addAll(projectNode.getAssignments().getExpressions());
    unnestNode.getFilter().ifPresent(expressionsBuilder::add);
    // Extract dereferences for pushdown
    Set<SubscriptExpression> dereferences = extractRowSubscripts(expressionsBuilder.build(), false, context.getSession(), typeAnalyzer, context.getSymbolAllocator().getTypes());
    // Only retain dereferences on replicate symbols
    dereferences = dereferences.stream().filter(expression -> unnestNode.getReplicateSymbols().contains(getBase(expression))).collect(toImmutableSet());
    if (dereferences.isEmpty()) {
        return Result.empty();
    }
    // Create new symbols for dereference expressions
    Assignments dereferenceAssignments = Assignments.of(dereferences, context.getSession(), context.getSymbolAllocator(), typeAnalyzer);
    // Rewrite project node assignments using new symbols for dereference expressions
    Map<Expression, SymbolReference> mappings = HashBiMap.create(dereferenceAssignments.getMap()).inverse().entrySet().stream().collect(toImmutableMap(Map.Entry::getKey, entry -> entry.getValue().toSymbolReference()));
    Assignments newAssignments = projectNode.getAssignments().rewrite(expression -> replaceExpression(expression, mappings));
    // Create a new ProjectNode (above the original source) adding dereference projections on replicated symbols
    ProjectNode source = new ProjectNode(context.getIdAllocator().getNextId(), unnestNode.getSource(), Assignments.builder().putIdentities(unnestNode.getSource().getOutputSymbols()).putAll(dereferenceAssignments).build());
    // Create projectNode with the new unnest node and assignments with replaced dereferences
    return Result.ofPlanNode(new ProjectNode(context.getIdAllocator().getNextId(), new UnnestNode(context.getIdAllocator().getNextId(), source, ImmutableList.<Symbol>builder().addAll(unnestNode.getReplicateSymbols()).addAll(dereferenceAssignments.getSymbols()).build(), unnestNode.getMappings(), unnestNode.getOrdinalitySymbol(), unnestNode.getJoinType(), unnestNode.getFilter().map(filter -> replaceExpression(filter, mappings))), newAssignments));
}

Example 45 with Assignments

use of io.trino.sql.planner.plan.Assignments in project trino by trinodb.

the class PushDownDereferencesThroughLimit method apply.

@Override
public Result apply(ProjectNode projectNode, Captures captures, Context context) {
    LimitNode limitNode = captures.get(CHILD);
    // Extract dereferences from project node assignments for pushdown
    Set<SubscriptExpression> dereferences = extractRowSubscripts(projectNode.getAssignments().getExpressions(), false, context.getSession(), typeAnalyzer, context.getSymbolAllocator().getTypes());
    // Exclude dereferences on symbols being used in tiesResolvingScheme and requiresPreSortedInputs
    Set<Symbol> excludedSymbols = ImmutableSet.<Symbol>builder().addAll(limitNode.getTiesResolvingScheme().map(OrderingScheme::getOrderBy).orElse(ImmutableList.of())).addAll(limitNode.getPreSortedInputs()).build();
    dereferences = dereferences.stream().filter(expression -> !excludedSymbols.contains(getBase(expression))).collect(toImmutableSet());
    if (dereferences.isEmpty()) {
        return Result.empty();
    }
    // Create new symbols for dereference expressions
    Assignments dereferenceAssignments = Assignments.of(dereferences, context.getSession(), context.getSymbolAllocator(), typeAnalyzer);
    // Rewrite project node assignments using new symbols for dereference expressions
    Map<Expression, SymbolReference> mappings = HashBiMap.create(dereferenceAssignments.getMap()).inverse().entrySet().stream().collect(toImmutableMap(Map.Entry::getKey, entry -> entry.getValue().toSymbolReference()));
    Assignments newAssignments = projectNode.getAssignments().rewrite(expression -> replaceExpression(expression, mappings));
    return Result.ofPlanNode(new ProjectNode(context.getIdAllocator().getNextId(), limitNode.replaceChildren(ImmutableList.of(new ProjectNode(context.getIdAllocator().getNextId(), limitNode.getSource(), Assignments.builder().putIdentities(limitNode.getSource().getOutputSymbols()).putAll(dereferenceAssignments).build()))), newAssignments));
}