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

Example 6 with LambdaExpression

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

the class TestEqualityInference method testExpressionsThatMayReturnNullOnNonNullInput.

@Test
public void testExpressionsThatMayReturnNullOnNonNullInput() {
    List<Expression> candidates = ImmutableList.of(// try_cast
    new Cast(nameReference("b"), toSqlType(BIGINT), true), functionResolution.functionCallBuilder(QualifiedName.of(TryFunction.NAME)).addArgument(new FunctionType(ImmutableList.of(), VARCHAR), new LambdaExpression(ImmutableList.of(), nameReference("b"))).build(), new NullIfExpression(nameReference("b"), number(1)), new IfExpression(nameReference("b"), number(1), new NullLiteral()), new InPredicate(nameReference("b"), new InListExpression(ImmutableList.of(new NullLiteral()))), new SearchedCaseExpression(ImmutableList.of(new WhenClause(new IsNotNullPredicate(nameReference("b")), new NullLiteral())), Optional.empty()), new SimpleCaseExpression(nameReference("b"), ImmutableList.of(new WhenClause(number(1), new NullLiteral())), Optional.empty()), new SubscriptExpression(new ArrayConstructor(ImmutableList.of(new NullLiteral())), nameReference("b")));
    for (Expression candidate : candidates) {
        EqualityInference inference = EqualityInference.newInstance(metadata, equals(nameReference("b"), nameReference("x")), equals(nameReference("a"), candidate));
        List<Expression> equalities = inference.generateEqualitiesPartitionedBy(symbols("b")).getScopeStraddlingEqualities();
        assertEquals(equalities.size(), 1);
        assertTrue(equalities.get(0).equals(equals(nameReference("x"), nameReference("b"))) || equalities.get(0).equals(equals(nameReference("b"), nameReference("x"))));
    }
}

Example 7 with LambdaExpression

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

the class TestDesugarTryExpressionRewriter method testTryExpressionDesugaringRewriter.

@Test
public void testTryExpressionDesugaringRewriter() {
    // 1 + try(2)
    Expression initial = new ArithmeticBinaryExpression(ADD, new DecimalLiteral("1"), new TryExpression(new DecimalLiteral("2")));
    Expression rewritten = DesugarTryExpressionRewriter.rewrite(initial, tester().getMetadata(), tester().getTypeAnalyzer(), tester().getSession(), new SymbolAllocator());
    // 1 + try_function(() -> 2)
    Expression expected = new ArithmeticBinaryExpression(ADD, new DecimalLiteral("1"), new TestingFunctionResolution().functionCallBuilder(QualifiedName.of(TryFunction.NAME)).addArgument(new FunctionType(ImmutableList.of(), createDecimalType(1)), new LambdaExpression(ImmutableList.of(), new DecimalLiteral("2"))).build());
    ExpressionVerifier verifier = new ExpressionVerifier(new SymbolAliases());
    assertTrue(verifier.process(rewritten, expected));
}

Example 8 with LambdaExpression

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

the class QueryPlanner method planAggregation.

private PlanBuilder planAggregation(PlanBuilder subPlan, List<List<Symbol>> groupingSets, Optional<Symbol> groupIdSymbol, List<FunctionCall> aggregates, Function<Expression, Symbol> coercions) {
    ImmutableList.Builder<AggregationAssignment> aggregateMappingBuilder = ImmutableList.builder();
    // deduplicate based on scope-aware equality
    for (FunctionCall function : scopeAwareDistinct(subPlan, aggregates)) {
        Symbol symbol = symbolAllocator.newSymbol(function, analysis.getType(function));
        // TODO: for ORDER BY arguments, rewrite them such that they match the actual arguments to the function. This is necessary to maintain the semantics of DISTINCT + ORDER BY,
        // which requires that ORDER BY be a subset of arguments
        // What can happen currently is that if the argument requires a coercion, the argument will take a different input that the ORDER BY clause, which is undefined behavior
        Aggregation aggregation = new Aggregation(analysis.getResolvedFunction(function), function.getArguments().stream().map(argument -> {
            if (argument instanceof LambdaExpression) {
                return subPlan.rewrite(argument);
            }
            return coercions.apply(argument).toSymbolReference();
        }).collect(toImmutableList()), function.isDistinct(), function.getFilter().map(coercions), function.getOrderBy().map(orderBy -> translateOrderingScheme(orderBy.getSortItems(), coercions)), Optional.empty());
        aggregateMappingBuilder.add(new AggregationAssignment(symbol, function, aggregation));
    }
    List<AggregationAssignment> aggregateMappings = aggregateMappingBuilder.build();
    ImmutableSet.Builder<Integer> globalGroupingSets = ImmutableSet.builder();
    for (int i = 0; i < groupingSets.size(); i++) {
        if (groupingSets.get(i).isEmpty()) {
            globalGroupingSets.add(i);
        }
    }
    ImmutableList.Builder<Symbol> groupingKeys = ImmutableList.builder();
    groupingSets.stream().flatMap(List::stream).distinct().forEach(groupingKeys::add);
    groupIdSymbol.ifPresent(groupingKeys::add);
    AggregationNode aggregationNode = new AggregationNode(idAllocator.getNextId(), subPlan.getRoot(), aggregateMappings.stream().collect(toImmutableMap(AggregationAssignment::getSymbol, AggregationAssignment::getRewritten)), groupingSets(groupingKeys.build(), groupingSets.size(), globalGroupingSets.build()), ImmutableList.of(), AggregationNode.Step.SINGLE, Optional.empty(), groupIdSymbol);
    return new PlanBuilder(subPlan.getTranslations().withAdditionalMappings(aggregateMappings.stream().collect(toImmutableMap(assignment -> scopeAwareKey(assignment.getAstExpression(), analysis, subPlan.getScope()), AggregationAssignment::getSymbol))), aggregationNode);
}

