Search in sources :

Example 26 with RexInputRef

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

the class EnumerableWindow method translateBound.

private Expression translateBound(RexToLixTranslator translator, ParameterExpression i_, Expression row_, Expression min_, Expression max_, Expression rows_, Group group, boolean lower, PhysType physType, Expression rowComparator, Expression keySelector, Expression keyComparator) {
    RexWindowBound bound = lower ? group.lowerBound : group.upperBound;
    if (bound.isUnbounded()) {
        return bound.isPreceding() ? min_ : max_;
    }
    if (group.isRows) {
        if (bound.isCurrentRow()) {
            return i_;
        }
        RexNode node = bound.getOffset();
        Expression offs = translator.translate(node);
        // Floating offset does not make sense since we refer to array index.
        // Nulls do not make sense as well.
        offs = RexToLixTranslator.convert(offs, int.class);
        Expression b = i_;
        if (bound.isFollowing()) {
            b = Expressions.add(b, offs);
        } else {
            b = Expressions.subtract(b, offs);
        }
        return b;
    }
    Expression searchLower = min_;
    Expression searchUpper = max_;
    if (bound.isCurrentRow()) {
        if (lower) {
            searchUpper = i_;
        } else {
            searchLower = i_;
        }
    }
    List<RelFieldCollation> fieldCollations = group.collation().getFieldCollations();
    if (bound.isCurrentRow() && fieldCollations.size() != 1) {
        return Expressions.call((lower ? BuiltInMethod.BINARY_SEARCH5_LOWER : BuiltInMethod.BINARY_SEARCH5_UPPER).method, rows_, row_, searchLower, searchUpper, keySelector, keyComparator);
    }
    assert fieldCollations.size() == 1 : "When using range window specification, ORDER BY should have" + " exactly one expression." + " Actual collation is " + group.collation();
    // isRange
    int orderKey = fieldCollations.get(0).getFieldIndex();
    RelDataType keyType = physType.getRowType().getFieldList().get(orderKey).getType();
    Type desiredKeyType = translator.typeFactory.getJavaClass(keyType);
    if (bound.getOffset() == null) {
        desiredKeyType = Primitive.box(desiredKeyType);
    }
    Expression val = translator.translate(new RexInputRef(orderKey, keyType), desiredKeyType);
    if (!bound.isCurrentRow()) {
        RexNode node = bound.getOffset();
        Expression offs = translator.translate(node);
        // TODO: support date + interval somehow
        if (bound.isFollowing()) {
            val = Expressions.add(val, offs);
        } else {
            val = Expressions.subtract(val, offs);
        }
    }
    return Expressions.call((lower ? BuiltInMethod.BINARY_SEARCH6_LOWER : BuiltInMethod.BINARY_SEARCH6_UPPER).method, rows_, val, searchLower, searchUpper, keySelector, keyComparator);
}
Also used : RelDataType(org.apache.calcite.rel.type.RelDataType) Type(java.lang.reflect.Type) BinaryExpression(org.apache.calcite.linq4j.tree.BinaryExpression) Expression(org.apache.calcite.linq4j.tree.Expression) ParameterExpression(org.apache.calcite.linq4j.tree.ParameterExpression) RexWindowBound(org.apache.calcite.rex.RexWindowBound) RelFieldCollation(org.apache.calcite.rel.RelFieldCollation) RexInputRef(org.apache.calcite.rex.RexInputRef) RelDataType(org.apache.calcite.rel.type.RelDataType) RexNode(org.apache.calcite.rex.RexNode)

Example 27 with RexInputRef

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

the class AggregateExpandDistinctAggregatesRule method onMatch.

