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

Example 1 with Assignments

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

the class InlineProjectIntoFilter method apply.

@Override
public Result apply(FilterNode node, Captures captures, Context context) {
    ProjectNode projectNode = captures.get(PROJECTION);
    List<Expression> filterConjuncts = extractConjuncts(node.getPredicate());
    Map<Boolean, List<Expression>> conjuncts = filterConjuncts.stream().collect(partitioningBy(SymbolReference.class::isInstance));
    List<Expression> simpleConjuncts = conjuncts.get(true);
    List<Expression> complexConjuncts = conjuncts.get(false);
    // Do not inline expression if the symbol is used multiple times in simple conjuncts.
    Set<Expression> simpleUniqueConjuncts = simpleConjuncts.stream().collect(Collectors.groupingBy(identity(), counting())).entrySet().stream().filter(entry -> entry.getValue() == 1).map(Map.Entry::getKey).collect(toImmutableSet());
    Set<Expression> complexConjunctSymbols = SymbolsExtractor.extractUnique(complexConjuncts).stream().map(Symbol::toSymbolReference).collect(toImmutableSet());
    // Do not inline expression if the symbol is used in complex conjuncts.
    Set<Expression> simpleConjunctsToInline = Sets.difference(simpleUniqueConjuncts, complexConjunctSymbols);
    if (simpleConjunctsToInline.isEmpty()) {
        return Result.empty();
    }
    ImmutableList.Builder<Expression> newConjuncts = ImmutableList.builder();
    Assignments.Builder newAssignments = Assignments.builder();
    Assignments.Builder postFilterAssignmentsBuilder = Assignments.builder();
    for (Expression conjunct : filterConjuncts) {
        if (simpleConjunctsToInline.contains(conjunct)) {
            Expression expression = projectNode.getAssignments().get(Symbol.from(conjunct));
            if (expression == null || expression instanceof SymbolReference) {
                // expression == null -> The symbol is not produced by the underlying projection (i.e. it is a correlation symbol).
                // expression instanceof SymbolReference -> Do not inline trivial projections.
                newConjuncts.add(conjunct);
            } else {
                newConjuncts.add(expression);
                newAssignments.putIdentities(SymbolsExtractor.extractUnique(expression));
                postFilterAssignmentsBuilder.put(Symbol.from(conjunct), TRUE_LITERAL);
            }
        } else {
            newConjuncts.add(conjunct);
        }
    }
    Assignments postFilterAssignments = postFilterAssignmentsBuilder.build();
    if (postFilterAssignments.isEmpty()) {
        return Result.empty();
    }
    Set<Symbol> postFilterSymbols = postFilterAssignments.getSymbols();
    // Remove inlined expressions from the underlying projection.
    newAssignments.putAll(projectNode.getAssignments().filter(symbol -> !postFilterSymbols.contains(symbol)));
    Map<Symbol, Expression> outputAssignments = new HashMap<>();
    outputAssignments.putAll(Assignments.identity(node.getOutputSymbols()).getMap());
    // Restore inlined symbols.
    outputAssignments.putAll(postFilterAssignments.getMap());
    return Result.ofPlanNode(new ProjectNode(context.getIdAllocator().getNextId(), new FilterNode(node.getId(), new ProjectNode(projectNode.getId(), projectNode.getSource(), newAssignments.build()), combineConjuncts(metadata, newConjuncts.build())), Assignments.builder().putAll(outputAssignments).build()));
}

Example 2 with Assignments

use of io.trino.sql.planner.plan.Assignments 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();
}

Example 3 with Assignments

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

the class PushDownDereferencesThroughWindow method apply.

