Examples with Captures io.prestosql.matching.Captures used on opensource projects

Example 1 with Captures

use of io.prestosql.matching.Captures in project hetu-core by openlookeng.

the class InlineProjections method apply.

@Override
public Result apply(ProjectNode parent, Captures captures, Context context) {
    ProjectNode child = captures.get(CHILD);
    Sets.SetView<Symbol> targets = extractInliningTargets(parent, child);
    if (targets.isEmpty()) {
        return Result.empty();
    }
    // inline the expressions
    Assignments assignments = child.getAssignments().filter(targets::contains);
    Map<Symbol, RowExpression> parentAssignments = parent.getAssignments().entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, entry -> (inlineReferences(entry.getValue(), assignments))));
    // 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<Symbol> inputs = child.getAssignments().entrySet().stream().filter(entry -> targets.contains(entry.getKey())).map(Map.Entry::getValue).flatMap(entry -> extractInputs(entry).stream()).collect(toSet());
    Assignments.Builder childAssignments = Assignments.builder();
    for (Map.Entry<Symbol, 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 (Symbol input : inputs) {
        if (allTranslated) {
            Type inputType = context.getSymbolAllocator().getSymbols().get(input);
            childAssignments.put(input, new VariableReferenceExpression(input.getName(), inputType));
        } else {
            childAssignments.put(input, castToRowExpression(toSymbolReference(input)));
        }
    }
    return Result.ofPlanNode(new ProjectNode(parent.getId(), new ProjectNode(child.getId(), child.getSource(), childAssignments.build()), Assignments.copyOf(parentAssignments)));
}

Example 2 with Captures

use of io.prestosql.matching.Captures in project hetu-core by openlookeng.

the class PushProjectionIntoTableScan method apply.

@Override
public Result apply(ProjectNode project, Captures captures, Context context) {
    TableScanNode tableScan = captures.get(TABLE_SCAN);
    List<ConnectorExpression> projections;
    try {
        projections = project.getAssignments().getExpressions().stream().map(expression -> ConnectorExpressionTranslator.translate(expression)).collect(toImmutableList());
    } catch (UnsupportedOperationException e) {
        // rewritten in terms of the new assignments for the columns passed in #2
        return Result.empty();
    }
    Map<String, ColumnHandle> assignments = tableScan.getAssignments().entrySet().stream().collect(toImmutableMap(entry -> entry.getKey().getName(), Map.Entry::getValue));
    Optional<ProjectionApplicationResult<TableHandle>> result = metadata.applyProjection(context.getSession(), tableScan.getTable(), projections, assignments);
    if (!result.isPresent()) {
        return Result.empty();
    }
    List<Symbol> newScanOutputs = new ArrayList<>();
    Map<Symbol, ColumnHandle> newScanAssignments = new HashMap<>();
    Map<String, Symbol> variableMappings = new HashMap<>();
    for (ProjectionApplicationResult.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);
    }
    // TODO: ensure newProjections.size == original projections.size
    List<RowExpression> newProjections = result.get().getProjections().stream().map(expression -> ConnectorExpressionTranslator.translate(expression, variableMappings, new LiteralEncoder(metadata))).collect(toImmutableList());
    Assignments.Builder newProjectionAssignments = Assignments.builder();
    for (int i = 0; i < project.getOutputSymbols().size(); i++) {
        newProjectionAssignments.put(project.getOutputSymbols().get(i), newProjections.get(i));
    }
    return Result.ofPlanNode(new ProjectNode(context.getIdAllocator().getNextId(), TableScanNode.newInstance(tableScan.getId(), result.get().getHandle(), newScanOutputs, newScanAssignments, tableScan.getStrategy(), tableScan.getReuseTableScanMappingId(), tableScan.getConsumerTableScanNodeCount(), tableScan.isForDelete()), newProjectionAssignments.build()));
}

Example 3 with Captures

use of io.prestosql.matching.Captures in project hetu-core by openlookeng.

the class PushProjectionThroughExchange method apply.

