Examples with RexLocalRef org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexLocalRef used on opensource projects

Example 21 with RexLocalRef

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexLocalRef in project calcite by apache.

the class RexProgramTest method createProg.

/**
 * Creates a program, depending on variant:
 *
 * <ol>
 * <li><code>select (x + y) + (x + 1) as a, (x + x) as b from t(x, y)</code>
 * <li><code>select (x + y) + (x + 1) as a, (x + (x + 1)) as b
 * from t(x, y)</code>
 * <li><code>select (x + y) + (x + 1) as a, (x + x) as b from t(x, y)
 * where ((x + y) &gt; 1) and ((x + y) &gt; 1)</code>
 * <li><code>select (x + y) + (x + 1) as a, (x + x) as b from t(x, y)
 * where not case
 *           when x + 1 &gt; 5 then true
 *           when y is null then null
 *           else false
 *           end</code>
 * </ol>
 */
private RexProgramBuilder createProg(int variant) {
    assert variant >= 0 && variant <= 4;
    List<RelDataType> types = Arrays.asList(typeFactory.createSqlType(SqlTypeName.INTEGER), typeFactory.createSqlType(SqlTypeName.INTEGER));
    List<String> names = Arrays.asList("x", "y");
    RelDataType inputRowType = typeFactory.createStructType(types, names);
    final RexProgramBuilder builder = new RexProgramBuilder(inputRowType, rexBuilder);
    // $t0 = x
    // $t1 = y
    // $t2 = $t0 + 1 (i.e. x + 1)
    final RexNode i0 = rexBuilder.makeInputRef(types.get(0), 0);
    final RexLiteral c1 = rexBuilder.makeExactLiteral(BigDecimal.ONE);
    final RexLiteral c5 = rexBuilder.makeExactLiteral(BigDecimal.valueOf(5L));
    RexLocalRef t2 = builder.addExpr(rexBuilder.makeCall(SqlStdOperatorTable.PLUS, i0, c1));
    // $t3 = 77 (not used)
    final RexLiteral c77 = rexBuilder.makeExactLiteral(BigDecimal.valueOf(77));
    RexLocalRef t3 = builder.addExpr(c77);
    Util.discard(t3);
    // $t4 = $t0 + $t1 (i.e. x + y)
    final RexNode i1 = rexBuilder.makeInputRef(types.get(1), 1);
    RexLocalRef t4 = builder.addExpr(rexBuilder.makeCall(SqlStdOperatorTable.PLUS, i0, i1));
    RexLocalRef t5;
    final RexLocalRef t1;
    switch(variant) {
        case 0:
        case 2:
            // $t5 = $t0 + $t0 (i.e. x + x)
            t5 = builder.addExpr(rexBuilder.makeCall(SqlStdOperatorTable.PLUS, i0, i0));
            t1 = null;
            break;
        case 1:
        case 3:
        case 4:
            // $tx = $t0 + 1
            t1 = builder.addExpr(rexBuilder.makeCall(SqlStdOperatorTable.PLUS, i0, c1));
            // $t5 = $t0 + $tx (i.e. x + (x + 1))
            t5 = builder.addExpr(rexBuilder.makeCall(SqlStdOperatorTable.PLUS, i0, t1));
            break;
        default:
            throw new AssertionError("unexpected variant " + variant);
    }
    // $t6 = $t4 + $t2 (i.e. (x + y) + (x + 1))
    RexLocalRef t6 = builder.addExpr(rexBuilder.makeCall(SqlStdOperatorTable.PLUS, t4, t2));
    builder.addProject(t6.getIndex(), "a");
    builder.addProject(t5.getIndex(), "b");
    final RexLocalRef t7;
    final RexLocalRef t8;
    switch(variant) {
        case 2:
            // $t7 = $t4 > $i0 (i.e. (x + y) > 0)
            t7 = builder.addExpr(rexBuilder.makeCall(SqlStdOperatorTable.GREATER_THAN, t4, i0));
            // $t8 = $t7 AND $t7
            t8 = builder.addExpr(and(t7, t7));
            builder.addCondition(t8);
            builder.addCondition(t7);
            break;
        case 3:
        case 4:
            // $t7 = 5
            t7 = builder.addExpr(c5);
            // $t8 = $t2 > $t7 (i.e. (x + 1) > 5)
            t8 = builder.addExpr(gt(t2, t7));
            // $t9 = true
            final RexLocalRef t9 = builder.addExpr(trueLiteral);
            // $t10 = $t1 is not null (i.e. y is not null)
            assert t1 != null;
            final RexLocalRef t10 = builder.addExpr(rexBuilder.makeCall(SqlStdOperatorTable.IS_NOT_NULL, t1));
            // $t11 = false
            final RexLocalRef t11 = builder.addExpr(falseLiteral);
            // $t12 = unknown
            final RexLocalRef t12 = builder.addExpr(unknownLiteral);
            // $t13 = case when $t8 then $t9 when $t10 then $t11 else $t12 end
            final RexLocalRef t13 = builder.addExpr(case_(t8, t9, t10, t11, t12));
            // $t14 = not $t13 (i.e. not case ... end)
            final RexLocalRef t14 = builder.addExpr(not(t13));
            // don't add 't14 is true' - that is implicit
            if (variant == 3) {
                builder.addCondition(t14);
            } else {
                // $t15 = $14 is true
                final RexLocalRef t15 = builder.addExpr(rexBuilder.makeCall(SqlStdOperatorTable.IS_TRUE, t14));
                builder.addCondition(t15);
            }
    }
    return builder;
}

