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

Example 16 with RexShuttle

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

the class AbstractMaterializedViewRule method shuttleReferences.

/**
 * Replaces all the possible subexpressions by input references
 * to the input node. If available, it uses the rewriting mapping
 * to change the position to reference. Takes the reference type
 * from the input node.
 */
private static RexNode shuttleReferences(final RexBuilder rexBuilder, final RexNode expr, final Multimap<String, Integer> exprsLineage, final RelNode node, final Mapping rewritingMapping) {
    try {
        RexShuttle visitor = new RexShuttle() {

            @Override
            public RexNode visitTableInputRef(RexTableInputRef ref) {
                Collection<Integer> c = exprsLineage.get(ref.toString());
                if (c.isEmpty()) {
                    // Cannot map expression
                    throw Util.FoundOne.NULL;
                }
                int pos = c.iterator().next();
                if (rewritingMapping != null) {
                    pos = rewritingMapping.getTargetOpt(pos);
                    if (pos == -1) {
                        // Cannot map expression
                        throw Util.FoundOne.NULL;
                    }
                }
                if (node != null) {
                    return rexBuilder.makeInputRef(node, pos);
                }
                return rexBuilder.makeInputRef(ref.getType(), pos);
            }

            @Override
            public RexNode visitInputRef(RexInputRef inputRef) {
                Collection<Integer> c = exprsLineage.get(inputRef.toString());
                if (c.isEmpty()) {
                    // Cannot map expression
                    throw Util.FoundOne.NULL;
                }
                int pos = c.iterator().next();
                if (rewritingMapping != null) {
                    pos = rewritingMapping.getTargetOpt(pos);
                    if (pos == -1) {
                        // Cannot map expression
                        throw Util.FoundOne.NULL;
                    }
                }
                if (node != null) {
                    return rexBuilder.makeInputRef(node, pos);
                }
                return rexBuilder.makeInputRef(inputRef.getType(), pos);
            }

            @Override
            public RexNode visitCall(final RexCall call) {
                Collection<Integer> c = exprsLineage.get(call.toString());
                if (c.isEmpty()) {
                    // Cannot map expression
                    return super.visitCall(call);
                }
                int pos = c.iterator().next();
                if (rewritingMapping != null) {
                    pos = rewritingMapping.getTargetOpt(pos);
                    if (pos == -1) {
                        // Cannot map expression
                        return super.visitCall(call);
                    }
                }
                if (node != null) {
                    return rexBuilder.makeInputRef(node, pos);
                }
                return rexBuilder.makeInputRef(call.getType(), pos);
            }
        };
        return visitor.apply(expr);
    } catch (Util.FoundOne ex) {
        Util.swallow(ex, null);
        return null;
    }
}

Example 17 with RexShuttle

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

the class ProjectWindowTransposeRule method findReference.

private ImmutableBitSet findReference(final LogicalProject project, final LogicalWindow window) {
    final int windowInputColumn = window.getInput().getRowType().getFieldCount();
    final ImmutableBitSet.Builder beReferred = ImmutableBitSet.builder();
    final RexShuttle referenceFinder = new RexShuttle() {

        @Override
        public RexNode visitInputRef(RexInputRef inputRef) {
            final int index = inputRef.getIndex();
            if (index < windowInputColumn) {
                beReferred.set(index);
            }
            return inputRef;
        }
    };
    // Reference in LogicalProject
    for (RexNode rexNode : project.getChildExps()) {
        rexNode.accept(referenceFinder);
    }
    // Reference in LogicalWindow
    for (Window.Group group : window.groups) {
        // Reference in Partition-By
        for (int index : group.keys) {
            if (index < windowInputColumn) {
                beReferred.set(index);
            }
        }
        // Reference in Order-By
        for (RelFieldCollation relFieldCollation : group.orderKeys.getFieldCollations()) {
            if (relFieldCollation.getFieldIndex() < windowInputColumn) {
                beReferred.set(relFieldCollation.getFieldIndex());
            }
        }
        // Reference in Window Functions
        for (Window.RexWinAggCall rexWinAggCall : group.aggCalls) {
            rexWinAggCall.accept(referenceFinder);
        }
    }
    return beReferred.build();
}