// ~ Methods ----------------------------------------------------------------
public void onMatch(RelOptRuleCall call) {
    final Aggregate aggregate = call.rel(0);
    if (!aggregate.containsDistinctCall()) {
        return;
    }
    // Find all of the agg expressions. We use a LinkedHashSet to ensure determinism.
    // find all aggregate calls without distinct
    int nonDistinctAggCallCount = 0;
    int filterCount = 0;
    int unsupportedNonDistinctAggCallCount = 0;
    final Set<Pair<List<Integer>, Integer>> argLists = new LinkedHashSet<>();
    for (AggregateCall aggCall : aggregate.getAggCallList()) {
        if (aggCall.filterArg >= 0) {
            ++filterCount;
        }
        if (!aggCall.isDistinct()) {
            ++nonDistinctAggCallCount;
            final SqlKind aggCallKind = aggCall.getAggregation().getKind();
            // We only support COUNT/SUM/MIN/MAX for the "single" count distinct optimization
            switch(aggCallKind) {
                case COUNT:
                case SUM:
                case SUM0:
                case MIN:
                case MAX:
                    break;
                default:
                    ++unsupportedNonDistinctAggCallCount;
            }
        } else {
            argLists.add(Pair.of(aggCall.getArgList(), aggCall.filterArg));
        }
    }
    final int distinctAggCallCount = aggregate.getAggCallList().size() - nonDistinctAggCallCount;
    Preconditions.checkState(argLists.size() > 0, "containsDistinctCall lied");
    // arguments then we can use a more efficient form.
    if (nonDistinctAggCallCount == 0 && argLists.size() == 1 && aggregate.getGroupType() == Group.SIMPLE) {
        final Pair<List<Integer>, Integer> pair = Iterables.getOnlyElement(argLists);
        final RelBuilder relBuilder = call.builder();
        convertMonopole(relBuilder, aggregate, pair.left, pair.right);
        call.transformTo(relBuilder.build());
        return;
    }
    if (useGroupingSets) {
        rewriteUsingGroupingSets(call, aggregate);
        return;
    }
    // we can generate multi-phase aggregates
    if (// one distinct aggregate
    distinctAggCallCount == 1 && // no filter
    filterCount == 0 && // sum/min/max/count in non-distinct aggregate
    unsupportedNonDistinctAggCallCount == 0 && nonDistinctAggCallCount > 0) {
        // one or more non-distinct aggregates
        final RelBuilder relBuilder = call.builder();
        convertSingletonDistinct(relBuilder, aggregate, argLists);
        call.transformTo(relBuilder.build());
        return;
    }
    // Create a list of the expressions which will yield the final result.
    // Initially, the expressions point to the input field.
    final List<RelDataTypeField> aggFields = aggregate.getRowType().getFieldList();
    final List<RexInputRef> refs = new ArrayList<>();
    final List<String> fieldNames = aggregate.getRowType().getFieldNames();
    final ImmutableBitSet groupSet = aggregate.getGroupSet();
    final int groupAndIndicatorCount = aggregate.getGroupCount() + aggregate.getIndicatorCount();
    for (int i : Util.range(groupAndIndicatorCount)) {
        refs.add(RexInputRef.of(i, aggFields));
    }
    // Aggregate the original relation, including any non-distinct aggregates.
    final List<AggregateCall> newAggCallList = new ArrayList<>();
    int i = -1;
    for (AggregateCall aggCall : aggregate.getAggCallList()) {
        ++i;
        if (aggCall.isDistinct()) {
            refs.add(null);
            continue;
        }
        refs.add(new RexInputRef(groupAndIndicatorCount + newAggCallList.size(), aggFields.get(groupAndIndicatorCount + i).getType()));
        newAggCallList.add(aggCall);
    }
    // In the case where there are no non-distinct aggregates (regardless of
    // whether there are group bys), there's no need to generate the
    // extra aggregate and join.
    final RelBuilder relBuilder = call.builder();
    relBuilder.push(aggregate.getInput());
    int n = 0;
    if (!newAggCallList.isEmpty()) {
        final RelBuilder.GroupKey groupKey = relBuilder.groupKey(groupSet, aggregate.getGroupSets());
        relBuilder.aggregate(groupKey, newAggCallList);
        ++n;
    }
    // set of operands.
    for (Pair<List<Integer>, Integer> argList : argLists) {
        doRewrite(relBuilder, aggregate, n++, argList.left, argList.right, refs);
    }
    relBuilder.project(refs, fieldNames);
    call.transformTo(relBuilder.build());
}
Also used : LinkedHashSet(java.util.LinkedHashSet) RelBuilder(org.apache.calcite.tools.RelBuilder) ImmutableBitSet(org.apache.calcite.util.ImmutableBitSet) ArrayList(java.util.ArrayList) SqlKind(org.apache.calcite.sql.SqlKind) AggregateCall(org.apache.calcite.rel.core.AggregateCall) RelDataTypeField(org.apache.calcite.rel.type.RelDataTypeField) RexInputRef(org.apache.calcite.rex.RexInputRef) ArrayList(java.util.ArrayList) ImmutableList(com.google.common.collect.ImmutableList) ImmutableIntList(org.apache.calcite.util.ImmutableIntList) List(java.util.List) Aggregate(org.apache.calcite.rel.core.Aggregate) LogicalAggregate(org.apache.calcite.rel.logical.LogicalAggregate) Pair(org.apache.calcite.util.Pair)

