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

Example 6 with IfExpression

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

the class TestPatternRecognitionNodeSerialization method testMeasureRoundtrip.

@Test
public void testMeasureRoundtrip() {
    ObjectMapperProvider provider = new ObjectMapperProvider();
    provider.setJsonSerializers(ImmutableMap.of(Expression.class, new ExpressionSerialization.ExpressionSerializer()));
    provider.setJsonDeserializers(ImmutableMap.of(Expression.class, new ExpressionSerialization.ExpressionDeserializer(new SqlParser()), Type.class, new TypeDeserializer(TESTING_TYPE_MANAGER)));
    JsonCodec<Measure> codec = new JsonCodecFactory(provider).jsonCodec(Measure.class);
    assertJsonRoundTrip(codec, new Measure(new ExpressionAndValuePointers(new NullLiteral(), ImmutableList.of(), ImmutableList.of(), ImmutableSet.of(), ImmutableSet.of()), BOOLEAN));
    assertJsonRoundTrip(codec, new Measure(new ExpressionAndValuePointers(new IfExpression(new ComparisonExpression(GREATER_THAN, new SymbolReference("match_number"), new SymbolReference("x")), new GenericLiteral("BIGINT", "10"), new ArithmeticUnaryExpression(MINUS, new SymbolReference("y"))), ImmutableList.of(new Symbol("match_number"), new Symbol("x"), new Symbol("y")), ImmutableList.of(new ScalarValuePointer(new LogicalIndexPointer(ImmutableSet.of(), true, true, 0, 0), new Symbol("input_symbol_a")), new ScalarValuePointer(new LogicalIndexPointer(ImmutableSet.of(new IrLabel("A")), false, true, 1, -1), new Symbol("input_symbol_a")), new ScalarValuePointer(new LogicalIndexPointer(ImmutableSet.of(new IrLabel("B")), false, true, 1, -1), new Symbol("input_symbol_b"))), ImmutableSet.of(), ImmutableSet.of(new Symbol("match_number"))), BIGINT));
}

Example 7 with IfExpression

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

the class DecorrelateUnnest method apply.

