Search in sources :

Example 11 with RexProgram

use of org.apache.calcite.rex.RexProgram 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);
}
Also used : PythonUtil(org.apache.flink.table.planner.plan.utils.PythonUtil) RexProgram(org.apache.calcite.rex.RexProgram) RexBuilder(org.apache.calcite.rex.RexBuilder) FlinkLogicalTableFunctionScan(org.apache.flink.table.planner.plan.nodes.logical.FlinkLogicalTableFunctionScan) FlinkLogicalCalc(org.apache.flink.table.planner.plan.nodes.logical.FlinkLogicalCalc) RelOptUtil(org.apache.calcite.plan.RelOptUtil) Collectors(java.util.stream.Collectors) RelOptRuleCall(org.apache.calcite.plan.RelOptRuleCall) FlinkLogicalRel(org.apache.flink.table.planner.plan.nodes.logical.FlinkLogicalRel) RelOptRule(org.apache.calcite.plan.RelOptRule) RexProgramBuilder(org.apache.calcite.rex.RexProgramBuilder) RexUtil(org.apache.calcite.rex.RexUtil) List(java.util.List) StreamPhysicalCorrelateRule(org.apache.flink.table.planner.plan.rules.physical.stream.StreamPhysicalCorrelateRule) RexNode(org.apache.calcite.rex.RexNode) JoinRelType(org.apache.calcite.rel.core.JoinRelType) FlinkLogicalCorrelate(org.apache.flink.table.planner.plan.nodes.logical.FlinkLogicalCorrelate) FlinkLogicalCorrelate(org.apache.flink.table.planner.plan.nodes.logical.FlinkLogicalCorrelate) FlinkLogicalCalc(org.apache.flink.table.planner.plan.nodes.logical.FlinkLogicalCalc) RexProgram(org.apache.calcite.rex.RexProgram) RexBuilder(org.apache.calcite.rex.RexBuilder) FlinkLogicalTableFunctionScan(org.apache.flink.table.planner.plan.nodes.logical.FlinkLogicalTableFunctionScan) RexProgramBuilder(org.apache.calcite.rex.RexProgramBuilder) RexNode(org.apache.calcite.rex.RexNode)

Example 12 with RexProgram

use of org.apache.calcite.rex.RexProgram 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);
}
Also used : RexShuttle(org.apache.calcite.rex.RexShuttle) FlinkLogicalCalc(org.apache.flink.table.planner.plan.nodes.logical.FlinkLogicalCalc) RexProgram(org.apache.calcite.rex.RexProgram) FlinkLogicalTableSourceScan(org.apache.flink.table.planner.plan.nodes.logical.FlinkLogicalTableSourceScan) FlinkLogicalWatermarkAssigner(org.apache.flink.table.planner.plan.nodes.logical.FlinkLogicalWatermarkAssigner) FlinkTypeFactory(org.apache.flink.table.planner.calcite.FlinkTypeFactory) RexInputRef(org.apache.calcite.rex.RexInputRef) RexBuilder(org.apache.calcite.rex.RexBuilder) RexProgramBuilder(org.apache.calcite.rex.RexProgramBuilder) RexNode(org.apache.calcite.rex.RexNode)

Example 13 with RexProgram