Example 22 with RexLocalRef

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexLocalRef in project flink by apache.

the class PushFilterPastChangelogNormalizeRule method projectWith.

/**
 * Returns a {@link StreamPhysicalCalc} which is a copy of {@param calc}, but with the
 * projections applied from {@param projectFromCalc}.
 */
private StreamPhysicalCalc projectWith(RelBuilder relBuilder, StreamPhysicalCalc projectFromCalc, StreamPhysicalCalc calc) {
    final RexProgramBuilder programBuilder = new RexProgramBuilder(calc.getRowType(), relBuilder.getRexBuilder());
    if (calc.getProgram().getCondition() != null) {
        programBuilder.addCondition(calc.getProgram().expandLocalRef(calc.getProgram().getCondition()));
    }
    for (Pair<RexLocalRef, String> projectRef : projectFromCalc.getProgram().getNamedProjects()) {
        final RexNode project = projectFromCalc.getProgram().expandLocalRef(projectRef.left);
        programBuilder.addProject(project, projectRef.right);
    }
    final RexProgram newProgram = programBuilder.getProgram();
    return (StreamPhysicalCalc) calc.copy(calc.getTraitSet(), calc.getInput(), newProgram);
}

Example 23 with RexLocalRef

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexLocalRef in project flink by apache.

the class RelTimeIndicatorConverter method mergeCalcToMaterializeTimeIndicators.

private RelNode mergeCalcToMaterializeTimeIndicators(FlinkLogicalCalc calc, Set<Integer> refIndices) {
    RexProgram program = calc.getProgram();
    RexProgramBuilder newProgramBuilder = new RexProgramBuilder(program.getInputRowType(), rexBuilder);
    for (int idx = 0; idx < program.getNamedProjects().size(); idx++) {
        Pair<RexLocalRef, String> pair = program.getNamedProjects().get(idx);
        RexNode project = program.expandLocalRef(pair.left);
        if (refIndices.contains(idx)) {
            project = materializeTimeIndicators(project);
        }
        newProgramBuilder.addProject(project, pair.right);
    }
    if (program.getCondition() != null) {
        newProgramBuilder.addCondition(program.expandLocalRef(program.getCondition()));
    }
    RexProgram newProgram = newProgramBuilder.getProgram();
    return FlinkLogicalCalc.create(calc.getInput(), newProgram);
}

Example 24 with RexLocalRef

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexLocalRef in project calcite by apache.

the class LogicalWindow method create.

/**
 * Creates a LogicalWindow by parsing a {@link RexProgram}.
 */
