Examples with RexNode org.apache.calcite.rex.RexNode used on opensource projects

Example 1 with RexNode

use of org.apache.calcite.rex.RexNode in project flink by apache.

the class FlinkAggregateExpandDistinctAggregatesRule method createSelectDistinct.

/**
	 * Given an {@link org.apache.calcite.rel.logical.LogicalAggregate}
	 * and the ordinals of the arguments to a
	 * particular call to an aggregate function, creates a 'select distinct'
	 * relational expression which projects the group columns and those
	 * arguments but nothing else.
	 *
	 * <p>For example, given
	 *
	 * <blockquote>
	 * <pre>select f0, count(distinct f1), count(distinct f2)
	 * from t group by f0</pre>
	 * </blockquote>
	 *
	 * and the argument list
	 *
	 * <blockquote>{2}</blockquote>
	 *
	 * returns
	 *
	 * <blockquote>
	 * <pre>select distinct f0, f2 from t</pre>
	 * </blockquote>
	 *
	 * '
	 *
	 * <p>The <code>sourceOf</code> map is populated with the source of each
	 * column; in this case sourceOf.get(0) = 0, and sourceOf.get(1) = 2.</p>
	 *
	 * @param relBuilder Relational expression builder
	 * @param aggregate Aggregate relational expression
	 * @param argList   Ordinals of columns to make distinct
	 * @param filterArg Ordinal of column to filter on, or -1
	 * @param sourceOf  Out parameter, is populated with a map of where each
	 *				  output field came from
	 * @return Aggregate relational expression which projects the required
	 * columns
	 */
private RelBuilder createSelectDistinct(RelBuilder relBuilder, Aggregate aggregate, List<Integer> argList, int filterArg, Map<Integer, Integer> sourceOf) {
    relBuilder.push(aggregate.getInput());
    final List<Pair<RexNode, String>> projects = new ArrayList<>();
    final List<RelDataTypeField> childFields = relBuilder.peek().getRowType().getFieldList();
    for (int i : aggregate.getGroupSet()) {
        sourceOf.put(i, projects.size());
        projects.add(RexInputRef.of2(i, childFields));
    }
    for (Integer arg : argList) {
        if (filterArg >= 0) {
            // Implement
            //   agg(DISTINCT arg) FILTER $f
            // by generating
            //   SELECT DISTINCT ... CASE WHEN $f THEN arg ELSE NULL END AS arg
            // and then applying
            //   agg(arg)
            // as usual.
            //
            // It works except for (rare) agg functions that need to see null
            // values.
            final RexBuilder rexBuilder = aggregate.getCluster().getRexBuilder();
            final RexInputRef filterRef = RexInputRef.of(filterArg, childFields);
            final Pair<RexNode, String> argRef = RexInputRef.of2(arg, childFields);
            RexNode condition = rexBuilder.makeCall(SqlStdOperatorTable.CASE, filterRef, argRef.left, rexBuilder.ensureType(argRef.left.getType(), rexBuilder.constantNull(), true));
            sourceOf.put(arg, projects.size());
            projects.add(Pair.of(condition, "i$" + argRef.right));
            continue;
        }
        if (sourceOf.get(arg) != null) {
            continue;
        }
        sourceOf.put(arg, projects.size());
        projects.add(RexInputRef.of2(arg, childFields));
    }
    relBuilder.project(Pair.left(projects), Pair.right(projects));
    // Get the distinct values of the GROUP BY fields and the arguments
    // to the agg functions.
    relBuilder.push(aggregate.copy(aggregate.getTraitSet(), relBuilder.build(), false, ImmutableBitSet.range(projects.size()), null, ImmutableList.<AggregateCall>of()));
    return relBuilder;
}

Example 2 with RexNode

use of org.apache.calcite.rex.RexNode in project flink by apache.

the class FlinkAggregateJoinTransposeRule method populateEquivalences.

private static void populateEquivalences(Map<Integer, BitSet> equivalence, RexNode predicate) {
    switch(predicate.getKind()) {
        case EQUALS:
            RexCall call = (RexCall) predicate;
            final List<RexNode> operands = call.getOperands();
            if (operands.get(0) instanceof RexInputRef) {
                final RexInputRef ref0 = (RexInputRef) operands.get(0);
                if (operands.get(1) instanceof RexInputRef) {
                    final RexInputRef ref1 = (RexInputRef) operands.get(1);
                    populateEquivalence(equivalence, ref0.getIndex(), ref1.getIndex());
                    populateEquivalence(equivalence, ref1.getIndex(), ref0.getIndex());
                }
            }
    }
}

Example 3 with RexNode

use of org.apache.calcite.rex.RexNode in project flink by apache.

the class FlinkRelDecorrelator method projectJoinOutputWithNullability.