Example 18 with RexShuttle

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

the class ProjectWatermarkAssignerTransposeRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    LogicalProject project = call.rel(0);
    LogicalWatermarkAssigner watermarkAssigner = call.rel(1);
    // NOTES: DON'T use the nestedSchema datatype to build the transposed project.
    NestedSchema nestedSchema = getUsedFieldsInTopLevelProjectAndWatermarkAssigner(project, watermarkAssigner);
    FlinkRelBuilder builder = (FlinkRelBuilder) call.builder().push(watermarkAssigner.getInput());
    List<RexInputRef> transposedProjects = new LinkedList<>();
    List<String> usedNames = new LinkedList<>();
    // add the used column RexInputRef and names into list
    for (NestedColumn column : nestedSchema.columns().values()) {
        // mark by hand
        column.setIndexOfLeafInNewSchema(transposedProjects.size());
        column.markLeaf();
        usedNames.add(column.name());
        transposedProjects.add(builder.field(column.indexInOriginSchema()));
    }
    // get the rowtime field index in the transposed project
    String rowTimeName = watermarkAssigner.getRowType().getFieldNames().get(watermarkAssigner.rowtimeFieldIndex());
    int indexOfRowTimeInTransposedProject;
    if (nestedSchema.columns().get(rowTimeName) == null) {
        // push the RexInputRef of the rowtime into the list
        int rowTimeIndexInInput = watermarkAssigner.rowtimeFieldIndex();
        indexOfRowTimeInTransposedProject = transposedProjects.size();
        transposedProjects.add(builder.field(rowTimeIndexInInput));
        usedNames.add(rowTimeName);
    } else {
        // find rowtime ref in the list and mark the location
        indexOfRowTimeInTransposedProject = nestedSchema.columns().get(rowTimeName).indexOfLeafInNewSchema();
    }
    // the rowtime column has no rowtime indicator
    builder.project(transposedProjects, usedNames);
    // rewrite the top level field reference
    RexNode newWatermarkExpr = watermarkAssigner.watermarkExpr().accept(new RexShuttle() {

        @Override
        public RexNode visitInputRef(RexInputRef inputRef) {
            String fieldName = watermarkAssigner.getRowType().getFieldNames().get(inputRef.getIndex());
            return builder.field(nestedSchema.columns().get(fieldName).indexOfLeafInNewSchema());
        }
    });
    builder.watermark(indexOfRowTimeInTransposedProject, newWatermarkExpr);
    List<RexNode> newProjects = NestedProjectionUtil.rewrite(project.getProjects(), nestedSchema, call.builder().getRexBuilder());
    RelNode newProject = builder.project(newProjects, project.getRowType().getFieldNames()).build();
    call.transformTo(newProject);
}

Example 19 with RexShuttle

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

the class LogicalWindow method create.

/**
 * Creates a LogicalWindow by parsing a {@link RexProgram}.
 */
