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

Example 6 with FrameBound

use of io.trino.sql.tree.FrameBound 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;
}