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

Example 41 with RexProgram

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexProgram 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 42 with RexProgram

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

the class FlinkSemiAntiJoinProjectTransposeRule method adjustCondition.

/**
 * Pulls the project above the semi/anti join and returns the resulting semi/anti join
 * condition. As a result, the semi/anti join condition should be modified such that references
 * to the LHS of a semi/anti join should now reference the children of the project that's on the
 * LHS.
 *
 * @param project LogicalProject on the LHS of the semi/anti join
 * @param join the semi/anti join
 * @return the modified semi/anti join condition
 */
private RexNode adjustCondition(LogicalProject project, Join join) {
    // create two RexPrograms -- the bottom one representing a
    // concatenation of the project and the RHS of the semi/anti join and the
    // top one representing the semi/anti join condition
    RexBuilder rexBuilder = project.getCluster().getRexBuilder();
    RelDataTypeFactory typeFactory = rexBuilder.getTypeFactory();
    RelNode rightChild = join.getRight();
    // for the bottom RexProgram, the input is a concatenation of the
    // child of the project and the RHS of the semi/anti join
    RelDataType bottomInputRowType = SqlValidatorUtil.deriveJoinRowType(project.getInput().getRowType(), rightChild.getRowType(), JoinRelType.INNER, typeFactory, null, join.getSystemFieldList());
    RexProgramBuilder bottomProgramBuilder = new RexProgramBuilder(bottomInputRowType, rexBuilder);
    // of the semi/anti join
    for (Pair<RexNode, String> pair : project.getNamedProjects()) {
        bottomProgramBuilder.addProject(pair.left, pair.right);
    }
    int nLeftFields = project.getInput().getRowType().getFieldCount();
    List<RelDataTypeField> rightFields = rightChild.getRowType().getFieldList();
    int nRightFields = rightFields.size();
    for (int i = 0; i < nRightFields; i++) {
        final RelDataTypeField field = rightFields.get(i);
        RexNode inputRef = rexBuilder.makeInputRef(field.getType(), i + nLeftFields);
        bottomProgramBuilder.addProject(inputRef, field.getName());
    }
    RexProgram bottomProgram = bottomProgramBuilder.getProgram();
    // input rowtype into the top program is the concatenation of the
    // project and the RHS of the semi/anti join
    RelDataType topInputRowType = SqlValidatorUtil.deriveJoinRowType(project.getRowType(), rightChild.getRowType(), JoinRelType.INNER, typeFactory, null, join.getSystemFieldList());
    RexProgramBuilder topProgramBuilder = new RexProgramBuilder(topInputRowType, rexBuilder);
    topProgramBuilder.addIdentity();
    topProgramBuilder.addCondition(join.getCondition());
    RexProgram topProgram = topProgramBuilder.getProgram();
    // merge the programs and expand out the local references to form
    // the new semi/anti join condition; it now references a concatenation of
    // the project's child and the RHS of the semi/anti join
    RexProgram mergedProgram = RexProgramBuilder.mergePrograms(topProgram, bottomProgram, rexBuilder);
    return mergedProgram.expandLocalRef(mergedProgram.getCondition());
}

Example 43 with RexProgram

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

the class PythonMapMergeRule method matches.

