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

Example 11 with UnionNode

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

the class RemoveEmptyUnionBranches method apply.

@Override
public Result apply(UnionNode node, Captures captures, Context context) {
    long emptyBranches = node.getSources().stream().filter(source -> isEmpty(source, context.getLookup())).count();
    if (emptyBranches == 0) {
        return Result.empty();
    }
    if (emptyBranches == node.getSources().size()) {
        return Result.ofPlanNode(new ValuesNode(node.getId(), node.getOutputSymbols(), ImmutableList.of()));
    }
    ImmutableList.Builder<PlanNode> newSourcesBuilder = ImmutableList.builder();
    ImmutableListMultimap.Builder<Symbol, Symbol> outputsToInputsBuilder = ImmutableListMultimap.builder();
    for (int i = 0; i < node.getSources().size(); i++) {
        PlanNode source = node.getSources().get(i);
        if (!isEmpty(source, context.getLookup())) {
            newSourcesBuilder.add(source);
            for (Symbol column : node.getOutputSymbols()) {
                outputsToInputsBuilder.put(column, node.getSymbolMapping().get(column).get(i));
            }
        }
    }
    List<PlanNode> newSources = newSourcesBuilder.build();
    ListMultimap<Symbol, Symbol> outputsToInputs = outputsToInputsBuilder.build();
    if (newSources.size() == 1) {
        Assignments.Builder assignments = Assignments.builder();
        outputsToInputs.entries().stream().forEach(entry -> assignments.put(entry.getKey(), entry.getValue().toSymbolReference()));
        return Result.ofPlanNode(new ProjectNode(node.getId(), newSources.get(0), assignments.build()));
    }
    return Result.ofPlanNode(new UnionNode(node.getId(), newSources, outputsToInputs, node.getOutputSymbols()));
}

Example 12 with UnionNode

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

the class TestTypeValidator method testInvalidUnion.

@Test
public void testInvalidUnion() {
    Symbol outputSymbol = symbolAllocator.newSymbol("output", DATE);
    ListMultimap<Symbol, Symbol> mappings = ImmutableListMultimap.<Symbol, Symbol>builder().put(outputSymbol, columnD).put(outputSymbol, // should be a symbol with DATE type
    columnA).build();
    PlanNode node = new UnionNode(newId(), ImmutableList.of(baseTableScan, baseTableScan), mappings, ImmutableList.copyOf(mappings.keySet()));
    assertThatThrownBy(() -> assertTypesValid(node)).isInstanceOf(IllegalArgumentException.class).hasMessageMatching("type of symbol 'output(_[0-9]+)?' is expected to be date, but the actual type is bigint");
}

Example 13 with UnionNode

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

the class TestTypeValidator method testValidUnion.

@Test
public void testValidUnion() {
    Symbol outputSymbol = symbolAllocator.newSymbol("output", DATE);
    ListMultimap<Symbol, Symbol> mappings = ImmutableListMultimap.<Symbol, Symbol>builder().put(outputSymbol, columnD).put(outputSymbol, columnD).build();
    PlanNode node = new UnionNode(newId(), ImmutableList.of(baseTableScan, baseTableScan), mappings, ImmutableList.copyOf(mappings.keySet()));
    assertTypesValid(node);
}

Example 14 with UnionNode

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

the class RelationPlanner method visitUnion.

@Override
protected RelationPlan visitUnion(Union node, Void context) {
    checkArgument(!node.getRelations().isEmpty(), "No relations specified for UNION");
    SetOperationPlan setOperationPlan = process(node);
    PlanNode planNode = new UnionNode(idAllocator.getNextId(), setOperationPlan.getSources(), setOperationPlan.getSymbolMapping(), ImmutableList.copyOf(setOperationPlan.getSymbolMapping().keySet()));
    if (node.isDistinct()) {
        planNode = distinct(planNode);
    }
    return new RelationPlan(planNode, analysis.getScope(node), planNode.getOutputSymbols(), outerContext);
}

Example 15 with UnionNode

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

the class QueryPlanner method planExpand.