public static RelNode create(RelOptCluster cluster, RelTraitSet traitSet, RelBuilder relBuilder, RelNode child, final RexProgram program) {
    final RelDataType outRowType = program.getOutputRowType();
    // Build a list of distinct groups, partitions and aggregate
    // functions.
    final Multimap<WindowKey, RexOver> windowMap = LinkedListMultimap.create();
    final int inputFieldCount = child.getRowType().getFieldCount();
    final Map<RexLiteral, RexInputRef> constantPool = new HashMap<>();
    final List<RexLiteral> constants = new ArrayList<>();
    // Identify constants in the expression tree and replace them with
    // references to newly generated constant pool.
    RexShuttle replaceConstants = new RexShuttle() {

        @Override
        public RexNode visitLiteral(RexLiteral literal) {
            RexInputRef ref = constantPool.get(literal);
            if (ref != null) {
                return ref;
            }
            constants.add(literal);
            ref = new RexInputRef(constantPool.size() + inputFieldCount, literal.getType());
            constantPool.put(literal, ref);
            return ref;
        }
    };
    // Build a list of groups, partitions, and aggregate functions. Each
    // aggregate function will add its arguments as outputs of the input
    // program.
    final Map<RexOver, RexOver> origToNewOver = new IdentityHashMap<>();
    for (RexNode agg : program.getExprList()) {
        if (agg instanceof RexOver) {
            final RexOver origOver = (RexOver) agg;
            final RexOver newOver = (RexOver) origOver.accept(replaceConstants);
            origToNewOver.put(origOver, newOver);
            addWindows(windowMap, newOver, inputFieldCount);
        }
    }
    final Map<RexOver, Window.RexWinAggCall> aggMap = new HashMap<>();
    List<Group> groups = new ArrayList<>();
    for (Map.Entry<WindowKey, Collection<RexOver>> entry : windowMap.asMap().entrySet()) {
        final WindowKey windowKey = entry.getKey();
        final List<RexWinAggCall> aggCalls = new ArrayList<>();
        for (RexOver over : entry.getValue()) {
            final RexWinAggCall aggCall = new RexWinAggCall(over.getAggOperator(), over.getType(), toInputRefs(over.operands), aggMap.size(), over.isDistinct());
            aggCalls.add(aggCall);
            aggMap.put(over, aggCall);
        }
        RexShuttle toInputRefs = new RexShuttle() {

            @Override
            public RexNode visitLocalRef(RexLocalRef localRef) {
                return new RexInputRef(localRef.getIndex(), localRef.getType());
            }
        };
        groups.add(new Group(windowKey.groupSet, windowKey.isRows, windowKey.lowerBound.accept(toInputRefs), windowKey.upperBound.accept(toInputRefs), windowKey.orderKeys, aggCalls));
    }
    // Figure out the type of the inputs to the output program.
    // They are: the inputs to this rel, followed by the outputs of
    // each window.
    final List<Window.RexWinAggCall> flattenedAggCallList = new ArrayList<>();
    final List<Map.Entry<String, RelDataType>> fieldList = new ArrayList<Map.Entry<String, RelDataType>>(child.getRowType().getFieldList());
    final int offset = fieldList.size();
    // Use better field names for agg calls that are projected.
    final Map<Integer, String> fieldNames = new HashMap<>();
    for (Ord<RexLocalRef> ref : Ord.zip(program.getProjectList())) {
        final int index = ref.e.getIndex();
        if (index >= offset) {
            fieldNames.put(index - offset, outRowType.getFieldNames().get(ref.i));
        }
    }
    for (Ord<Group> window : Ord.zip(groups)) {
        for (Ord<RexWinAggCall> over : Ord.zip(window.e.aggCalls)) {
            // Add the k-th over expression of
            // the i-th window to the output of the program.
            String name = fieldNames.get(over.i);
            if (name == null || name.startsWith("$")) {
                name = "w" + window.i + "$o" + over.i;
            }
            fieldList.add(Pair.of(name, over.e.getType()));
            flattenedAggCallList.add(over.e);
        }
    }
    final RelDataType intermediateRowType = cluster.getTypeFactory().createStructType(fieldList);
    // The output program is the windowed agg's program, combined with
    // the output calc (if it exists).
    RexShuttle shuttle = new RexShuttle() {

        public RexNode visitOver(RexOver over) {
            // Look up the aggCall which this expr was translated to.
            final Window.RexWinAggCall aggCall = aggMap.get(origToNewOver.get(over));
            assert aggCall != null;
            assert RelOptUtil.eq("over", over.getType(), "aggCall", aggCall.getType(), Litmus.THROW);
            // Find the index of the aggCall among all partitions of all
            // groups.
            final int aggCallIndex = flattenedAggCallList.indexOf(aggCall);
            assert aggCallIndex >= 0;
            // Replace expression with a reference to the window slot.
            final int index = inputFieldCount + aggCallIndex;
            assert RelOptUtil.eq("over", over.getType(), "intermed", intermediateRowType.getFieldList().get(index).getType(), Litmus.THROW);
            return new RexInputRef(index, over.getType());
        }

        public RexNode visitLocalRef(RexLocalRef localRef) {
            final int index = localRef.getIndex();
            if (index < inputFieldCount) {
                // Reference to input field.
                return localRef;
            }
            return new RexLocalRef(flattenedAggCallList.size() + index, localRef.getType());
        }
    };
    final LogicalWindow window = LogicalWindow.create(traitSet, child, constants, intermediateRowType, groups);
    // The order that the "over" calls occur in the groups and
    // partitions may not match the order in which they occurred in the
    // original expression.
    // Add a project to permute them.
    final List<RexNode> rexNodesWindow = new ArrayList<>();
    for (RexNode rexNode : program.getExprList()) {
        rexNodesWindow.add(rexNode.accept(shuttle));
    }
    final List<RexNode> refToWindow = toInputRefs(rexNodesWindow);
    final List<RexNode> projectList = new ArrayList<>();
    for (RexLocalRef inputRef : program.getProjectList()) {
        final int index = inputRef.getIndex();
        final RexInputRef ref = (RexInputRef) refToWindow.get(index);
        projectList.add(ref);
    }
    return relBuilder.push(window).project(projectList, outRowType.getFieldNames()).build();
}

