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

Example 1 with Assignment

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

the class UnwrapSingleColumnRowInApply method apply.

@Override
public Result apply(ApplyNode node, Captures captures, Context context) {
    Assignments.Builder inputAssignments = Assignments.builder().putIdentities(node.getInput().getOutputSymbols());
    Assignments.Builder nestedPlanAssignments = Assignments.builder().putIdentities(node.getSubquery().getOutputSymbols());
    boolean applied = false;
    Assignments.Builder applyAssignments = Assignments.builder();
    for (Map.Entry<Symbol, Expression> assignment : node.getSubqueryAssignments().entrySet()) {
        Symbol output = assignment.getKey();
        Expression expression = assignment.getValue();
        Optional<Unwrapping> unwrapped = Optional.empty();
        if (expression instanceof InPredicate) {
            InPredicate predicate = (InPredicate) expression;
            unwrapped = unwrapSingleColumnRow(context, predicate.getValue(), predicate.getValueList(), (value, list) -> new InPredicate(value.toSymbolReference(), list.toSymbolReference()));
        } else if (expression instanceof QuantifiedComparisonExpression) {
            QuantifiedComparisonExpression comparison = (QuantifiedComparisonExpression) expression;
            unwrapped = unwrapSingleColumnRow(context, comparison.getValue(), comparison.getSubquery(), (value, list) -> new QuantifiedComparisonExpression(comparison.getOperator(), comparison.getQuantifier(), value.toSymbolReference(), list.toSymbolReference()));
        }
        if (unwrapped.isPresent()) {
            applied = true;
            Unwrapping unwrapping = unwrapped.get();
            inputAssignments.add(unwrapping.getInputAssignment());
            nestedPlanAssignments.add(unwrapping.getNestedPlanAssignment());
            applyAssignments.put(output, unwrapping.getExpression());
        } else {
            applyAssignments.put(assignment);
        }
    }
    if (!applied) {
        return Result.empty();
    }
    return Result.ofPlanNode(new ProjectNode(context.getIdAllocator().getNextId(), new ApplyNode(node.getId(), new ProjectNode(context.getIdAllocator().getNextId(), node.getInput(), inputAssignments.build()), new ProjectNode(context.getIdAllocator().getNextId(), node.getSubquery(), nestedPlanAssignments.build()), applyAssignments.build(), node.getCorrelation(), node.getOriginSubquery()), Assignments.identity(node.getOutputSymbols())));
}

Example 2 with Assignment

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

the class UnwrapSingleColumnRowInApply method unwrapSingleColumnRow.

private Optional<Unwrapping> unwrapSingleColumnRow(Context context, Expression value, Expression list, BiFunction<Symbol, Symbol, Expression> function) {
    Type type = typeAnalyzer.getType(context.getSession(), context.getSymbolAllocator().getTypes(), value);
    if (type instanceof RowType) {
        RowType rowType = (RowType) type;
        if (rowType.getFields().size() == 1) {
            Type elementType = rowType.getTypeParameters().get(0);
            Symbol valueSymbol = context.getSymbolAllocator().newSymbol("input", elementType);
            Symbol listSymbol = context.getSymbolAllocator().newSymbol("subquery", elementType);
            Assignment inputAssignment = new Assignment(valueSymbol, new SubscriptExpression(value, new LongLiteral("1")));
            Assignment nestedPlanAssignment = new Assignment(listSymbol, new SubscriptExpression(list, new LongLiteral("1")));
            Expression comparison = function.apply(valueSymbol, listSymbol);
            return Optional.of(new Unwrapping(comparison, inputAssignment, nestedPlanAssignment));
        }
    }
    return Optional.empty();
}