use of org.apache.calcite.rex.RexProgram 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);
}
Also used : Mappings(org.apache.calcite.util.mapping.Mappings) RexProgram(org.apache.calcite.rex.RexProgram) StreamPhysicalExchange(org.apache.flink.table.planner.plan.nodes.physical.stream.StreamPhysicalExchange) Tuple2(org.apache.flink.api.java.tuple.Tuple2) HashMap(java.util.HashMap) FlinkRelDistribution(org.apache.flink.table.planner.plan.trait.FlinkRelDistribution) RelOptUtil(org.apache.calcite.plan.RelOptUtil) Function(java.util.function.Function) ArrayList(java.util.ArrayList) RexUtil(org.apache.calcite.rex.RexUtil) StreamPhysicalChangelogNormalize(org.apache.flink.table.planner.plan.nodes.physical.stream.StreamPhysicalChangelogNormalize) RexNode(org.apache.calcite.rex.RexNode) Map(java.util.Map) RelTraitSet(org.apache.calcite.plan.RelTraitSet) FlinkRelDistributionTraitDef(org.apache.flink.table.planner.plan.trait.FlinkRelDistributionTraitDef) RelDataType(org.apache.calcite.rel.type.RelDataType) RexBuilder(org.apache.calcite.rex.RexBuilder) StreamPhysicalCalc(org.apache.flink.table.planner.plan.nodes.physical.stream.StreamPhysicalCalc) RelRule(org.apache.calcite.plan.RelRule) RelNode(org.apache.calcite.rel.RelNode) RelOptRuleCall(org.apache.calcite.plan.RelOptRuleCall) RowTypeUtils(org.apache.flink.table.planner.typeutils.RowTypeUtils) RexInputRef(org.apache.calcite.rex.RexInputRef) RelOptRule(org.apache.calcite.plan.RelOptRule) RexProgramBuilder(org.apache.calcite.rex.RexProgramBuilder) List(java.util.List) Preconditions.checkArgument(org.apache.flink.util.Preconditions.checkArgument) RexLocalRef(org.apache.calcite.rex.RexLocalRef) RelDistribution(org.apache.calcite.rel.RelDistribution) RelDataTypeField(org.apache.calcite.rel.type.RelDataTypeField) StreamPhysicalWatermarkAssigner(org.apache.flink.table.planner.plan.nodes.physical.stream.StreamPhysicalWatermarkAssigner) RexShuttle(org.apache.calcite.rex.RexShuttle) Collections(java.util.Collections) RexShuttle(org.apache.calcite.rex.RexShuttle) Mappings(org.apache.calcite.util.mapping.Mappings) RelNode(org.apache.calcite.rel.RelNode) RexProgram(org.apache.calcite.rex.RexProgram) RexInputRef(org.apache.calcite.rex.RexInputRef) RexProgramBuilder(org.apache.calcite.rex.RexProgramBuilder) RexNode(org.apache.calcite.rex.RexNode)

Example 14 with RexProgram

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

the class WatermarkAssignerChangelogNormalizeTransposeRule method pushDownTransformedWatermarkAndCalc.

private RelNode pushDownTransformedWatermarkAndCalc(RexProgram newPushDownProgram, StreamPhysicalWatermarkAssigner watermark, StreamPhysicalExchange exchange, StreamPhysicalChangelogNormalize changelogNormalize, StreamPhysicalCalc calc) {
    final RelNode pushDownCalc = calc.copy(// Clears distribution on new Calc
    calc.getTraitSet().plus(FlinkRelDistribution.DEFAULT()), exchange.getInput(), newPushDownProgram);
    final Mappings.TargetMapping mappingOfPushDownCalc = buildMapping(newPushDownProgram);
    final RelDistribution newDistribution = exchange.getDistribution().apply(mappingOfPushDownCalc);
    final RelNode newChangelogNormalize = buildTreeInOrder(pushDownCalc, Tuple2.of(watermark, watermark.getTraitSet().plus(FlinkRelDistribution.DEFAULT())), // mapping of Calc
    Tuple2.of(exchange, exchange.getTraitSet().plus(newDistribution)), Tuple2.of(changelogNormalize, changelogNormalize.getTraitSet().plus(newDistribution)));
    final List<String> newInputFieldNames = newChangelogNormalize.getRowType().getFieldNames();
    final RexProgramBuilder topProgramBuilder = new RexProgramBuilder(newChangelogNormalize.getRowType(), changelogNormalize.getCluster().getRexBuilder());
    for (int fieldIdx = 0; fieldIdx < calc.getRowType().getFieldCount(); fieldIdx++) {
        topProgramBuilder.addProject(RexInputRef.of(fieldIdx, newChangelogNormalize.getRowType()), newInputFieldNames.get(fieldIdx));
    }
    final RexProgram topProgram = topProgramBuilder.getProgram();
    return calc.copy(calc.getTraitSet(), newChangelogNormalize, topProgram);
}
Also used : RelNode(org.apache.calcite.rel.RelNode) Mappings(org.apache.calcite.util.mapping.Mappings) RexProgram(org.apache.calcite.rex.RexProgram) RexProgramBuilder(org.apache.calcite.rex.RexProgramBuilder) FlinkRelDistribution(org.apache.flink.table.planner.plan.trait.FlinkRelDistribution) RelDistribution(org.apache.calcite.rel.RelDistribution)

Example 15 with RexProgram

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

the class PythonCorrelateSplitRule method createTopCalc.

