Examples with RexDefaultVisitor org.apache.flink.table.planner.plan.utils.RexDefaultVisitor used on opensource projects

Example 1 with RexDefaultVisitor

use of org.apache.flink.table.planner.plan.utils.RexDefaultVisitor in project flink by apache.

the class PythonCorrelateSplitRule method createNewFieldNames.

private List<String> createNewFieldNames(RelDataType rowType, RexBuilder rexBuilder, int primitiveFieldCount, ArrayBuffer<RexNode> extractedRexNodes, List<RexNode> calcProjects) {
    for (int i = 0; i < primitiveFieldCount; i++) {
        calcProjects.add(RexInputRef.of(i, rowType));
    }
    // change RexCorrelVariable to RexInputRef.
    RexDefaultVisitor<RexNode> visitor = new RexDefaultVisitor<RexNode>() {

        @Override
        public RexNode visitFieldAccess(RexFieldAccess fieldAccess) {
            RexNode expr = fieldAccess.getReferenceExpr();
            if (expr instanceof RexCorrelVariable) {
                RelDataTypeField field = fieldAccess.getField();
                return new RexInputRef(field.getIndex(), field.getType());
            } else {
                return rexBuilder.makeFieldAccess(expr.accept(this), fieldAccess.getField().getIndex());
            }
        }

        @Override
        public RexNode visitNode(RexNode rexNode) {
            return rexNode;
        }
    };
    // add the fields of the extracted rex calls.
    Iterator<RexNode> iterator = extractedRexNodes.iterator();
    while (iterator.hasNext()) {
        RexNode rexNode = iterator.next();
        if (rexNode instanceof RexCall) {
            RexCall rexCall = (RexCall) rexNode;
            List<RexNode> newProjects = rexCall.getOperands().stream().map(x -> x.accept(visitor)).collect(Collectors.toList());
            RexCall newRexCall = rexCall.clone(rexCall.getType(), newProjects);
            calcProjects.add(newRexCall);
        } else {
            calcProjects.add(rexNode);
        }
    }
    List<String> nameList = new LinkedList<>();
    for (int i = 0; i < primitiveFieldCount; i++) {
        nameList.add(rowType.getFieldNames().get(i));
    }
    Iterator<Object> indicesIterator = extractedRexNodes.indices().iterator();
    while (indicesIterator.hasNext()) {
        nameList.add("f" + indicesIterator.next());
    }
    return SqlValidatorUtil.uniquify(nameList, rexBuilder.getTypeFactory().getTypeSystem().isSchemaCaseSensitive());
}