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

Example 1 with UnionNode

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

the class PushProjectionThroughUnion method apply.

@Override
public Result apply(ProjectNode parent, Captures captures, Context context) {
    UnionNode source = captures.get(CHILD);
    // OutputLayout of the resultant Union, will be same as the layout of the Project
    List<Symbol> outputLayout = parent.getOutputSymbols();
    // Mapping from the output symbol to ordered list of symbols from each of the sources
    ImmutableListMultimap.Builder<Symbol, Symbol> mappings = ImmutableListMultimap.builder();
    // sources for the resultant UnionNode
    ImmutableList.Builder<PlanNode> outputSources = ImmutableList.builder();
    for (int i = 0; i < source.getSources().size(); i++) {
        // Map: output of union -> input of this source to the union
        Map<Symbol, SymbolReference> outputToInput = source.sourceSymbolMap(i);
        // assignments for the new ProjectNode
        Assignments.Builder assignments = Assignments.builder();
        // mapping from current ProjectNode to new ProjectNode, used to identify the output layout
        Map<Symbol, Symbol> projectSymbolMapping = new HashMap<>();
        // Translate the assignments in the ProjectNode using symbols of the source of the UnionNode
        for (Map.Entry<Symbol, Expression> entry : parent.getAssignments().entrySet()) {
            Expression translatedExpression = inlineSymbols(outputToInput, entry.getValue());
            Type type = context.getSymbolAllocator().getTypes().get(entry.getKey());
            Symbol symbol = context.getSymbolAllocator().newSymbol(translatedExpression, type);
            assignments.put(symbol, translatedExpression);
            projectSymbolMapping.put(entry.getKey(), symbol);
        }
        outputSources.add(new ProjectNode(context.getIdAllocator().getNextId(), source.getSources().get(i), assignments.build()));
        outputLayout.forEach(symbol -> mappings.put(symbol, projectSymbolMapping.get(symbol)));
    }
    return Result.ofPlanNode(new UnionNode(parent.getId(), outputSources.build(), mappings.build(), ImmutableList.copyOf(mappings.build().keySet())));
}

Example 2 with UnionNode

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

the class SetOperationMerge method mergeFirstSource.

/**
 * Only merge first source node. This method is assumed to be used for EXCEPT, which is a non-associative set operation.
 * Provides a correct plan transformation also for associative set operations: UNION and INTERSECT.
 *
 * @return Merged plan node if applied.
 */
public Optional<SetOperationNode> mergeFirstSource() {
    Lookup lookup = context.getLookup();
    List<PlanNode> sources = node.getSources().stream().flatMap(lookup::resolveGroup).collect(Collectors.toList());
    PlanNode child = sources.get(0);
    // Determine if set operations can be merged and whether the resulting set operation is quantified DISTINCT or ALL
    Optional<Boolean> mergedQuantifier = mergedQuantifierIsDistinct(node, child);
    if (mergedQuantifier.isEmpty()) {
        return Optional.empty();
    }
    ImmutableListMultimap.Builder<Symbol, Symbol> newMappingsBuilder = ImmutableListMultimap.builder();
    // Merge all sources of the first source.
    addMergedMappings((SetOperationNode) child, 0, newMappingsBuilder);
    // Keep remaining as it is
    for (int i = 1; i < sources.size(); i++) {
        PlanNode source = sources.get(i);
        addOriginalMappings(source, i, newMappingsBuilder);
    }
    if (node instanceof UnionNode) {
        return Optional.of(new UnionNode(node.getId(), newSources, newMappingsBuilder.build(), node.getOutputSymbols()));
    }
    if (node instanceof IntersectNode) {
        return Optional.of(new IntersectNode(node.getId(), newSources, newMappingsBuilder.build(), node.getOutputSymbols(), mergedQuantifier.get()));
    }
    if (node instanceof ExceptNode) {
        return Optional.of(new ExceptNode(node.getId(), newSources, newMappingsBuilder.build(), node.getOutputSymbols(), mergedQuantifier.get()));
    }
    throw new IllegalArgumentException("unexpected node type: " + node.getClass().getSimpleName());
}

Example 3 with UnionNode

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

the class PushTableWriteThroughUnion method apply.

@Override
public Result apply(TableWriterNode writerNode, Captures captures, Context context) {
    UnionNode unionNode = captures.get(CHILD);
    ImmutableList.Builder<PlanNode> rewrittenSources = ImmutableList.builder();
    List<Map<Symbol, Symbol>> sourceMappings = new ArrayList<>();
    for (int source = 0; source < unionNode.getSources().size(); source++) {
        rewrittenSources.add(rewriteSource(writerNode, unionNode, source, sourceMappings, context));
    }
    ImmutableListMultimap.Builder<Symbol, Symbol> unionMappings = ImmutableListMultimap.builder();
    sourceMappings.forEach(mappings -> mappings.forEach(unionMappings::put));
    return Result.ofPlanNode(new UnionNode(context.getIdAllocator().getNextId(), rewrittenSources.build(), unionMappings.build(), ImmutableList.copyOf(unionMappings.build().keySet())));
}

Example 4 with UnionNode

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

the class PushTopNThroughUnion method apply.

@Override
public Result apply(TopNNode topNNode, Captures captures, Context context) {
    UnionNode unionNode = captures.get(CHILD);
    ImmutableList.Builder<PlanNode> sources = ImmutableList.builder();
    for (PlanNode source : unionNode.getSources()) {
        SymbolMapper.Builder symbolMapper = SymbolMapper.builder();
        Set<Symbol> sourceOutputSymbols = ImmutableSet.copyOf(source.getOutputSymbols());
        for (Symbol unionOutput : unionNode.getOutputSymbols()) {
            Set<Symbol> inputSymbols = ImmutableSet.copyOf(unionNode.getSymbolMapping().get(unionOutput));
            Symbol unionInput = getLast(intersection(inputSymbols, sourceOutputSymbols));
            symbolMapper.put(unionOutput, unionInput);
        }
        sources.add(symbolMapper.build().map(topNNode, source, context.getIdAllocator().getNextId()));
    }
    return Result.ofPlanNode(new UnionNode(unionNode.getId(), sources.build(), unionNode.getSymbolMapping(), unionNode.getOutputSymbols()));
}

Example 5 with UnionNode

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

the class TestEffectivePredicateExtractor method testUnion.

@Test
public void testUnion() {
    ImmutableListMultimap<Symbol, Symbol> symbolMapping = ImmutableListMultimap.of(A, B, A, C, A, E);
    PlanNode node = new UnionNode(newId(), ImmutableList.of(filter(baseTableScan, greaterThan(AE, bigintLiteral(10))), filter(baseTableScan, and(greaterThan(AE, bigintLiteral(10)), lessThan(AE, bigintLiteral(100)))), filter(baseTableScan, and(greaterThan(AE, bigintLiteral(10)), lessThan(AE, bigintLiteral(100))))), symbolMapping, ImmutableList.copyOf(symbolMapping.keySet()));
    Expression effectivePredicate = effectivePredicateExtractor.extract(SESSION, node, TypeProvider.empty(), typeAnalyzer);
    // Only the common conjuncts can be inferred through a Union
    assertEquals(normalizeConjuncts(effectivePredicate), normalizeConjuncts(greaterThan(AE, bigintLiteral(10))));
}