@Override
public Result apply(CorrelatedJoinNode correlatedJoinNode, Captures captures, Context context) {
    // determine shape of the subquery
    PlanNode searchRoot = correlatedJoinNode.getSubquery();
    // 1. find EnforceSingleRowNode in the subquery
    Optional<EnforceSingleRowNode> enforceSingleRow = PlanNodeSearcher.searchFrom(searchRoot, context.getLookup()).where(EnforceSingleRowNode.class::isInstance).recurseOnlyWhen(planNode -> false).findFirst();
    if (enforceSingleRow.isPresent()) {
        searchRoot = enforceSingleRow.get().getSource();
    }
    // 2. find correlated UnnestNode in the subquery
    Optional<UnnestNode> subqueryUnnest = PlanNodeSearcher.searchFrom(searchRoot, context.getLookup()).where(node -> isSupportedUnnest(node, correlatedJoinNode.getCorrelation(), context.getLookup())).recurseOnlyWhen(node -> node instanceof ProjectNode || (node instanceof LimitNode && ((LimitNode) node).getCount() > 0) || (node instanceof TopNNode && ((TopNNode) node).getCount() > 0)).findFirst();
    if (subqueryUnnest.isEmpty()) {
        return Result.empty();
    }
    UnnestNode unnestNode = subqueryUnnest.get();
    // assign unique id to input rows
    Symbol uniqueSymbol = context.getSymbolAllocator().newSymbol("unique", BIGINT);
    PlanNode input = new AssignUniqueId(context.getIdAllocator().getNextId(), correlatedJoinNode.getInput(), uniqueSymbol);
    // pre-project unnest symbols if they were pre-projected in subquery
    // The correlated UnnestNode either unnests correlation symbols directly, or unnests symbols produced by a projection that uses only correlation symbols.
    // Here, any underlying projection that was a source of the correlated UnnestNode, is appended as a source of the rewritten UnnestNode.
    // If the projection is not necessary for UnnestNode (i.e. it does not produce any unnest symbols), it should be pruned afterwards.
    PlanNode unnestSource = context.getLookup().resolve(unnestNode.getSource());
    if (unnestSource instanceof ProjectNode) {
        ProjectNode sourceProjection = (ProjectNode) unnestSource;
        input = new ProjectNode(sourceProjection.getId(), input, Assignments.builder().putIdentities(input.getOutputSymbols()).putAll(sourceProjection.getAssignments()).build());
    }
    // determine join type for rewritten UnnestNode
    Type unnestJoinType = LEFT;
    if (enforceSingleRow.isEmpty() && correlatedJoinNode.getType() == CorrelatedJoinNode.Type.INNER && unnestNode.getJoinType() == INNER) {
        unnestJoinType = INNER;
    }
    // make sure that the rewritten node is with ordinality, which might be necessary to restore inner unnest semantics after rewrite.
    Symbol ordinalitySymbol = unnestNode.getOrdinalitySymbol().orElseGet(() -> context.getSymbolAllocator().newSymbol("ordinality", BIGINT));
    // rewrite correlated join to UnnestNode.
    UnnestNode rewrittenUnnest = new UnnestNode(context.getIdAllocator().getNextId(), input, input.getOutputSymbols(), unnestNode.getMappings(), Optional.of(ordinalitySymbol), unnestJoinType, Optional.empty());
    // restore all nodes from the subquery
    PlanNode rewrittenPlan = Rewriter.rewriteNodeSequence(correlatedJoinNode.getSubquery(), input.getOutputSymbols(), ordinalitySymbol, uniqueSymbol, rewrittenUnnest, context.getSession(), metadata, context.getLookup(), context.getIdAllocator(), context.getSymbolAllocator());
    // between unnested rows and synthetic rows added by left unnest.
    if (unnestNode.getJoinType() == INNER && rewrittenUnnest.getJoinType() == LEFT) {
        Assignments.Builder assignments = Assignments.builder().putIdentities(correlatedJoinNode.getInput().getOutputSymbols());
        for (Symbol subquerySymbol : correlatedJoinNode.getSubquery().getOutputSymbols()) {
            assignments.put(subquerySymbol, new IfExpression(new IsNullPredicate(ordinalitySymbol.toSymbolReference()), new Cast(new NullLiteral(), toSqlType(context.getSymbolAllocator().getTypes().get(subquerySymbol))), subquerySymbol.toSymbolReference()));
        }
        rewrittenPlan = new ProjectNode(context.getIdAllocator().getNextId(), rewrittenPlan, assignments.build());
    }
    // restrict outputs
    return Result.ofPlanNode(restrictOutputs(context.getIdAllocator(), rewrittenPlan, ImmutableSet.copyOf(correlatedJoinNode.getOutputSymbols())).orElse(rewrittenPlan));
}

Example 8 with IfExpression

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

the class AstBuilder method visitFunctionCall.

