Examples with RexProgramBuilder org.apache.calcite.rex.RexProgramBuilder used on opensource projects

Example 6 with RexProgramBuilder

use of org.apache.calcite.rex.RexProgramBuilder in project samza by apache.

the class RexToJavaCompiler method compile.

/**
 * Compiles a relational expression to a instance of {@link Expression}
 *
 * for e.g.
 *    Query : select id from profile
 *      where profile table has relational schema with id(NUMBER) and name(VARCHAR) columns.
 *    This query will result in the following relational plan
 *      LogicalProject(id=[$1])
 *        LogicalTableScan(table=[[profile]])
 *
 *    And the corresponding expressions are
 *       inputs : EnumerableTableScan (Which is the output of LogicalTableScan)
 *       nodes : [$1] Which essentially means take pick the first column from the input
 *
 *    This function converts the LogicalProject expression "[$1]" with input RexNode which is an output of TableScan
 *    to a java code that implements the interface {@link Expression}
 *
 * @param inputs Input relations/time-varying relations for this row expression
 * @param nodes relational expressions that needs to be converted to java code.
 * @return compiled expression of type {@link org.apache.samza.sql.data.Expression}
 */
public org.apache.samza.sql.data.Expression compile(List<RelNode> inputs, List<RexNode> nodes) {
    /*
     *  In case there are multiple input relations, we build a single input row type combining types of all the inputs.
     */
    final RelDataTypeFactory.FieldInfoBuilder fieldBuilder = rexBuilder.getTypeFactory().builder();
    for (RelNode input : inputs) {
        fieldBuilder.addAll(input.getRowType().getFieldList());
    }
    final RelDataType inputRowType = fieldBuilder.build();
    final RexProgramBuilder programBuilder = new RexProgramBuilder(inputRowType, rexBuilder);
    for (RexNode node : nodes) {
        programBuilder.addProject(node, null);
    }
    final RexProgram program = programBuilder.getProgram();
    final BlockBuilder builder = new BlockBuilder();
    final ParameterExpression sqlContext = Expressions.parameter(SamzaSqlExecutionContext.class, "sqlContext");
    final ParameterExpression context = Expressions.parameter(Context.class, "context");
    final ParameterExpression root = DataContext.ROOT;
    final ParameterExpression inputValues = Expressions.parameter(Object[].class, "inputValues");
    final ParameterExpression outputValues = Expressions.parameter(Object[].class, "outputValues");
    final JavaTypeFactoryImpl javaTypeFactory = new SamzaSqlJavaTypeFactoryImpl(rexBuilder.getTypeFactory().getTypeSystem());
    // public void execute(Object[] inputValues, Object[] outputValues)
    final RexToLixTranslator.InputGetter inputGetter = new RexToLixTranslator.InputGetterImpl(ImmutableList.of(Pair.of(Expressions.variable(Object[].class, "inputValues"), PhysTypeImpl.of(javaTypeFactory, inputRowType, JavaRowFormat.ARRAY, false))));
    final List<org.apache.calcite.linq4j.tree.Expression> list = RexToLixTranslator.translateProjects(program, javaTypeFactory, SqlConformanceEnum.DEFAULT, builder, null, DataContext.ROOT, inputGetter, null);
    for (int i = 0; i < list.size(); i++) {
        builder.add(Expressions.statement(Expressions.assign(Expressions.arrayIndex(outputValues, Expressions.constant(i)), list.get(i))));
    }
    return createSamzaExpressionFromCalcite(sqlContext, context, root, inputValues, outputValues, builder.toBlock());
}

Example 7 with RexProgramBuilder

use of org.apache.calcite.rex.RexProgramBuilder in project flink by apache.

the class CalcPythonCorrelateTransposeRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    FlinkLogicalCorrelate correlate = call.rel(0);
    FlinkLogicalCalc right = call.rel(2);
    RexBuilder rexBuilder = call.builder().getRexBuilder();
    FlinkLogicalCalc mergedCalc = StreamPhysicalCorrelateRule.getMergedCalc(right);
    FlinkLogicalTableFunctionScan tableScan = StreamPhysicalCorrelateRule.getTableScan(mergedCalc);
    RexProgram mergedCalcProgram = mergedCalc.getProgram();
    InputRefRewriter inputRefRewriter = new InputRefRewriter(correlate.getRowType().getFieldCount() - mergedCalc.getRowType().getFieldCount());
    List<RexNode> correlateFilters = RelOptUtil.conjunctions(mergedCalcProgram.expandLocalRef(mergedCalcProgram.getCondition())).stream().map(x -> x.accept(inputRefRewriter)).collect(Collectors.toList());
    FlinkLogicalCorrelate newCorrelate = new FlinkLogicalCorrelate(correlate.getCluster(), correlate.getTraitSet(), correlate.getLeft(), tableScan, correlate.getCorrelationId(), correlate.getRequiredColumns(), correlate.getJoinType());
    RexNode topCalcCondition = RexUtil.composeConjunction(rexBuilder, correlateFilters);
    RexProgram rexProgram = new RexProgramBuilder(newCorrelate.getRowType(), rexBuilder).getProgram();
    FlinkLogicalCalc newTopCalc = new FlinkLogicalCalc(newCorrelate.getCluster(), newCorrelate.getTraitSet(), newCorrelate, RexProgram.create(newCorrelate.getRowType(), rexProgram.getExprList(), topCalcCondition, newCorrelate.getRowType(), rexBuilder));
    call.transformTo(newTopCalc);
}

Example 8 with RexProgramBuilder

use of org.apache.calcite.rex.RexProgramBuilder in project flink by apache.

the class PushWatermarkIntoTableSourceScanAcrossCalcRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    FlinkLogicalWatermarkAssigner watermarkAssigner = call.rel(0);
    FlinkLogicalCalc calc = call.rel(1);
    RexProgram originProgram = calc.getProgram();
    List<RexNode> projectList = originProgram.getProjectList().stream().map(originProgram::expandLocalRef).collect(Collectors.toList());
    // get watermark expression
    RexNode rowTimeColumn = projectList.get(watermarkAssigner.rowtimeFieldIndex());
    RexNode newWatermarkExpr = watermarkAssigner.watermarkExpr().accept(new RexShuttle() {

        @Override
        public RexNode visitInputRef(RexInputRef inputRef) {
            return projectList.get(inputRef.getIndex());
        }
    });
    // push watermark assigner into the scan
    FlinkLogicalTableSourceScan newScan = getNewScan(watermarkAssigner, newWatermarkExpr, call.rel(2), ShortcutUtils.unwrapContext(calc).getTableConfig(), // useWatermarkAssignerRowType
    false);
    FlinkTypeFactory typeFactory = ShortcutUtils.unwrapTypeFactory(calc);
    RexBuilder builder = call.builder().getRexBuilder();
    // cast timestamp/timestamp_ltz type to rowtime type.
    RexNode newRowTimeColumn = builder.makeReinterpretCast(typeFactory.createRowtimeIndicatorType(rowTimeColumn.getType().isNullable(), rowTimeColumn.getType().getSqlTypeName() == SqlTypeName.TIMESTAMP_WITH_LOCAL_TIME_ZONE), rowTimeColumn, null);
    // build new calc program
    RexProgramBuilder programBuilder = new RexProgramBuilder(newScan.getRowType(), builder);
    List<String> outputFieldNames = originProgram.getOutputRowType().getFieldNames();
    for (int i = 0; i < projectList.size(); i++) {
        if (i == watermarkAssigner.rowtimeFieldIndex()) {
            // replace the origin computed column to keep type consistent
            programBuilder.addProject(newRowTimeColumn, outputFieldNames.get(i));
        } else {
            programBuilder.addProject(projectList.get(i), outputFieldNames.get(i));
        }
    }
    if (originProgram.getCondition() != null) {
        programBuilder.addCondition(originProgram.expandLocalRef(originProgram.getCondition()));
    }
    FlinkLogicalCalc newCalc = FlinkLogicalCalc.create(newScan, programBuilder.getProgram());
    call.transformTo(newCalc);
}

Example 9 with RexProgramBuilder

use of org.apache.calcite.rex.RexProgramBuilder in project flink by apache.

