| /* |
| * Licensed to the Apache Software Foundation (ASF) under one or more |
| * contributor license agreements. See the NOTICE file distributed with |
| * this work for additional information regarding copyright ownership. |
| * The ASF licenses this file to you under the Apache License, Version 2.0 |
| * (the "License"); you may not use this file except in compliance with |
| * the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| package org.apache.calcite.rel.rules; |
| |
| import org.apache.calcite.plan.Contexts; |
| import org.apache.calcite.plan.RelOptRuleCall; |
| import org.apache.calcite.plan.RelRule; |
| import org.apache.calcite.rel.RelCollations; |
| import org.apache.calcite.rel.core.Aggregate; |
| import org.apache.calcite.rel.core.Aggregate.Group; |
| import org.apache.calcite.rel.core.AggregateCall; |
| import org.apache.calcite.rel.core.JoinRelType; |
| import org.apache.calcite.rel.core.RelFactories; |
| import org.apache.calcite.rel.logical.LogicalAggregate; |
| import org.apache.calcite.rel.type.RelDataTypeField; |
| import org.apache.calcite.rex.RexBuilder; |
| import org.apache.calcite.rex.RexInputRef; |
| import org.apache.calcite.rex.RexNode; |
| import org.apache.calcite.runtime.PairList; |
| import org.apache.calcite.sql.SqlAggFunction; |
| import org.apache.calcite.sql.SqlKind; |
| import org.apache.calcite.sql.fun.SqlStdOperatorTable; |
| import org.apache.calcite.sql.fun.SqlSumEmptyIsZeroAggFunction; |
| import org.apache.calcite.tools.RelBuilder; |
| import org.apache.calcite.tools.RelBuilderFactory; |
| import org.apache.calcite.util.ImmutableBitSet; |
| import org.apache.calcite.util.ImmutableIntList; |
| import org.apache.calcite.util.Optionality; |
| import org.apache.calcite.util.Pair; |
| import org.apache.calcite.util.Util; |
| |
| import com.google.common.collect.ImmutableList; |
| import com.google.common.collect.Iterables; |
| import com.google.common.collect.Lists; |
| |
| import org.checkerframework.checker.nullness.qual.Nullable; |
| import org.immutables.value.Value; |
| |
| import java.util.ArrayList; |
| import java.util.Collection; |
| import java.util.HashMap; |
| import java.util.HashSet; |
| import java.util.LinkedHashMap; |
| import java.util.LinkedHashSet; |
| import java.util.List; |
| import java.util.Map; |
| import java.util.NavigableSet; |
| import java.util.Set; |
| import java.util.TreeSet; |
| import java.util.stream.Collectors; |
| import java.util.stream.Stream; |
| |
| import static com.google.common.base.Preconditions.checkArgument; |
| import static com.google.common.base.Preconditions.checkState; |
| |
| import static java.util.Objects.requireNonNull; |
| |
| /** |
| * Planner rule that expands distinct aggregates |
| * (such as {@code COUNT(DISTINCT x)}) from a |
| * {@link org.apache.calcite.rel.core.Aggregate}. |
| * |
| * <p>How this is done depends upon the arguments to the function. If all |
| * functions have the same argument |
| * (e.g. {@code COUNT(DISTINCT x), SUM(DISTINCT x)} both have the argument |
| * {@code x}) then one extra {@link org.apache.calcite.rel.core.Aggregate} is |
| * sufficient. |
| * |
| * <p>If there are multiple arguments |
| * (e.g. {@code COUNT(DISTINCT x), COUNT(DISTINCT y)}) |
| * the rule creates separate {@code Aggregate}s and combines using a |
| * {@link org.apache.calcite.rel.core.Join}. |
| * |
| * @see CoreRules#AGGREGATE_EXPAND_DISTINCT_AGGREGATES |
| * @see CoreRules#AGGREGATE_EXPAND_DISTINCT_AGGREGATES_TO_JOIN |
| */ |
| @Value.Enclosing |
| public final class AggregateExpandDistinctAggregatesRule |
| extends RelRule<AggregateExpandDistinctAggregatesRule.Config> |
| implements TransformationRule { |
| |
| /** Creates an AggregateExpandDistinctAggregatesRule. */ |
| AggregateExpandDistinctAggregatesRule(Config config) { |
| super(config); |
| } |
| |
| @Deprecated // to be removed before 2.0 |
| public AggregateExpandDistinctAggregatesRule( |
| Class<? extends Aggregate> clazz, |
| boolean useGroupingSets, |
| RelBuilderFactory relBuilderFactory) { |
| this(Config.DEFAULT.withRelBuilderFactory(relBuilderFactory) |
| .withOperandSupplier(b -> |
| b.operand(clazz).anyInputs()) |
| .as(Config.class) |
| .withUsingGroupingSets(useGroupingSets)); |
| } |
| |
| @Deprecated // to be removed before 2.0 |
| public AggregateExpandDistinctAggregatesRule( |
| Class<? extends LogicalAggregate> clazz, |
| boolean useGroupingSets, |
| RelFactories.JoinFactory joinFactory) { |
| this(clazz, useGroupingSets, RelBuilder.proto(Contexts.of(joinFactory))); |
| } |
| |
| @Deprecated // to be removed before 2.0 |
| public AggregateExpandDistinctAggregatesRule( |
| Class<? extends LogicalAggregate> clazz, |
| RelFactories.JoinFactory joinFactory) { |
| this(clazz, false, RelBuilder.proto(Contexts.of(joinFactory))); |
| } |
| |
| //~ Methods ---------------------------------------------------------------- |
| |
| @Override public void onMatch(RelOptRuleCall call) { |
| final Aggregate aggregate = call.rel(0); |
| if (!aggregate.containsDistinctCall()) { |
| return; |
| } |
| |
| if (!config.isUsingGroupingSets() |
| && aggregate.groupSets.size() > 1) { |
| // Grouping sets are not handled correctly |
| // when generating joins. |
| return; |
| } |
| |
| // Find all of the agg expressions. We use a LinkedHashSet to ensure determinism. |
| final List<AggregateCall> aggCalls = aggregate.getAggCallList(); |
| // Find all aggregate calls with distinct |
| final List<AggregateCall> distinctAggCalls = aggCalls.stream() |
| .filter(AggregateCall::isDistinct).collect(Collectors.toList()); |
| // Find all aggregate calls without distinct |
| final List<AggregateCall> nonDistinctAggCalls = aggCalls.stream() |
| .filter(aggCall -> !aggCall.isDistinct()).collect(Collectors.toList()); |
| final long filterCount = aggCalls.stream() |
| .filter(aggCall -> aggCall.filterArg >= 0).count(); |
| final long unsupportedNonDistinctAggCallCount = nonDistinctAggCalls.stream() |
| .filter(aggCall -> { |
| 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: |
| return false; |
| default: |
| return true; |
| } |
| }).count(); |
| // Argument list of distinct agg calls. |
| final Set<Pair<List<Integer>, Integer>> distinctCallArgLists = distinctAggCalls.stream() |
| .map(aggCall -> Pair.of(aggCall.getArgList(), aggCall.filterArg)) |
| .collect(Collectors.toCollection(LinkedHashSet::new)); |
| |
| checkState(!distinctCallArgLists.isEmpty(), "containsDistinctCall lied"); |
| |
| // If all of the agg expressions are distinct and have the same |
| // arguments then we can use a more efficient form. |
| |
| // MAX, MIN, BIT_AND, BIT_OR always ignore distinct attribute, |
| // when they are mixed in with other distinct agg calls, |
| // we can still use this promotion. |
| |
| // Treat the agg expression with Optionality.IGNORED as distinct and |
| // re-statistic the non-distinct agg call count and the distinct agg |
| // call arguments. |
| final List<AggregateCall> nonDistinctAggCallsOfIgnoredOptionality = |
| nonDistinctAggCalls.stream().filter(aggCall -> |
| aggCall.getAggregation().getDistinctOptionality() == Optionality.IGNORED) |
| .collect(Collectors.toList()); |
| // Different with distinctCallArgLists, this list also contains args that come from |
| // agg call which can ignore the distinct constraint. |
| final Set<Pair<List<Integer>, Integer>> distinctCallArgLists2 = |
| Stream.of(distinctAggCalls, nonDistinctAggCallsOfIgnoredOptionality) |
| .flatMap(Collection::stream) |
| .map(aggCall -> Pair.of(aggCall.getArgList(), aggCall.filterArg)) |
| .collect(Collectors.toCollection(LinkedHashSet::new)); |
| |
| if ((nonDistinctAggCalls.size() - nonDistinctAggCallsOfIgnoredOptionality.size()) == 0 |
| && distinctCallArgLists2.size() == 1 |
| && aggregate.getGroupType() == Group.SIMPLE) { |
| final Pair<List<Integer>, Integer> pair = |
| Iterables.getOnlyElement(distinctCallArgLists2); |
| final RelBuilder relBuilder = call.builder(); |
| convertMonopole(relBuilder, aggregate, pair.left, pair.right); |
| call.transformTo(relBuilder.build()); |
| return; |
| } |
| |
| if (config.isUsingGroupingSets()) { |
| rewriteUsingGroupingSets(call, aggregate); |
| return; |
| } |
| |
| // If only one distinct aggregate and one or more non-distinct aggregates, |
| // we can generate multi-phase aggregates |
| if (distinctAggCalls.size() == 1 // one distinct aggregate |
| && filterCount == 0 // no filter |
| && unsupportedNonDistinctAggCallCount == 0 // sum/min/max/count in non-distinct aggregate |
| && nonDistinctAggCalls.size() > 0) { // one or more non-distinct aggregates |
| final RelBuilder relBuilder = call.builder(); |
| convertSingletonDistinct(relBuilder, aggregate, distinctCallArgLists); |
| 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<@Nullable RexInputRef> refs = new ArrayList<>(); |
| final List<String> fieldNames = aggregate.getRowType().getFieldNames(); |
| final ImmutableBitSet groupSet = aggregate.getGroupSet(); |
| final int groupCount = aggregate.getGroupCount(); |
| for (int i : Util.range(groupCount)) { |
| 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( |
| groupCount + newAggCallList.size(), |
| aggFields.get(groupCount + 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; |
| } |
| |
| // For each set of operands, find and rewrite all calls which have that |
| // set of operands. |
| for (Pair<List<Integer>, Integer> argList : distinctCallArgLists) { |
| doRewrite(relBuilder, aggregate, n++, argList.left, argList.right, refs); |
| } |
| // It is assumed doRewrite above replaces nulls in refs |
| @SuppressWarnings("assignment.type.incompatible") |
| List<RexInputRef> nonNullRefs = refs; |
| relBuilder.project(nonNullRefs, fieldNames); |
| call.transformTo(relBuilder.build()); |
| } |
| |
| /** |
| * Converts an aggregate with one distinct aggregate and one or more |
| * non-distinct aggregates to multi-phase aggregates (see reference example |
| * below). |
| * |
| * @param relBuilder Contains the input relational expression |
| * @param aggregate Original aggregate |
| * @param argLists Arguments and filters to the distinct aggregate function |
| * |
| */ |
| private static RelBuilder convertSingletonDistinct(RelBuilder relBuilder, |
| Aggregate aggregate, Set<Pair<List<Integer>, Integer>> argLists) { |
| |
| // In this case, we are assuming that there is a single distinct function. |
| // So make sure that argLists is of size one. |
| checkArgument(argLists.size() == 1); |
| |
| // For example, |
| // SELECT deptno, COUNT(*), SUM(bonus), MIN(DISTINCT sal) |
| // FROM emp |
| // GROUP BY deptno |
| // |
| // becomes |
| // |
| // SELECT deptno, SUM(cnt), SUM(bonus), MIN(sal) |
| // FROM ( |
| // SELECT deptno, COUNT(*) as cnt, SUM(bonus), sal |
| // FROM EMP |
| // GROUP BY deptno, sal) // Aggregate B |
| // GROUP BY deptno // Aggregate A |
| relBuilder.push(aggregate.getInput()); |
| |
| final List<AggregateCall> originalAggCalls = aggregate.getAggCallList(); |
| final ImmutableBitSet originalGroupSet = aggregate.getGroupSet(); |
| |
| // Add the distinct aggregate column(s) to the group-by columns, |
| // if not already a part of the group-by |
| final NavigableSet<Integer> bottomGroups = new TreeSet<>(aggregate.getGroupSet().asList()); |
| for (AggregateCall aggCall : originalAggCalls) { |
| if (aggCall.isDistinct()) { |
| bottomGroups.addAll(aggCall.getArgList()); |
| break; // since we only have single distinct call |
| } |
| } |
| final ImmutableBitSet bottomGroupSet = ImmutableBitSet.of(bottomGroups); |
| |
| // Generate the intermediate aggregate B, the one on the bottom that converts |
| // a distinct call to group by call. |
| // Bottom aggregate is the same as the original aggregate, except that |
| // the bottom aggregate has converted the DISTINCT aggregate to a group by clause. |
| final List<AggregateCall> bottomAggregateCalls = new ArrayList<>(); |
| for (AggregateCall aggCall : originalAggCalls) { |
| // Project the column corresponding to the distinct aggregate. Project |
| // as-is all the non-distinct aggregates |
| if (!aggCall.isDistinct()) { |
| final AggregateCall newCall = |
| AggregateCall.create(aggCall.getAggregation(), false, |
| aggCall.isApproximate(), aggCall.ignoreNulls(), aggCall.rexList, |
| aggCall.getArgList(), -1, aggCall.distinctKeys, |
| aggCall.collation, bottomGroupSet.cardinality(), |
| relBuilder.peek(), null, aggCall.name); |
| bottomAggregateCalls.add(newCall); |
| } |
| } |
| // Generate the aggregate B (see the reference example above) |
| relBuilder.push( |
| aggregate.copy(aggregate.getTraitSet(), relBuilder.build(), |
| bottomGroupSet, null, bottomAggregateCalls)); |
| |
| // Add aggregate A (see the reference example above), the top aggregate |
| // to handle the rest of the aggregation that the bottom aggregate hasn't handled |
| final List<AggregateCall> topAggregateCalls = new ArrayList<>(); |
| // Use the remapped arguments for the (non)distinct aggregate calls |
| int nonDistinctAggCallProcessedSoFar = 0; |
| for (AggregateCall aggCall : originalAggCalls) { |
| final AggregateCall newCall; |
| if (aggCall.isDistinct()) { |
| List<Integer> newArgList = new ArrayList<>(); |
| for (int arg : aggCall.getArgList()) { |
| newArgList.add(bottomGroups.headSet(arg, false).size()); |
| } |
| newCall = |
| AggregateCall.create(aggCall.getAggregation(), |
| false, |
| aggCall.isApproximate(), |
| aggCall.ignoreNulls(), |
| aggCall.rexList, |
| newArgList, |
| -1, |
| aggCall.distinctKeys, |
| aggCall.collation, |
| originalGroupSet.cardinality(), |
| relBuilder.peek(), |
| aggCall.getType(), |
| aggCall.name); |
| } else { |
| // If aggregate B had a COUNT aggregate call the corresponding aggregate at |
| // aggregate A must be SUM. For other aggregates, it remains the same. |
| final int arg = bottomGroups.size() + nonDistinctAggCallProcessedSoFar; |
| final List<Integer> newArgs = ImmutableList.of(arg); |
| if (aggCall.getAggregation().getKind() == SqlKind.COUNT) { |
| newCall = |
| AggregateCall.create(new SqlSumEmptyIsZeroAggFunction(), false, |
| aggCall.isApproximate(), aggCall.ignoreNulls(), |
| aggCall.rexList, newArgs, -1, aggCall.distinctKeys, |
| aggCall.collation, originalGroupSet.cardinality(), |
| relBuilder.peek(), null, aggCall.getName()); |
| } else { |
| newCall = |
| AggregateCall.create(aggCall.getAggregation(), false, |
| aggCall.isApproximate(), aggCall.ignoreNulls(), |
| aggCall.rexList, newArgs, -1, aggCall.distinctKeys, |
| aggCall.collation, originalGroupSet.cardinality(), |
| relBuilder.peek(), null, aggCall.name); |
| } |
| nonDistinctAggCallProcessedSoFar++; |
| } |
| |
| topAggregateCalls.add(newCall); |
| } |
| |
| // Populate the group-by keys with the remapped arguments for aggregate A |
| // The top groupset is basically an identity (first X fields of aggregate B's |
| // output), minus the distinct aggCall's input. |
| final Set<Integer> topGroupSet = new HashSet<>(); |
| int groupSetToAdd = 0; |
| for (int bottomGroup : bottomGroups) { |
| if (originalGroupSet.get(bottomGroup)) { |
| topGroupSet.add(groupSetToAdd); |
| } |
| groupSetToAdd++; |
| } |
| relBuilder.push( |
| aggregate.copy(aggregate.getTraitSet(), relBuilder.build(), |
| ImmutableBitSet.of(topGroupSet), null, topAggregateCalls)); |
| |
| // Add projection node for case: SUM of COUNT(*): |
| // Type of the SUM may be larger than type of COUNT. |
| // CAST to original type must be added. |
| relBuilder.convert(aggregate.getRowType(), true); |
| |
| return relBuilder; |
| } |
| |
| private static void rewriteUsingGroupingSets(RelOptRuleCall call, |
| Aggregate aggregate) { |
| final Set<ImmutableBitSet> groupSetTreeSet = |
| new TreeSet<>(ImmutableBitSet.ORDERING); |
| // GroupSet to distinct filter arg map, |
| // filterArg will be -1 for non-distinct agg call. |
| |
| // Using `Set` here because it's possible that two agg calls |
| // have different filterArgs but same groupSet. |
| final Map<ImmutableBitSet, Set<Integer>> distinctFilterArgMap = new HashMap<>(); |
| for (AggregateCall aggCall : aggregate.getAggCallList()) { |
| ImmutableBitSet groupSet; |
| int filterArg; |
| if (!aggCall.isDistinct()) { |
| filterArg = -1; |
| groupSet = aggregate.getGroupSet(); |
| groupSetTreeSet.add(aggregate.getGroupSet()); |
| } else { |
| filterArg = aggCall.filterArg; |
| groupSet = |
| ImmutableBitSet.of(aggCall.getArgList()) |
| .setIf(filterArg, filterArg >= 0) |
| .union(aggregate.getGroupSet()); |
| groupSetTreeSet.add(groupSet); |
| } |
| Set<Integer> filterList = distinctFilterArgMap |
| .computeIfAbsent(groupSet, g -> new HashSet<>()); |
| filterList.add(filterArg); |
| } |
| |
| final ImmutableList<ImmutableBitSet> groupSets = |
| ImmutableList.copyOf(groupSetTreeSet); |
| final ImmutableBitSet fullGroupSet = ImmutableBitSet.union(groupSets); |
| |
| final List<AggregateCall> distinctAggCalls = new ArrayList<>(); |
| for (Pair<AggregateCall, String> aggCall : aggregate.getNamedAggCalls()) { |
| if (!aggCall.left.isDistinct()) { |
| AggregateCall newAggCall = |
| aggCall.left.adaptTo(aggregate.getInput(), |
| aggCall.left.getArgList(), aggCall.left.filterArg, |
| aggregate.getGroupCount(), fullGroupSet.cardinality()); |
| distinctAggCalls.add(newAggCall.withName(aggCall.right)); |
| } |
| } |
| |
| final RelBuilder relBuilder = call.builder(); |
| relBuilder.push(aggregate.getInput()); |
| final int groupCount = fullGroupSet.cardinality(); |
| |
| // Get the base ordinal of filter args for different groupSets. |
| final Map<Pair<ImmutableBitSet, Integer>, Integer> filters = new LinkedHashMap<>(); |
| int z = groupCount + distinctAggCalls.size(); |
| for (ImmutableBitSet groupSet : groupSets) { |
| Set<Integer> filterArgList = distinctFilterArgMap.get(groupSet); |
| for (Integer filterArg : requireNonNull(filterArgList, "filterArgList")) { |
| filters.put(Pair.of(groupSet, filterArg), z); |
| z += 1; |
| } |
| } |
| |
| distinctAggCalls.add( |
| AggregateCall.create(SqlStdOperatorTable.GROUPING, false, false, false, |
| ImmutableList.of(), ImmutableIntList.copyOf(fullGroupSet), -1, |
| null, RelCollations.EMPTY, |
| groupSets.size(), relBuilder.peek(), null, "$g")); |
| |
| relBuilder.aggregate( |
| relBuilder.groupKey(fullGroupSet, groupSets), |
| distinctAggCalls); |
| |
| // GROUPING returns an integer (0 or 1). Add a project to convert those |
| // values to BOOLEAN. |
| if (!filters.isEmpty()) { |
| final List<RexNode> nodes = new ArrayList<>(relBuilder.fields()); |
| final RexNode nodeZ = nodes.remove(nodes.size() - 1); |
| for (Map.Entry<Pair<ImmutableBitSet, Integer>, Integer> entry : filters.entrySet()) { |
| final long v = groupValue(fullGroupSet.asList(), entry.getKey().left); |
| int distinctFilterArg = remap(fullGroupSet, entry.getKey().right); |
| RexNode expr = relBuilder.equals(nodeZ, relBuilder.literal(v)); |
| if (distinctFilterArg > -1) { |
| // 'AND' the filter of the distinct aggregate call and the group value. |
| expr = |
| relBuilder.and(expr, |
| relBuilder.call(SqlStdOperatorTable.IS_TRUE, |
| relBuilder.field(distinctFilterArg))); |
| } |
| // "f" means filter. |
| nodes.add( |
| relBuilder.alias(expr, |
| "$g_" + v + (distinctFilterArg < 0 ? "" : "_f_" + distinctFilterArg))); |
| } |
| relBuilder.project(nodes); |
| } |
| |
| int x = groupCount; |
| final ImmutableBitSet groupSet = aggregate.getGroupSet(); |
| final List<AggregateCall> newCalls = new ArrayList<>(); |
| for (AggregateCall aggCall : aggregate.getAggCallList()) { |
| final int newFilterArg; |
| final List<Integer> newArgList; |
| final SqlAggFunction aggregation; |
| if (!aggCall.isDistinct()) { |
| aggregation = SqlStdOperatorTable.MIN; |
| newArgList = ImmutableIntList.of(x++); |
| newFilterArg = |
| requireNonNull(filters.get(Pair.of(groupSet, -1)), |
| "filters.get(Pair.of(groupSet, -1))"); |
| } else { |
| aggregation = aggCall.getAggregation(); |
| newArgList = remap(fullGroupSet, aggCall.getArgList()); |
| final ImmutableBitSet newGroupSet = ImmutableBitSet.of(aggCall.getArgList()) |
| .setIf(aggCall.filterArg, aggCall.filterArg >= 0) |
| .union(groupSet); |
| newFilterArg = |
| requireNonNull(filters.get(Pair.of(newGroupSet, aggCall.filterArg)), |
| "filters.get(of(newGroupSet, aggCall.filterArg))"); |
| } |
| final AggregateCall newCall = |
| AggregateCall.create(aggregation, false, |
| aggCall.isApproximate(), aggCall.ignoreNulls(), |
| aggCall.rexList, newArgList, newFilterArg, |
| aggCall.distinctKeys, aggCall.collation, |
| aggregate.getGroupCount(), relBuilder.peek(), null, aggCall.name); |
| newCalls.add(newCall); |
| } |
| |
| relBuilder.aggregate( |
| relBuilder.groupKey( |
| remap(fullGroupSet, groupSet), |
| remap(fullGroupSet, aggregate.getGroupSets())), |
| newCalls); |
| relBuilder.convert(aggregate.getRowType(), true); |
| call.transformTo(relBuilder.build()); |
| } |
| |
| /** Returns the value that "GROUPING(fullGroupSet)" will return for |
| * "groupSet". |
| * |
| * <p>It is important that {@code fullGroupSet} is not an |
| * {@link ImmutableBitSet}; the order of the bits matters. */ |
| static long groupValue(Collection<Integer> fullGroupSet, |
| ImmutableBitSet groupSet) { |
| long v = 0; |
| long x = 1L << (fullGroupSet.size() - 1); |
| assert ImmutableBitSet.of(fullGroupSet).contains(groupSet); |
| for (int i : fullGroupSet) { |
| if (!groupSet.get(i)) { |
| v |= x; |
| } |
| x >>= 1; |
| } |
| return v; |
| } |
| |
| static ImmutableBitSet remap(ImmutableBitSet groupSet, |
| ImmutableBitSet bitSet) { |
| final ImmutableBitSet.Builder builder = ImmutableBitSet.builder(); |
| for (Integer bit : bitSet) { |
| builder.set(remap(groupSet, bit)); |
| } |
| return builder.build(); |
| } |
| |
| static ImmutableList<ImmutableBitSet> remap(ImmutableBitSet groupSet, |
| Iterable<ImmutableBitSet> bitSets) { |
| final ImmutableList.Builder<ImmutableBitSet> builder = |
| ImmutableList.builder(); |
| for (ImmutableBitSet bitSet : bitSets) { |
| builder.add(remap(groupSet, bitSet)); |
| } |
| return builder.build(); |
| } |
| |
| private static List<Integer> remap(ImmutableBitSet groupSet, |
| List<Integer> argList) { |
| ImmutableIntList list = ImmutableIntList.of(); |
| for (int arg : argList) { |
| list = list.append(remap(groupSet, arg)); |
| } |
| return list; |
| } |
| |
| private static int remap(ImmutableBitSet groupSet, int arg) { |
| return arg < 0 ? -1 : groupSet.indexOf(arg); |
| } |
| |
| /** |
| * Converts an aggregate relational expression that contains just one |
| * distinct aggregate function (or perhaps several over the same arguments) |
| * and no non-distinct aggregate functions. |
| */ |
| private static RelBuilder convertMonopole(RelBuilder relBuilder, Aggregate aggregate, |
| List<Integer> argList, int filterArg) { |
| // For example, |
| // SELECT deptno, COUNT(DISTINCT sal), SUM(DISTINCT sal) |
| // FROM emp |
| // GROUP BY deptno |
| // |
| // becomes |
| // |
| // SELECT deptno, COUNT(distinct_sal), SUM(distinct_sal) |
| // FROM ( |
| // SELECT DISTINCT deptno, sal AS distinct_sal |
| // FROM EMP GROUP BY deptno) |
| // GROUP BY deptno |
| |
| // 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); |
| |
| // Create an aggregate on top, with the new aggregate list. |
| final List<AggregateCall> newAggCalls = |
| Lists.newArrayList(aggregate.getAggCallList()); |
| rewriteAggCalls(newAggCalls, argList, sourceOf); |
| final int cardinality = aggregate.getGroupSet().cardinality(); |
| relBuilder.push( |
| aggregate.copy(aggregate.getTraitSet(), relBuilder.build(), |
| ImmutableBitSet.range(cardinality), null, newAggCalls)); |
| return relBuilder; |
| } |
| |
| /** |
| * 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 |
| */ |
| private static void doRewrite(RelBuilder relBuilder, Aggregate aggregate, int n, |
| List<Integer> argList, int filterArg, List<@Nullable 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 groupCount = aggregate.getGroupCount(); |
| int i = groupCount - 1; |
| for (AggregateCall aggCall : aggCalls) { |
| ++i; |
| |
| // Ignore agg calls which are not distinct or have the wrong set |
| // arguments. If we're rewriting aggs whose args are {sal}, we will |
| // rewrite COUNT(DISTINCT sal) and SUM(DISTINCT sal) but ignore |
| // 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 (Integer arg : aggCall.getArgList()) { |
| newArgs.add(requireNonNull(sourceOf.get(arg), () -> "sourceOf.get(" + arg + ")")); |
| } |
| final int newFilterArg = |
| aggCall.filterArg < 0 ? -1 |
| : requireNonNull(sourceOf.get(aggCall.filterArg), |
| () -> "sourceOf.get(" + aggCall.filterArg + ")"); |
| final AggregateCall newAggCall = |
| AggregateCall.create(aggCall.getAggregation(), false, |
| aggCall.isApproximate(), aggCall.ignoreNulls(), aggCall.rexList, |
| newArgs, newFilterArg, aggCall.distinctKeys, aggCall.collation, |
| aggCall.getType(), aggCall.getName()); |
| assert refs.get(i) == null; |
| if (leftFields == null) { |
| refs.set(i, |
| new RexInputRef(groupCount + aggCallList.size(), |
| newAggCall.getType())); |
| } else { |
| refs.set(i, |
| new RexInputRef(leftFields.size() + groupCount |
| + 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; |
| |
| relBuilder.push( |
| aggregate.copy(aggregate.getTraitSet(), relBuilder.build(), |
| newGroupSet, newGroupingSets, aggCallList)); |
| |
| // If there's no left child yet, no need to create the join |
| if (leftFields == null) { |
| 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 = new ArrayList<>(); |
| for (i = 0; i < groupCount; ++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); |
| } |
| |
| private static void rewriteAggCalls( |
| List<AggregateCall> newAggCalls, |
| List<Integer> argList, |
| Map<Integer, Integer> sourceOf) { |
| // Rewrite the agg calls. 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)". |
| for (int i = 0; i < newAggCalls.size(); i++) { |
| final AggregateCall aggCall = newAggCalls.get(i); |
| |
| // Ignore agg calls which are not distinct or have the wrong set |
| // arguments. If we're rewriting aggregates whose args are {sal}, we will |
| // rewrite COUNT(DISTINCT sal) and SUM(DISTINCT sal) but ignore |
| // COUNT(DISTINCT gender) or SUM(sal). |
| if (!aggCall.isDistinct() |
| && aggCall.getAggregation().getDistinctOptionality() != Optionality.IGNORED) { |
| 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( |
| requireNonNull(sourceOf.get(arg), |
| () -> "sourceOf.get(" + arg + ")")); |
| } |
| final AggregateCall newAggCall = |
| AggregateCall.create(aggCall.getAggregation(), false, |
| aggCall.isApproximate(), aggCall.ignoreNulls(), |
| aggCall.rexList, newArgs, -1, |
| aggCall.distinctKeys, aggCall.collation, |
| aggCall.getType(), aggCall.getName()); |
| newAggCalls.set(i, newAggCall); |
| } |
| } |
| |
| /** |
| * Given an {@link org.apache.calcite.rel.core.Aggregate} |
| * 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> |
| * |
| * <p>and the argument list |
| * |
| * <blockquote>{2}</blockquote> |
| * |
| * <p>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. |
| * |
| * @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 static RelBuilder createSelectDistinct(RelBuilder relBuilder, |
| Aggregate aggregate, List<Integer> argList, int filterArg, |
| Map<Integer, Integer> sourceOf) { |
| relBuilder.push(aggregate.getInput()); |
| final PairList<RexNode, String> projects = PairList.of(); |
| final List<RelDataTypeField> childFields = |
| relBuilder.peek().getRowType().getFieldList(); |
| for (int i : aggregate.getGroupSet()) { |
| sourceOf.put(i, projects.size()); |
| RexInputRef.add2(projects, 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.makeNullLiteral(argRef.left.getType())); |
| sourceOf.put(arg, projects.size()); |
| projects.add(condition, "i$" + argRef.right); |
| continue; |
| } |
| if (sourceOf.get(arg) != null) { |
| continue; |
| } |
| sourceOf.put(arg, projects.size()); |
| RexInputRef.add2(projects, arg, childFields); |
| } |
| relBuilder.project(projects.leftList(), projects.rightList()); |
| |
| // Get the distinct values of the GROUP BY fields and the arguments |
| // to the agg functions. |
| relBuilder.push( |
| aggregate.copy(aggregate.getTraitSet(), relBuilder.build(), |
| ImmutableBitSet.range(projects.size()), null, ImmutableList.of())); |
| return relBuilder; |
| } |
| |
| /** Rule configuration. */ |
| @Value.Immutable |
| public interface Config extends RelRule.Config { |
| Config DEFAULT = ImmutableAggregateExpandDistinctAggregatesRule.Config.of() |
| .withOperandSupplier(b -> |
| b.operand(LogicalAggregate.class).anyInputs()); |
| |
| Config JOIN = DEFAULT.withUsingGroupingSets(false); |
| |
| @Override default AggregateExpandDistinctAggregatesRule toRule() { |
| return new AggregateExpandDistinctAggregatesRule(this); |
| } |
| |
| /** Whether to use GROUPING SETS, default true. */ |
| @Value.Default default boolean isUsingGroupingSets() { |
| return true; |
| } |
| |
| /** Sets {@link #isUsingGroupingSets()}. */ |
| Config withUsingGroupingSets(boolean usingGroupingSets); |
| } |
| } |