public static RelNode create(RelOptCluster cluster, RelTraitSet traitSet, RelBuilder relBuilder, RelNode child, final RexProgram program) {
    final RelDataType outRowType = program.getOutputRowType();
    // Build a list of distinct groups, partitions and aggregate
    // functions.
    final Multimap<WindowKey, RexOver> windowMap = LinkedListMultimap.create();
    final int inputFieldCount = child.getRowType().getFieldCount();
    final Map<RexLiteral, RexInputRef> constantPool = new HashMap<>();
    final List<RexLiteral> constants = new ArrayList<>();
    // Identify constants in the expression tree and replace them with
    // references to newly generated constant pool.
    RexShuttle replaceConstants = new RexShuttle() {

        @Override
        public RexNode visitLiteral(RexLiteral literal) {
            RexInputRef ref = constantPool.get(literal);
            if (ref != null) {
                return ref;
            }
            constants.add(literal);
            ref = new RexInputRef(constantPool.size() + inputFieldCount, literal.getType());
            constantPool.put(literal, ref);
            return ref;
        }
    };
    // Build a list of groups, partitions, and aggregate functions. Each
    // aggregate function will add its arguments as outputs of the input
    // program.
    final Map<RexOver, RexOver> origToNewOver = new IdentityHashMap<>();
    for (RexNode agg : program.getExprList()) {
        if (agg instanceof RexOver) {
            final RexOver origOver = (RexOver) agg;
            final RexOver newOver = (RexOver) origOver.accept(replaceConstants);
            origToNewOver.put(origOver, newOver);
            addWindows(windowMap, newOver, inputFieldCount);
        }
    }
    final Map<RexOver, Window.RexWinAggCall> aggMap = new HashMap<>();
    List<Group> groups = new ArrayList<>();
    for (Map.Entry<WindowKey, Collection<RexOver>> entry : windowMap.asMap().entrySet()) {
        final WindowKey windowKey = entry.getKey();
        final List<RexWinAggCall> aggCalls = new ArrayList<>();
        for (RexOver over : entry.getValue()) {
            final RexWinAggCall aggCall = new RexWinAggCall(over.getAggOperator(), over.getType(), toInputRefs(over.operands), aggMap.size(), over.isDistinct());
            aggCalls.add(aggCall);
            aggMap.put(over, aggCall);
        }
        RexShuttle toInputRefs = new RexShuttle() {

            @Override
            public RexNode visitLocalRef(RexLocalRef localRef) {
                return new RexInputRef(localRef.getIndex(), localRef.getType());
            }
        };
        groups.add(new Group(windowKey.groupSet, windowKey.isRows, windowKey.lowerBound.accept(toInputRefs), windowKey.upperBound.accept(toInputRefs), windowKey.orderKeys, aggCalls));
    }
    // Figure out the type of the inputs to the output program.
    // They are: the inputs to this rel, followed by the outputs of
    // each window.
    final List<Window.RexWinAggCall> flattenedAggCallList = new ArrayList<>();
    final List<Map.Entry<String, RelDataType>> fieldList = new ArrayList<Map.Entry<String, RelDataType>>(child.getRowType().getFieldList());
    final int offset = fieldList.size();
    // Use better field names for agg calls that are projected.
    final Map<Integer, String> fieldNames = new HashMap<>();
    for (Ord<RexLocalRef> ref : Ord.zip(program.getProjectList())) {
        final int index = ref.e.getIndex();
        if (index >= offset) {
            fieldNames.put(index - offset, outRowType.getFieldNames().get(ref.i));
        }
    }
    for (Ord<Group> window : Ord.zip(groups)) {
        for (Ord<RexWinAggCall> over : Ord.zip(window.e.aggCalls)) {
            // Add the k-th over expression of
            // the i-th window to the output of the program.
            String name = fieldNames.get(over.i);
            if (name == null || name.startsWith("$")) {
                name = "w" + window.i + "$o" + over.i;
            }
            fieldList.add(Pair.of(name, over.e.getType()));
            flattenedAggCallList.add(over.e);
        }
    }
    final RelDataType intermediateRowType = cluster.getTypeFactory().createStructType(fieldList);
    // The output program is the windowed agg's program, combined with
    // the output calc (if it exists).
    RexShuttle shuttle = new RexShuttle() {

        public RexNode visitOver(RexOver over) {
            // Look up the aggCall which this expr was translated to.
            final Window.RexWinAggCall aggCall = aggMap.get(origToNewOver.get(over));
            assert aggCall != null;
            assert RelOptUtil.eq("over", over.getType(), "aggCall", aggCall.getType(), Litmus.THROW);
            // Find the index of the aggCall among all partitions of all
            // groups.
            final int aggCallIndex = flattenedAggCallList.indexOf(aggCall);
            assert aggCallIndex >= 0;
            // Replace expression with a reference to the window slot.
            final int index = inputFieldCount + aggCallIndex;
            assert RelOptUtil.eq("over", over.getType(), "intermed", intermediateRowType.getFieldList().get(index).getType(), Litmus.THROW);
            return new RexInputRef(index, over.getType());
        }

        public RexNode visitLocalRef(RexLocalRef localRef) {
            final int index = localRef.getIndex();
            if (index < inputFieldCount) {
                // Reference to input field.
                return localRef;
            }
            return new RexLocalRef(flattenedAggCallList.size() + index, localRef.getType());
        }
    };
    final LogicalWindow window = LogicalWindow.create(traitSet, child, constants, intermediateRowType, groups);
    // The order that the "over" calls occur in the groups and
    // partitions may not match the order in which they occurred in the
    // original expression.
    // Add a project to permute them.
    final List<RexNode> rexNodesWindow = new ArrayList<>();
    for (RexNode rexNode : program.getExprList()) {
        rexNodesWindow.add(rexNode.accept(shuttle));
    }
    final List<RexNode> refToWindow = toInputRefs(rexNodesWindow);
    final List<RexNode> projectList = new ArrayList<>();
    for (RexLocalRef inputRef : program.getProjectList()) {
        final int index = inputRef.getIndex();
        final RexInputRef ref = (RexInputRef) refToWindow.get(index);
        projectList.add(ref);
    }
    return relBuilder.push(window).project(projectList, outRowType.getFieldNames()).build();
}

Example 25 with RexLocalRef

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexLocalRef in project calcite by apache.

the class Calc method accept.

public RelNode accept(RexShuttle shuttle) {
    List<RexNode> oldExprs = program.getExprList();
    List<RexNode> exprs = shuttle.apply(oldExprs);
    List<RexLocalRef> oldProjects = program.getProjectList();
    List<RexLocalRef> projects = shuttle.apply(oldProjects);
    RexLocalRef oldCondition = program.getCondition();
    RexNode condition;
    if (oldCondition != null) {
        condition = shuttle.apply(oldCondition);
        assert condition instanceof RexLocalRef : "Invalid condition after rewrite. Expected RexLocalRef, got " + condition;
    } else {
        condition = null;
    }
    if (exprs == oldExprs && projects == oldProjects && condition == oldCondition) {
        return this;
    }
    return copy(traitSet, getInput(), new RexProgram(program.getInputRowType(), exprs, projects, (RexLocalRef) condition, program.getOutputRowType()));
}