@Override
public Result apply(ProjectNode projectNode, Captures captures, Context context) {
    WindowNode windowNode = captures.get(CHILD);
    // Extract dereferences for pushdown
    Set<SubscriptExpression> dereferences = extractRowSubscripts(ImmutableList.<Expression>builder().addAll(projectNode.getAssignments().getExpressions()).addAll(windowNode.getWindowFunctions().values().stream().flatMap(function -> function.getArguments().stream()).collect(toImmutableList())).build(), false, context.getSession(), typeAnalyzer, context.getSymbolAllocator().getTypes());
    WindowNode.Specification specification = windowNode.getSpecification();
    dereferences = dereferences.stream().filter(expression -> {
        Symbol symbol = getBase(expression);
        // Exclude partitionBy, orderBy and synthesized symbols
        return !specification.getPartitionBy().contains(symbol) && !specification.getOrderingScheme().map(OrderingScheme::getOrderBy).orElse(ImmutableList.of()).contains(symbol) && !windowNode.getCreatedSymbols().contains(symbol);
    }).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(), new WindowNode(windowNode.getId(), new ProjectNode(context.getIdAllocator().getNextId(), windowNode.getSource(), Assignments.builder().putIdentities(windowNode.getSource().getOutputSymbols()).putAll(dereferenceAssignments).build()), windowNode.getSpecification(), // Replace dereference expressions in functions
    windowNode.getWindowFunctions().entrySet().stream().collect(toImmutableMap(Map.Entry::getKey, entry -> {
        WindowNode.Function oldFunction = entry.getValue();
        return new WindowNode.Function(oldFunction.getResolvedFunction(), oldFunction.getArguments().stream().map(expression -> replaceExpression(expression, mappings)).collect(toImmutableList()), oldFunction.getFrame(), oldFunction.isIgnoreNulls());
    })), windowNode.getHashSymbol(), windowNode.getPrePartitionedInputs(), windowNode.getPreSortedOrderPrefix()), newAssignments));
}

Example 4 with Assignments

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

the class PushDownProjectionsFromPatternRecognition method apply.

@Override
public Result apply(PatternRecognitionNode node, Captures captures, Context context) {
    Assignments.Builder assignments = Assignments.builder();
    Map<IrLabel, ExpressionAndValuePointers> rewrittenVariableDefinitions = rewriteVariableDefinitions(node.getVariableDefinitions(), assignments, context);
    Map<Symbol, Measure> rewrittenMeasureDefinitions = rewriteMeasureDefinitions(node.getMeasures(), assignments, context);
    if (assignments.build().isEmpty()) {
        return Result.empty();
    }
    assignments.putIdentities(node.getSource().getOutputSymbols());
    ProjectNode projectNode = new ProjectNode(context.getIdAllocator().getNextId(), node.getSource(), assignments.build());
    PatternRecognitionNode patternRecognitionNode = new PatternRecognitionNode(node.getId(), projectNode, node.getSpecification(), node.getHashSymbol(), node.getPrePartitionedInputs(), node.getPreSortedOrderPrefix(), node.getWindowFunctions(), rewrittenMeasureDefinitions, node.getCommonBaseFrame(), node.getRowsPerMatch(), node.getSkipToLabel(), node.getSkipToPosition(), node.isInitial(), node.getPattern(), node.getSubsets(), rewrittenVariableDefinitions);
    return Result.ofPlanNode(restrictOutputs(context.getIdAllocator(), patternRecognitionNode, ImmutableSet.copyOf(node.getOutputSymbols())).orElse(patternRecognitionNode));
}

Example 5 with Assignments

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

the class PushJoinIntoTableScan method apply.

