use of org.apache.beam.vendor.calcite.v1_28_0.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);
}
use of org.apache.beam.vendor.calcite.v1_28_0.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);
}
use of org.apache.beam.vendor.calcite.v1_28_0.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);
}
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(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);
}
use of org.apache.beam.vendor.calcite.v1_28_0.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));
}
Aggregations