Example 28 with RexInputRef

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

the class AggregateExpandDistinctAggregatesRule method doRewrite.

/**
 * Converts all distinct aggregate calls to a given set of arguments.
 *
 * <p>This method is called several times, one for each set of arguments.
 * Each time it is called, it generates a JOIN to a new SELECT DISTINCT
 * relational expression, and modifies the set of top-level calls.
 *
 * @param aggregate Original aggregate
 * @param n         Ordinal of this in a join. {@code relBuilder} contains the
 *                  input relational expression (either the original
 *                  aggregate, the output from the previous call to this
 *                  method. {@code n} is 0 if we're converting the
 *                  first distinct aggregate in a query with no non-distinct
 *                  aggregates)
 * @param argList   Arguments to the distinct aggregate function
 * @param filterArg Argument that filters input to aggregate function, or -1
 * @param refs      Array of expressions which will be the projected by the
 *                  result of this rule. Those relating to this arg list will
 *                  be modified  @return Relational expression
 */
private void doRewrite(RelBuilder relBuilder, Aggregate aggregate, int n, List<Integer> argList, int filterArg, List<RexInputRef> refs) {
    final RexBuilder rexBuilder = aggregate.getCluster().getRexBuilder();
    final List<RelDataTypeField> leftFields;
    if (n == 0) {
        leftFields = null;
    } else {
        leftFields = relBuilder.peek().getRowType().getFieldList();
    }
    // Aggregate(
    // child,
    // {COUNT(DISTINCT 1), SUM(DISTINCT 1), SUM(2)})
    // 
    // becomes
    // 
    // Aggregate(
    // Join(
    // child,
    // Aggregate(child, < all columns > {}),
    // INNER,
    // <f2 = f5>))
    // 
    // E.g.
    // SELECT deptno, SUM(DISTINCT sal), COUNT(DISTINCT gender), MAX(age)
    // FROM Emps
    // GROUP BY deptno
    // 
    // becomes
    // 
    // SELECT e.deptno, adsal.sum_sal, adgender.count_gender, e.max_age
    // FROM (
    // SELECT deptno, MAX(age) as max_age
    // FROM Emps GROUP BY deptno) AS e
    // JOIN (
    // SELECT deptno, COUNT(gender) AS count_gender FROM (
    // SELECT DISTINCT deptno, gender FROM Emps) AS dgender
    // GROUP BY deptno) AS adgender
    // ON e.deptno = adgender.deptno
    // JOIN (
    // SELECT deptno, SUM(sal) AS sum_sal FROM (
    // SELECT DISTINCT deptno, sal FROM Emps) AS dsal
    // GROUP BY deptno) AS adsal
    // ON e.deptno = adsal.deptno
    // GROUP BY e.deptno
    // 
    // Note that if a query contains no non-distinct aggregates, then the
    // very first join/group by is omitted.  In the example above, if
    // MAX(age) is removed, then the sub-select of "e" is not needed, and
    // instead the two other group by's are joined to one another.
    // Project the columns of the GROUP BY plus the arguments
    // to the agg function.
    final Map<Integer, Integer> sourceOf = new HashMap<>();
    createSelectDistinct(relBuilder, aggregate, argList, filterArg, sourceOf);
    // Now compute the aggregate functions on top of the distinct dataset.
    // Each distinct agg becomes a non-distinct call to the corresponding
    // field from the right; for example,
    // "COUNT(DISTINCT e.sal)"
    // becomes
    // "COUNT(distinct_e.sal)".
    final List<AggregateCall> aggCallList = new ArrayList<>();
    final List<AggregateCall> aggCalls = aggregate.getAggCallList();
    final int groupAndIndicatorCount = aggregate.getGroupCount() + aggregate.getIndicatorCount();
    int i = groupAndIndicatorCount - 1;
    for (AggregateCall aggCall : aggCalls) {
        ++i;
        // COUNT(DISTINCT gender) or SUM(sal).
        if (!aggCall.isDistinct()) {
            continue;
        }
        if (!aggCall.getArgList().equals(argList)) {
            continue;
        }
        // Re-map arguments.
        final int argCount = aggCall.getArgList().size();
        final List<Integer> newArgs = new ArrayList<>(argCount);
        for (int j = 0; j < argCount; j++) {
            final Integer arg = aggCall.getArgList().get(j);
            newArgs.add(sourceOf.get(arg));
        }
        final int newFilterArg = aggCall.filterArg >= 0 ? sourceOf.get(aggCall.filterArg) : -1;
        final AggregateCall newAggCall = AggregateCall.create(aggCall.getAggregation(), false, aggCall.isApproximate(), newArgs, newFilterArg, aggCall.getType(), aggCall.getName());
        assert refs.get(i) == null;
        if (n == 0) {
            refs.set(i, new RexInputRef(groupAndIndicatorCount + aggCallList.size(), newAggCall.getType()));
        } else {
            refs.set(i, new RexInputRef(leftFields.size() + groupAndIndicatorCount + aggCallList.size(), newAggCall.getType()));
        }
        aggCallList.add(newAggCall);
    }
    final Map<Integer, Integer> map = new HashMap<>();
    for (Integer key : aggregate.getGroupSet()) {
        map.put(key, map.size());
    }
    final ImmutableBitSet newGroupSet = aggregate.getGroupSet().permute(map);
    assert newGroupSet.equals(ImmutableBitSet.range(aggregate.getGroupSet().cardinality()));
    ImmutableList<ImmutableBitSet> newGroupingSets = null;
    if (aggregate.indicator) {
        newGroupingSets = ImmutableBitSet.ORDERING.immutableSortedCopy(ImmutableBitSet.permute(aggregate.getGroupSets(), map));
    }
    relBuilder.push(aggregate.copy(aggregate.getTraitSet(), relBuilder.build(), aggregate.indicator, newGroupSet, newGroupingSets, aggCallList));
    // If there's no left child yet, no need to create the join
    if (n == 0) {
        return;
    }
    // Create the join condition. It is of the form
    // 'left.f0 = right.f0 and left.f1 = right.f1 and ...'
    // where {f0, f1, ...} are the GROUP BY fields.
    final List<RelDataTypeField> distinctFields = relBuilder.peek().getRowType().getFieldList();
    final List<RexNode> conditions = Lists.newArrayList();
    for (i = 0; i < groupAndIndicatorCount; ++i) {
        // null values form its own group
        // use "is not distinct from" so that the join condition
        // allows null values to match.
        conditions.add(rexBuilder.makeCall(SqlStdOperatorTable.IS_NOT_DISTINCT_FROM, RexInputRef.of(i, leftFields), new RexInputRef(leftFields.size() + i, distinctFields.get(i).getType())));
    }
    // Join in the new 'select distinct' relation.
    relBuilder.join(JoinRelType.INNER, conditions);
}
Also used : ImmutableBitSet(org.apache.calcite.util.ImmutableBitSet) HashMap(java.util.HashMap) LinkedHashMap(java.util.LinkedHashMap) ArrayList(java.util.ArrayList) AggregateCall(org.apache.calcite.rel.core.AggregateCall) RelDataTypeField(org.apache.calcite.rel.type.RelDataTypeField) RexBuilder(org.apache.calcite.rex.RexBuilder) RexInputRef(org.apache.calcite.rex.RexInputRef) RexNode(org.apache.calcite.rex.RexNode)