@Override
public Node visitFunctionCall(SqlBaseParser.FunctionCallContext context) {
    Optional<Expression> filter = visitIfPresent(context.filter(), Expression.class);
    Optional<Window> window = visitIfPresent(context.over(), Window.class);
    Optional<OrderBy> orderBy = Optional.empty();
    if (context.ORDER() != null) {
        orderBy = Optional.of(new OrderBy(visit(context.sortItem(), SortItem.class)));
    }
    QualifiedName name = getQualifiedName(context.qualifiedName());
    boolean distinct = isDistinct(context.setQuantifier());
    SqlBaseParser.NullTreatmentContext nullTreatment = context.nullTreatment();
    SqlBaseParser.ProcessingModeContext processingMode = context.processingMode();
    if (name.toString().equalsIgnoreCase("if")) {
        check(context.expression().size() == 2 || context.expression().size() == 3, "Invalid number of arguments for 'if' function", context);
        check(!window.isPresent(), "OVER clause not valid for 'if' function", context);
        check(!distinct, "DISTINCT not valid for 'if' function", context);
        check(nullTreatment == null, "Null treatment clause not valid for 'if' function", context);
        check(processingMode == null, "Running or final semantics not valid for 'if' function", context);
        check(!filter.isPresent(), "FILTER not valid for 'if' function", context);
        Expression elseExpression = null;
        if (context.expression().size() == 3) {
            elseExpression = (Expression) visit(context.expression(2));
        }
        return new IfExpression(getLocation(context), (Expression) visit(context.expression(0)), (Expression) visit(context.expression(1)), elseExpression);
    }
    if (name.toString().equalsIgnoreCase("nullif")) {
        check(context.expression().size() == 2, "Invalid number of arguments for 'nullif' function", context);
        check(!window.isPresent(), "OVER clause not valid for 'nullif' function", context);
        check(!distinct, "DISTINCT not valid for 'nullif' function", context);
        check(nullTreatment == null, "Null treatment clause not valid for 'nullif' function", context);
        check(processingMode == null, "Running or final semantics not valid for 'nullif' function", context);
        check(!filter.isPresent(), "FILTER not valid for 'nullif' function", context);
        return new NullIfExpression(getLocation(context), (Expression) visit(context.expression(0)), (Expression) visit(context.expression(1)));
    }
    if (name.toString().equalsIgnoreCase("coalesce")) {
        check(context.expression().size() >= 2, "The 'coalesce' function must have at least two arguments", context);
        check(!window.isPresent(), "OVER clause not valid for 'coalesce' function", context);
        check(!distinct, "DISTINCT not valid for 'coalesce' function", context);
        check(nullTreatment == null, "Null treatment clause not valid for 'coalesce' function", context);
        check(processingMode == null, "Running or final semantics not valid for 'coalesce' function", context);
        check(!filter.isPresent(), "FILTER not valid for 'coalesce' function", context);
        return new CoalesceExpression(getLocation(context), visit(context.expression(), Expression.class));
    }
    if (name.toString().equalsIgnoreCase("try")) {
        check(context.expression().size() == 1, "The 'try' function must have exactly one argument", context);
        check(!window.isPresent(), "OVER clause not valid for 'try' function", context);
        check(!distinct, "DISTINCT not valid for 'try' function", context);
        check(nullTreatment == null, "Null treatment clause not valid for 'try' function", context);
        check(processingMode == null, "Running or final semantics not valid for 'try' function", context);
        check(!filter.isPresent(), "FILTER not valid for 'try' function", context);
        return new TryExpression(getLocation(context), (Expression) visit(getOnlyElement(context.expression())));
    }
    if (name.toString().equalsIgnoreCase("format")) {
        check(context.expression().size() >= 2, "The 'format' function must have at least two arguments", context);
        check(!window.isPresent(), "OVER clause not valid for 'format' function", context);
        check(!distinct, "DISTINCT not valid for 'format' function", context);
        check(nullTreatment == null, "Null treatment clause not valid for 'format' function", context);
        check(processingMode == null, "Running or final semantics not valid for 'format' function", context);
        check(!filter.isPresent(), "FILTER not valid for 'format' function", context);
        return new Format(getLocation(context), visit(context.expression(), Expression.class));
    }
    if (name.toString().equalsIgnoreCase("$internal$bind")) {
        check(context.expression().size() >= 1, "The '$internal$bind' function must have at least one arguments", context);
        check(!window.isPresent(), "OVER clause not valid for '$internal$bind' function", context);
        check(!distinct, "DISTINCT not valid for '$internal$bind' function", context);
        check(nullTreatment == null, "Null treatment clause not valid for '$internal$bind' function", context);
        check(processingMode == null, "Running or final semantics not valid for '$internal$bind' function", context);
        check(!filter.isPresent(), "FILTER not valid for '$internal$bind' function", context);
        int numValues = context.expression().size() - 1;
        List<Expression> arguments = context.expression().stream().map(this::visit).map(Expression.class::cast).collect(toImmutableList());
        return new BindExpression(getLocation(context), arguments.subList(0, numValues), arguments.get(numValues));
    }
    Optional<NullTreatment> nulls = Optional.empty();
    if (nullTreatment != null) {
        if (nullTreatment.IGNORE() != null) {
            nulls = Optional.of(NullTreatment.IGNORE);
        } else if (nullTreatment.RESPECT() != null) {
            nulls = Optional.of(NullTreatment.RESPECT);
        }
    }
    Optional<ProcessingMode> mode = Optional.empty();
    if (processingMode != null) {
        if (processingMode.RUNNING() != null) {
            mode = Optional.of(new ProcessingMode(getLocation(processingMode), RUNNING));
        } else if (processingMode.FINAL() != null) {
            mode = Optional.of(new ProcessingMode(getLocation(processingMode), FINAL));
        }
    }
    List<Expression> arguments = visit(context.expression(), Expression.class);
    if (context.label != null) {
        arguments = ImmutableList.of(new DereferenceExpression(getLocation(context.label), (Identifier) visit(context.label)));
    }
    return new FunctionCall(Optional.of(getLocation(context)), name, window, filter, orderBy, distinct, nulls, mode, arguments);
}

