SqlSplittableAggFunction.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.sql;
import org.apache.calcite.rel.RelCollations;
import org.apache.calcite.rel.core.AggregateCall;
import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.rel.type.RelDataTypeFactory;
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.rex.RexUtil;
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
import org.apache.calcite.sql.parser.SqlParserPos;
import org.apache.calcite.sql.type.SqlTypeName;
import org.apache.calcite.util.ImmutableIntList;
import org.apache.calcite.util.mapping.Mappings;
import com.google.common.collect.ImmutableList;
import org.checkerframework.checker.nullness.qual.Nullable;
import java.math.BigDecimal;
import java.util.ArrayList;
import java.util.List;
/**
* Aggregate function that can be split into partial aggregates.
*
* <p>For example, {@code COUNT(x)} can be split into {@code COUNT(x)} on
* subsets followed by {@code SUM} to combine those counts.
*/
public interface SqlSplittableAggFunction extends SqlSingletonAggFunction {
AggregateCall split(AggregateCall aggregateCall,
Mappings.TargetMapping mapping);
/** Called to generate an aggregate for the other side of the join
* than the side aggregate call's arguments come from. Returns null if
* no aggregate is required. */
@Nullable AggregateCall other(RelDataTypeFactory typeFactory, AggregateCall e);
/** Generates an aggregate call to merge sub-totals.
*
* <p>Most implementations will add a single aggregate call to
* {@code aggCalls}, and return a {@link RexInputRef} that points to it.
*
* @param rexBuilder Rex builder
* @param extra Place to define extra input expressions
* @param offset Offset due to grouping columns (and indicator columns if
* applicable)
* @param inputRowType Input row type
* @param aggregateCall Source aggregate call
* @param leftSubTotal Ordinal of the sub-total coming from the left side of
* the join, or -1 if there is no such sub-total
* @param rightSubTotal Ordinal of the sub-total coming from the right side
* of the join, or -1 if there is no such sub-total
*
* @return Aggregate call
*/
AggregateCall topSplit(RexBuilder rexBuilder, Registry<RexNode> extra,
int offset, RelDataType inputRowType, AggregateCall aggregateCall,
int leftSubTotal, int rightSubTotal);
/**
* Merge top and bottom aggregate calls into a single aggregate call,
* if they are legit to merge.
*
* <p>SUM of SUM becomes SUM; SUM of COUNT becomes COUNT;
* MAX of MAX becomes MAX; MIN of MIN becomes MIN.
* AVG of AVG would not match, nor would COUNT of COUNT.
*
* @param top top aggregate call
* @param bottom bottom aggregate call
* @return Merged aggregate call, null if fails to merge aggregate calls
*/
@Nullable AggregateCall merge(AggregateCall top, AggregateCall bottom);
/** Collection in which one can register an element. Registering may return
* a reference to an existing element.
*
* @param <E> element type */
interface Registry<E> {
int register(E e);
}
/** Splitting strategy for {@code COUNT}.
*
* <p>COUNT splits into itself followed by SUM. (Actually
* SUM0, because the total needs to be 0, not null, if there are 0 rows.)
* This rule works for any number of arguments to COUNT, including COUNT(*).
*/
class CountSplitter implements SqlSplittableAggFunction {
public static final CountSplitter INSTANCE = new CountSplitter();
@Override public AggregateCall split(AggregateCall aggregateCall,
Mappings.TargetMapping mapping) {
return aggregateCall.transform(mapping);
}
@Override public @Nullable AggregateCall other(RelDataTypeFactory typeFactory,
AggregateCall e) {
final RelDataType type = typeFactory.createSqlType(SqlTypeName.BIGINT);
return AggregateCall.create(e.getParserPosition(), SqlStdOperatorTable.COUNT, false, false,
false, ImmutableList.of(), ImmutableIntList.of(), -1, null,
RelCollations.EMPTY, type, null);
}
@Override public AggregateCall topSplit(RexBuilder rexBuilder,
Registry<RexNode> extra, int offset, RelDataType inputRowType,
AggregateCall aggregateCall, int leftSubTotal, int rightSubTotal) {
final List<RexNode> merges = new ArrayList<>();
final SqlParserPos pos = aggregateCall.getParserPosition();
if (leftSubTotal >= 0) {
merges.add(
rexBuilder.makeInputRef(aggregateCall.type, leftSubTotal));
}
if (rightSubTotal >= 0) {
merges.add(
rexBuilder.makeInputRef(aggregateCall.type, rightSubTotal));
}
RexNode node;
switch (merges.size()) {
case 1:
node = merges.get(0);
break;
case 2:
node = rexBuilder.makeCall(pos, SqlStdOperatorTable.MULTIPLY, merges);
break;
default:
throw new AssertionError("unexpected count " + merges);
}
int ordinal = extra.register(node);
return AggregateCall.create(pos, SqlStdOperatorTable.SUM0, false, false,
false, aggregateCall.rexList, ImmutableList.of(ordinal), -1,
aggregateCall.distinctKeys, aggregateCall.collation,
aggregateCall.type, aggregateCall.name);
}
/**
* {@inheritDoc}
*
* <p>{@code COUNT(*)}, and {@code COUNT} applied to all NOT NULL arguments,
* become {@code 1}; otherwise
* {@code CASE WHEN arg0 IS NOT NULL THEN 1 ELSE 0 END}.
*/
@Override public RexNode singleton(RexBuilder rexBuilder, RelDataType inputRowType,
AggregateCall aggregateCall) {
final List<RexNode> predicates = new ArrayList<>();
for (Integer arg : aggregateCall.getArgList()) {
final RelDataType type = inputRowType.getFieldList().get(arg).getType();
if (type.isNullable()) {
predicates.add(
rexBuilder.makeCall(SqlStdOperatorTable.IS_NOT_NULL,
rexBuilder.makeInputRef(type, arg)));
}
}
final RexNode predicate =
RexUtil.composeConjunction(rexBuilder, predicates, true);
final RexNode rexOne =
rexBuilder.makeExactLiteral(BigDecimal.ONE, aggregateCall.getType());
if (predicate == null) {
return rexOne;
} else {
return rexBuilder.makeCall(SqlStdOperatorTable.CASE, predicate, rexOne,
rexBuilder.makeExactLiteral(BigDecimal.ZERO, aggregateCall.getType()));
}
}
@Override public @Nullable AggregateCall merge(AggregateCall top, AggregateCall bottom) {
if (bottom.getAggregation().getKind() == SqlKind.COUNT
&& (top.getAggregation().getKind() == SqlKind.SUM
|| top.getAggregation().getKind() == SqlKind.SUM0)) {
return AggregateCall.create(top.getParserPosition(), bottom.getAggregation(),
bottom.isDistinct(), bottom.isApproximate(), false,
bottom.rexList, bottom.getArgList(), bottom.filterArg,
bottom.distinctKeys, bottom.getCollation(),
bottom.getType(), top.getName());
} else {
return null;
}
}
}
/** Aggregate function that splits into two applications of itself.
*
* <p>Examples are MIN and MAX. */
class SelfSplitter implements SqlSplittableAggFunction {
public static final SelfSplitter INSTANCE = new SelfSplitter();
@Override public RexNode singleton(RexBuilder rexBuilder,
RelDataType inputRowType, AggregateCall aggregateCall) {
final int arg = aggregateCall.getArgList().get(0);
final RelDataTypeField field = inputRowType.getFieldList().get(arg);
return rexBuilder.makeInputRef(field.getType(), arg);
}
@Override public AggregateCall split(AggregateCall aggregateCall,
Mappings.TargetMapping mapping) {
return aggregateCall.transform(mapping);
}
@Override public @Nullable AggregateCall other(RelDataTypeFactory typeFactory,
AggregateCall e) {
return null; // no aggregate function required on other side
}
@Override public AggregateCall topSplit(RexBuilder rexBuilder,
Registry<RexNode> extra, int offset, RelDataType inputRowType,
AggregateCall aggregateCall, int leftSubTotal, int rightSubTotal) {
assert (leftSubTotal >= 0) != (rightSubTotal >= 0);
assert aggregateCall.collation.getFieldCollations().isEmpty();
final int arg = leftSubTotal >= 0 ? leftSubTotal : rightSubTotal;
return aggregateCall.withArgList(ImmutableIntList.of(arg));
}
@Override public @Nullable AggregateCall merge(AggregateCall top, AggregateCall bottom) {
if (top.getAggregation().getKind() == bottom.getAggregation().getKind()) {
return AggregateCall.create(top.getParserPosition(), bottom.getAggregation(),
bottom.isDistinct(), bottom.isApproximate(), false,
bottom.rexList, bottom.getArgList(), bottom.filterArg,
bottom.distinctKeys, bottom.getCollation(),
bottom.getType(), top.getName());
} else {
return null;
}
}
}
/** Common splitting strategy for {@code SUM} and {@code SUM0} functions. */
abstract class AbstractSumSplitter implements SqlSplittableAggFunction {
@Override public RexNode singleton(RexBuilder rexBuilder,
RelDataType inputRowType, AggregateCall aggregateCall) {
final int arg = aggregateCall.getArgList().get(0);
final RelDataTypeField field = inputRowType.getFieldList().get(arg);
final RelDataType fieldType = field.getType();
return rexBuilder.makeInputRef(fieldType, arg);
}
@Override public AggregateCall split(AggregateCall aggregateCall,
Mappings.TargetMapping mapping) {
return aggregateCall.transform(mapping);
}
@Override public @Nullable AggregateCall other(RelDataTypeFactory typeFactory,
AggregateCall e) {
final RelDataType type = typeFactory.createSqlType(SqlTypeName.BIGINT);
return AggregateCall.create(e.getParserPosition(), SqlStdOperatorTable.COUNT, false, false,
false, ImmutableList.of(), ImmutableIntList.of(), -1, null,
RelCollations.EMPTY, type, null);
}
@Override public AggregateCall topSplit(RexBuilder rexBuilder,
Registry<RexNode> extra, int offset, RelDataType inputRowType,
AggregateCall aggregateCall, int leftSubTotal, int rightSubTotal) {
final List<RexNode> merges = new ArrayList<>();
final SqlParserPos pos = aggregateCall.getParserPosition();
final List<RelDataTypeField> fieldList = inputRowType.getFieldList();
if (leftSubTotal >= 0) {
final RelDataType type = fieldList.get(leftSubTotal).getType();
merges.add(rexBuilder.makeInputRef(type, leftSubTotal));
}
if (rightSubTotal >= 0) {
final RelDataType type = fieldList.get(rightSubTotal).getType();
merges.add(rexBuilder.makeInputRef(type, rightSubTotal));
}
RexNode node;
switch (merges.size()) {
case 1:
node = merges.get(0);
break;
case 2:
node = rexBuilder.makeCall(pos, SqlStdOperatorTable.MULTIPLY, merges);
node = rexBuilder.makeAbstractCast(pos, aggregateCall.type, node, false);
break;
default:
throw new AssertionError("unexpected count " + merges);
}
int ordinal = extra.register(node);
return AggregateCall.create(pos, getMergeAggFunctionOfTopSplit(), false, false,
false, aggregateCall.rexList, ImmutableList.of(ordinal), -1,
aggregateCall.distinctKeys, aggregateCall.collation,
aggregateCall.type, aggregateCall.name);
}
@Override public @Nullable AggregateCall merge(AggregateCall top, AggregateCall bottom) {
SqlKind topKind = top.getAggregation().getKind();
if (topKind == bottom.getAggregation().getKind()
&& (topKind == SqlKind.SUM
|| topKind == SqlKind.SUM0)) {
return AggregateCall.create(top.getParserPosition(), bottom.getAggregation(),
bottom.isDistinct(), bottom.isApproximate(), false,
bottom.rexList, bottom.getArgList(), bottom.filterArg,
bottom.distinctKeys, bottom.getCollation(),
top.getType(), top.getName());
} else {
return null;
}
}
protected abstract SqlAggFunction getMergeAggFunctionOfTopSplit();
}
/** Splitting strategy for {@code SUM} function. */
class SumSplitter extends AbstractSumSplitter {
public static final SumSplitter INSTANCE = new SumSplitter();
@Override public SqlAggFunction getMergeAggFunctionOfTopSplit() {
return SqlStdOperatorTable.SUM;
}
}
/** Splitting strategy for {@code SUM0} function. */
class Sum0Splitter extends AbstractSumSplitter {
public static final Sum0Splitter INSTANCE = new Sum0Splitter();
@Override public SqlAggFunction getMergeAggFunctionOfTopSplit() {
return SqlStdOperatorTable.SUM0;
}
@Override public RexNode singleton(RexBuilder rexBuilder,
RelDataType inputRowType, AggregateCall aggregateCall) {
final int arg = aggregateCall.getArgList().get(0);
final RelDataType type = inputRowType.getFieldList().get(arg).getType();
final RelDataTypeFactory typeFactory = rexBuilder.getTypeFactory();
final RelDataType type1 = typeFactory.getTypeSystem().deriveSumType(typeFactory, type);
final RexNode inputRef = rexBuilder.makeInputRef(type1, arg);
if (type.isNullable()) {
return rexBuilder.makeCall(SqlStdOperatorTable.COALESCE, inputRef,
rexBuilder.makeExactLiteral(BigDecimal.ZERO, type));
} else {
return inputRef;
}
}
}
}