Examples with Expression io.trino.sql.tree.Expression used on opensource projects

Example 1 with Expression

use of io.trino.sql.tree.Expression in project trino by trinodb.

the class FilterStatsCalculator method extractCorrelatedGroups.

private static List<List<Expression>> extractCorrelatedGroups(List<Expression> terms, double filterConjunctionIndependenceFactor) {
    if (filterConjunctionIndependenceFactor == 1) {
        // Allows the filters to be estimated as if there is no correlation between any of the terms
        return ImmutableList.of(terms);
    }
    ListMultimap<Expression, Symbol> expressionUniqueSymbols = ArrayListMultimap.create();
    terms.forEach(expression -> expressionUniqueSymbols.putAll(expression, extractUnique(expression)));
    // Partition symbols into disjoint sets such that the symbols belonging to different disjoint sets
    // do not appear together in any expression.
    DisjointSet<Symbol> symbolsPartitioner = new DisjointSet<>();
    for (Expression term : terms) {
        List<Symbol> expressionSymbols = expressionUniqueSymbols.get(term);
        if (expressionSymbols.isEmpty()) {
            continue;
        }
        // Ensure that symbol is added to DisjointSet when there is only one symbol in the list
        symbolsPartitioner.find(expressionSymbols.get(0));
        for (int i = 1; i < expressionSymbols.size(); i++) {
            symbolsPartitioner.findAndUnion(expressionSymbols.get(0), expressionSymbols.get(i));
        }
    }
    // Use disjoint sets of symbols to partition the given list of expressions
    List<Set<Symbol>> symbolPartitions = ImmutableList.copyOf(symbolsPartitioner.getEquivalentClasses());
    checkState(symbolPartitions.size() <= terms.size(), "symbolPartitions size exceeds number of expressions");
    ListMultimap<Integer, Expression> expressionPartitions = ArrayListMultimap.create();
    for (Expression term : terms) {
        List<Symbol> expressionSymbols = expressionUniqueSymbols.get(term);
        int expressionPartitionId;
        if (expressionSymbols.isEmpty()) {
            // For expressions with no symbols
            expressionPartitionId = symbolPartitions.size();
        } else {
            // Lookup any symbol to find the partition id
            Symbol symbol = expressionSymbols.get(0);
            expressionPartitionId = IntStream.range(0, symbolPartitions.size()).filter(partition -> symbolPartitions.get(partition).contains(symbol)).findFirst().orElseThrow();
        }
        expressionPartitions.put(expressionPartitionId, term);
    }
    return expressionPartitions.keySet().stream().map(expressionPartitions::get).collect(toImmutableList());
}

Example 2 with Expression

use of io.trino.sql.tree.Expression in project trino by trinodb.

the class ExtractSpatialJoins method tryCreateSpatialJoin.