Example 9 with IfExpression

use of io.trino.sql.tree.IfExpression 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 10 with IfExpression

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

the class QueryPlanner method planFrameOffset.

private FrameOffsetPlanAndSymbol planFrameOffset(PlanBuilder subPlan, Optional<Symbol> frameOffset) {
    if (frameOffset.isEmpty()) {
        return new FrameOffsetPlanAndSymbol(subPlan, Optional.empty());
    }
    Symbol offsetSymbol = frameOffset.get();
    Type offsetType = symbolAllocator.getTypes().get(offsetSymbol);
    // Append filter to validate offset values. They mustn't be negative or null.
    Expression zeroOffset = zeroOfType(offsetType);
    ResolvedFunction fail = plannerContext.getMetadata().resolveFunction(session, QualifiedName.of("fail"), fromTypes(VARCHAR));
    Expression predicate = new IfExpression(new ComparisonExpression(GREATER_THAN_OR_EQUAL, offsetSymbol.toSymbolReference(), zeroOffset), TRUE_LITERAL, new Cast(new FunctionCall(fail.toQualifiedName(), ImmutableList.of(new Cast(new StringLiteral("Window frame offset value must not be negative or null"), toSqlType(VARCHAR)))), toSqlType(BOOLEAN)));
    subPlan = subPlan.withNewRoot(new FilterNode(idAllocator.getNextId(), subPlan.getRoot(), predicate));
    if (offsetType.equals(BIGINT)) {
        return new FrameOffsetPlanAndSymbol(subPlan, Optional.of(offsetSymbol));
    }
    Expression offsetToBigint;
    if (offsetType instanceof DecimalType && !((DecimalType) offsetType).isShort()) {
        String maxBigint = Long.toString(Long.MAX_VALUE);
        int maxBigintPrecision = maxBigint.length();
        int actualPrecision = ((DecimalType) offsetType).getPrecision();
        if (actualPrecision < maxBigintPrecision) {
            offsetToBigint = new Cast(offsetSymbol.toSymbolReference(), toSqlType(BIGINT));
        } else if (actualPrecision > maxBigintPrecision) {
            // If the offset value exceeds max bigint, it implies that the frame bound falls beyond the partition bound.
            // In such case, the frame bound is set to the partition bound. Passing max bigint as the offset value has
            // the same effect. The offset value can be truncated to max bigint for the purpose of cast.
            offsetToBigint = new GenericLiteral("BIGINT", maxBigint);
        } else {
            offsetToBigint = new IfExpression(new ComparisonExpression(LESS_THAN_OR_EQUAL, offsetSymbol.toSymbolReference(), new DecimalLiteral(maxBigint)), new Cast(offsetSymbol.toSymbolReference(), toSqlType(BIGINT)), new GenericLiteral("BIGINT", maxBigint));
        }
    } else {
        offsetToBigint = new Cast(offsetSymbol.toSymbolReference(), toSqlType(BIGINT), false, typeCoercion.isTypeOnlyCoercion(offsetType, BIGINT));
    }
    Symbol coercedOffsetSymbol = symbolAllocator.newSymbol(offsetToBigint, BIGINT);
    subPlan = subPlan.withNewRoot(new ProjectNode(idAllocator.getNextId(), subPlan.getRoot(), Assignments.builder().putIdentities(subPlan.getRoot().getOutputSymbols()).put(coercedOffsetSymbol, offsetToBigint).build()));
    return new FrameOffsetPlanAndSymbol(subPlan, Optional.of(coercedOffsetSymbol));
}