This is an automated email from the ASF dual-hosted git repository. jhyde pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/calcite.git
commit 76f26b8319e66c30c6cbfdfba6d5b2839cdedfa2 Author: Julian Hyde <[email protected]> AuthorDate: Wed Sep 2 14:51:53 2020 -0700 [CALCITE-4220] In SqlToRelConverter, use RelBuilder for creating Aggregate Move the code added by "[CALCITE-1824] GROUP_ID returns wrong result" from SqlToRelConverter to RelBuilder. Now GROUP_ID is handled correctly if created via RelBuilder. Fix a bug where, in a query with multiple GROUPING SETS, RelBuilder incorrectly simplifed an Aggregate to a Project. That simplification is valid only with a single groupng set. --- .../sql/validate/AggregatingSelectScope.java | 25 +- .../apache/calcite/sql2rel/SqlToRelConverter.java | 121 +--------- .../org/apache/calcite/tools/PigRelBuilder.java | 15 +- .../java/org/apache/calcite/tools/RelBuilder.java | 268 +++++++++++++++++---- .../org/apache/calcite/test/RelBuilderTest.java | 38 +++ .../org/apache/calcite/test/RelOptRulesTest.java | 7 +- .../apache/calcite/test/SqlToRelConverterTest.java | 7 + .../org/apache/calcite/test/SqlToRelTestBase.java | 3 +- .../apache/calcite/test/SqlToRelConverterTest.xml | 20 ++ core/src/test/resources/sql/agg.iq | 15 ++ .../org/apache/calcite/piglet/PigRelBuilder.java | 5 +- 11 files changed, 333 insertions(+), 191 deletions(-) diff --git a/core/src/main/java/org/apache/calcite/sql/validate/AggregatingSelectScope.java b/core/src/main/java/org/apache/calcite/sql/validate/AggregatingSelectScope.java index b61fec5..279772d 100644 --- a/core/src/main/java/org/apache/calcite/sql/validate/AggregatingSelectScope.java +++ b/core/src/main/java/org/apache/calcite/sql/validate/AggregatingSelectScope.java @@ -30,11 +30,11 @@ import org.apache.calcite.util.Pair; import com.google.common.base.Suppliers; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; -import com.google.common.collect.Maps; +import com.google.common.collect.ImmutableSortedMultiset; +import java.util.ArrayList; import java.util.List; import java.util.Map; -import java.util.SortedMap; import java.util.function.Supplier; import static org.apache.calcite.sql.SqlUtil.stripAs; @@ -96,21 +96,18 @@ public class AggregatingSelectScope } } - final SortedMap<ImmutableBitSet, Integer> flatGroupSetCount = - Maps.newTreeMap(ImmutableBitSet.COMPARATOR); + final List<ImmutableBitSet> flatGroupSets = new ArrayList<>(); for (List<ImmutableBitSet> groupSet : Linq4j.product(builder.build())) { - final ImmutableBitSet set = ImmutableBitSet.union(groupSet); - flatGroupSetCount.put(set, flatGroupSetCount.getOrDefault(set, 0) + 1); + flatGroupSets.add(ImmutableBitSet.union(groupSet)); } // For GROUP BY (), we need a singleton grouping set. - if (flatGroupSetCount.isEmpty()) { - flatGroupSetCount.put(ImmutableBitSet.of(), 1); + if (flatGroupSets.isEmpty()) { + flatGroupSets.add(ImmutableBitSet.of()); } return new Resolved(groupAnalyzer.extraExprs, groupAnalyzer.groupExprs, - flatGroupSetCount.keySet(), flatGroupSetCount, - groupAnalyzer.groupExprProjection); + flatGroupSets, groupAnalyzer.groupExprProjection); } finally { groupAnalyzer = null; } @@ -224,23 +221,21 @@ public class AggregatingSelectScope /** Information about an aggregating scope that can only be determined * after validation has occurred. Therefore it cannot be populated when * the scope is created. */ + @SuppressWarnings("UnstableApiUsage") public static class Resolved { public final ImmutableList<SqlNode> extraExprList; public final ImmutableList<SqlNode> groupExprList; public final ImmutableBitSet groupSet; - public final ImmutableList<ImmutableBitSet> groupSets; - public final Map<ImmutableBitSet, Integer> groupSetCount; + public final ImmutableSortedMultiset<ImmutableBitSet> groupSets; public final Map<Integer, Integer> groupExprProjection; Resolved(List<SqlNode> extraExprList, List<SqlNode> groupExprList, Iterable<ImmutableBitSet> groupSets, - Map<ImmutableBitSet, Integer> groupSetCount, Map<Integer, Integer> groupExprProjection) { this.extraExprList = ImmutableList.copyOf(extraExprList); this.groupExprList = ImmutableList.copyOf(groupExprList); this.groupSet = ImmutableBitSet.range(groupExprList.size()); - this.groupSets = ImmutableList.copyOf(groupSets); - this.groupSetCount = ImmutableMap.copyOf(groupSetCount); + this.groupSets = ImmutableSortedMultiset.copyOf(groupSets); this.groupExprProjection = ImmutableMap.copyOf(groupExprProjection); } diff --git a/core/src/main/java/org/apache/calcite/sql2rel/SqlToRelConverter.java b/core/src/main/java/org/apache/calcite/sql2rel/SqlToRelConverter.java index af7c390..b941e72 100644 --- a/core/src/main/java/org/apache/calcite/sql2rel/SqlToRelConverter.java +++ b/core/src/main/java/org/apache/calcite/sql2rel/SqlToRelConverter.java @@ -181,7 +181,6 @@ import com.google.common.collect.ImmutableList.Builder; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Iterables; -import com.google.common.collect.Sets; import org.slf4j.Logger; @@ -217,6 +216,7 @@ import static org.apache.calcite.sql.SqlUtil.stripAs; * <p>The public entry points are: {@link #convertQuery}, * {@link #convertExpression(SqlNode)}. */ +@SuppressWarnings("UnstableApiUsage") public class SqlToRelConverter { //~ Static fields/initializers --------------------------------------------- @@ -3020,12 +3020,10 @@ public class SqlToRelConverter { bb.scope.getMonotonicity(groupItem)); } - final RelNode relNode = aggConverter.containsGroupId() - ? rewriteAggregateWithGroupId(bb, r, aggConverter) - : createAggregate(bb, r.groupSet, r.groupSets, - aggConverter.getAggCalls()); - - bb.setRoot(relNode, false); + // Add the aggregator + bb.setRoot( + createAggregate(bb, r.groupSet, r.groupSets.asList(), + aggConverter.getAggCalls()), false); bb.mapRootRelToFieldProjection.put(bb.root, r.groupExprProjection); // Replace sub-queries in having here and modify having to use @@ -3100,104 +3098,6 @@ public class SqlToRelConverter { } /** - * The {@code GROUP_ID()} function is used to distinguish duplicate groups. - * However, as Aggregate normalizes group sets to canonical form (i.e., - * flatten, sorting, redundancy removal), this information is lost in RelNode. - * Therefore, it is impossible to implement the function in runtime. - * - * To fill this gap, an aggregation query that contains {@code GROUP_ID()} function - * will generally be rewritten into UNION when converting to RelNode. - * - * Also see the discussion in JIRA - * <a href="https://issues.apache.org/jira/browse/CALCITE-1824";>[CALCITE-1824] - * GROUP_ID returns wrong result</a>. - */ - private RelNode rewriteAggregateWithGroupId(Blackboard bb, - AggregatingSelectScope.Resolved r, AggConverter converter) { - final List<AggregateCall> aggregateCalls = converter.getAggCalls(); - final ImmutableBitSet groupSet = r.groupSet; - final Map<ImmutableBitSet, Integer> groupSetCount = r.groupSetCount; - - final List<String> fieldNamesIfNoRewrite = createAggregate(bb, groupSet, - r.groupSets, aggregateCalls).getRowType().getFieldNames(); - - // If n duplicates exist for a particular grouping, the {@code GROUP_ID()} - // function produces values in the range 0 to n-1. For each value, - // we need to figure out the corresponding group sets. - // - // For example, "... GROUPING SETS (a, a, b, c, c, c, c)" - // (i) The max value of the GROUP_ID() function returns is 3 - // (ii) GROUPING SETS (a, b, c) produces value 0, - // GROUPING SETS (a, c) produces value 1, - // GROUPING SETS (c) produces value 2 - // GROUPING SETS (c) produces value 3 - final Map<Integer, Set<ImmutableBitSet>> groupIdToGroupSets = new HashMap<>(); - int maxGroupId = 0; - for (Map.Entry<ImmutableBitSet, Integer> entry: groupSetCount.entrySet()) { - int groupId = entry.getValue() - 1; - if (groupId > maxGroupId) { - maxGroupId = groupId; - } - for (int i = 0; i <= groupId; i++) { - groupIdToGroupSets.computeIfAbsent(i, - k -> Sets.newTreeSet(ImmutableBitSet.COMPARATOR)) - .add(entry.getKey()); - } - } - - // AggregateCall list without GROUP_ID function - final List<AggregateCall> aggregateCallsWithoutGroupId = new ArrayList<>(); - for (AggregateCall aggregateCall : aggregateCalls) { - if (aggregateCall.getAggregation().kind != SqlKind.GROUP_ID) { - aggregateCallsWithoutGroupId.add(aggregateCall); - } - } - final List<RelNode> projects = new ArrayList<>(); - // For each group id value , we first construct an Aggregate without - // GROUP_ID() function call, and then create a Project node on top of it. - // The Project adds literal value for group id in right position. - for (int groupId = 0; groupId <= maxGroupId; groupId++) { - // Create the Aggregate node without GROUP_ID() call - final ImmutableList<ImmutableBitSet> groupSets = - ImmutableList.copyOf(groupIdToGroupSets.get(groupId)); - final RelNode aggregate = createAggregate(bb, groupSet, - groupSets, aggregateCallsWithoutGroupId); - - // RexLiteral for each GROUP_ID, note the type should be BIGINT - final RelDataType groupIdType = typeFactory.createSqlType(SqlTypeName.BIGINT); - final RexNode groupIdLiteral = rexBuilder.makeExactLiteral( - BigDecimal.valueOf(groupId), groupIdType); - - relBuilder.push(aggregate); - final List<RexNode> selectList = new ArrayList<>(); - final int groupExprLength = r.groupExprList.size(); - // Project fields in group by expressions - for (int i = 0; i < groupExprLength; i++) { - selectList.add(relBuilder.field(i)); - } - // Project fields in aggregate calls - int groupIdCount = 0; - for (int i = 0; i < aggregateCalls.size(); i++) { - if (aggregateCalls.get(i).getAggregation().kind == SqlKind.GROUP_ID) { - selectList.add(groupIdLiteral); - groupIdCount++; - } else { - int ordinal = groupExprLength + i - groupIdCount; - selectList.add(relBuilder.field(ordinal)); - } - } - final RelNode project = relBuilder.project( - selectList, fieldNamesIfNoRewrite).build(); - projects.add(project); - } - // Skip to create Union when there is only one child, i.e., no duplicate group set. - if (projects.size() == 1) { - return projects.get(0); - } - return LogicalUnion.create(projects, true); - } - - /** * Creates an Aggregate. * * <p>In case the aggregate rel changes the order in which it projects @@ -3217,7 +3117,11 @@ public class SqlToRelConverter { */ protected RelNode createAggregate(Blackboard bb, ImmutableBitSet groupSet, ImmutableList<ImmutableBitSet> groupSets, List<AggregateCall> aggCalls) { - return LogicalAggregate.create(bb.root, ImmutableList.of(), groupSet, groupSets, aggCalls); + relBuilder.push(bb.root); + final RelBuilder.GroupKey groupKey = + relBuilder.groupKey(groupSet, (Iterable<ImmutableBitSet>) groupSets); + return relBuilder.aggregate(groupKey, aggCalls) + .build(); } public RexDynamicParam convertDynamicParam( @@ -5493,11 +5397,6 @@ public class SqlToRelConverter { return aggCalls; } - private boolean containsGroupId() { - return aggCalls.stream().anyMatch( - agg -> agg.getAggregation().kind == SqlKind.GROUP_ID); - } - public RelDataTypeFactory getTypeFactory() { return typeFactory; } diff --git a/core/src/main/java/org/apache/calcite/tools/PigRelBuilder.java b/core/src/main/java/org/apache/calcite/tools/PigRelBuilder.java index 0c46267..f41f151 100644 --- a/core/src/main/java/org/apache/calcite/tools/PigRelBuilder.java +++ b/core/src/main/java/org/apache/calcite/tools/PigRelBuilder.java @@ -122,13 +122,12 @@ public class PigRelBuilder extends RelBuilder { public PigRelBuilder group(GroupOption option, Partitioner partitioner, int parallel, Iterable<? extends GroupKey> groupKeys) { - @SuppressWarnings("unchecked") final List<GroupKeyImpl> groupKeyList = - ImmutableList.copyOf((Iterable) groupKeys); + final List<GroupKey> groupKeyList = ImmutableList.copyOf(groupKeys); validateGroupList(groupKeyList); - final int groupCount = groupKeyList.get(0).nodes.size(); + final int groupCount = groupKeyList.get(0).groupKeyCount(); final int n = groupKeyList.size(); - for (Ord<GroupKeyImpl> groupKey : Ord.reverse(groupKeyList)) { + for (Ord<GroupKey> groupKey : Ord.reverse(groupKeyList)) { RelNode r = null; if (groupKey.i < n - 1) { r = build(); @@ -152,13 +151,13 @@ public class PigRelBuilder extends RelBuilder { return this; } - protected void validateGroupList(List<GroupKeyImpl> groupKeyList) { + protected void validateGroupList(List<GroupKey> groupKeyList) { if (groupKeyList.isEmpty()) { throw new IllegalArgumentException("must have at least one group"); } - final int groupCount = groupKeyList.get(0).nodes.size(); - for (GroupKeyImpl groupKey : groupKeyList) { - if (groupKey.nodes.size() != groupCount) { + final int groupCount = groupKeyList.get(0).groupKeyCount(); + for (GroupKey groupKey : groupKeyList) { + if (groupKey.groupKeyCount() != groupCount) { throw new IllegalArgumentException("group key size mismatch"); } } diff --git a/core/src/main/java/org/apache/calcite/tools/RelBuilder.java b/core/src/main/java/org/apache/calcite/tools/RelBuilder.java index af33cc4..bcd860b 100644 --- a/core/src/main/java/org/apache/calcite/tools/RelBuilder.java +++ b/core/src/main/java/org/apache/calcite/tools/RelBuilder.java @@ -105,8 +105,11 @@ import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; +import com.google.common.collect.ImmutableSortedMultiset; import com.google.common.collect.Iterables; import com.google.common.collect.Lists; +import com.google.common.collect.Multiset; +import com.google.common.collect.Sets; import java.math.BigDecimal; import java.util.AbstractList; @@ -1299,8 +1302,7 @@ public class RelBuilder { /** Creates a {@link Project} of all original fields, plus the given list of * expressions. */ public RelBuilder projectPlus(Iterable<RexNode> nodes) { - final ImmutableList.Builder<RexNode> builder = ImmutableList.builder(); - return project(builder.addAll(fields()).addAll(nodes).build()); + return project(Iterables.concat(fields(), nodes)); } /** Creates a {@link Project} of all original fields, except the given @@ -1652,8 +1654,7 @@ public class RelBuilder { .collect(Collectors.toList())); } - /** Creates an {@link Aggregate} with multiple - * calls. */ + /** Creates an {@link Aggregate} with multiple calls. */ public RelBuilder aggregate(GroupKey groupKey, Iterable<AggCall> aggCalls) { final Registrar registrar = new Registrar(fields(), peek().getRowType().getFieldNames()); @@ -1673,22 +1674,26 @@ public class RelBuilder { } if (registrar.extraNodes.size() == fields().size()) { final Boolean unique = mq.areColumnsUnique(peek(), groupSet); - if (unique != null && unique) { + if (unique != null && unique + && !config.aggregateUnique() + && groupKey_.isSimple()) { // Rel is already unique. return project(fields(groupSet)); } } final Double maxRowCount = mq.getMaxRowCount(peek()); - if (maxRowCount != null && maxRowCount <= 1D) { + if (maxRowCount != null && maxRowCount <= 1D + && !config.aggregateUnique() + && groupKey_.isSimple()) { // If there is at most one row, rel is already unique. return project(fields(groupSet)); } } + ImmutableList<ImmutableBitSet> groupSets; if (groupKey_.nodeLists != null) { final int sizeBefore = registrar.extraNodes.size(); - final SortedSet<ImmutableBitSet> groupSetSet = - new TreeSet<>(ImmutableBitSet.ORDERING); + final List<ImmutableBitSet> groupSetList = new ArrayList<>(); for (ImmutableList<RexNode> nodeList : groupKey_.nodeLists) { final ImmutableBitSet groupSet2 = ImmutableBitSet.of(registrar.registerExpressions(nodeList)); @@ -1696,9 +1701,16 @@ public class RelBuilder { throw new IllegalArgumentException("group set element " + nodeList + " must be a subset of group key"); } - groupSetSet.add(groupSet2); + groupSetList.add(groupSet2); + } + final ImmutableSortedMultiset<ImmutableBitSet> groupSetMultiset = + ImmutableSortedMultiset.copyOf(ImmutableBitSet.COMPARATOR, + groupSetList); + if (Iterables.any(aggCalls, RelBuilder::isGroupId)) { + return rewriteAggregateWithGroupId(groupSet, groupSetMultiset, + ImmutableList.copyOf(aggCalls)); } - groupSets = ImmutableList.copyOf(groupSetSet); + groupSets = ImmutableList.copyOf(groupSetMultiset.elementSet()); if (registrar.extraNodes.size() > sizeBefore) { throw new IllegalArgumentException( "group sets contained expressions not in group key: " @@ -1708,14 +1720,9 @@ public class RelBuilder { } else { groupSets = ImmutableList.of(groupSet); } + for (AggCall aggCall : aggCalls) { - if (aggCall instanceof AggCallImpl) { - final AggCallImpl aggCall1 = (AggCallImpl) aggCall; - registrar.registerExpressions(aggCall1.operands); - if (aggCall1.filter != null) { - registrar.registerExpression(aggCall1.filter); - } - } + ((AggCallPlus) aggCall).register(registrar); } project(registrar.extraNodes); rename(registrar.names); @@ -1723,35 +1730,8 @@ public class RelBuilder { RelNode r = frame.rel; final List<AggregateCall> aggregateCalls = new ArrayList<>(); for (AggCall aggCall : aggCalls) { - final AggregateCall aggregateCall; - if (aggCall instanceof AggCallImpl) { - final AggCallImpl aggCall1 = (AggCallImpl) aggCall; - final List<Integer> args = - registrar.registerExpressions(aggCall1.operands); - final int filterArg = aggCall1.filter == null ? -1 - : registrar.registerExpression(aggCall1.filter); - if (aggCall1.distinct && !aggCall1.aggFunction.isQuantifierAllowed()) { - throw new IllegalArgumentException("DISTINCT not allowed"); - } - if (aggCall1.filter != null && !aggCall1.aggFunction.allowsFilter()) { - throw new IllegalArgumentException("FILTER not allowed"); - } - RelCollation collation = - RelCollations.of(aggCall1.orderKeys - .stream() - .map(orderKey -> - collation(orderKey, RelFieldCollation.Direction.ASCENDING, - null, Collections.emptyList())) - .collect(Collectors.toList())); - aggregateCall = - AggregateCall.create(aggCall1.aggFunction, aggCall1.distinct, - aggCall1.approximate, - aggCall1.ignoreNulls, args, filterArg, collation, - groupSet.cardinality(), r, null, aggCall1.alias); - } else { - aggregateCall = ((AggCallImpl2) aggCall).aggregateCall; - } - aggregateCalls.add(aggregateCall); + aggregateCalls.add( + ((AggCallPlus) aggCall).aggregateCall(registrar, groupSet, r)); } assert ImmutableBitSet.ORDERING.isStrictlyOrdered(groupSets) : groupSets; @@ -1880,6 +1860,95 @@ public class RelBuilder { return this; } + /** + * The {@code GROUP_ID()} function is used to distinguish duplicate groups. + * However, as Aggregate normalizes group sets to canonical form (i.e., + * flatten, sorting, redundancy removal), this information is lost in RelNode. + * Therefore, it is impossible to implement the function in runtime. + * + * <p>To fill this gap, an aggregation query that contains {@code GROUP_ID()} + * function will generally be rewritten into UNION when converting to RelNode. + * + * <p>Also see the discussion in + * <a href="https://issues.apache.org/jira/browse/CALCITE-1824";>[CALCITE-1824] + * GROUP_ID returns wrong result</a>. + */ + private RelBuilder rewriteAggregateWithGroupId(ImmutableBitSet groupSet, + ImmutableSortedMultiset<ImmutableBitSet> groupSets, + List<AggCall> aggregateCalls) { + final List<String> fieldNamesIfNoRewrite = + Aggregate.deriveRowType(getTypeFactory(), peek().getRowType(), false, + groupSet, groupSets.asList(), + aggregateCalls.stream().map(c -> ((AggCallPlus) c).aggregateCall()) + .collect(Util.toImmutableList())).getFieldNames(); + + // If n duplicates exist for a particular grouping, the {@code GROUP_ID()} + // function produces values in the range 0 to n-1. For each value, + // we need to figure out the corresponding group sets. + // + // For example, "... GROUPING SETS (a, a, b, c, c, c, c)" + // (i) The max value of the GROUP_ID() function returns is 3 + // (ii) GROUPING SETS (a, b, c) produces value 0, + // GROUPING SETS (a, c) produces value 1, + // GROUPING SETS (c) produces value 2 + // GROUPING SETS (c) produces value 3 + final Map<Integer, Set<ImmutableBitSet>> groupIdToGroupSets = new HashMap<>(); + int maxGroupId = 0; + for (Multiset.Entry<ImmutableBitSet> entry: groupSets.entrySet()) { + int groupId = entry.getCount() - 1; + if (groupId > maxGroupId) { + maxGroupId = groupId; + } + for (int i = 0; i <= groupId; i++) { + groupIdToGroupSets.computeIfAbsent(i, + k -> Sets.newTreeSet(ImmutableBitSet.COMPARATOR)) + .add(entry.getElement()); + } + } + + // AggregateCall list without GROUP_ID function + final List<AggCall> aggregateCallsWithoutGroupId = + new ArrayList<>(aggregateCalls); + aggregateCallsWithoutGroupId.removeIf(RelBuilder::isGroupId); + + // For each group id value, we first construct an Aggregate without + // GROUP_ID() function call, and then create a Project node on top of it. + // The Project adds literal value for group id in right position. + final Frame frame = stack.pop(); + for (int groupId = 0; groupId <= maxGroupId; groupId++) { + // Create the Aggregate node without GROUP_ID() call + stack.push(frame); + aggregate(groupKey(groupSet, groupIdToGroupSets.get(groupId)), + aggregateCallsWithoutGroupId); + + final List<RexNode> selectList = new ArrayList<>(); + final int groupExprLength = groupSet.cardinality(); + // Project fields in group by expressions + for (int i = 0; i < groupExprLength; i++) { + selectList.add(field(i)); + } + // Project fields in aggregate calls + int groupIdCount = 0; + for (int i = 0; i < aggregateCalls.size(); i++) { + if (isGroupId(aggregateCalls.get(i))) { + selectList.add( + getRexBuilder().makeExactLiteral(BigDecimal.valueOf(groupId), + getTypeFactory().createSqlType(SqlTypeName.BIGINT))); + groupIdCount++; + } else { + selectList.add(field(groupExprLength + i - groupIdCount)); + } + } + project(selectList, fieldNamesIfNoRewrite); + } + + return union(true, maxGroupId + 1); + } + + private static boolean isGroupId(AggCall c) { + return ((AggCallPlus) c).op().kind == SqlKind.GROUP_ID; + } + private RelBuilder setOp(boolean all, SqlKind kind, int n) { List<RelNode> inputs = new LinkedList<>(); for (int i = 0; i < n; i++) { @@ -2787,6 +2856,24 @@ public class RelBuilder { AggCall distinct(); } + /** Internal methods shared by all implementations of {@link AggCall}. */ + private interface AggCallPlus extends AggCall { + /** Returns the aggregate function. */ + SqlAggFunction op(); + + /** Returns an {@link AggregateCall} that is approximately equivalent + * to this {@code AggCall} and is good for certain things, such as deriving + * field names. */ + AggregateCall aggregateCall(); + + /** Converts this {@code AggCall} to a good {@link AggregateCall}. */ + AggregateCall aggregateCall(Registrar registrar, ImmutableBitSet groupSet, + RelNode r); + + /** Registers expressions in operands and filters. */ + void register(Registrar registrar); + } + /** Information necessary to create the GROUP BY clause of an Aggregate. * * @see RelBuilder#groupKey */ @@ -2795,13 +2882,16 @@ public class RelBuilder { * * <p>Used to assign field names in the {@code group} operation. */ GroupKey alias(String alias); + + /** Returns the number of columns in the group key. */ + int groupKeyCount(); } /** Implementation of {@link RelBuilder.GroupKey}. */ - public static class GroupKeyImpl implements GroupKey { - public final ImmutableList<RexNode> nodes; - public final ImmutableList<ImmutableList<RexNode>> nodeLists; - public final String alias; + static class GroupKeyImpl implements GroupKey { + final ImmutableList<RexNode> nodes; + final ImmutableList<ImmutableList<RexNode>> nodeLists; + final String alias; GroupKeyImpl(ImmutableList<RexNode> nodes, ImmutableList<ImmutableList<RexNode>> nodeLists, String alias) { @@ -2814,15 +2904,23 @@ public class RelBuilder { return alias == null ? nodes.toString() : nodes + " as " + alias; } + @Override public int groupKeyCount() { + return nodes.size(); + } + public GroupKey alias(String alias) { return Objects.equals(this.alias, alias) ? this : new GroupKeyImpl(nodes, nodeLists, alias); } + + boolean isSimple() { + return nodeLists == null || nodeLists.size() == 1; + } } /** Implementation of {@link AggCall}. */ - private class AggCallImpl implements AggCall { + private class AggCallImpl implements AggCallPlus { private final SqlAggFunction aggFunction; private final boolean distinct; private final boolean approximate; @@ -2879,6 +2977,46 @@ public class RelBuilder { return b.toString(); } + @Override public SqlAggFunction op() { + return aggFunction; + } + + @Override public AggregateCall aggregateCall() { + return AggregateCall.create(aggFunction, distinct, approximate, + ignoreNulls, ImmutableList.of(), -1, null, null, alias); + } + + @Override public AggregateCall aggregateCall(Registrar registrar, + ImmutableBitSet groupSet, RelNode r) { + final List<Integer> args = + registrar.registerExpressions(this.operands); + final int filterArg = this.filter == null ? -1 + : registrar.registerExpression(this.filter); + if (this.distinct && !this.aggFunction.isQuantifierAllowed()) { + throw new IllegalArgumentException("DISTINCT not allowed"); + } + if (this.filter != null && !this.aggFunction.allowsFilter()) { + throw new IllegalArgumentException("FILTER not allowed"); + } + RelCollation collation = + RelCollations.of(this.orderKeys + .stream() + .map(orderKey -> + collation(orderKey, RelFieldCollation.Direction.ASCENDING, + null, Collections.emptyList())) + .collect(Collectors.toList())); + return AggregateCall.create(aggFunction, distinct, approximate, + ignoreNulls, args, filterArg, collation, groupSet.cardinality(), r, + null, alias); + } + + @Override public void register(Registrar registrar) { + registrar.registerExpressions(operands); + if (filter != null) { + registrar.registerExpression(filter); + } + } + public AggCall sort(Iterable<RexNode> orderKeys) { final ImmutableList<RexNode> orderKeyList = ImmutableList.copyOf(orderKeys); @@ -2934,7 +3072,7 @@ public class RelBuilder { /** Implementation of {@link AggCall} that wraps an * {@link AggregateCall}. */ - private static class AggCallImpl2 implements AggCall { + private static class AggCallImpl2 implements AggCallPlus { private final AggregateCall aggregateCall; AggCallImpl2(AggregateCall aggregateCall) { @@ -2945,6 +3083,23 @@ public class RelBuilder { return aggregateCall.toString(); } + @Override public SqlAggFunction op() { + return aggregateCall.getAggregation(); + } + + @Override public AggregateCall aggregateCall() { + return aggregateCall; + } + + @Override public AggregateCall aggregateCall(Registrar registrar, + ImmutableBitSet groupSet, RelNode r) { + return aggregateCall; + } + + @Override public void register(Registrar registrar) { + // nothing to do + } + public AggCall sort(Iterable<RexNode> orderKeys) { throw new UnsupportedOperationException(); } @@ -3231,6 +3386,15 @@ public class RelBuilder { /** Sets {@link #simplify}. */ Config withSimplify(boolean simplify); + + /** Whether to create an Aggregate even if we know that the input is + * already unique; default false. */ + @ImmutableBeans.Property + @ImmutableBeans.BooleanDefault(false) + boolean aggregateUnique(); + + /** Sets {@link #aggregateUnique()}. */ + Config withAggregateUnique(boolean aggregateUnique); } /** Creates a {@link RelBuilder.Config}. diff --git a/core/src/test/java/org/apache/calcite/test/RelBuilderTest.java b/core/src/test/java/org/apache/calcite/test/RelBuilderTest.java index 610905e..ea74078 100644 --- a/core/src/test/java/org/apache/calcite/test/RelBuilderTest.java +++ b/core/src/test/java/org/apache/calcite/test/RelBuilderTest.java @@ -1569,6 +1569,44 @@ public class RelBuilderTest { } } + @Test void testAggregateOneRow() { + final Function<RelBuilder, RelNode> f = builder -> + builder.values(new String[] {"a", "b"}, 1, 2) + .aggregate(builder.groupKey(1)) + .build(); + final String plan = "LogicalProject(b=[$1])\n" + + " LogicalValues(tuples=[[{ 1, 2 }]])\n"; + assertThat(f.apply(createBuilder()), hasTree(plan)); + + final String plan2 = "LogicalAggregate(group=[{1}])\n" + + " LogicalValues(tuples=[[{ 1, 2 }]])\n"; + assertThat(f.apply(createBuilder(c -> c.withAggregateUnique(true))), + hasTree(plan2)); + } + + /** Tests that we do not convert an Aggregate to a Project if there are + * multiple group sets. */ + @Test void testAggregateGroupingSetsOneRow() { + final Function<RelBuilder, RelNode> f = builder -> { + final List<Integer> list01 = Arrays.asList(0, 1); + final List<Integer> list0 = Collections.singletonList(0); + final List<Integer> list1 = Collections.singletonList(1); + return builder.values(new String[] {"a", "b"}, 1, 2) + .aggregate( + builder.groupKey(builder.fields(list01), + ImmutableList.of(builder.fields(list0), + builder.fields(list1), + builder.fields(list01)))) + .build(); + }; + final String plan = "" + + "LogicalAggregate(group=[{0, 1}], groups=[[{0, 1}, {0}, {1}]])\n" + + " LogicalValues(tuples=[[{ 1, 2 }]])\n"; + assertThat(f.apply(createBuilder()), hasTree(plan)); + assertThat(f.apply(createBuilder(c -> c.withAggregateUnique(true))), + hasTree(plan)); + } + @Test void testDistinct() { // Equivalent SQL: // SELECT DISTINCT deptno diff --git a/core/src/test/java/org/apache/calcite/test/RelOptRulesTest.java b/core/src/test/java/org/apache/calcite/test/RelOptRulesTest.java index 42a3c87..b049154 100644 --- a/core/src/test/java/org/apache/calcite/test/RelOptRulesTest.java +++ b/core/src/test/java/org/apache/calcite/test/RelOptRulesTest.java @@ -241,6 +241,7 @@ class RelOptRulesTest extends RelOptTestBase { + " group by empno, deptno))\n" + "or deptno < 40 + 60"; checkSubQuery(sql) + .withRelBuilderConfig(b -> b.withAggregateUnique(true)) .withRule(CoreRules.FILTER_REDUCE_EXPRESSIONS) .check(); } @@ -4488,6 +4489,7 @@ class RelOptRulesTest extends RelOptTestBase { + "left outer join sales.dept as d on e.empno < d.deptno\n" + "group by e.empno,d.deptno"; sql(sql) + .withRelBuilderConfig(b -> b.withAggregateUnique(true)) .withPreRule(CoreRules.AGGREGATE_PROJECT_MERGE) .withRule(CoreRules.AGGREGATE_JOIN_TRANSPOSE_EXTENDED) .checkUnchanged(); @@ -4577,6 +4579,7 @@ class RelOptRulesTest extends RelOptTestBase { + "join sales.dept as d on e.empno < d.deptno\n" + "group by e.empno,d.deptno"; sql(sql) + .withRelBuilderConfig(b -> b.withAggregateUnique(true)) .withPreRule(CoreRules.AGGREGATE_PROJECT_MERGE) .withRule(CoreRules.AGGREGATE_JOIN_TRANSPOSE_EXTENDED) .checkUnchanged(); @@ -5022,6 +5025,7 @@ class RelOptRulesTest extends RelOptTestBase { @Test void testAggregateRemove2() { final String sql = "select distinct empno, deptno from sales.emp\n"; sql(sql) + .withRelBuilderConfig(b -> b.withAggregateUnique(true)) .withRule(CoreRules.AGGREGATE_REMOVE, CoreRules.PROJECT_MERGE) .check(); @@ -5988,7 +5992,8 @@ class RelOptRulesTest extends RelOptTestBase { @Test void testWhereInCorrelated() { final String sql = "select sal from emp where empno IN (\n" + " select deptno from dept where emp.job = dept.name)"; - checkSubQuery(sql).withLateDecorrelation(true).check(); + checkSubQuery(sql).withLateDecorrelation(true) + .check(); } @Test void testWhereExpressionInCorrelated() { diff --git a/core/src/test/java/org/apache/calcite/test/SqlToRelConverterTest.java b/core/src/test/java/org/apache/calcite/test/SqlToRelConverterTest.java index 44d46ef..0f76773 100644 --- a/core/src/test/java/org/apache/calcite/test/SqlToRelConverterTest.java +++ b/core/src/test/java/org/apache/calcite/test/SqlToRelConverterTest.java @@ -537,6 +537,13 @@ class SqlToRelConverterTest extends SqlToRelTestBase { sql(sql).ok(); } + @Test void testGroupingSetsRepeated() { + final String sql = "select deptno, group_id()\n" + + "from emp\n" + + "group by grouping sets (deptno, (), deptno)"; + sql(sql).ok(); + } + @Test void testGroupingSetsWith() { final String sql = "with t(a, b, c, d) as (values (1, 2, 3, 4))\n" + "select 1 from t\n" diff --git a/core/src/test/java/org/apache/calcite/test/SqlToRelTestBase.java b/core/src/test/java/org/apache/calcite/test/SqlToRelTestBase.java index 68bef38..e74b67d 100644 --- a/core/src/test/java/org/apache/calcite/test/SqlToRelTestBase.java +++ b/core/src/test/java/org/apache/calcite/test/SqlToRelTestBase.java @@ -115,7 +115,8 @@ public abstract class SqlToRelTestBase { c.withTrimUnusedFields(true) .withExpand(true) .addRelBuilderConfigTransform(b -> - b.withPruneInputOfAggregate(false))); + b.withAggregateUnique(true) + .withPruneInputOfAggregate(false))); } protected Tester getTesterWithDynamicTable() { diff --git a/core/src/test/resources/org/apache/calcite/test/SqlToRelConverterTest.xml b/core/src/test/resources/org/apache/calcite/test/SqlToRelConverterTest.xml index adb27d2..46fdbd0 100644 --- a/core/src/test/resources/org/apache/calcite/test/SqlToRelConverterTest.xml +++ b/core/src/test/resources/org/apache/calcite/test/SqlToRelConverterTest.xml @@ -1406,6 +1406,26 @@ LogicalProject(EXPR$0=[JSON_TYPE(ROW($2.TYPE, $2.DESC, ROW($2.OTHERS.A, $2.OTHER ]]> </Resource> </TestCase> + <TestCase name="testGroupingSetsRepeated"> + <Resource name="sql"> + <![CDATA[select deptno, group_id() +from emp +group by grouping sets (deptno, (), deptno)]]> + </Resource> + <Resource name="plan"> + <![CDATA[ +LogicalUnion(all=[true]) + LogicalProject(DEPTNO=[$0], EXPR$1=[0:BIGINT]) + LogicalAggregate(group=[{0}], groups=[[{0}, {}]]) + LogicalProject(DEPTNO=[$7]) + LogicalTableScan(table=[[CATALOG, SALES, EMP]]) + LogicalProject(DEPTNO=[$0], EXPR$1=[1:BIGINT]) + LogicalAggregate(group=[{0}]) + LogicalProject(DEPTNO=[$7]) + LogicalTableScan(table=[[CATALOG, SALES, EMP]]) +]]> + </Resource> + </TestCase> <TestCase name="testOrder"> <Resource name="sql"> <![CDATA[select empno from emp order by empno, empno desc]]> diff --git a/core/src/test/resources/sql/agg.iq b/core/src/test/resources/sql/agg.iq index c3573e6..16cb93f 100644 --- a/core/src/test/resources/sql/agg.iq +++ b/core/src/test/resources/sql/agg.iq @@ -351,6 +351,21 @@ select deptno + 1, count(*) as c from emps group by grouping sets ((), (deptno + !ok +# GROUPING SETS on single-row relation returns multiple rows +select 1 as c +from (values ('a', 'b')) as t (a, b) +group by grouping sets ((a), (b), (b, a)); ++---+ +| C | ++---+ +| 1 | +| 1 | +| 1 | ++---+ +(3 rows) + +!ok + # CUBE select deptno + 1, count(*) as c from emp group by cube(deptno, gender); +--------+---+ diff --git a/piglet/src/main/java/org/apache/calcite/piglet/PigRelBuilder.java b/piglet/src/main/java/org/apache/calcite/piglet/PigRelBuilder.java index 6af2f24..93afa88 100644 --- a/piglet/src/main/java/org/apache/calcite/piglet/PigRelBuilder.java +++ b/piglet/src/main/java/org/apache/calcite/piglet/PigRelBuilder.java @@ -384,9 +384,8 @@ public class PigRelBuilder extends RelBuilder { * @return This builder */ public RelBuilder cogroup(Iterable<? extends GroupKey> groupKeys) { - @SuppressWarnings("unchecked") final List<GroupKeyImpl> groupKeyList = - ImmutableList.copyOf((Iterable) groupKeys); - final int groupCount = groupKeyList.get(0).nodes.size(); + final List<GroupKey> groupKeyList = ImmutableList.copyOf(groupKeys); + final int groupCount = groupKeyList.get(0).groupKeyCount(); // Pull out all relations needed for the group final int numRels = groupKeyList.size();