Example 29 with RexInputRef

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

the class AggregateJoinTransposeRule method onMatch.

public void onMatch(RelOptRuleCall call) {
    final Aggregate aggregate = call.rel(0);
    final Join join = call.rel(1);
    final RexBuilder rexBuilder = aggregate.getCluster().getRexBuilder();
    final RelBuilder relBuilder = call.builder();
    // If any aggregate call has a filter, bail out
    for (AggregateCall aggregateCall : aggregate.getAggCallList()) {
        if (aggregateCall.getAggregation().unwrap(SqlSplittableAggFunction.class) == null) {
            return;
        }
        if (aggregateCall.filterArg >= 0) {
            return;
        }
    }
    // aggregate operator
    if (join.getJoinType() != JoinRelType.INNER) {
        return;
    }
    if (!allowFunctions && !aggregate.getAggCallList().isEmpty()) {
        return;
    }
    // Do the columns used by the join appear in the output of the aggregate?
    final ImmutableBitSet aggregateColumns = aggregate.getGroupSet();
    final RelMetadataQuery mq = call.getMetadataQuery();
    final ImmutableBitSet keyColumns = keyColumns(aggregateColumns, mq.getPulledUpPredicates(join).pulledUpPredicates);
    final ImmutableBitSet joinColumns = RelOptUtil.InputFinder.bits(join.getCondition());
    final boolean allColumnsInAggregate = keyColumns.contains(joinColumns);
    final ImmutableBitSet belowAggregateColumns = aggregateColumns.union(joinColumns);
    // Split join condition
    final List<Integer> leftKeys = Lists.newArrayList();
    final List<Integer> rightKeys = Lists.newArrayList();
    final List<Boolean> filterNulls = Lists.newArrayList();
    RexNode nonEquiConj = RelOptUtil.splitJoinCondition(join.getLeft(), join.getRight(), join.getCondition(), leftKeys, rightKeys, filterNulls);
    // If it contains non-equi join conditions, we bail out
    if (!nonEquiConj.isAlwaysTrue()) {
        return;
    }
    // Push each aggregate function down to each side that contains all of its
    // arguments. Note that COUNT(*), because it has no arguments, can go to
    // both sides.
    final Map<Integer, Integer> map = new HashMap<>();
    final List<Side> sides = new ArrayList<>();
    int uniqueCount = 0;
    int offset = 0;
    int belowOffset = 0;
    for (int s = 0; s < 2; s++) {
        final Side side = new Side();
        final RelNode joinInput = join.getInput(s);
        int fieldCount = joinInput.getRowType().getFieldCount();
        final ImmutableBitSet fieldSet = ImmutableBitSet.range(offset, offset + fieldCount);
        final ImmutableBitSet belowAggregateKeyNotShifted = belowAggregateColumns.intersect(fieldSet);
        for (Ord<Integer> c : Ord.zip(belowAggregateKeyNotShifted)) {
            map.put(c.e, belowOffset + c.i);
        }
        final Mappings.TargetMapping mapping = s == 0 ? Mappings.createIdentity(fieldCount) : Mappings.createShiftMapping(fieldCount + offset, 0, offset, fieldCount);
        final ImmutableBitSet belowAggregateKey = belowAggregateKeyNotShifted.shift(-offset);
        final boolean unique;
        if (!allowFunctions) {
            assert aggregate.getAggCallList().isEmpty();
            // If there are no functions, it doesn't matter as much whether we
            // aggregate the inputs before the join, because there will not be
            // any functions experiencing a cartesian product effect.
            // 
            // But finding out whether the input is already unique requires a call
            // to areColumnsUnique that currently (until [CALCITE-1048] "Make
            // metadata more robust" is fixed) places a heavy load on
            // the metadata system.
            // 
            // So we choose to imagine the the input is already unique, which is
            // untrue but harmless.
            // 
            Util.discard(Bug.CALCITE_1048_FIXED);
            unique = true;
        } else {
            final Boolean unique0 = mq.areColumnsUnique(joinInput, belowAggregateKey);
            unique = unique0 != null && unique0;
        }
        if (unique) {
            ++uniqueCount;
            side.aggregate = false;
            relBuilder.push(joinInput);
            final List<RexNode> projects = new ArrayList<>();
            for (Integer i : belowAggregateKey) {
                projects.add(relBuilder.field(i));
            }
            for (Ord<AggregateCall> aggCall : Ord.zip(aggregate.getAggCallList())) {
                final SqlAggFunction aggregation = aggCall.e.getAggregation();
                final SqlSplittableAggFunction splitter = Preconditions.checkNotNull(aggregation.unwrap(SqlSplittableAggFunction.class));
                if (!aggCall.e.getArgList().isEmpty() && fieldSet.contains(ImmutableBitSet.of(aggCall.e.getArgList()))) {
                    final RexNode singleton = splitter.singleton(rexBuilder, joinInput.getRowType(), aggCall.e.transform(mapping));
                    if (singleton instanceof RexInputRef) {
                        side.split.put(aggCall.i, ((RexInputRef) singleton).getIndex());
                    } else {
                        projects.add(singleton);
                        side.split.put(aggCall.i, projects.size() - 1);
                    }
                }
            }
            relBuilder.project(projects);
            side.newInput = relBuilder.build();
        } else {
            side.aggregate = true;
            List<AggregateCall> belowAggCalls = new ArrayList<>();
            final SqlSplittableAggFunction.Registry<AggregateCall> belowAggCallRegistry = registry(belowAggCalls);
            final int oldGroupKeyCount = aggregate.getGroupCount();
            final int newGroupKeyCount = belowAggregateKey.cardinality();
            for (Ord<AggregateCall> aggCall : Ord.zip(aggregate.getAggCallList())) {
                final SqlAggFunction aggregation = aggCall.e.getAggregation();
                final SqlSplittableAggFunction splitter = Preconditions.checkNotNull(aggregation.unwrap(SqlSplittableAggFunction.class));
                final AggregateCall call1;
                if (fieldSet.contains(ImmutableBitSet.of(aggCall.e.getArgList()))) {
                    final AggregateCall splitCall = splitter.split(aggCall.e, mapping);
                    call1 = splitCall.adaptTo(joinInput, splitCall.getArgList(), splitCall.filterArg, oldGroupKeyCount, newGroupKeyCount);
                } else {
                    call1 = splitter.other(rexBuilder.getTypeFactory(), aggCall.e);
                }
                if (call1 != null) {
                    side.split.put(aggCall.i, belowAggregateKey.cardinality() + belowAggCallRegistry.register(call1));
                }
            }
            side.newInput = relBuilder.push(joinInput).aggregate(relBuilder.groupKey(belowAggregateKey, null), belowAggCalls).build();
        }
        offset += fieldCount;
        belowOffset += side.newInput.getRowType().getFieldCount();
        sides.add(side);
    }
    if (uniqueCount == 2) {
        // invocation of this rule; if we continue we might loop forever.
        return;
    }
    // Update condition
    final Mapping mapping = (Mapping) Mappings.target(new Function<Integer, Integer>() {

        public Integer apply(Integer a0) {
            return map.get(a0);
        }
    }, join.getRowType().getFieldCount(), belowOffset);
    final RexNode newCondition = RexUtil.apply(mapping, join.getCondition());
    // Create new join
    relBuilder.push(sides.get(0).newInput).push(sides.get(1).newInput).join(join.getJoinType(), newCondition);
    // Aggregate above to sum up the sub-totals
    final List<AggregateCall> newAggCalls = new ArrayList<>();
    final int groupIndicatorCount = aggregate.getGroupCount() + aggregate.getIndicatorCount();
    final int newLeftWidth = sides.get(0).newInput.getRowType().getFieldCount();
    final List<RexNode> projects = new ArrayList<>(rexBuilder.identityProjects(relBuilder.peek().getRowType()));
    for (Ord<AggregateCall> aggCall : Ord.zip(aggregate.getAggCallList())) {
        final SqlAggFunction aggregation = aggCall.e.getAggregation();
        final SqlSplittableAggFunction splitter = Preconditions.checkNotNull(aggregation.unwrap(SqlSplittableAggFunction.class));
        final Integer leftSubTotal = sides.get(0).split.get(aggCall.i);
        final Integer rightSubTotal = sides.get(1).split.get(aggCall.i);
        newAggCalls.add(splitter.topSplit(rexBuilder, registry(projects), groupIndicatorCount, relBuilder.peek().getRowType(), aggCall.e, leftSubTotal == null ? -1 : leftSubTotal, rightSubTotal == null ? -1 : rightSubTotal + newLeftWidth));
    }
    relBuilder.project(projects);
    boolean aggConvertedToProjects = false;
    if (allColumnsInAggregate) {
        // let's see if we can convert aggregate into projects
        List<RexNode> projects2 = new ArrayList<>();
        for (int key : Mappings.apply(mapping, aggregate.getGroupSet())) {
            projects2.add(relBuilder.field(key));
        }
        for (AggregateCall newAggCall : newAggCalls) {
            final SqlSplittableAggFunction splitter = newAggCall.getAggregation().unwrap(SqlSplittableAggFunction.class);
            if (splitter != null) {
                final RelDataType rowType = relBuilder.peek().getRowType();
                projects2.add(splitter.singleton(rexBuilder, rowType, newAggCall));
            }
        }
        if (projects2.size() == aggregate.getGroupSet().cardinality() + newAggCalls.size()) {
            // We successfully converted agg calls into projects.
            relBuilder.project(projects2);
            aggConvertedToProjects = true;
        }
    }
    if (!aggConvertedToProjects) {
        relBuilder.aggregate(relBuilder.groupKey(Mappings.apply(mapping, aggregate.getGroupSet()), Mappings.apply2(mapping, aggregate.getGroupSets())), newAggCalls);
    }
    call.transformTo(relBuilder.build());
}
Also used : RelMetadataQuery(org.apache.calcite.rel.metadata.RelMetadataQuery) ImmutableBitSet(org.apache.calcite.util.ImmutableBitSet) HashMap(java.util.HashMap) ArrayList(java.util.ArrayList) Mapping(org.apache.calcite.util.mapping.Mapping) RelDataType(org.apache.calcite.rel.type.RelDataType) Function(com.google.common.base.Function) SqlSplittableAggFunction(org.apache.calcite.sql.SqlSplittableAggFunction) SqlAggFunction(org.apache.calcite.sql.SqlAggFunction) RexBuilder(org.apache.calcite.rex.RexBuilder) SqlSplittableAggFunction(org.apache.calcite.sql.SqlSplittableAggFunction) RelBuilder(org.apache.calcite.tools.RelBuilder) Join(org.apache.calcite.rel.core.Join) LogicalJoin(org.apache.calcite.rel.logical.LogicalJoin) SqlAggFunction(org.apache.calcite.sql.SqlAggFunction) AggregateCall(org.apache.calcite.rel.core.AggregateCall) RelNode(org.apache.calcite.rel.RelNode) Mappings(org.apache.calcite.util.mapping.Mappings) RexInputRef(org.apache.calcite.rex.RexInputRef) Aggregate(org.apache.calcite.rel.core.Aggregate) LogicalAggregate(org.apache.calcite.rel.logical.LogicalAggregate) RexNode(org.apache.calcite.rex.RexNode)

