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

Example 1 with UnnestNode

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

the class ExtractSpatialJoins method addPartitioningNodes.

private static PlanNode addPartitioningNodes(PlannerContext plannerContext, Context context, PlanNode node, Symbol partitionSymbol, KdbTree kdbTree, Expression geometry, Optional<Expression> radius) {
    Assignments.Builder projections = Assignments.builder();
    for (Symbol outputSymbol : node.getOutputSymbols()) {
        projections.putIdentity(outputSymbol);
    }
    TypeSignature typeSignature = new TypeSignature(KDB_TREE_TYPENAME);
    FunctionCallBuilder spatialPartitionsCall = FunctionCallBuilder.resolve(context.getSession(), plannerContext.getMetadata()).setName(QualifiedName.of("spatial_partitions")).addArgument(typeSignature, new Cast(new StringLiteral(KdbTreeUtils.toJson(kdbTree)), toSqlType(plannerContext.getTypeManager().getType(typeSignature)))).addArgument(GEOMETRY_TYPE_SIGNATURE, geometry);
    radius.ifPresent(value -> spatialPartitionsCall.addArgument(DOUBLE, value));
    FunctionCall partitioningFunction = spatialPartitionsCall.build();
    Symbol partitionsSymbol = context.getSymbolAllocator().newSymbol(partitioningFunction, new ArrayType(INTEGER));
    projections.put(partitionsSymbol, partitioningFunction);
    return new UnnestNode(context.getIdAllocator().getNextId(), new ProjectNode(context.getIdAllocator().getNextId(), node, projections.build()), node.getOutputSymbols(), ImmutableList.of(new UnnestNode.Mapping(partitionsSymbol, ImmutableList.of(partitionSymbol))), Optional.empty(), INNER, Optional.empty());
}

Example 2 with UnnestNode

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

the class DecorrelateInnerUnnestWithGlobalAggregation method apply.

@Override
public Result apply(CorrelatedJoinNode correlatedJoinNode, Captures captures, Context context) {
    // find global aggregation in subquery
    List<PlanNode> globalAggregations = PlanNodeSearcher.searchFrom(correlatedJoinNode.getSubquery(), context.getLookup()).where(DecorrelateInnerUnnestWithGlobalAggregation::isGlobalAggregation).recurseOnlyWhen(node -> node instanceof ProjectNode || isGlobalAggregation(node)).findAll();
    if (globalAggregations.isEmpty()) {
        return Result.empty();
    }
    // if there are multiple global aggregations, the one that is closest to the source is the "reducing" aggregation, because it reduces multiple input rows to single output row
    AggregationNode reducingAggregation = (AggregationNode) globalAggregations.get(globalAggregations.size() - 1);
    // find unnest in subquery
    Optional<UnnestNode> subqueryUnnest = PlanNodeSearcher.searchFrom(reducingAggregation.getSource(), context.getLookup()).where(node -> isSupportedUnnest(node, correlatedJoinNode.getCorrelation(), context.getLookup())).recurseOnlyWhen(node -> node instanceof ProjectNode || isGroupedAggregation(node)).findFirst();
    if (subqueryUnnest.isEmpty()) {
        return Result.empty();
    }
    UnnestNode unnestNode = subqueryUnnest.get();
    // assign unique id to input rows to restore semantics of aggregations after rewrite
    PlanNode input = new AssignUniqueId(context.getIdAllocator().getNextId(), correlatedJoinNode.getInput(), context.getSymbolAllocator().newSymbol("unique", BIGINT));
    // pre-project unnest symbols if they were pre-projected in subquery
    // The correlated UnnestNode either unnests correlation symbols directly, or unnests symbols produced by a projection that uses only correlation symbols.
    // Here, any underlying projection that was a source of the correlated UnnestNode, is appended as a source of the rewritten UnnestNode.
    // If the projection is not necessary for UnnestNode (i.e. it does not produce any unnest symbols), it should be pruned afterwards.
    PlanNode unnestSource = context.getLookup().resolve(unnestNode.getSource());
    if (unnestSource instanceof ProjectNode) {
        ProjectNode sourceProjection = (ProjectNode) unnestSource;
        input = new ProjectNode(sourceProjection.getId(), input, Assignments.builder().putIdentities(input.getOutputSymbols()).putAll(sourceProjection.getAssignments()).build());
    }
    // rewrite correlated join to UnnestNode
    Symbol ordinalitySymbol = unnestNode.getOrdinalitySymbol().orElseGet(() -> context.getSymbolAllocator().newSymbol("ordinality", BIGINT));
    UnnestNode rewrittenUnnest = new UnnestNode(context.getIdAllocator().getNextId(), input, input.getOutputSymbols(), unnestNode.getMappings(), Optional.of(ordinalitySymbol), LEFT, Optional.empty());
    // append mask symbol based on ordinality to distinguish between the unnested rows and synthetic null rows
    Symbol mask = context.getSymbolAllocator().newSymbol("mask", BOOLEAN);
    ProjectNode sourceWithMask = new ProjectNode(context.getIdAllocator().getNextId(), rewrittenUnnest, Assignments.builder().putIdentities(rewrittenUnnest.getOutputSymbols()).put(mask, new IsNotNullPredicate(ordinalitySymbol.toSymbolReference())).build());
    // restore all projections, grouped aggregations and global aggregations from the subquery
    PlanNode result = rewriteNodeSequence(context.getLookup().resolve(correlatedJoinNode.getSubquery()), input.getOutputSymbols(), mask, sourceWithMask, reducingAggregation.getId(), unnestNode.getId(), context.getSymbolAllocator(), context.getIdAllocator(), context.getLookup());
    // restrict outputs
    return Result.ofPlanNode(restrictOutputs(context.getIdAllocator(), result, ImmutableSet.copyOf(correlatedJoinNode.getOutputSymbols())).orElse(result));
}