@Override
public Result apply(ProjectNode project, Captures captures, Context context) {
    ExchangeNode exchange = captures.get(CHILD);
    Set<Symbol> partitioningColumns = exchange.getPartitioningScheme().getPartitioning().getColumns();
    ImmutableList.Builder<PlanNode> newSourceBuilder = ImmutableList.builder();
    ImmutableList.Builder<List<Symbol>> inputsBuilder = ImmutableList.builder();
    for (int i = 0; i < exchange.getSources().size(); i++) {
        Map<Symbol, VariableReferenceExpression> outputToInputMap = extractExchangeOutputToInput(exchange, i, context.getSymbolAllocator().getTypes());
        Assignments.Builder projections = Assignments.builder();
        ImmutableList.Builder<Symbol> inputs = ImmutableList.builder();
        // Need to retain the partition keys for the exchange
        partitioningColumns.stream().map(outputToInputMap::get).forEach(nameReference -> {
            Symbol symbol = new Symbol(nameReference.getName());
            projections.put(symbol, nameReference);
            inputs.add(symbol);
        });
        if (exchange.getPartitioningScheme().getHashColumn().isPresent()) {
            // Need to retain the hash symbol for the exchange
            projections.put(exchange.getPartitioningScheme().getHashColumn().get(), toVariableReference(exchange.getPartitioningScheme().getHashColumn().get(), context.getSymbolAllocator().getTypes()));
            inputs.add(exchange.getPartitioningScheme().getHashColumn().get());
        }
        if (exchange.getOrderingScheme().isPresent()) {
            // need to retain ordering columns for the exchange
            exchange.getOrderingScheme().get().getOrderBy().stream().filter(symbol -> !partitioningColumns.contains(symbol)).map(outputToInputMap::get).forEach(nameReference -> {
                Symbol symbol = new Symbol(nameReference.getName());
                projections.put(symbol, nameReference);
                inputs.add(symbol);
            });
        }
        for (Map.Entry<Symbol, RowExpression> projection : project.getAssignments().entrySet()) {
            checkArgument(!isExpression(projection.getValue()), "Cannot contain OriginalExpression after AddExchange");
            Map<VariableReferenceExpression, VariableReferenceExpression> variableOutputToInputMap = new LinkedHashMap<>();
            outputToInputMap.forEach(((symbol, variable) -> variableOutputToInputMap.put(new VariableReferenceExpression(symbol.getName(), context.getSymbolAllocator().getTypes().get(symbol)), variable)));
            RowExpression translatedExpression = RowExpressionVariableInliner.inlineVariables(variableOutputToInputMap, projection.getValue());
            Symbol symbol = context.getSymbolAllocator().newSymbol(translatedExpression);
            projections.put(symbol, translatedExpression);
            inputs.add(symbol);
        }
        newSourceBuilder.add(new ProjectNode(context.getIdAllocator().getNextId(), exchange.getSources().get(i), projections.build()));
        inputsBuilder.add(inputs.build());
    }
    // Construct the output symbols in the same order as the sources
    ImmutableList.Builder<Symbol> outputBuilder = ImmutableList.builder();
    partitioningColumns.forEach(outputBuilder::add);
    exchange.getPartitioningScheme().getHashColumn().ifPresent(outputBuilder::add);
    if (exchange.getOrderingScheme().isPresent()) {
        exchange.getOrderingScheme().get().getOrderBy().stream().filter(symbol -> !partitioningColumns.contains(symbol)).forEach(outputBuilder::add);
    }
    for (Map.Entry<Symbol, 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.getId(), exchange.getType(), exchange.getScope(), partitioningScheme, newSourceBuilder.build(), inputsBuilder.build(), exchange.getOrderingScheme(), AggregationNode.AggregationType.HASH);
    // we need to strip unnecessary symbols (hash, partitioning columns).
    return Result.ofPlanNode(restrictOutputs(context.getIdAllocator(), result, ImmutableSet.copyOf(project.getOutputSymbols()), true, context.getSymbolAllocator().getTypes()).orElse(result));
}

Example 4 with Captures

use of io.prestosql.matching.Captures in project hetu-core by openlookeng.

the class TransformCorrelatedLateralJoinToJoin method apply.

@Override
public Result apply(LateralJoinNode lateralJoinNode, Captures captures, Context context) {
    PlanNode subquery = lateralJoinNode.getSubquery();
    PlanNodeDecorrelator planNodeDecorrelator = new PlanNodeDecorrelator(context.getIdAllocator(), context.getLookup());
    Optional<DecorrelatedNode> decorrelatedNodeOptional = planNodeDecorrelator.decorrelateFilters(subquery, lateralJoinNode.getCorrelation());
    return decorrelatedNodeOptional.map(decorrelatedNode -> {
        Expression joinFilter = combineConjuncts(decorrelatedNode.getCorrelatedPredicates().orElse(TRUE_LITERAL), lateralJoinNode.getFilter());
        return Result.ofPlanNode(new JoinNode(context.getIdAllocator().getNextId(), lateralJoinNode.getType().toJoinNodeType(), lateralJoinNode.getInput(), decorrelatedNode.getNode(), ImmutableList.of(), lateralJoinNode.getOutputSymbols(), joinFilter.equals(TRUE_LITERAL) ? Optional.empty() : Optional.of(castToRowExpression(joinFilter)), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), ImmutableMap.of()));
    }).orElseGet(Result::empty);
}

Example 5 with Captures

use of io.prestosql.matching.Captures in project hetu-core by openlookeng.

the class SingleDistinctAggregationToGroupBy method apply.

@Override
public Result apply(AggregationNode aggregation, Captures captures, Context context) {
    List<Set<Expression>> argumentSets = extractArgumentSets(aggregation).collect(Collectors.toList());
    Set<Symbol> symbols = Iterables.getOnlyElement(argumentSets).stream().map(SymbolUtils::from).collect(Collectors.toSet());
    return Result.ofPlanNode(new AggregationNode(aggregation.getId(), new AggregationNode(context.getIdAllocator().getNextId(), aggregation.getSource(), ImmutableMap.of(), singleGroupingSet(ImmutableList.<Symbol>builder().addAll(aggregation.getGroupingKeys()).addAll(symbols).build()), ImmutableList.of(), SINGLE, Optional.empty(), Optional.empty(), aggregation.getAggregationType(), aggregation.getFinalizeSymbol()), // 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.getHashSymbol(), aggregation.getGroupIdSymbol(), aggregation.getAggregationType(), aggregation.getFinalizeSymbol()));
}