Example 9 with LambdaExpression

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

the class QueryPlanner method aggregate.

private PlanBuilder aggregate(PlanBuilder subPlan, QuerySpecification node) {
    if (!analysis.isAggregation(node)) {
        return subPlan;
    }
    ImmutableList.Builder<Expression> inputBuilder = ImmutableList.builder();
    analysis.getAggregates(node).stream().map(FunctionCall::getArguments).flatMap(List::stream).filter(// lambda expression is generated at execution time
    expression -> !(expression instanceof LambdaExpression)).forEach(inputBuilder::add);
    analysis.getAggregates(node).stream().map(FunctionCall::getOrderBy).map(NodeUtils::getSortItemsFromOrderBy).flatMap(List::stream).map(SortItem::getSortKey).forEach(inputBuilder::add);
    // filter expressions need to be projected first
    analysis.getAggregates(node).stream().map(FunctionCall::getFilter).filter(Optional::isPresent).map(Optional::get).forEach(inputBuilder::add);
    GroupingSetAnalysis groupingSetAnalysis = analysis.getGroupingSets(node);
    inputBuilder.addAll(groupingSetAnalysis.getComplexExpressions());
    List<Expression> inputs = inputBuilder.build();
    subPlan = subqueryPlanner.handleSubqueries(subPlan, inputs, analysis.getSubqueries(node));
    subPlan = subPlan.appendProjections(inputs, symbolAllocator, idAllocator);
    // Add projection to coerce inputs to their site-specific types.
    // This is important because the same lexical expression may need to be coerced
    // in different ways if it's referenced by multiple arguments to the window function.
    // For example, given v::integer,
    // avg(v)
    // Needs to be rewritten as
    // avg(CAST(v AS double))
    PlanAndMappings coercions = coerce(subPlan, inputs, analysis, idAllocator, symbolAllocator, typeCoercion);
    subPlan = coercions.getSubPlan();
    GroupingSetsPlan groupingSets = planGroupingSets(subPlan, node, groupingSetAnalysis);
    subPlan = planAggregation(groupingSets.getSubPlan(), groupingSets.getGroupingSets(), groupingSets.getGroupIdSymbol(), analysis.getAggregates(node), coercions::get);
    return planGroupingOperations(subPlan, node, groupingSets.getGroupIdSymbol(), groupingSets.getColumnOnlyGroupingSets());
}

Example 10 with LambdaExpression

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

the class QueryPlanner method planWindowFunctions.

