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

Example 81 with RexInputRef

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

the class CommonExecPythonCorrelate method extractPythonTableFunctionInfo.

private Tuple2<int[], PythonFunctionInfo> extractPythonTableFunctionInfo() {
    LinkedHashMap<RexNode, Integer> inputNodes = new LinkedHashMap<>();
    PythonFunctionInfo pythonTableFunctionInfo = CommonPythonUtil.createPythonFunctionInfo(invocation, inputNodes);
    int[] udtfInputOffsets = inputNodes.keySet().stream().filter(x -> x instanceof RexInputRef).map(x -> ((RexInputRef) x).getIndex()).mapToInt(i -> i).toArray();
    return Tuple2.of(udtfInputOffsets, pythonTableFunctionInfo);
}

Example 82 with RexInputRef

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

the class FlinkRelMdCollation method project.

/**
 * Helper method to determine a {@link Project}'s collation.
 */
public static List<RelCollation> project(RelMetadataQuery mq, RelNode input, List<? extends RexNode> projects) {
    final SortedSet<RelCollation> collations = new TreeSet<>();
    final List<RelCollation> inputCollations = mq.collations(input);
    if (inputCollations == null || inputCollations.isEmpty()) {
        return com.google.common.collect.ImmutableList.of();
    }
    final com.google.common.collect.Multimap<Integer, Integer> targets = com.google.common.collect.LinkedListMultimap.create();
    final Map<Integer, SqlMonotonicity> targetsWithMonotonicity = new HashMap<>();
    for (Ord<RexNode> project : Ord.<RexNode>zip(projects)) {
        if (project.e instanceof RexInputRef) {
            targets.put(((RexInputRef) project.e).getIndex(), project.i);
        } else if (project.e instanceof RexCall) {
            final RexCall call = (RexCall) project.e;
            final RexCallBinding binding = RexCallBinding.create(input.getCluster().getTypeFactory(), call, inputCollations);
            targetsWithMonotonicity.put(project.i, call.getOperator().getMonotonicity(binding));
        }
    }
    final List<RelFieldCollation> fieldCollations = new ArrayList<>();
    loop: for (RelCollation ic : inputCollations) {
        if (ic.getFieldCollations().isEmpty()) {
            continue;
        }
        fieldCollations.clear();
        for (RelFieldCollation ifc : ic.getFieldCollations()) {
            final Collection<Integer> integers = targets.get(ifc.getFieldIndex());
            if (integers.isEmpty()) {
                // cannot do this collation
                continue loop;
            }
            fieldCollations.add(ifc.withFieldIndex(integers.iterator().next()));
        }
        assert !fieldCollations.isEmpty();
        collations.add(RelCollations.of(fieldCollations));
    }
    final List<RelFieldCollation> fieldCollationsForRexCalls = new ArrayList<>();
    for (Map.Entry<Integer, SqlMonotonicity> entry : targetsWithMonotonicity.entrySet()) {
        final SqlMonotonicity value = entry.getValue();
        switch(value) {
            case NOT_MONOTONIC:
            case CONSTANT:
                break;
            default:
                fieldCollationsForRexCalls.add(new RelFieldCollation(entry.getKey(), RelFieldCollation.Direction.of(value)));
                break;
        }
    }
    if (!fieldCollationsForRexCalls.isEmpty()) {
        collations.add(RelCollations.of(fieldCollationsForRexCalls));
    }
    return com.google.common.collect.ImmutableList.copyOf(collations);
}

Example 83 with RexInputRef

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

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

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