AggregateExpandWithinDistinctRule.java
/*
* 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.linq4j.Ord;
import org.apache.calcite.plan.RelOptRuleCall;
import org.apache.calcite.plan.RelRule;
import org.apache.calcite.rel.core.Aggregate;
import org.apache.calcite.rel.core.AggregateCall;
import org.apache.calcite.rel.logical.LogicalAggregate;
import org.apache.calcite.rex.RexNode;
import org.apache.calcite.sql.SqlKind;
import org.apache.calcite.sql.fun.SqlInternalOperators;
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
import org.apache.calcite.tools.RelBuilder;
import org.apache.calcite.util.ImmutableBitSet;
import org.apache.calcite.util.ImmutableIntList;
import org.apache.calcite.util.Util;
import org.apache.calcite.util.mapping.IntPair;
import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Multimap;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.immutables.value.Value;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeSet;
import java.util.function.IntPredicate;
import java.util.stream.Collectors;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static org.apache.calcite.rel.rules.AggregateExpandDistinctAggregatesRule.groupValue;
import static org.apache.calcite.rel.rules.AggregateExpandDistinctAggregatesRule.remap;
import static java.util.Objects.requireNonNull;
/**
* Planner rule that rewrites an {@link Aggregate} that contains
* {@code WITHIN DISTINCT} aggregate functions.
*
* <p>For example,
* <blockquote>
* <pre>SELECT o.paymentType,
* COUNT(*) as "count",
* COUNT(*) WITHIN DISTINCT (o.orderId) AS orderCount,
* SUM(o.shipping) WITHIN DISTINCT (o.orderId) as sumShipping,
* SUM(i.units) as sumUnits
* FROM Orders AS o
* JOIN OrderItems AS i USING (orderId)
* GROUP BY o.paymentType</pre>
* </blockquote>
*
* <p>becomes
*
* <blockquote>
* <pre>
* SELECT paymentType,
* COUNT(*) as "count",
* COUNT(*) FILTER (WHERE g = 0) AS orderCount,
* SUM(minShipping) FILTER (WHERE g = 0) AS sumShipping,
* SUM(sumUnits) FILTER (WHERE g = 1) as sumUnits
* FROM (
* SELECT o.paymentType,
* GROUPING(o.orderId) AS g,
* SUM(o.shipping) AS sumShipping,
* MIN(o.shipping) AS minShipping,
* SUM(i.units) AS sumUnits
* FROM Orders AS o
* JOIN OrderItems ON o.orderId = i.orderId
* GROUP BY GROUPING SETS ((o.paymentType), (o.paymentType, o.orderId)))
* GROUP BY o.paymentType</pre>
* </blockquote>
*
* <p>By the way, note that {@code COUNT(*) WITHIN DISTINCT (o.orderId)}
* is identical to {@code COUNT(DISTINCT o.orderId)}.
* {@code WITHIN DISTINCT} is a generalization of aggregate(DISTINCT).
* So, it is perhaps not surprising that the rewrite to {@code GROUPING SETS}
* is similar.
*
* <p>If there are multiple arguments
* (e.g. {@code SUM(a) WITHIN DISTINCT (x), SUM(a) WITHIN DISTINCT (y)})
* the rule creates separate {@code GROUPING SET}s.
*/
@Value.Enclosing
public class AggregateExpandWithinDistinctRule
extends RelRule<AggregateExpandWithinDistinctRule.Config> {
/** Creates an AggregateExpandWithinDistinctRule. */
protected AggregateExpandWithinDistinctRule(Config config) {
super(config);
}
private static boolean hasWithinDistinct(Aggregate aggregate) {
return aggregate.getAggCallList().stream()
.anyMatch(c -> c.distinctKeys != null)
// Wait until AggregateReduceFunctionsRule has dealt with AVG etc.
&& aggregate.getAggCallList().stream()
.noneMatch(CoreRules.AGGREGATE_REDUCE_FUNCTIONS::canReduce)
// Don't think we can handle GROUPING SETS yet
&& aggregate.getGroupType() == Aggregate.Group.SIMPLE;
}
//~ Methods ----------------------------------------------------------------
@Override public void onMatch(RelOptRuleCall call) {
final Aggregate aggregate = call.rel(0);
// Throughout this method, assume we are working on the following SQL:
//
// SELECT deptno, SUM(sal), SUM(sal) WITHIN DISTINCT (job)
// FROM emp
// GROUP BY deptno
//
// or in algebra,
//
// Aggregate($0, SUM($2), SUM($2) WITHIN DISTINCT ($4))
// Scan(emp)
//
// We plan to generate the following:
//
// SELECT deptno, SUM(sal), SUM(sal) WITHIN DISTINCT (job)
// FROM (
// SELECT deptno, GROUPING(deptno, job), SUM(sal), MIN(sal)
// FROM emp
// GROUP BY GROUPING SETS ((deptno), (deptno, job)))
// GROUP BY deptno
//
// This rewrite also handles DISTINCT aggregates. We treat
//
// SUM(DISTINCT sal)
// SUM(DISTINCT sal) WITHIN DISTINCT (job)
//
// as if the user had written
//
// SUM(sal) WITHIN DISTINCT (sal)
//
final List<AggregateCall> aggCallList =
aggregate.getAggCallList()
.stream()
.map(c -> unDistinct(c, aggregate.getInput()::fieldIsNullable))
.collect(toImmutableList());
// Find all within-distinct expressions.
final Multimap<ImmutableBitSet, AggregateCall> argLists =
ArrayListMultimap.create();
// A bit set that represents an aggregate with no WITHIN DISTINCT.
// Different from "WITHIN DISTINCT ()".
final ImmutableBitSet notDistinct =
ImmutableBitSet.of(aggregate.getInput().getRowType().getFieldCount());
for (AggregateCall aggCall : aggCallList) {
ImmutableBitSet distinctKeys = aggCall.distinctKeys;
if (distinctKeys == null) {
distinctKeys = notDistinct;
} else {
if (distinctKeys.intersects(aggregate.getGroupSet())) {
// Remove group keys. E.g.
// sum(x) within distinct (y, z) ... group by y
// can be simplified to
// sum(x) within distinct (z) ... group by y
// Note that this assumes a single grouping set for the original agg.
distinctKeys = distinctKeys.rebuild()
.removeAll(aggregate.getGroupSet()).build();
}
}
argLists.put(distinctKeys, aggCall);
}
// Compute the set of all grouping sets that will be used in the output
// query. For each WITHIN DISTINCT aggregate call, we will need a grouping
// set that is the union of the aggregate call's unique keys and the input
// query's overall grouping. Redundant grouping sets can be reused for
// multiple aggregate calls.
final Set<ImmutableBitSet> groupSetTreeSet =
new TreeSet<>(ImmutableBitSet.ORDERING);
for (ImmutableBitSet key : argLists.keySet()) {
groupSetTreeSet.add(
(key == notDistinct)
? aggregate.getGroupSet()
: ImmutableBitSet.of(key).union(aggregate.getGroupSet()));
}
final ImmutableList<ImmutableBitSet> groupSets =
ImmutableList.copyOf(groupSetTreeSet);
final boolean hasMultipleGroupSets = groupSets.size() > 1;
final ImmutableBitSet fullGroupSet = ImmutableBitSet.union(groupSets);
final Set<Integer> fullGroupOrderedSet = new LinkedHashSet<>();
fullGroupOrderedSet.addAll(aggregate.getGroupSet().asSet());
fullGroupOrderedSet.addAll(fullGroupSet.asSet());
final ImmutableIntList fullGroupList =
ImmutableIntList.copyOf(fullGroupOrderedSet);
// Build the inner query
//
// SELECT deptno, SUM(sal) AS sum_sal, MIN(sal) AS min_sal,
// MAX(sal) AS max_sal, GROUPING(deptno, job) AS g
// FROM emp
// GROUP BY GROUPING SETS ((deptno), (deptno, job))
//
// or in algebra,
//
// Aggregate([($0), ($0, $4)], SUM($2), MIN($2), MAX($2), GROUPING($0, $4))
// Scan(emp)
final RelBuilder b = call.builder();
b.push(aggregate.getInput());
final List<RelBuilder.AggCall> aggCalls = new ArrayList<>();
// Helper class for building the inner query.
// CHECKSTYLE: IGNORE 1
class Registrar {
final int g = fullGroupSet.cardinality();
/** Map of input fields (below the original aggregation) and filter args
* to inner query aggregate calls. */
final Map<IntPair, Integer> args = new HashMap<>();
/** Map of aggregate calls from the original aggregation to inner query
* aggregate calls. */
final Map<Integer, Integer> aggs = new HashMap<>();
/** Map of aggregate calls from the original aggregation to inner-query
* {@code COUNT(*)} calls, which are only needed for filters in the outer
* aggregate when the original aggregate call does not ignore null
* inputs. */
final Map<Integer, Integer> counts = new HashMap<>();
List<Integer> fields(List<Integer> fields, int filterArg) {
return Util.transform(fields, f -> this.field(f, filterArg));
}
int field(int field, int filterArg) {
return requireNonNull(args.get(IntPair.of(field, filterArg)));
}
/** Computes an aggregate call argument's values for a
* {@code WITHIN DISTINCT} aggregate call.
*
* <p>For example, to compute
* {@code SUM(x) WITHIN DISTINCT (y) GROUP BY (z)},
* the inner aggregate must first group {@code x} by {@code (y, z)}
* — using {@code MIN} to select the (hopefully) unique value of
* {@code x} for each {@code (y, z)} group. Actually summing over the
* grouped {@code x} values must occur in an outer aggregate.
*
* @param field Index of an input field that's used in a
* {@code WITHIN DISTINCT} aggregate call
* @param filterArg Filter arg used in the original aggregate call, or
* {@code -1} if there is no filter. We use the same filter in
* the inner query.
* @return Index of the inner query aggregate call representing the
* grouped field, which can be referenced in the outer query
* aggregate call
*/
int register(int field, int filterArg) {
return args.computeIfAbsent(IntPair.of(field, filterArg), j -> {
final int ordinal = g + aggCalls.size();
RelBuilder.AggCall groupedField =
b.aggregateCall(SqlStdOperatorTable.MIN, b.field(field));
aggCalls.add(
filterArg < 0
? groupedField
: groupedField.filter(b.field(filterArg)));
if (config.throwIfNotUnique()) {
groupedField =
b.aggregateCall(SqlStdOperatorTable.MAX, b.field(field));
aggCalls.add(
filterArg < 0
? groupedField
: groupedField.filter(b.field(filterArg)));
}
return ordinal;
});
}
/** Registers an aggregate call that is <em>not</em> a
* {@code WITHIN DISTINCT} call.
*
* <p>Unlike the case handled by {@link #register(int, int)} above,
* aggregate calls without any distinct keys do not need a second round
* of aggregation in the outer query, so they can be computed "as-is" in
* the inner query.
*
* @param i Index of the aggregate call in the original aggregation
* @param aggregateCall Original aggregate call
* @return Index of the aggregate call in the computed inner query
*/
int registerAgg(int i, RelBuilder.AggCall aggregateCall) {
final int ordinal = g + aggCalls.size();
aggs.put(i, ordinal);
aggCalls.add(aggregateCall);
return ordinal;
}
int getAgg(int i) {
return requireNonNull(aggs.get(i));
}
/** Registers an extra {@code COUNT} aggregate call when it's needed to
* filter out null inputs in the outer aggregate.
*
* <p>This should only be called for aggregate calls with filters. It's
* possible that the filter would eliminate all input rows to the
* {@code MIN} call in the inner query, so calls in the outer
* aggregate may need to be aware of this. See usage of
* {@link AggregateExpandWithinDistinctRule#mustBeCounted(AggregateCall)}.
*
* @param filterArg The original aggregate call's filter; must be
* non-negative
* @return Index of the {@code COUNT} call in the computed inner query
*/
int registerCount(int filterArg) {
assert filterArg >= 0;
return counts.computeIfAbsent(filterArg, i -> {
final int ordinal = g + aggCalls.size();
aggCalls.add(b.aggregateCall(SqlStdOperatorTable.COUNT)
.filter(b.field(filterArg)));
return ordinal;
});
}
int getCount(int filterArg) {
return requireNonNull(counts.get(filterArg));
}
}
final Registrar registrar = new Registrar();
Ord.forEach(aggCallList, (c, i) -> {
if (c.distinctKeys == null) {
RelBuilder.AggCall aggCall =
b.aggregateCall(c.getParserPosition(), c.getAggregation(), b.fields(c.getArgList()));
if (c.hasFilter()) {
aggCall = aggCall.filter(b.field(c.filterArg));
}
if (c.hasCollation()) {
aggCall = aggCall.sort(b.fields(c.getCollation()));
}
registrar.registerAgg(i, aggCall);
} else {
for (int inputIdx : c.getArgList()) {
registrar.register(inputIdx, c.filterArg);
}
if (mustBeCounted(c)) {
registrar.registerCount(c.filterArg);
}
}
});
// Add an additional GROUPING() aggregate call so we can select only the
// relevant inner-aggregate rows from the outer aggregate. If there is only
// 1 grouping set (i.e. every aggregate call has the same distinct keys),
// no GROUPING() call is necessary.
final int grouping =
hasMultipleGroupSets
? registrar.registerAgg(-1,
b.aggregateCall(
SqlStdOperatorTable.GROUPING,
b.fields(fullGroupList)))
: -1;
b.aggregate(b.groupKey(fullGroupSet, groupSets), aggCalls);
// Build the outer query
//
// SELECT deptno,
// MIN(sum_sal) FILTER (g = 0),
// SUM(min_sal) FILTER (g = 1)
// FROM ( ... )
// GROUP BY deptno
//
// or in algebra,
//
// Aggregate($0, SUM($2 WHERE $4 = 0), SUM($3 WHERE $4 = 1))
// Aggregate([($0), ($0, $2)], SUM($2), MIN($2), GROUPING($0, $4))
// Scan(emp)
//
// If throwIfNotUnique, the "SUM(min_sal) FILTER (g = 1)" term above becomes
//
// SUM(min_sal) FILTER (
// $THROW_UNLESS(g <> 1 OR min_sal IS NOT DISTINCT FROM max_sal,
// 'more than one distinct value in agg UNIQUE_VALUE')
// AND g = 1)
aggCalls.clear();
Ord.forEach(aggCallList, (c, i) -> {
final List<RexNode> filters = new ArrayList<>();
RexNode groupFilter = null;
if (hasMultipleGroupSets) {
groupFilter =
b.equals(
b.field(grouping),
b.literal(
groupValue(fullGroupList,
union(aggregate.getGroupSet(), c.distinctKeys))));
filters.add(groupFilter);
}
RelBuilder.AggCall aggCall;
if (c.distinctKeys == null) {
aggCall =
b.aggregateCall(c.getParserPosition(), SqlStdOperatorTable.MIN,
b.field(registrar.getAgg(i)));
} else {
// The inputs to this aggregate are outputs from MIN() calls from the
// inner agg, and MIN() returns null iff it has no non-null inputs,
// which can only happen if an original aggregate's filter causes all
// non-null input rows to be discarded for a particular group in the
// inner aggregate. In this case, it should be ignored by the outer
// aggregate as well. In case the aggregate call does not naturally
// ignore null inputs, we add a filter based on a COUNT() in the inner
// aggregate.
aggCall =
b.aggregateCall(c.getParserPosition(), c.getAggregation(),
b.fields(registrar.fields(c.getArgList(), c.filterArg)));
if (mustBeCounted(c)) {
filters.add(
b.greaterThan(b.field(registrar.getCount(c.filterArg)),
b.literal(0)));
}
if (config.throwIfNotUnique()) {
for (int j : c.getArgList()) {
RexNode isUniqueCondition =
b.isNotDistinctFrom(
b.field(registrar.field(j, c.filterArg)),
b.field(registrar.field(j, c.filterArg) + 1));
if (groupFilter != null) {
isUniqueCondition = b.or(b.not(groupFilter), isUniqueCondition);
}
String message = "more than one distinct value in agg UNIQUE_VALUE";
filters.add(
b.call(SqlInternalOperators.THROW_UNLESS, isUniqueCondition,
b.literal(message)));
}
}
}
if (!filters.isEmpty()) {
aggCall = aggCall.filter(b.and(filters));
}
aggCalls.add(aggCall);
});
b.aggregate(
b.groupKey(
remap(fullGroupSet, aggregate.getGroupSet()),
remap(fullGroupSet, aggregate.getGroupSets())),
aggCalls);
b.convert(aggregate.getRowType(), false);
call.transformTo(b.build());
}
private static boolean mustBeCounted(AggregateCall aggCall) {
// Always count filtered inner aggregates to be safe.
//
// It's possible that, for some aggregate calls (namely, those that
// completely ignore null inputs), we could neglect counting the
// grouped-and-filtered rows of the inner aggregate and filtering the empty
// ones out from the outer aggregate, since those empty groups would produce
// null values as the result of MIN and thus be ignored by the outer
// aggregate anyway.
//
// Note that using "aggCall.ignoreNulls()" is not sufficient to determine
// when it's safe to do this, since for COUNT the value of ignoreNulls()
// should generally be true even though COUNT(*) will never ignore anything.
return aggCall.hasFilter();
}
/** Converts a {@code DISTINCT} aggregate call into an equivalent one with
* {@code WITHIN DISTINCT}.
*
* <p>Examples:
* <ul>
* <li>{@code SUM(DISTINCT x)} →
* {@code SUM(x) WITHIN DISTINCT (x)} has distinct key (x);
* <li>{@code SUM(DISTINCT x)} WITHIN DISTINCT (y) →
* {@code SUM(x) WITHIN DISTINCT (x)} has distinct key (x);
* <li>{@code SUM(x)} WITHIN DISTINCT (y, z) has distinct key (y, z);
* <li>{@code SUM(x)} has no distinct key.
* </ul>
*/
private static AggregateCall unDistinct(AggregateCall aggregateCall,
IntPredicate isNullable) {
if (aggregateCall.isDistinct()) {
final List<Integer> newArgList = aggregateCall.getArgList()
.stream()
.filter(i ->
aggregateCall.getAggregation().getKind() != SqlKind.COUNT
|| aggregateCall.hasFilter()
|| isNullable.test(i))
.collect(Collectors.toList());
return aggregateCall.withDistinct(false)
.withDistinctKeys(ImmutableBitSet.of(aggregateCall.getArgList()))
.withArgList(newArgList);
}
return aggregateCall;
}
private static ImmutableBitSet union(ImmutableBitSet s0,
@Nullable ImmutableBitSet s1) {
return s1 == null ? s0 : s0.union(s1);
}
/** Rule configuration. */
@Value.Immutable
public interface Config extends RelRule.Config {
Config DEFAULT = ImmutableAggregateExpandWithinDistinctRule.Config.of()
.withOperandSupplier(b -> b.operand(LogicalAggregate.class)
.predicate(AggregateExpandWithinDistinctRule::hasWithinDistinct)
.anyInputs());
@Override default AggregateExpandWithinDistinctRule toRule() {
return new AggregateExpandWithinDistinctRule(this);
}
/** Whether the code generated by the rule should throw if the arguments
* are not functionally dependent.
*
* <p>For example, if implementing {@code SUM(sal) WITHIN DISTINCT job)} ...
* {@code GROUP BY deptno},
* suppose that within department 10, (job, sal) has the values
* ('CLERK', 100), ('CLERK', 120), ('MANAGER', 150), ('MANAGER', 150). If
* {@code throwIfNotUnique} is true, the query would throw because of the
* values [100, 120]; if false, the query would sum the distinct values
* [100, 120, 150]. */
@Value.Default default boolean throwIfNotUnique() {
return true;
}
/** Sets {@link #throwIfNotUnique()}. */
Config withThrowIfNotUnique(boolean throwIfNotUnique);
}
}