private PlanBuilder planWindowFunctions(Node node, PlanBuilder subPlan, List<FunctionCall> windowFunctions) {
    if (windowFunctions.isEmpty()) {
        return subPlan;
    }
    for (FunctionCall windowFunction : scopeAwareDistinct(subPlan, windowFunctions)) {
        checkArgument(windowFunction.getFilter().isEmpty(), "Window functions cannot have filter");
        ResolvedWindow window = analysis.getWindow(windowFunction);
        checkState(window != null, "no resolved window for: " + windowFunction);
        // Pre-project inputs.
        // Predefined window parts (specified in WINDOW clause) can only use source symbols, and no output symbols.
        // It matters in case when this window planning takes place in ORDER BY clause, where both source and output
        // symbols are visible.
        // This issue is solved by analyzing window definitions in the source scope. After analysis, the expressions
        // are recorded as belonging to the source scope, and consequentially source symbols will be used to plan them.
        ImmutableList.Builder<Expression> inputsBuilder = ImmutableList.<Expression>builder().addAll(windowFunction.getArguments().stream().filter(// lambda expression is generated at execution time
        argument -> !(argument instanceof LambdaExpression)).collect(Collectors.toList())).addAll(window.getPartitionBy()).addAll(getSortItemsFromOrderBy(window.getOrderBy()).stream().map(SortItem::getSortKey).iterator());
        if (window.getFrame().isPresent()) {
            WindowFrame frame = window.getFrame().get();
            frame.getStart().getValue().ifPresent(inputsBuilder::add);
            if (frame.getEnd().isPresent()) {
                frame.getEnd().get().getValue().ifPresent(inputsBuilder::add);
            }
        }
        List<Expression> inputs = inputsBuilder.build();
        subPlan = subqueryPlanner.handleSubqueries(subPlan, inputs, analysis.getSubqueries(node));
        subPlan = subPlan.appendProjections(inputs, symbolAllocator, idAllocator);
        // Add projection to coerce inputs to their site-specific types.
        // This is important because the same lexical expression may need to be coerced
        // in different ways if it's referenced by multiple arguments to the window function.
        // For example, given v::integer,
        // avg(v) OVER (ORDER BY v)
        // Needs to be rewritten as
        // avg(CAST(v AS double)) OVER (ORDER BY v)
        PlanAndMappings coercions = coerce(subPlan, inputs, analysis, idAllocator, symbolAllocator, typeCoercion);
        subPlan = coercions.getSubPlan();
        // For frame of type RANGE, append casts and functions necessary for frame bound calculations
        Optional<Symbol> frameStart = Optional.empty();
        Optional<Symbol> frameEnd = Optional.empty();
        Optional<Symbol> sortKeyCoercedForFrameStartComparison = Optional.empty();
        Optional<Symbol> sortKeyCoercedForFrameEndComparison = Optional.empty();
        if (window.getFrame().isPresent() && window.getFrame().get().getType() == RANGE) {
            Optional<Expression> startValue = window.getFrame().get().getStart().getValue();
            Optional<Expression> endValue = window.getFrame().get().getEnd().flatMap(FrameBound::getValue);
            // record sortKey coercions for reuse
            Map<Type, Symbol> sortKeyCoercions = new HashMap<>();
            // process frame start
            FrameBoundPlanAndSymbols plan = planFrameBound(subPlan, coercions, startValue, window, sortKeyCoercions);
            subPlan = plan.getSubPlan();
            frameStart = plan.getFrameBoundSymbol();
            sortKeyCoercedForFrameStartComparison = plan.getSortKeyCoercedForFrameBoundComparison();
            // process frame end
            plan = planFrameBound(subPlan, coercions, endValue, window, sortKeyCoercions);
            subPlan = plan.getSubPlan();
            frameEnd = plan.getFrameBoundSymbol();
            sortKeyCoercedForFrameEndComparison = plan.getSortKeyCoercedForFrameBoundComparison();
        } else if (window.getFrame().isPresent() && (window.getFrame().get().getType() == ROWS || window.getFrame().get().getType() == GROUPS)) {
            Optional<Expression> startValue = window.getFrame().get().getStart().getValue();
            Optional<Expression> endValue = window.getFrame().get().getEnd().flatMap(FrameBound::getValue);
            // process frame start
            FrameOffsetPlanAndSymbol plan = planFrameOffset(subPlan, startValue.map(coercions::get));
            subPlan = plan.getSubPlan();
            frameStart = plan.getFrameOffsetSymbol();
            // process frame end
            plan = planFrameOffset(subPlan, endValue.map(coercions::get));
            subPlan = plan.getSubPlan();
            frameEnd = plan.getFrameOffsetSymbol();
        } else if (window.getFrame().isPresent()) {
            throw new IllegalArgumentException("unexpected window frame type: " + window.getFrame().get().getType());
        }
        if (window.getFrame().isPresent() && window.getFrame().get().getPattern().isPresent()) {
            WindowFrame frame = window.getFrame().get();
            subPlan = subqueryPlanner.handleSubqueries(subPlan, extractPatternRecognitionExpressions(frame.getVariableDefinitions(), frame.getMeasures()), analysis.getSubqueries(node));
            subPlan = planPatternRecognition(subPlan, windowFunction, window, coercions, frameEnd);
        } else {
            subPlan = planWindow(subPlan, windowFunction, window, coercions, frameStart, sortKeyCoercedForFrameStartComparison, frameEnd, sortKeyCoercedForFrameEndComparison);
        }
    }
    return subPlan;
}