Example 3 with UnnestNode

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

the class DecorrelateLeftUnnestWithGlobalAggregation method isSupportedUnnest.

/**
 * This rule supports decorrelation of UnnestNode meeting certain conditions:
 * - the UnnestNode should be based on correlation symbols, that is: either unnest correlation symbols directly,
 * or unnest symbols produced by a projection that uses only correlation symbols.
 * - the UnnestNode should not have any replicate symbols,
 * - the UnnestNode should be of type LEFT,
 * - the UnnestNode should not have a filter.
 */
private static boolean isSupportedUnnest(PlanNode node, List<Symbol> correlation, Lookup lookup) {
    if (!(node instanceof UnnestNode)) {
        return false;
    }
    UnnestNode unnestNode = (UnnestNode) node;
    List<Symbol> unnestSymbols = unnestNode.getMappings().stream().map(UnnestNode.Mapping::getInput).collect(toImmutableList());
    PlanNode unnestSource = lookup.resolve(unnestNode.getSource());
    boolean basedOnCorrelation = ImmutableSet.copyOf(correlation).containsAll(unnestSymbols) || unnestSource instanceof ProjectNode && ImmutableSet.copyOf(correlation).containsAll(SymbolsExtractor.extractUnique(((ProjectNode) unnestSource).getAssignments().getExpressions()));
    return isScalar(unnestNode.getSource(), lookup) && unnestNode.getReplicateSymbols().isEmpty() && basedOnCorrelation && unnestNode.getJoinType() == LEFT && (unnestNode.getFilter().isEmpty() || unnestNode.getFilter().get().equals(TRUE_LITERAL));
}

Example 4 with UnnestNode

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

the class DecorrelateUnnest method isSupportedUnnest.

/**
 * This rule supports decorrelation of UnnestNode meeting certain conditions:
 * - the UnnestNode should be based on correlation symbols, that is: either unnest correlation symbols directly,
 * or unnest symbols produced by a projection that uses only correlation symbols.
 * - the UnnestNode should not have any replicate symbols,
 * - the UnnestNode should be of type INNER or LEFT,
 * - the UnnestNode should not have a filter.
 */
private static boolean isSupportedUnnest(PlanNode node, List<Symbol> correlation, Lookup lookup) {
    if (!(node instanceof UnnestNode)) {
        return false;
    }
    UnnestNode unnestNode = (UnnestNode) node;
    List<Symbol> unnestSymbols = unnestNode.getMappings().stream().map(UnnestNode.Mapping::getInput).collect(toImmutableList());
    PlanNode unnestSource = lookup.resolve(unnestNode.getSource());
    boolean basedOnCorrelation = ImmutableSet.copyOf(correlation).containsAll(unnestSymbols) || unnestSource instanceof ProjectNode && ImmutableSet.copyOf(correlation).containsAll(SymbolsExtractor.extractUnique(((ProjectNode) unnestSource).getAssignments().getExpressions()));
    return isScalar(unnestNode.getSource(), lookup) && unnestNode.getReplicateSymbols().isEmpty() && basedOnCorrelation && (unnestNode.getJoinType() == INNER || unnestNode.getJoinType() == LEFT) && (unnestNode.getFilter().isEmpty() || unnestNode.getFilter().get().equals(TRUE_LITERAL));
}

Example 5 with UnnestNode

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

the class RelationPlanner method planUnnest.

private RelationPlan planUnnest(PlanBuilder subPlan, Unnest node, List<Symbol> replicatedColumns, Optional<Expression> filter, Join.Type type, Scope outputScope) {
    subPlan = subPlan.appendProjections(node.getExpressions(), symbolAllocator, idAllocator);
    Map<Field, Symbol> allocations = analysis.getOutputDescriptor(node).getVisibleFields().stream().collect(toImmutableMap(Function.identity(), symbolAllocator::newSymbol));
    UnnestAnalysis unnestAnalysis = analysis.getUnnest(node);
    ImmutableList.Builder<UnnestNode.Mapping> mappings = ImmutableList.builder();
    for (Expression expression : node.getExpressions()) {
        Symbol input = subPlan.translate(expression);
        List<Symbol> outputs = unnestAnalysis.getMappings().get(NodeRef.of(expression)).stream().map(allocations::get).collect(toImmutableList());
        mappings.add(new UnnestNode.Mapping(input, outputs));
    }
    UnnestNode unnestNode = new UnnestNode(idAllocator.getNextId(), subPlan.getRoot(), replicatedColumns, mappings.build(), unnestAnalysis.getOrdinalityField().map(allocations::get), JoinNode.Type.typeConvert(type), filter);
    // Currently, it works out because, by construction, the order of the output symbols in the UnnestNode will match the order of the fields in the Join node.
    return new RelationPlan(unnestNode, outputScope, unnestNode.getOutputSymbols(), outerContext);
}