/**
	 * Pulls project above the join from its RHS input. Enforces nullability
	 * for join output.
	 *
	 * @param join             Join
	 * @param project          Original project as the right-hand input of the join
	 * @param nullIndicatorPos Position of null indicator
	 * @return the subtree with the new LogicalProject at the root
	 */
private RelNode projectJoinOutputWithNullability(LogicalJoin join, LogicalProject project, int nullIndicatorPos) {
    final RelDataTypeFactory typeFactory = join.getCluster().getTypeFactory();
    final RelNode left = join.getLeft();
    final JoinRelType joinType = join.getJoinType();
    RexInputRef nullIndicator = new RexInputRef(nullIndicatorPos, typeFactory.createTypeWithNullability(join.getRowType().getFieldList().get(nullIndicatorPos).getType(), true));
    // now create the new project
    List<Pair<RexNode, String>> newProjExprs = Lists.newArrayList();
    // project everything from the LHS and then those from the original
    // projRel
    List<RelDataTypeField> leftInputFields = left.getRowType().getFieldList();
    for (int i = 0; i < leftInputFields.size(); i++) {
        newProjExprs.add(RexInputRef.of2(i, leftInputFields));
    }
    // Marked where the projected expr is coming from so that the types will
    // become nullable for the original projections which are now coming out
    // of the nullable side of the OJ.
    boolean projectPulledAboveLeftCorrelator = joinType.generatesNullsOnRight();
    for (Pair<RexNode, String> pair : project.getNamedProjects()) {
        RexNode newProjExpr = removeCorrelationExpr(pair.left, projectPulledAboveLeftCorrelator, nullIndicator);
        newProjExprs.add(Pair.of(newProjExpr, pair.right));
    }
    return RelOptUtil.createProject(join, newProjExprs, false);
}

Example 4 with RexNode

use of org.apache.calcite.rex.RexNode in project flink by apache.

the class FlinkRelDecorrelator method decorrelateRel.

/**
	 * Rewrites a {@link LogicalAggregate}.
	 *
	 * @param rel Aggregate to rewrite
	 */
public Frame decorrelateRel(LogicalAggregate rel) {
    if (rel.getGroupType() != Aggregate.Group.SIMPLE) {
        throw new AssertionError(Bug.CALCITE_461_FIXED);
    }
    // Aggregate itself should not reference cor vars.
    assert !cm.mapRefRelToCorVar.containsKey(rel);
    final RelNode oldInput = rel.getInput();
    final Frame frame = getInvoke(oldInput, rel);
    if (frame == null) {
        // If input has not been rewritten, do not rewrite this rel.
        return null;
    }
    final RelNode newInput = frame.r;
    // map from newInput
    Map<Integer, Integer> mapNewInputToProjOutputPos = Maps.newHashMap();
    final int oldGroupKeyCount = rel.getGroupSet().cardinality();
    // Project projects the original expressions,
    // plus any correlated variables the input wants to pass along.
    final List<Pair<RexNode, String>> projects = Lists.newArrayList();
    List<RelDataTypeField> newInputOutput = newInput.getRowType().getFieldList();
    int newPos = 0;
    // oldInput has the original group by keys in the front.
    final NavigableMap<Integer, RexLiteral> omittedConstants = new TreeMap<>();
    for (int i = 0; i < oldGroupKeyCount; i++) {
        final RexLiteral constant = projectedLiteral(newInput, i);
        if (constant != null) {
            // Exclude constants. Aggregate({true}) occurs because Aggregate({})
            // would generate 1 row even when applied to an empty table.
            omittedConstants.put(i, constant);
            continue;
        }
        int newInputPos = frame.oldToNewOutputPos.get(i);
        projects.add(RexInputRef.of2(newInputPos, newInputOutput));
        mapNewInputToProjOutputPos.put(newInputPos, newPos);
        newPos++;
    }
    final SortedMap<Correlation, Integer> mapCorVarToOutputPos = new TreeMap<>();
    if (!frame.corVarOutputPos.isEmpty()) {
        // position oldGroupKeyCount.
        for (Map.Entry<Correlation, Integer> entry : frame.corVarOutputPos.entrySet()) {
            projects.add(RexInputRef.of2(entry.getValue(), newInputOutput));
            mapCorVarToOutputPos.put(entry.getKey(), newPos);
            mapNewInputToProjOutputPos.put(entry.getValue(), newPos);
            newPos++;
        }
    }
    // add the remaining fields
    final int newGroupKeyCount = newPos;
    for (int i = 0; i < newInputOutput.size(); i++) {
        if (!mapNewInputToProjOutputPos.containsKey(i)) {
            projects.add(RexInputRef.of2(i, newInputOutput));
            mapNewInputToProjOutputPos.put(i, newPos);
            newPos++;
        }
    }
    assert newPos == newInputOutput.size();
    // This Project will be what the old input maps to,
    // replacing any previous mapping from old input).
    RelNode newProject = RelOptUtil.createProject(newInput, projects, false);
    // update mappings:
    // oldInput ----> newInput
    //
    //                newProject
    //                   |
    // oldInput ----> newInput
    //
    // is transformed to
    //
    // oldInput ----> newProject
    //                   |
    //                newInput
    Map<Integer, Integer> combinedMap = Maps.newHashMap();
    for (Integer oldInputPos : frame.oldToNewOutputPos.keySet()) {
        combinedMap.put(oldInputPos, mapNewInputToProjOutputPos.get(frame.oldToNewOutputPos.get(oldInputPos)));
    }
    register(oldInput, newProject, combinedMap, mapCorVarToOutputPos);
    // now it's time to rewrite the Aggregate
    final ImmutableBitSet newGroupSet = ImmutableBitSet.range(newGroupKeyCount);
    List<AggregateCall> newAggCalls = Lists.newArrayList();
    List<AggregateCall> oldAggCalls = rel.getAggCallList();
    int oldInputOutputFieldCount = rel.getGroupSet().cardinality();
    int newInputOutputFieldCount = newGroupSet.cardinality();
    int i = -1;
    for (AggregateCall oldAggCall : oldAggCalls) {
        ++i;
        List<Integer> oldAggArgs = oldAggCall.getArgList();
        List<Integer> aggArgs = Lists.newArrayList();
        // for the argument.
        for (int oldPos : oldAggArgs) {
            aggArgs.add(combinedMap.get(oldPos));
        }
        final int filterArg = oldAggCall.filterArg < 0 ? oldAggCall.filterArg : combinedMap.get(oldAggCall.filterArg);
        newAggCalls.add(oldAggCall.adaptTo(newProject, aggArgs, filterArg, oldGroupKeyCount, newGroupKeyCount));
        // The old to new output position mapping will be the same as that
        // of newProject, plus any aggregates that the oldAgg produces.
        combinedMap.put(oldInputOutputFieldCount + i, newInputOutputFieldCount + i);
    }
    relBuilder.push(LogicalAggregate.create(newProject, false, newGroupSet, null, newAggCalls));
    if (!omittedConstants.isEmpty()) {
        final List<RexNode> postProjects = new ArrayList<>(relBuilder.fields());
        for (Map.Entry<Integer, RexLiteral> entry : omittedConstants.descendingMap().entrySet()) {
            postProjects.add(entry.getKey() + frame.corVarOutputPos.size(), entry.getValue());
        }
        relBuilder.project(postProjects);
    }
    // located at the same position as the input newProject.
    return register(rel, relBuilder.build(), combinedMap, mapCorVarToOutputPos);
}

