Search in sources :

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;
    }
}
Also used : RexCall(org.apache.calcite.rex.RexCall) RexShuttle(org.apache.calcite.rex.RexShuttle) RexInputRef(org.apache.calcite.rex.RexInputRef) RexUtil(org.apache.calcite.rex.RexUtil) RelOptUtil(org.apache.calcite.plan.RelOptUtil) Util(org.apache.calcite.util.Util) RexTableInputRef(org.apache.calcite.rex.RexTableInputRef) RelReferentialConstraint(org.apache.calcite.rel.RelReferentialConstraint)

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();
}
Also used : Window(org.apache.calcite.rel.core.Window) LogicalWindow(org.apache.calcite.rel.logical.LogicalWindow) RexShuttle(org.apache.calcite.rex.RexShuttle) ImmutableBitSet(org.apache.calcite.util.ImmutableBitSet) RelFieldCollation(org.apache.calcite.rel.RelFieldCollation) RexInputRef(org.apache.calcite.rex.RexInputRef) RexNode(org.apache.calcite.rex.RexNode)

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);
}
Also used : RexShuttle(org.apache.calcite.rex.RexShuttle) NestedColumn(org.apache.flink.table.planner.plan.utils.NestedColumn) LinkedList(java.util.LinkedList) RelNode(org.apache.calcite.rel.RelNode) LogicalWatermarkAssigner(org.apache.flink.table.planner.plan.nodes.calcite.LogicalWatermarkAssigner) FlinkRelBuilder(org.apache.flink.table.planner.calcite.FlinkRelBuilder) RexInputRef(org.apache.calcite.rex.RexInputRef) LogicalProject(org.apache.calcite.rel.logical.LogicalProject) NestedSchema(org.apache.flink.table.planner.plan.utils.NestedSchema) RexNode(org.apache.calcite.rex.RexNode)

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();
}
Also used : RexLiteral(org.apache.calcite.rex.RexLiteral) HashMap(java.util.HashMap) IdentityHashMap(java.util.IdentityHashMap) IdentityHashMap(java.util.IdentityHashMap) ArrayList(java.util.ArrayList) RelDataType(org.apache.calcite.rel.type.RelDataType) RexOver(org.apache.calcite.rex.RexOver) RexWindow(org.apache.calcite.rex.RexWindow) Window(org.apache.calcite.rel.core.Window) RexShuttle(org.apache.calcite.rex.RexShuttle) RexInputRef(org.apache.calcite.rex.RexInputRef) Collection(java.util.Collection) RexLocalRef(org.apache.calcite.rex.RexLocalRef) HashMap(java.util.HashMap) Map(java.util.Map) IdentityHashMap(java.util.IdentityHashMap) RexNode(org.apache.calcite.rex.RexNode)

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;
    }
}
Also used : RexShuttle(org.apache.calcite.rex.RexShuttle) SqlCall(org.apache.calcite.sql.SqlCall) ImmutableList(com.google.common.collect.ImmutableList) SqlKind(org.apache.calcite.sql.SqlKind) RexFieldCollation(org.apache.calcite.rex.RexFieldCollation) SqlWindow(org.apache.calcite.sql.SqlWindow) SqlNodeList(org.apache.calcite.sql.SqlNodeList) SqlLiteral(org.apache.calcite.sql.SqlLiteral) SqlNode(org.apache.calcite.sql.SqlNode) RexNode(org.apache.calcite.rex.RexNode)

Aggregations

RexShuttle (org.apache.calcite.rex.RexShuttle)20 RexNode (org.apache.calcite.rex.RexNode)16 RexInputRef (org.apache.calcite.rex.RexInputRef)13 ArrayList (java.util.ArrayList)9 RelNode (org.apache.calcite.rel.RelNode)8 RexCall (org.apache.calcite.rex.RexCall)7 RexBuilder (org.apache.calcite.rex.RexBuilder)6 RelOptUtil (org.apache.calcite.plan.RelOptUtil)5 List (java.util.List)4 Map (java.util.Map)4 RexLiteral (org.apache.calcite.rex.RexLiteral)4 RexLocalRef (org.apache.calcite.rex.RexLocalRef)4 RexProgram (org.apache.calcite.rex.RexProgram)4 RexProgramBuilder (org.apache.calcite.rex.RexProgramBuilder)4 Collections (java.util.Collections)3 LinkedList (java.util.LinkedList)3 RelDataTypeField (org.apache.calcite.rel.type.RelDataTypeField)3 RelBuilder (org.apache.calcite.tools.RelBuilder)3 HashMap (java.util.HashMap)2 Function (java.util.function.Function)2