the class WatermarkAssignerChangelogNormalizeTransposeRule method createTransformedProgramWithAllShuffleKeys.

private RexProgram createTransformedProgramWithAllShuffleKeys(RexProgram program, List<Integer> projectsOutShuffleKeys, RexBuilder rexBuilder) {
    RelDataType oldInputRowType = program.getInputRowType();
    List<String> visitedProjectNames = new ArrayList<>();
    RexProgramBuilder newProgramBuilder = new RexProgramBuilder(oldInputRowType, rexBuilder);
    program.getNamedProjects().forEach(pair -> {
        newProgramBuilder.addProject(program.expandLocalRef(pair.left), pair.right);
        visitedProjectNames.add(pair.right);
    });
    List<RelDataTypeField> oldFieldList = oldInputRowType.getFieldList();
    for (Integer projectsOutShuffleKey : projectsOutShuffleKeys) {
        RelDataTypeField oldField = oldFieldList.get(projectsOutShuffleKey);
        String oldFieldName = oldField.getName();
        String newProjectName = RowTypeUtils.getUniqueName(oldFieldName, visitedProjectNames);
        newProgramBuilder.addProject(new RexInputRef(projectsOutShuffleKey, oldField.getType()), newProjectName);
        visitedProjectNames.add(newProjectName);
    }
    if (program.getCondition() != null) {
        newProgramBuilder.addCondition(program.expandLocalRef(program.getCondition()));
    }
    return newProgramBuilder.getProgram();
}

Example 10 with RexProgramBuilder

use of org.apache.calcite.rex.RexProgramBuilder in project flink by apache.

the class WatermarkAssignerChangelogNormalizeTransposeRule method pushDownTransformedWatermark.

private RelNode pushDownTransformedWatermark(StreamPhysicalWatermarkAssigner watermark, StreamPhysicalCalc calc, StreamPhysicalChangelogNormalize changelogNormalize, StreamPhysicalExchange exchange, Mappings.TargetMapping calcMapping, RexBuilder rexBuilder) {
    Mappings.TargetMapping inversedMapping = calcMapping.inverse();
    final int newRowTimeFieldIndex = inversedMapping.getTargetOpt(watermark.rowtimeFieldIndex());
    // Updates watermark properties after push down before Calc
    // 1. rewrites watermark expression
    // 2. clears distribution
    // 3. updates row time field index
    RexNode newWatermarkExpr = watermark.watermarkExpr();
    if (watermark.watermarkExpr() != null) {
        newWatermarkExpr = RexUtil.apply(inversedMapping, watermark.watermarkExpr());
    }
    final RelNode newWatermark = watermark.copy(watermark.getTraitSet().plus(FlinkRelDistribution.DEFAULT()), exchange.getInput(), newRowTimeFieldIndex, newWatermarkExpr);
    final RelNode newChangelogNormalize = buildTreeInOrder(newWatermark, Tuple2.of(exchange, exchange.getTraitSet()), Tuple2.of(changelogNormalize, changelogNormalize.getTraitSet()));
    // Rewrites Calc program because the field type of row time
    // field is changed after watermark pushed down
    final RexProgram oldProgram = calc.getProgram();
    final RexProgramBuilder programBuilder = new RexProgramBuilder(newChangelogNormalize.getRowType(), rexBuilder);
    final Function<RexNode, RexNode> rexShuttle = e -> e.accept(new RexShuttle() {

        @Override
        public RexNode visitInputRef(RexInputRef inputRef) {
            if (inputRef.getIndex() == newRowTimeFieldIndex) {
                return RexInputRef.of(newRowTimeFieldIndex, newChangelogNormalize.getRowType());
            } else {
                return inputRef;
            }
        }
    });
    oldProgram.getNamedProjects().forEach(pair -> programBuilder.addProject(rexShuttle.apply(oldProgram.expandLocalRef(pair.left)), pair.right));
    if (oldProgram.getCondition() != null) {
        programBuilder.addCondition(rexShuttle.apply(oldProgram.expandLocalRef(oldProgram.getCondition())));
    }
    final RexProgram newProgram = programBuilder.getProgram();
    return calc.copy(calc.getTraitSet(), newChangelogNormalize, newProgram);
}