Example 5 with RexNode

use of org.apache.calcite.rex.RexNode in project flink by apache.

the class FlinkRelDecorrelator method decorrelateRel.

/**
	 * Rewrite LogicalProject.
	 *
	 * @param rel the project rel to rewrite
	 */
public Frame decorrelateRel(LogicalProject rel) {
    //
    // Rewrite logic:
    //
    // 1. Pass along any correlated variables coming from the input.
    //
    final RelNode oldInput = rel.getInput();
    Frame frame = getInvoke(oldInput, rel);
    if (frame == null) {
        // If input has not been rewritten, do not rewrite this rel.
        return null;
    }
    final List<RexNode> oldProjects = rel.getProjects();
    final List<RelDataTypeField> relOutput = rel.getRowType().getFieldList();
    // LogicalProject projects the original expressions,
    // plus any correlated variables the input wants to pass along.
    final List<Pair<RexNode, String>> projects = Lists.newArrayList();
    // and produce the correlated variables in the new output.
    if (cm.mapRefRelToCorVar.containsKey(rel)) {
        decorrelateInputWithValueGenerator(rel);
        // The old input should be mapped to the LogicalJoin created by
        // rewriteInputWithValueGenerator().
        frame = map.get(oldInput);
    }
    // LogicalProject projects the original expressions
    final Map<Integer, Integer> mapOldToNewOutputPos = Maps.newHashMap();
    int newPos;
    for (newPos = 0; newPos < oldProjects.size(); newPos++) {
        projects.add(newPos, Pair.of(decorrelateExpr(oldProjects.get(newPos)), relOutput.get(newPos).getName()));
        mapOldToNewOutputPos.put(newPos, newPos);
    }
    // Project any correlated variables the input wants to pass along.
    final SortedMap<Correlation, Integer> mapCorVarToOutputPos = new TreeMap<>();
    for (Map.Entry<Correlation, Integer> entry : frame.corVarOutputPos.entrySet()) {
        projects.add(RexInputRef.of2(entry.getValue(), frame.r.getRowType().getFieldList()));
        mapCorVarToOutputPos.put(entry.getKey(), newPos);
        newPos++;
    }
    RelNode newProject = RelOptUtil.createProject(frame.r, projects, false);
    return register(rel, newProject, mapOldToNewOutputPos, mapCorVarToOutputPos);
}