private FlinkLogicalCalc createTopCalc(int primitiveLeftFieldCount, RexBuilder rexBuilder, ArrayBuffer<RexNode> extractedRexNodes, RelDataType calcRowType, FlinkLogicalCorrelate newCorrelate) {
    RexProgram rexProgram = new RexProgramBuilder(newCorrelate.getRowType(), rexBuilder).getProgram();
    int offset = extractedRexNodes.size() + primitiveLeftFieldCount;
    // extract correlate output RexNode.
    List<RexNode> newTopCalcProjects = rexProgram.getExprList().stream().filter(x -> x instanceof RexInputRef).filter(x -> {
        int index = ((RexInputRef) x).getIndex();
        return index < primitiveLeftFieldCount || index >= offset;
    }).collect(Collectors.toList());
    return new FlinkLogicalCalc(newCorrelate.getCluster(), newCorrelate.getTraitSet(), newCorrelate, RexProgram.create(newCorrelate.getRowType(), newTopCalcProjects, null, calcRowType, rexBuilder));
}
Also used : RexFieldAccess(org.apache.calcite.rex.RexFieldAccess) RexProgram(org.apache.calcite.rex.RexProgram) RexUtil(org.apache.calcite.rex.RexUtil) SqlValidatorUtil(org.apache.calcite.sql.validate.SqlValidatorUtil) RexNode(org.apache.calcite.rex.RexNode) LinkedList(java.util.LinkedList) ArrayBuffer(scala.collection.mutable.ArrayBuffer) PythonUtil(org.apache.flink.table.planner.plan.utils.PythonUtil) RelDataType(org.apache.calcite.rel.type.RelDataType) RexDefaultVisitor(org.apache.flink.table.planner.plan.utils.RexDefaultVisitor) RexBuilder(org.apache.calcite.rex.RexBuilder) Iterator(scala.collection.Iterator) FlinkLogicalTableFunctionScan(org.apache.flink.table.planner.plan.nodes.logical.FlinkLogicalTableFunctionScan) FlinkLogicalCalc(org.apache.flink.table.planner.plan.nodes.logical.FlinkLogicalCalc) RelNode(org.apache.calcite.rel.RelNode) Collectors(java.util.stream.Collectors) RelOptRuleCall(org.apache.calcite.plan.RelOptRuleCall) RexInputRef(org.apache.calcite.rex.RexInputRef) RelOptRule(org.apache.calcite.plan.RelOptRule) RexProgramBuilder(org.apache.calcite.rex.RexProgramBuilder) List(java.util.List) StreamPhysicalCorrelateRule(org.apache.flink.table.planner.plan.rules.physical.stream.StreamPhysicalCorrelateRule) RelDataTypeField(org.apache.calcite.rel.type.RelDataTypeField) RexCorrelVariable(org.apache.calcite.rex.RexCorrelVariable) HepRelVertex(org.apache.calcite.plan.hep.HepRelVertex) FlinkLogicalCorrelate(org.apache.flink.table.planner.plan.nodes.logical.FlinkLogicalCorrelate) RexCall(org.apache.calcite.rex.RexCall) FlinkLogicalCalc(org.apache.flink.table.planner.plan.nodes.logical.FlinkLogicalCalc) RexProgram(org.apache.calcite.rex.RexProgram) RexInputRef(org.apache.calcite.rex.RexInputRef) RexProgramBuilder(org.apache.calcite.rex.RexProgramBuilder) RexNode(org.apache.calcite.rex.RexNode)

Aggregations

RexProgram (org.apache.calcite.rex.RexProgram)46 RexNode (org.apache.calcite.rex.RexNode)29 RexProgramBuilder (org.apache.calcite.rex.RexProgramBuilder)26 RexBuilder (org.apache.calcite.rex.RexBuilder)21 RelNode (org.apache.calcite.rel.RelNode)17 RelDataType (org.apache.calcite.rel.type.RelDataType)13 ArrayList (java.util.ArrayList)12 RexLocalRef (org.apache.calcite.rex.RexLocalRef)12 LogicalCalc (org.apache.calcite.rel.logical.LogicalCalc)11 FlinkLogicalCalc (org.apache.flink.table.planner.plan.nodes.logical.FlinkLogicalCalc)10 RelDataTypeField (org.apache.calcite.rel.type.RelDataTypeField)9 List (java.util.List)7 RexInputRef (org.apache.calcite.rex.RexInputRef)7 Collectors (java.util.stream.Collectors)5 RelDataTypeFactory (org.apache.calcite.rel.type.RelDataTypeFactory)5 RexCall (org.apache.calcite.rex.RexCall)5 RexShuttle (org.apache.calcite.rex.RexShuttle)5 Collections (java.util.Collections)4 RelOptUtil (org.apache.calcite.plan.RelOptUtil)4 FlinkLogicalCorrelate (org.apache.flink.table.planner.plan.nodes.logical.FlinkLogicalCorrelate)4