Examples with Symbol io.trino.sql.planner.Symbol used on opensource projects

Example 36 with Symbol

use of io.trino.sql.planner.Symbol in project trino by trinodb.

the class RewriteSpatialPartitioningAggregation method apply.

@Override
public Result apply(AggregationNode node, Captures captures, Context context) {
    ResolvedFunction spatialPartitioningFunction = plannerContext.getMetadata().resolveFunction(context.getSession(), QualifiedName.of(NAME), fromTypeSignatures(GEOMETRY_TYPE_SIGNATURE, INTEGER.getTypeSignature()));
    ResolvedFunction stEnvelopeFunction = plannerContext.getMetadata().resolveFunction(context.getSession(), QualifiedName.of("ST_Envelope"), fromTypeSignatures(GEOMETRY_TYPE_SIGNATURE));
    ImmutableMap.Builder<Symbol, Aggregation> aggregations = ImmutableMap.builder();
    Symbol partitionCountSymbol = context.getSymbolAllocator().newSymbol("partition_count", INTEGER);
    ImmutableMap.Builder<Symbol, Expression> envelopeAssignments = ImmutableMap.builder();
    for (Map.Entry<Symbol, Aggregation> entry : node.getAggregations().entrySet()) {
        Aggregation aggregation = entry.getValue();
        String name = aggregation.getResolvedFunction().getSignature().getName();
        if (name.equals(NAME) && aggregation.getArguments().size() == 1) {
            Expression geometry = getOnlyElement(aggregation.getArguments());
            Symbol envelopeSymbol = context.getSymbolAllocator().newSymbol("envelope", plannerContext.getTypeManager().getType(GEOMETRY_TYPE_SIGNATURE));
            if (isStEnvelopeFunctionCall(geometry, stEnvelopeFunction)) {
                envelopeAssignments.put(envelopeSymbol, geometry);
            } else {
                envelopeAssignments.put(envelopeSymbol, FunctionCallBuilder.resolve(context.getSession(), plannerContext.getMetadata()).setName(QualifiedName.of("ST_Envelope")).addArgument(GEOMETRY_TYPE_SIGNATURE, geometry).build());
            }
            aggregations.put(entry.getKey(), new Aggregation(spatialPartitioningFunction, ImmutableList.of(envelopeSymbol.toSymbolReference(), partitionCountSymbol.toSymbolReference()), false, Optional.empty(), Optional.empty(), aggregation.getMask()));
        } else {
            aggregations.put(entry);
        }
    }
    return Result.ofPlanNode(new AggregationNode(node.getId(), new ProjectNode(context.getIdAllocator().getNextId(), node.getSource(), Assignments.builder().putIdentities(node.getSource().getOutputSymbols()).put(partitionCountSymbol, new LongLiteral(Integer.toString(getHashPartitionCount(context.getSession())))).putAll(envelopeAssignments.buildOrThrow()).build()), aggregations.buildOrThrow(), node.getGroupingSets(), node.getPreGroupedSymbols(), node.getStep(), node.getHashSymbol(), node.getGroupIdSymbol()));
}

Example 37 with Symbol

use of io.trino.sql.planner.Symbol 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 38 with Symbol

use of io.trino.sql.planner.Symbol in project trino by trinodb.

the class SetOperationMerge method addMergedMappings.

private void addMergedMappings(SetOperationNode child, int childIndex, ImmutableListMultimap.Builder<Symbol, Symbol> newMappingsBuilder) {
    newSources.addAll(child.getSources());
    for (Map.Entry<Symbol, Collection<Symbol>> mapping : node.getSymbolMapping().asMap().entrySet()) {
        Symbol input = Iterables.get(mapping.getValue(), childIndex);
        newMappingsBuilder.putAll(mapping.getKey(), child.getSymbolMapping().get(input));
    }
}

Example 39 with Symbol

use of io.trino.sql.planner.Symbol in project trino by trinodb.

the class SetOperationNodeTranslator method appendCounts.

private WindowNode appendCounts(UnionNode sourceNode, List<Symbol> originalColumns, List<Symbol> markers, List<Symbol> countOutputs, Symbol rowNumberSymbol) {
    ImmutableMap.Builder<Symbol, WindowNode.Function> functions = ImmutableMap.builder();
    WindowNode.Frame defaultFrame = new WindowNode.Frame(ROWS, UNBOUNDED_PRECEDING, Optional.empty(), Optional.empty(), UNBOUNDED_FOLLOWING, Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty());
    for (int i = 0; i < markers.size(); i++) {
        Symbol output = countOutputs.get(i);
        functions.put(output, new WindowNode.Function(countFunction, ImmutableList.of(markers.get(i).toSymbolReference()), defaultFrame, false));
    }
    functions.put(rowNumberSymbol, new WindowNode.Function(rowNumberFunction, ImmutableList.of(), defaultFrame, false));
    return new WindowNode(idAllocator.getNextId(), sourceNode, new Specification(originalColumns, Optional.empty()), functions.buildOrThrow(), Optional.empty(), ImmutableSet.of(), 0);
}

Example 40 with Symbol

use of io.trino.sql.planner.Symbol in project trino by trinodb.

the class SetOperationNodeTranslator method appendMarkers.

private static PlanNode appendMarkers(PlanNodeIdAllocator idAllocator, SymbolAllocator symbolAllocator, PlanNode source, int markerIndex, List<Symbol> markers, Map<Symbol, SymbolReference> projections) {
    Assignments.Builder assignments = Assignments.builder();
    // add existing intersect symbols to projection
    for (Map.Entry<Symbol, SymbolReference> entry : projections.entrySet()) {
        Symbol symbol = symbolAllocator.newSymbol(entry.getKey().getName(), symbolAllocator.getTypes().get(entry.getKey()));
        assignments.put(symbol, entry.getValue());
    }
    // add extra marker fields to the projection
    for (int i = 0; i < markers.size(); ++i) {
        Expression expression = (i == markerIndex) ? TRUE_LITERAL : new Cast(new NullLiteral(), toSqlType(BOOLEAN));
        assignments.put(symbolAllocator.newSymbol(markers.get(i).getName(), BOOLEAN), expression);
    }
    return new ProjectNode(idAllocator.getNextId(), source, assignments.build());
}