public RelationPlan planExpand(Query query) {
    checkArgument(analysis.isExpandableQuery(query), "query is not registered as expandable");
    Union union = (Union) query.getQueryBody();
    ImmutableList.Builder<NodeAndMappings> recursionSteps = ImmutableList.builder();
    // plan anchor relation
    Relation anchorNode = union.getRelations().get(0);
    RelationPlan anchorPlan = new RelationPlanner(analysis, symbolAllocator, idAllocator, lambdaDeclarationToSymbolMap, plannerContext, outerContext, session, recursiveSubqueries).process(anchorNode, null);
    // prune anchor plan outputs to contain only the symbols exposed in the scope
    NodeAndMappings prunedAnchorPlan = pruneInvisibleFields(anchorPlan, idAllocator);
    // if the anchor plan has duplicate output symbols, add projection on top to make the symbols unique
    // This is necessary to successfully unroll recursion: the recursion step relation must follow
    // the same layout while it might not have duplicate outputs where the anchor plan did
    NodeAndMappings disambiguatedAnchorPlan = disambiguateOutputs(prunedAnchorPlan, symbolAllocator, idAllocator);
    anchorPlan = new RelationPlan(disambiguatedAnchorPlan.getNode(), analysis.getScope(query), disambiguatedAnchorPlan.getFields(), outerContext);
    recursionSteps.add(copy(anchorPlan.getRoot(), anchorPlan.getFieldMappings()));
    // plan recursion step
    Relation recursionStepRelation = union.getRelations().get(1);
    RelationPlan recursionStepPlan = new RelationPlanner(analysis, symbolAllocator, idAllocator, lambdaDeclarationToSymbolMap, plannerContext, outerContext, session, ImmutableMap.of(NodeRef.of(analysis.getRecursiveReference(query)), anchorPlan)).process(recursionStepRelation, null);
    // coerce recursion step outputs and prune them to contain only the symbols exposed in the scope
    NodeAndMappings coercedRecursionStep;
    List<Type> types = analysis.getRelationCoercion(recursionStepRelation);
    if (types == null) {
        coercedRecursionStep = pruneInvisibleFields(recursionStepPlan, idAllocator);
    } else {
        coercedRecursionStep = coerce(recursionStepPlan, types, symbolAllocator, idAllocator);
    }
    NodeAndMappings replacementSpot = new NodeAndMappings(anchorPlan.getRoot(), anchorPlan.getFieldMappings());
    PlanNode recursionStep = coercedRecursionStep.getNode();
    List<Symbol> mappings = coercedRecursionStep.getFields();
    // unroll recursion
    int maxRecursionDepth = getMaxRecursionDepth(session);
    for (int i = 0; i < maxRecursionDepth; i++) {
        recursionSteps.add(copy(recursionStep, mappings));
        NodeAndMappings replacement = copy(recursionStep, mappings);
        // if the recursion step plan has duplicate output symbols, add projection on top to make the symbols unique
        // This is necessary to successfully unroll recursion: the relation on the next recursion step must follow
        // the same layout while it might not have duplicate outputs where the plan for this step did
        replacement = disambiguateOutputs(replacement, symbolAllocator, idAllocator);
        recursionStep = replace(recursionStep, replacementSpot, replacement);
        replacementSpot = replacement;
    }
    // after the last recursion step, check if the recursion converged. the last step is expected to return empty result
    // 1. append window to count rows
    NodeAndMappings checkConvergenceStep = copy(recursionStep, mappings);
    Symbol countSymbol = symbolAllocator.newSymbol("count", BIGINT);
    ResolvedFunction function = plannerContext.getMetadata().resolveFunction(session, QualifiedName.of("count"), ImmutableList.of());
    WindowNode.Function countFunction = new WindowNode.Function(function, ImmutableList.of(), DEFAULT_FRAME, false);
    WindowNode windowNode = new WindowNode(idAllocator.getNextId(), checkConvergenceStep.getNode(), new WindowNode.Specification(ImmutableList.of(), Optional.empty()), ImmutableMap.of(countSymbol, countFunction), Optional.empty(), ImmutableSet.of(), 0);
    // 2. append filter to fail on non-empty result
    ResolvedFunction fail = plannerContext.getMetadata().resolveFunction(session, QualifiedName.of("fail"), fromTypes(VARCHAR));
    String recursionLimitExceededMessage = format("Recursion depth limit exceeded (%s). Use 'max_recursion_depth' session property to modify the limit.", maxRecursionDepth);
    Expression predicate = new IfExpression(new ComparisonExpression(GREATER_THAN_OR_EQUAL, countSymbol.toSymbolReference(), new GenericLiteral("BIGINT", "0")), new Cast(new FunctionCall(fail.toQualifiedName(), ImmutableList.of(new Cast(new StringLiteral(recursionLimitExceededMessage), toSqlType(VARCHAR)))), toSqlType(BOOLEAN)), TRUE_LITERAL);
    FilterNode filterNode = new FilterNode(idAllocator.getNextId(), windowNode, predicate);
    recursionSteps.add(new NodeAndMappings(filterNode, checkConvergenceStep.getFields()));
    // union all the recursion steps
    List<NodeAndMappings> recursionStepsToUnion = recursionSteps.build();
    List<Symbol> unionOutputSymbols = anchorPlan.getFieldMappings().stream().map(symbol -> symbolAllocator.newSymbol(symbol, "_expanded")).collect(toImmutableList());
    ImmutableListMultimap.Builder<Symbol, Symbol> unionSymbolMapping = ImmutableListMultimap.builder();
    for (NodeAndMappings plan : recursionStepsToUnion) {
        for (int i = 0; i < unionOutputSymbols.size(); i++) {
            unionSymbolMapping.put(unionOutputSymbols.get(i), plan.getFields().get(i));
        }
    }
    List<PlanNode> nodesToUnion = recursionStepsToUnion.stream().map(NodeAndMappings::getNode).collect(toImmutableList());
    PlanNode result = new UnionNode(idAllocator.getNextId(), nodesToUnion, unionSymbolMapping.build(), unionOutputSymbols);
    if (union.isDistinct()) {
        result = new AggregationNode(idAllocator.getNextId(), result, ImmutableMap.of(), singleGroupingSet(result.getOutputSymbols()), ImmutableList.of(), AggregationNode.Step.SINGLE, Optional.empty(), Optional.empty());
    }
    return new RelationPlan(result, anchorPlan.getScope(), unionOutputSymbols, outerContext);
}