private static Result tryCreateSpatialJoin(Context context, JoinNode joinNode, Expression filter, PlanNodeId nodeId, List<Symbol> outputSymbols, FunctionCall spatialFunction, Optional<Expression> radius, PlannerContext plannerContext, SplitManager splitManager, PageSourceManager pageSourceManager, TypeAnalyzer typeAnalyzer) {
    // TODO Add support for distributed left spatial joins
    Optional<String> spatialPartitioningTableName = joinNode.getType() == INNER ? getSpatialPartitioningTableName(context.getSession()) : Optional.empty();
    Optional<KdbTree> kdbTree = spatialPartitioningTableName.map(tableName -> loadKdbTree(tableName, context.getSession(), plannerContext.getMetadata(), splitManager, pageSourceManager));
    List<Expression> arguments = spatialFunction.getArguments();
    verify(arguments.size() == 2);
    Expression firstArgument = arguments.get(0);
    Expression secondArgument = arguments.get(1);
    Type sphericalGeographyType = plannerContext.getTypeManager().getType(SPHERICAL_GEOGRAPHY_TYPE_SIGNATURE);
    if (typeAnalyzer.getType(context.getSession(), context.getSymbolAllocator().getTypes(), firstArgument).equals(sphericalGeographyType) || typeAnalyzer.getType(context.getSession(), context.getSymbolAllocator().getTypes(), secondArgument).equals(sphericalGeographyType)) {
        return Result.empty();
    }
    Set<Symbol> firstSymbols = extractUnique(firstArgument);
    Set<Symbol> secondSymbols = extractUnique(secondArgument);
    if (firstSymbols.isEmpty() || secondSymbols.isEmpty()) {
        return Result.empty();
    }
    Optional<Symbol> newFirstSymbol = newGeometrySymbol(context, firstArgument, plannerContext.getTypeManager());
    Optional<Symbol> newSecondSymbol = newGeometrySymbol(context, secondArgument, plannerContext.getTypeManager());
    PlanNode leftNode = joinNode.getLeft();
    PlanNode rightNode = joinNode.getRight();
    PlanNode newLeftNode;
    PlanNode newRightNode;
    // Check if the order of arguments of the spatial function matches the order of join sides
    int alignment = checkAlignment(joinNode, firstSymbols, secondSymbols);
    if (alignment > 0) {
        newLeftNode = newFirstSymbol.map(symbol -> addProjection(context, leftNode, symbol, firstArgument)).orElse(leftNode);
        newRightNode = newSecondSymbol.map(symbol -> addProjection(context, rightNode, symbol, secondArgument)).orElse(rightNode);
    } else if (alignment < 0) {
        newLeftNode = newSecondSymbol.map(symbol -> addProjection(context, leftNode, symbol, secondArgument)).orElse(leftNode);
        newRightNode = newFirstSymbol.map(symbol -> addProjection(context, rightNode, symbol, firstArgument)).orElse(rightNode);
    } else {
        return Result.empty();
    }
    Expression newFirstArgument = toExpression(newFirstSymbol, firstArgument);
    Expression newSecondArgument = toExpression(newSecondSymbol, secondArgument);
    Optional<Symbol> leftPartitionSymbol = Optional.empty();
    Optional<Symbol> rightPartitionSymbol = Optional.empty();
    if (kdbTree.isPresent()) {
        leftPartitionSymbol = Optional.of(context.getSymbolAllocator().newSymbol("pid", INTEGER));
        rightPartitionSymbol = Optional.of(context.getSymbolAllocator().newSymbol("pid", INTEGER));
        if (alignment > 0) {
            newLeftNode = addPartitioningNodes(plannerContext, context, newLeftNode, leftPartitionSymbol.get(), kdbTree.get(), newFirstArgument, Optional.empty());
            newRightNode = addPartitioningNodes(plannerContext, context, newRightNode, rightPartitionSymbol.get(), kdbTree.get(), newSecondArgument, radius);
        } else {
            newLeftNode = addPartitioningNodes(plannerContext, context, newLeftNode, leftPartitionSymbol.get(), kdbTree.get(), newSecondArgument, Optional.empty());
            newRightNode = addPartitioningNodes(plannerContext, context, newRightNode, rightPartitionSymbol.get(), kdbTree.get(), newFirstArgument, radius);
        }
    }
    Expression newSpatialFunction = FunctionCallBuilder.resolve(context.getSession(), plannerContext.getMetadata()).setName(spatialFunction.getName()).addArgument(GEOMETRY_TYPE_SIGNATURE, newFirstArgument).addArgument(GEOMETRY_TYPE_SIGNATURE, newSecondArgument).build();
    Expression newFilter = replaceExpression(filter, ImmutableMap.of(spatialFunction, newSpatialFunction));
    return Result.ofPlanNode(new SpatialJoinNode(nodeId, SpatialJoinNode.Type.fromJoinNodeType(joinNode.getType()), newLeftNode, newRightNode, outputSymbols, newFilter, leftPartitionSymbol, rightPartitionSymbol, kdbTree.map(KdbTreeUtils::toJson)));
}

Example 3 with Expression

use of io.trino.sql.tree.Expression in project trino by trinodb.

the class ImplementFilteredAggregations method apply.

@Override
public Result apply(AggregationNode aggregationNode, Captures captures, Context context) {
    Assignments.Builder newAssignments = Assignments.builder();
    ImmutableMap.Builder<Symbol, Aggregation> aggregations = ImmutableMap.builder();
    ImmutableList.Builder<Expression> maskSymbols = ImmutableList.builder();
    boolean aggregateWithoutFilterOrMaskPresent = false;
    for (Map.Entry<Symbol, Aggregation> entry : aggregationNode.getAggregations().entrySet()) {
        Symbol output = entry.getKey();
        // strip the filters
        Aggregation aggregation = entry.getValue();
        Optional<Symbol> mask = aggregation.getMask();
        if (aggregation.getFilter().isPresent()) {
            Symbol filter = aggregation.getFilter().get();
            if (mask.isPresent()) {
                Symbol newMask = context.getSymbolAllocator().newSymbol("mask", BOOLEAN);
                Expression expression = and(mask.get().toSymbolReference(), filter.toSymbolReference());
                newAssignments.put(newMask, expression);
                mask = Optional.of(newMask);
                maskSymbols.add(newMask.toSymbolReference());
            } else {
                mask = Optional.of(filter);
                maskSymbols.add(filter.toSymbolReference());
            }
        } else if (mask.isPresent()) {
            maskSymbols.add(mask.get().toSymbolReference());
        } else {
            aggregateWithoutFilterOrMaskPresent = true;
        }
        aggregations.put(output, new Aggregation(aggregation.getResolvedFunction(), aggregation.getArguments(), aggregation.isDistinct(), Optional.empty(), aggregation.getOrderingScheme(), mask));
    }
    Expression predicate = TRUE_LITERAL;
    if (!aggregationNode.hasNonEmptyGroupingSet() && !aggregateWithoutFilterOrMaskPresent) {
        predicate = combineDisjunctsWithDefault(metadata, maskSymbols.build(), TRUE_LITERAL);
    }
    // identity projection for all existing inputs
    newAssignments.putIdentities(aggregationNode.getSource().getOutputSymbols());
    return Result.ofPlanNode(new AggregationNode(context.getIdAllocator().getNextId(), new FilterNode(context.getIdAllocator().getNextId(), new ProjectNode(context.getIdAllocator().getNextId(), aggregationNode.getSource(), newAssignments.build()), predicate), aggregations.buildOrThrow(), aggregationNode.getGroupingSets(), ImmutableList.of(), aggregationNode.getStep(), aggregationNode.getHashSymbol(), aggregationNode.getGroupIdSymbol()));
}

