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

Example 6 with UnionNode

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

the class TestCostCalculator method testUnion.

@Test
public void testUnion() {
    TableScanNode ts1 = tableScan("ts1", "orderkey");
    TableScanNode ts2 = tableScan("ts2", "orderkey_0");
    ImmutableListMultimap.Builder<Symbol, Symbol> outputMappings = ImmutableListMultimap.builder();
    outputMappings.put(new Symbol("orderkey_1"), new Symbol("orderkey"));
    outputMappings.put(new Symbol("orderkey_1"), new Symbol("orderkey_0"));
    UnionNode union = new UnionNode(new PlanNodeId("union"), ImmutableList.of(ts1, ts2), outputMappings.build(), ImmutableList.of(new Symbol("orderkey_1")));
    Map<String, PlanNodeStatsEstimate> stats = ImmutableMap.of("ts1", statsEstimate(ts1, 4000), "ts2", statsEstimate(ts2, 1000), "union", statsEstimate(ts1, 5000));
    Map<String, PlanCostEstimate> costs = ImmutableMap.of("ts1", cpuCost(1000), "ts2", cpuCost(1000));
    Map<String, Type> types = ImmutableMap.of("orderkey", BIGINT, "orderkey_0", BIGINT, "orderkey_1", BIGINT);
    assertCost(union, costs, stats, types).cpu(2000).memory(0).network(0);
    assertCostEstimatedExchanges(union, costs, stats, types).cpu(2000).memory(0).network(5000 * IS_NULL_OVERHEAD);
}

Example 7 with UnionNode

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

the class PushLimitThroughUnion method apply.

@Override
public Result apply(LimitNode parent, Captures captures, Context context) {
    UnionNode unionNode = captures.get(CHILD);
    ImmutableList.Builder<PlanNode> builder = ImmutableList.builder();
    boolean shouldApply = false;
    for (PlanNode source : unionNode.getSources()) {
        // This check is to ensure that we don't fire the optimizer if it was previously applied.
        if (isAtMost(source, context.getLookup(), parent.getCount())) {
            builder.add(source);
        } else {
            shouldApply = true;
            builder.add(new LimitNode(context.getIdAllocator().getNextId(), source, parent.getCount(), true));
        }
    }
    if (!shouldApply) {
        return Result.empty();
    }
    return Result.ofPlanNode(parent.replaceChildren(ImmutableList.of(unionNode.replaceChildren(builder.build()))));
}

Example 8 with UnionNode

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

the class SetOperationMerge method merge.

/**
 * Merge all matching source nodes. This method is assumed to be used only for associative set operations: UNION and INTERSECT.
 *
 * @return Merged plan node if applied.
 */
public Optional<SetOperationNode> merge() {
    checkState(node instanceof UnionNode || node instanceof IntersectNode, "unexpected node type: %s", node.getClass().getSimpleName());
    Lookup lookup = context.getLookup();
    List<PlanNode> sources = node.getSources().stream().flatMap(lookup::resolveGroup).collect(Collectors.toList());
    ImmutableListMultimap.Builder<Symbol, Symbol> newMappingsBuilder = ImmutableListMultimap.builder();
    boolean resultIsDistinct = false;
    boolean rewritten = false;
    for (int i = 0; i < sources.size(); i++) {
        PlanNode source = sources.get(i);
        // Determine if set operations can be merged and whether the resulting set operation is quantified DISTINCT or ALL
        Optional<Boolean> mergedQuantifier = mergedQuantifierIsDistinct(node, source);
        if (mergedQuantifier.isPresent()) {
            addMergedMappings((SetOperationNode) source, i, newMappingsBuilder);
            resultIsDistinct |= mergedQuantifier.get();
            rewritten = true;
        } else {
            // Keep mapping as it is
            addOriginalMappings(source, i, newMappingsBuilder);
        }
    }
    if (!rewritten) {
        return Optional.empty();
    }
    if (node instanceof UnionNode) {
        return Optional.of(new UnionNode(node.getId(), newSources, newMappingsBuilder.build(), node.getOutputSymbols()));
    }
    return Optional.of(new IntersectNode(node.getId(), newSources, newMappingsBuilder.build(), node.getOutputSymbols(), resultIsDistinct));
}

Example 9 with UnionNode

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

the class SetOperationNodeTranslator method makeSetContainmentPlanForAll.

public TranslationResult makeSetContainmentPlanForAll(SetOperationNode node) {
    checkArgument(!(node instanceof UnionNode), "Cannot simplify a UnionNode");
    List<Symbol> markers = allocateSymbols(node.getSources().size(), MARKER, BOOLEAN);
    // identity projection for all the fields in each of the sources plus marker columns
    List<PlanNode> withMarkers = appendMarkers(markers, node.getSources(), node);
    // add a union over all the rewritten sources
    List<Symbol> outputs = node.getOutputSymbols();
    UnionNode union = union(withMarkers, ImmutableList.copyOf(concat(outputs, markers)));
    // add counts and row number
    List<Symbol> countOutputs = allocateSymbols(markers.size(), "count", BIGINT);
    Symbol rowNumberSymbol = symbolAllocator.newSymbol("row_number", BIGINT);
    WindowNode window = appendCounts(union, outputs, markers, countOutputs, rowNumberSymbol);
    // prune markers
    ProjectNode project = new ProjectNode(idAllocator.getNextId(), window, Assignments.identity(ImmutableList.copyOf(concat(outputs, countOutputs, ImmutableList.of(rowNumberSymbol)))));
    return new TranslationResult(project, countOutputs, Optional.of(rowNumberSymbol));
}

Example 10 with UnionNode

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

the class SetOperationNodeTranslator method makeSetContainmentPlanForDistinct.

public TranslationResult makeSetContainmentPlanForDistinct(SetOperationNode node) {
    checkArgument(!(node instanceof UnionNode), "Cannot simplify a UnionNode");
    List<Symbol> markers = allocateSymbols(node.getSources().size(), MARKER, BOOLEAN);
    // identity projection for all the fields in each of the sources plus marker columns
    List<PlanNode> withMarkers = appendMarkers(markers, node.getSources(), node);
    // add a union over all the rewritten sources. The outputs of the union have the same name as the
    // original intersect node
    List<Symbol> outputs = node.getOutputSymbols();
    UnionNode union = union(withMarkers, ImmutableList.copyOf(concat(outputs, markers)));
    // add count aggregations
    List<Symbol> aggregationOutputs = allocateSymbols(markers.size(), "count", BIGINT);
    AggregationNode aggregation = computeCounts(union, outputs, markers, aggregationOutputs);
    return new TranslationResult(aggregation, aggregationOutputs);
}