@Override
public boolean matches(RelOptRuleCall call) {
    FlinkLogicalCalc topCalc = call.rel(0);
    FlinkLogicalCalc middleCalc = call.rel(1);
    FlinkLogicalCalc bottomCalc = call.rel(2);
    RexProgram topProgram = topCalc.getProgram();
    List<RexNode> topProjects = topProgram.getProjectList().stream().map(topProgram::expandLocalRef).collect(Collectors.toList());
    if (topProjects.size() != 1 || !PythonUtil.isPythonCall(topProjects.get(0), null) || !PythonUtil.takesRowAsInput((RexCall) topProjects.get(0))) {
        return false;
    }
    RexProgram bottomProgram = bottomCalc.getProgram();
    List<RexNode> bottomProjects = bottomProgram.getProjectList().stream().map(bottomProgram::expandLocalRef).collect(Collectors.toList());
    if (bottomProjects.size() != 1 || !PythonUtil.isPythonCall(bottomProjects.get(0), null)) {
        return false;
    }
    // Only Python Functions with same Python function kind can be merged together.
    if (PythonUtil.isPythonCall(topProjects.get(0), PythonFunctionKind.GENERAL) ^ PythonUtil.isPythonCall(bottomProjects.get(0), PythonFunctionKind.GENERAL)) {
        return false;
    }
    RexProgram middleProgram = middleCalc.getProgram();
    if (topProgram.getCondition() != null || middleProgram.getCondition() != null || bottomProgram.getCondition() != null) {
        return false;
    }
    List<RexNode> middleProjects = middleProgram.getProjectList().stream().map(middleProgram::expandLocalRef).collect(Collectors.toList());
    int inputRowFieldCount = middleProgram.getInputRowType().getFieldList().get(0).getValue().getFieldList().size();
    return isFlattenCalc(middleProjects, inputRowFieldCount) && isTopCalcTakesWholeMiddleCalcAsInputs((RexCall) topProjects.get(0), middleProjects.size());
}

Example 44 with RexProgram

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

the class PythonMapMergeRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    FlinkLogicalCalc topCalc = call.rel(0);
    FlinkLogicalCalc middleCalc = call.rel(1);
    FlinkLogicalCalc bottomCalc = call.rel(2);
    RexProgram topProgram = topCalc.getProgram();
    List<RexCall> topProjects = topProgram.getProjectList().stream().map(topProgram::expandLocalRef).map(x -> (RexCall) x).collect(Collectors.toList());
    RexCall topPythonCall = topProjects.get(0);
    // merge topCalc and middleCalc
    RexCall newPythonCall = topPythonCall.clone(topPythonCall.getType(), Collections.singletonList(RexInputRef.of(0, bottomCalc.getRowType())));
    List<RexCall> topMiddleMergedProjects = Collections.singletonList(newPythonCall);
    FlinkLogicalCalc topMiddleMergedCalc = new FlinkLogicalCalc(middleCalc.getCluster(), middleCalc.getTraitSet(), bottomCalc, RexProgram.create(bottomCalc.getRowType(), topMiddleMergedProjects, null, Collections.singletonList("f0"), call.builder().getRexBuilder()));
    // merge bottomCalc
    RexBuilder rexBuilder = call.builder().getRexBuilder();
    RexProgram mergedProgram = RexProgramBuilder.mergePrograms(topMiddleMergedCalc.getProgram(), bottomCalc.getProgram(), rexBuilder);
    Calc newCalc = topMiddleMergedCalc.copy(topMiddleMergedCalc.getTraitSet(), bottomCalc.getInput(), mergedProgram);
    call.transformTo(newCalc);
}

Example 45 with RexProgram

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

the class WatermarkAssignerChangelogNormalizeTransposeRule method pushDownTransformedWatermarkAndCalc.

private RelNode pushDownTransformedWatermarkAndCalc(StreamPhysicalWatermarkAssigner watermark, StreamPhysicalCalc calc, StreamPhysicalChangelogNormalize changelogNormalize, StreamPhysicalExchange exchange, List<Integer> completeShuffleKeys, Mappings.TargetMapping calcMapping) {
    final List<Integer> projectedOutShuffleKeys = new ArrayList<>();
    for (Integer key : completeShuffleKeys) {
        int targetIdx = calcMapping.getTargetOpt(key);
        if (targetIdx < 0) {
            projectedOutShuffleKeys.add(key);
        }
    }
    // Creates a new Program which contains all shuffle keys
    final RexBuilder rexBuilder = calc.getCluster().getRexBuilder();
    final RexProgram newPushDownProgram = createTransformedProgramWithAllShuffleKeys(calc.getProgram(), projectedOutShuffleKeys, rexBuilder);
    if (newPushDownProgram.isPermutation()) {
        // permutation of its inputs
        return pushDownTransformedWatermark(watermark, calc, changelogNormalize, exchange, calcMapping, rexBuilder);
    } else {
        // 2. Adds a top Calc to remove new added shuffle keys
        return pushDownTransformedWatermarkAndCalc(newPushDownProgram, watermark, exchange, changelogNormalize, calc);
    }
}