Example 4 with Expression

use of io.trino.sql.tree.Expression in project trino by trinodb.

the class ImplementIntersectDistinctAsUnion method apply.

@Override
public Result apply(IntersectNode node, Captures captures, Context context) {
    SetOperationNodeTranslator translator = new SetOperationNodeTranslator(context.getSession(), metadata, context.getSymbolAllocator(), context.getIdAllocator());
    SetOperationNodeTranslator.TranslationResult result = translator.makeSetContainmentPlanForDistinct(node);
    // intersect predicate: the row must be present in every source
    Expression predicate = and(result.getCountSymbols().stream().map(symbol -> new ComparisonExpression(GREATER_THAN_OR_EQUAL, symbol.toSymbolReference(), new GenericLiteral("BIGINT", "1"))).collect(toImmutableList()));
    return Result.ofPlanNode(new ProjectNode(context.getIdAllocator().getNextId(), new FilterNode(context.getIdAllocator().getNextId(), result.getPlanNode(), predicate), Assignments.identity(node.getOutputSymbols())));
}

Example 5 with Expression

use of io.trino.sql.tree.Expression in project trino by trinodb.

the class InlineProjectIntoFilter method apply.

@Override
public Result apply(FilterNode node, Captures captures, Context context) {
    ProjectNode projectNode = captures.get(PROJECTION);
    List<Expression> filterConjuncts = extractConjuncts(node.getPredicate());
    Map<Boolean, List<Expression>> conjuncts = filterConjuncts.stream().collect(partitioningBy(SymbolReference.class::isInstance));
    List<Expression> simpleConjuncts = conjuncts.get(true);
    List<Expression> complexConjuncts = conjuncts.get(false);
    // Do not inline expression if the symbol is used multiple times in simple conjuncts.
    Set<Expression> simpleUniqueConjuncts = simpleConjuncts.stream().collect(Collectors.groupingBy(identity(), counting())).entrySet().stream().filter(entry -> entry.getValue() == 1).map(Map.Entry::getKey).collect(toImmutableSet());
    Set<Expression> complexConjunctSymbols = SymbolsExtractor.extractUnique(complexConjuncts).stream().map(Symbol::toSymbolReference).collect(toImmutableSet());
    // Do not inline expression if the symbol is used in complex conjuncts.
    Set<Expression> simpleConjunctsToInline = Sets.difference(simpleUniqueConjuncts, complexConjunctSymbols);
    if (simpleConjunctsToInline.isEmpty()) {
        return Result.empty();
    }
    ImmutableList.Builder<Expression> newConjuncts = ImmutableList.builder();
    Assignments.Builder newAssignments = Assignments.builder();
    Assignments.Builder postFilterAssignmentsBuilder = Assignments.builder();
    for (Expression conjunct : filterConjuncts) {
        if (simpleConjunctsToInline.contains(conjunct)) {
            Expression expression = projectNode.getAssignments().get(Symbol.from(conjunct));
            if (expression == null || expression instanceof SymbolReference) {
                // expression == null -> The symbol is not produced by the underlying projection (i.e. it is a correlation symbol).
                // expression instanceof SymbolReference -> Do not inline trivial projections.
                newConjuncts.add(conjunct);
            } else {
                newConjuncts.add(expression);
                newAssignments.putIdentities(SymbolsExtractor.extractUnique(expression));
                postFilterAssignmentsBuilder.put(Symbol.from(conjunct), TRUE_LITERAL);
            }
        } else {
            newConjuncts.add(conjunct);
        }
    }
    Assignments postFilterAssignments = postFilterAssignmentsBuilder.build();
    if (postFilterAssignments.isEmpty()) {
        return Result.empty();
    }
    Set<Symbol> postFilterSymbols = postFilterAssignments.getSymbols();
    // Remove inlined expressions from the underlying projection.
    newAssignments.putAll(projectNode.getAssignments().filter(symbol -> !postFilterSymbols.contains(symbol)));
    Map<Symbol, Expression> outputAssignments = new HashMap<>();
    outputAssignments.putAll(Assignments.identity(node.getOutputSymbols()).getMap());
    // Restore inlined symbols.
    outputAssignments.putAll(postFilterAssignments.getMap());
    return Result.ofPlanNode(new ProjectNode(context.getIdAllocator().getNextId(), new FilterNode(node.getId(), new ProjectNode(projectNode.getId(), projectNode.getSource(), newAssignments.build()), combineConjuncts(metadata, newConjuncts.build())), Assignments.builder().putAll(outputAssignments).build()));
}