@Override
public Result apply(JoinNode joinNode, Captures captures, Context context) {
    if (joinNode.isCrossJoin()) {
        return Result.empty();
    }
    TableScanNode left = captures.get(LEFT_TABLE_SCAN);
    TableScanNode right = captures.get(RIGHT_TABLE_SCAN);
    verify(!left.isUpdateTarget() && !right.isUpdateTarget(), "Unexpected Join over for-update table scan");
    Expression effectiveFilter = getEffectiveFilter(joinNode);
    FilterSplitResult filterSplitResult = splitFilter(effectiveFilter, left.getOutputSymbols(), right.getOutputSymbols(), context);
    if (!filterSplitResult.getRemainingFilter().equals(BooleanLiteral.TRUE_LITERAL)) {
        // TODO add extra filter node above join
        return Result.empty();
    }
    if (left.getEnforcedConstraint().isNone() || right.getEnforcedConstraint().isNone()) {
        // enforced constraint harder below.
        return Result.empty();
    }
    Map<String, ColumnHandle> leftAssignments = left.getAssignments().entrySet().stream().collect(toImmutableMap(entry -> entry.getKey().getName(), Map.Entry::getValue));
    Map<String, ColumnHandle> rightAssignments = right.getAssignments().entrySet().stream().collect(toImmutableMap(entry -> entry.getKey().getName(), Map.Entry::getValue));
    /*
         * We are (lazily) computing estimated statistics for join node and left and right table
         * and passing those to connector via applyJoin.
         *
         * There are a couple reasons for this approach:
         * - the engine knows how to estimate join and connector may not
         * - the engine may have cached stats for the table scans (within context.getStatsProvider()), so can be able to provide information more inexpensively
         * - in the future, the engine may be able to provide stats for table scan even in case when connector no longer can (see https://github.com/trinodb/trino/issues/6998)
         * - the pushdown feasibility assessment logic may be different (or configured differently) for different connectors/catalogs.
         */
    JoinStatistics joinStatistics = getJoinStatistics(joinNode, left, right, context);
    Optional<JoinApplicationResult<TableHandle>> joinApplicationResult = metadata.applyJoin(context.getSession(), getJoinType(joinNode), left.getTable(), right.getTable(), filterSplitResult.getPushableConditions(), // TODO we could pass only subset of assignments here, those which are needed to resolve filterSplitResult.getPushableConditions
    leftAssignments, rightAssignments, joinStatistics);
    if (joinApplicationResult.isEmpty()) {
        return Result.empty();
    }
    TableHandle handle = joinApplicationResult.get().getTableHandle();
    Map<ColumnHandle, ColumnHandle> leftColumnHandlesMapping = joinApplicationResult.get().getLeftColumnHandles();
    Map<ColumnHandle, ColumnHandle> rightColumnHandlesMapping = joinApplicationResult.get().getRightColumnHandles();
    ImmutableMap.Builder<Symbol, ColumnHandle> assignmentsBuilder = ImmutableMap.builder();
    assignmentsBuilder.putAll(left.getAssignments().entrySet().stream().collect(toImmutableMap(Map.Entry::getKey, entry -> leftColumnHandlesMapping.get(entry.getValue()))));
    assignmentsBuilder.putAll(right.getAssignments().entrySet().stream().collect(toImmutableMap(Map.Entry::getKey, entry -> rightColumnHandlesMapping.get(entry.getValue()))));
    Map<Symbol, ColumnHandle> assignments = assignmentsBuilder.buildOrThrow();
    // convert enforced constraint
    JoinNode.Type joinType = joinNode.getType();
    TupleDomain<ColumnHandle> leftConstraint = deriveConstraint(left.getEnforcedConstraint(), leftColumnHandlesMapping, joinType == RIGHT || joinType == FULL);
    TupleDomain<ColumnHandle> rightConstraint = deriveConstraint(right.getEnforcedConstraint(), rightColumnHandlesMapping, joinType == LEFT || joinType == FULL);
    TupleDomain<ColumnHandle> newEnforcedConstraint = TupleDomain.withColumnDomains(ImmutableMap.<ColumnHandle, Domain>builder().putAll(leftConstraint.getDomains().orElseThrow()).putAll(rightConstraint.getDomains().orElseThrow()).buildOrThrow());
    return Result.ofPlanNode(new ProjectNode(context.getIdAllocator().getNextId(), new TableScanNode(joinNode.getId(), handle, ImmutableList.copyOf(assignments.keySet()), assignments, newEnforcedConstraint, deriveTableStatisticsForPushdown(context.getStatsProvider(), context.getSession(), joinApplicationResult.get().isPrecalculateStatistics(), joinNode), false, Optional.empty()), Assignments.identity(joinNode.getOutputSymbols())));
}