Example 30 with RexInputRef

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

the class AggregateProjectMergeRule method apply.

public static RelNode apply(RelOptRuleCall call, Aggregate aggregate, Project project) {
    final List<Integer> newKeys = Lists.newArrayList();
    final Map<Integer, Integer> map = new HashMap<>();
    for (int key : aggregate.getGroupSet()) {
        final RexNode rex = project.getProjects().get(key);
        if (rex instanceof RexInputRef) {
            final int newKey = ((RexInputRef) rex).getIndex();
            newKeys.add(newKey);
            map.put(key, newKey);
        } else {
            // Cannot handle "GROUP BY expression"
            return null;
        }
    }
    final ImmutableBitSet newGroupSet = aggregate.getGroupSet().permute(map);
    ImmutableList<ImmutableBitSet> newGroupingSets = null;
    if (aggregate.getGroupType() != Group.SIMPLE) {
        newGroupingSets = ImmutableBitSet.ORDERING.immutableSortedCopy(ImmutableBitSet.permute(aggregate.getGroupSets(), map));
    }
    final ImmutableList.Builder<AggregateCall> aggCalls = ImmutableList.builder();
    for (AggregateCall aggregateCall : aggregate.getAggCallList()) {
        final ImmutableList.Builder<Integer> newArgs = ImmutableList.builder();
        for (int arg : aggregateCall.getArgList()) {
            final RexNode rex = project.getProjects().get(arg);
            if (rex instanceof RexInputRef) {
                newArgs.add(((RexInputRef) rex).getIndex());
            } else {
                // Cannot handle "AGG(expression)"
                return null;
            }
        }
        final int newFilterArg;
        if (aggregateCall.filterArg >= 0) {
            final RexNode rex = project.getProjects().get(aggregateCall.filterArg);
            if (!(rex instanceof RexInputRef)) {
                return null;
            }
            newFilterArg = ((RexInputRef) rex).getIndex();
        } else {
            newFilterArg = -1;
        }
        aggCalls.add(aggregateCall.copy(newArgs.build(), newFilterArg));
    }
    final Aggregate newAggregate = aggregate.copy(aggregate.getTraitSet(), project.getInput(), aggregate.indicator, newGroupSet, newGroupingSets, aggCalls.build());
    // Add a project if the group set is not in the same order or
    // contains duplicates.
    final RelBuilder relBuilder = call.builder();
    relBuilder.push(newAggregate);
    if (!newKeys.equals(newGroupSet.asList())) {
        final List<Integer> posList = Lists.newArrayList();
        for (int newKey : newKeys) {
            posList.add(newGroupSet.indexOf(newKey));
        }
        if (aggregate.indicator) {
            for (int newKey : newKeys) {
                posList.add(aggregate.getGroupCount() + newGroupSet.indexOf(newKey));
            }
        }
        for (int i = newAggregate.getGroupCount() + newAggregate.getIndicatorCount(); i < newAggregate.getRowType().getFieldCount(); i++) {
            posList.add(i);
        }
        relBuilder.project(relBuilder.fields(posList));
    }
    return relBuilder.build();
}
Also used : RelBuilder(org.apache.calcite.tools.RelBuilder) ImmutableBitSet(org.apache.calcite.util.ImmutableBitSet) HashMap(java.util.HashMap) ImmutableList(com.google.common.collect.ImmutableList) AggregateCall(org.apache.calcite.rel.core.AggregateCall) RexInputRef(org.apache.calcite.rex.RexInputRef) Aggregate(org.apache.calcite.rel.core.Aggregate) RexNode(org.apache.calcite.rex.RexNode)

Aggregations

RexInputRef (org.apache.calcite.rex.RexInputRef)241 RexNode (org.apache.calcite.rex.RexNode)200 ArrayList (java.util.ArrayList)105 RelNode (org.apache.calcite.rel.RelNode)85 RelDataTypeField (org.apache.calcite.rel.type.RelDataTypeField)80 RexCall (org.apache.calcite.rex.RexCall)67 RelDataType (org.apache.calcite.rel.type.RelDataType)63 RexBuilder (org.apache.calcite.rex.RexBuilder)54 ImmutableBitSet (org.apache.calcite.util.ImmutableBitSet)52 HashMap (java.util.HashMap)47 AggregateCall (org.apache.calcite.rel.core.AggregateCall)36 List (java.util.List)35 HashSet (java.util.HashSet)32 Pair (org.apache.calcite.util.Pair)32 RexLiteral (org.apache.calcite.rex.RexLiteral)29 Map (java.util.Map)24 RelOptUtil (org.apache.calcite.plan.RelOptUtil)24 Set (java.util.Set)20 ImmutableList (com.google.common.collect.ImmutableList)19 LinkedHashMap (java.util.LinkedHashMap)19