Example 20 with RexShuttle

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

the class SqlToRelConverter method convertOver.

private RexNode convertOver(Blackboard bb, SqlNode node) {
    SqlCall call = (SqlCall) node;
    SqlCall aggCall = call.operand(0);
    SqlNode windowOrRef = call.operand(1);
    final SqlWindow window = validator.resolveWindow(windowOrRef, bb.scope, true);
    // ROW_NUMBER() expects specific kind of framing.
    if (aggCall.getKind() == SqlKind.ROW_NUMBER) {
        window.setLowerBound(SqlWindow.createUnboundedPreceding(SqlParserPos.ZERO));
        window.setUpperBound(SqlWindow.createCurrentRow(SqlParserPos.ZERO));
        window.setRows(SqlLiteral.createBoolean(true, SqlParserPos.ZERO));
    }
    final SqlNodeList partitionList = window.getPartitionList();
    final ImmutableList.Builder<RexNode> partitionKeys = ImmutableList.builder();
    for (SqlNode partition : partitionList) {
        partitionKeys.add(bb.convertExpression(partition));
    }
    RexNode lowerBound = bb.convertExpression(window.getLowerBound());
    RexNode upperBound = bb.convertExpression(window.getUpperBound());
    SqlNodeList orderList = window.getOrderList();
    if ((orderList.size() == 0) && !window.isRows()) {
        // A logical range requires an ORDER BY clause. Use the implicit
        // ordering of this relation. There must be one, otherwise it would
        // have failed validation.
        orderList = bb.scope.getOrderList();
        if (orderList == null) {
            throw new AssertionError("Relation should have sort key for implicit ORDER BY");
        }
    }
    final ImmutableList.Builder<RexFieldCollation> orderKeys = ImmutableList.builder();
    final Set<SqlKind> flags = EnumSet.noneOf(SqlKind.class);
    for (SqlNode order : orderList) {
        flags.clear();
        RexNode e = bb.convertSortExpression(order, flags);
        orderKeys.add(new RexFieldCollation(e, flags));
    }
    try {
        Preconditions.checkArgument(bb.window == null, "already in window agg mode");
        bb.window = window;
        RexNode rexAgg = exprConverter.convertCall(bb, aggCall);
        rexAgg = rexBuilder.ensureType(validator.getValidatedNodeType(call), rexAgg, false);
        // Walk over the tree and apply 'over' to all agg functions. This is
        // necessary because the returned expression is not necessarily a call
        // to an agg function. For example, AVG(x) becomes SUM(x) / COUNT(x).
        final SqlLiteral q = aggCall.getFunctionQuantifier();
        final boolean isDistinct = q != null && q.getValue() == SqlSelectKeyword.DISTINCT;
        final RexShuttle visitor = new HistogramShuttle(partitionKeys.build(), orderKeys.build(), RexWindowBound.create(window.getLowerBound(), lowerBound), RexWindowBound.create(window.getUpperBound(), upperBound), window, isDistinct);
        RexNode overNode = rexAgg.accept(visitor);
        return overNode;
    } finally {
        bb.window = null;
    }
}