diff --git a/core/src/main/java/org/apache/calcite/util/Bug.java b/core/src/main/java/org/apache/calcite/util/Bug.java index 7fdbb4c082f..80683183022 100644 --- a/core/src/main/java/org/apache/calcite/util/Bug.java +++ b/core/src/main/java/org/apache/calcite/util/Bug.java @@ -210,6 +210,12 @@ public abstract class Bug { * Support FORMAT in CAST from Numeric and BYTES to String (Enabled in BigQuery) is fixed. */ public static final boolean CALCITE_6270_FIXED = false; + /** Whether + * + * [CALCITE-6340] RelBuilder drops set conventions when aggregating over duplicate + * projected fields is fixed. */ + public static final boolean CALCITE_6340_FIXED = false; + /** * Use this to flag temporary code. */ 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 d38f7fcda33..d24b28d853d 100644 --- a/core/src/test/java/org/apache/calcite/test/RelBuilderTest.java +++ b/core/src/test/java/org/apache/calcite/test/RelBuilderTest.java @@ -26,6 +26,7 @@ import org.apache.calcite.plan.RelTraitSet; import org.apache.calcite.rel.RelCollations; import org.apache.calcite.rel.RelDistributions; +import org.apache.calcite.rel.RelFieldCollation; import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.core.AggregateCall; import org.apache.calcite.rel.core.Correlate; @@ -78,6 +79,7 @@ import org.apache.calcite.tools.RelRunners; import org.apache.calcite.tools.RuleSet; import org.apache.calcite.tools.RuleSets; +import org.apache.calcite.util.Bug; import org.apache.calcite.util.Holder; import org.apache.calcite.util.ImmutableBitSet; import org.apache.calcite.util.Pair; @@ -1453,55 +1455,155 @@ private RexNode caseCall(RelBuilder b, RexNode ref, RexNode... nodes) { /** Test case for * * [CALCITE-6340] RelBuilder drops set conventions when aggregating over duplicate - * projected fields.. + * projected fields */ - @Test void testPruneProjectInputOfAggregatePreservesConventionAndCollations() { + @Test void testPruneProjectInputOfAggregatePreservesConventionAndCollationsWhenEmpty() { final RelBuilder builder = createBuilder(config -> config.withPruneInputOfAggregate(true)); RelNode node = builder .scan("EMP") - .sort(builder.nullsLast(builder.desc(builder.field(0))), - builder.field(1)) + .sort(builder.nullsLast(builder.desc(builder.field(1))), + builder.field(0)) .project(builder.alias(builder.field(0), "a"), builder.alias(builder.field(1), "b"), builder.alias(builder.field(0), "c"), builder.alias(builder.field(1), "d")) .build(); - RelTraitSet desiredTraits = builder.getCluster().traitSet() + final RelTraitSet desiredTraits = builder.getCluster().traitSet() .replace(EnumerableConvention.INSTANCE); - - RuleSet prepareRules = + final RuleSet prepareRules = RuleSets.ofList(EnumerableRules.ENUMERABLE_PROJECT_RULE, EnumerableRules.ENUMERABLE_SORT_RULE, EnumerableRules.ENUMERABLE_TABLE_SCAN_RULE); - Program program = Programs.of(prepareRules); + final Program program = Programs.of(prepareRules); + + // turn the logical plan into a physical plan so that a convention can be set node = - program.run(node.getCluster().getPlanner(), node, desiredTraits, ImmutableList.of(), - ImmutableList.of()); + program.run(node.getCluster().getPlanner(), node, desiredTraits, ImmutableList.of(), ImmutableList.of()); + // collations are lost as the sort is on column [1, 0], but we group on 0, convention stays node = builder.push(node) .aggregate( builder.groupKey(0), builder.aggregateCall( - SqlStdOperatorTable.SUM, builder.field(0))) + SqlStdOperatorTable.SUM, builder.field(0))) .build(); - RelTraitSet relTraitSet = node.getInput(0).getTraitSet(); - assertTrue(relTraitSet.contains(EnumerableConvention.INSTANCE)); + final RelTraitSet expectedTraitSet = builder.getCluster().traitSet() + .replace(EnumerableConvention.INSTANCE); - // The below check on the collation currently fails because of - // https://issues.apache.org/jira/browse/CALCITE-6391. Uncomment once this issue is fixed. + if (Bug.CALCITE_6270_FIXED) { + assertThat(node.getInput(0).getTraitSet(), is(expectedTraitSet)); + } else { + assertThat(node.getInput(0).getTraitSet().get(0), is(expectedTraitSet.get(0))); + } + } - //final RelCollation collation1 = RelCollations.of(new RelFieldCollation(1, - // RelFieldCollation.Direction.DESCENDING, RelFieldCollation.NullDirection.LAST), - // new RelFieldCollation(0)); - // assertTrue(relTraitSet.getTrait(1).satisfies(collation1)); + /** Test case for + * + * [CALCITE-6340] RelBuilder drops set conventions when aggregating over duplicate + * projected fields + */ + @Test void testPruneProjectInputOfAggregatePreservesConventionAndSingletonCollation() { + final RelBuilder builder = createBuilder(config -> config.withPruneInputOfAggregate(true)); + + RelNode node = builder + .scan("EMP") + .sort(builder.nullsLast(builder.desc(builder.field(1)))) + .project(builder.alias(builder.field(0), "a"), + builder.alias(builder.field(1), "b"), + builder.alias(builder.field(0), "c"), + builder.alias(builder.field(1), "d")) + .build(); + + final RelTraitSet desiredTraits = builder.getCluster().traitSet() + .replace(EnumerableConvention.INSTANCE); + + final RuleSet prepareRules = + RuleSets.ofList(EnumerableRules.ENUMERABLE_PROJECT_RULE, + EnumerableRules.ENUMERABLE_SORT_RULE, + EnumerableRules.ENUMERABLE_TABLE_SCAN_RULE); + final Program program = Programs.of(prepareRules); + + // turn the logical plan into a physical plan so that a convention can be set + node = + program.run(node.getCluster().getPlanner(), node, desiredTraits, ImmutableList.of(), ImmutableList.of()); + + node = builder.push(node) + .aggregate( + builder.groupKey(1), builder.aggregateCall( + SqlStdOperatorTable.SUM, builder.field(1))) + .build(); + + final RelTraitSet expectedTraitSet = builder.getCluster().traitSet() + .replace(EnumerableConvention.INSTANCE) + .replace( + RelCollations.of( + new RelFieldCollation(0, + RelFieldCollation.Direction.DESCENDING, RelFieldCollation.NullDirection.LAST))); + + if (Bug.CALCITE_6270_FIXED) { + assertThat(node.getInput(0).getTraitSet(), is(expectedTraitSet)); + } else { + assertThat(node.getInput(0).getTraitSet().get(0), is(expectedTraitSet.get(0))); + } } /** Test case for * * [CALCITE-6340] RelBuilder drops set conventions when aggregating over duplicate - * projected fields.. + * projected fields + */ + @Test void testPruneProjectInputOfAggregatePreservesConventionAndCompositeCollation() { + final RelBuilder builder = createBuilder(config -> config.withPruneInputOfAggregate(true)); + + RelNode node = builder + .scan("EMP") + .sort(builder.nullsLast(builder.desc(builder.field(1))), + builder.field(0)) + .project(builder.alias(builder.field(0), "a"), + builder.alias(builder.field(1), "b"), + builder.alias(builder.field(0), "c"), + builder.alias(builder.field(1), "d")) + .build(); + + final RelTraitSet desiredTraits = builder.getCluster().traitSet() + .replace(EnumerableConvention.INSTANCE); + + final RuleSet prepareRules = + RuleSets.ofList(EnumerableRules.ENUMERABLE_PROJECT_RULE, + EnumerableRules.ENUMERABLE_SORT_RULE, + EnumerableRules.ENUMERABLE_TABLE_SCAN_RULE); + final Program program = Programs.of(prepareRules); + + // turn the logical plan into a physical plan so that a convention can be set + node = + program.run(node.getCluster().getPlanner(), node, desiredTraits, ImmutableList.of(), ImmutableList.of()); + + node = builder.push(node) + .aggregate( + builder.groupKey(1), builder.aggregateCall( + SqlStdOperatorTable.SUM, builder.field(1))) + .build(); + + final RelTraitSet expectedTraitSet = builder.getCluster().traitSet() + .replace(EnumerableConvention.INSTANCE) + .replace( + RelCollations.of( + new RelFieldCollation(0, + RelFieldCollation.Direction.DESCENDING, RelFieldCollation.NullDirection.LAST))); + + if (Bug.CALCITE_6270_FIXED) { + assertThat(node.getInput(0).getTraitSet(), is(expectedTraitSet)); + } else { + assertThat(node.getInput(0).getTraitSet().get(0), is(expectedTraitSet.get(0))); + } + } + + /** Test case for + * + * [CALCITE-6340] RelBuilder drops set conventions when aggregating over duplicate + * projected fields */ @Test void testPruneProjectInputOfAggregatePreservesConventionAndDistribution() { final RelBuilder builder = createBuilder(config -> config.withPruneInputOfAggregate(true)); @@ -1512,46 +1614,58 @@ private RexNode caseCall(RelBuilder b, RexNode ref, RexNode... nodes) { builder.alias(builder.field(0), "b"), builder.alias(builder.field(1), "c")) .build(); - RelTraitSet desiredTraits = builder.getCluster().traitSet() - .replace(EnumerableConvention.INSTANCE); - RuleSet prepareRules = + final RelTraitSet desiredTraits = builder.getCluster().traitSet() + .replace(EnumerableConvention.INSTANCE); + final RuleSet prepareRules = RuleSets.ofList(EnumerableRules.ENUMERABLE_PROJECT_RULE, EnumerableRules.ENUMERABLE_TABLE_SCAN_RULE); - Program program = Programs.of(prepareRules); + final Program program = Programs.of(prepareRules); + + // turn the logical plan into a physical plan so that a distribution can be set node = - program.run(node.getCluster().getPlanner(), node, desiredTraits, ImmutableList.of(), - ImmutableList.of()); + program.run(node.getCluster().getPlanner(), node, desiredTraits, ImmutableList.of(), ImmutableList.of()); + // setting the distribution drops the collations node = node.copy(desiredTraits.plus(RelDistributions.BROADCAST_DISTRIBUTED), node.getInputs()); node = builder.push(node) .aggregate( builder.groupKey(0), builder.aggregateCall( - SqlStdOperatorTable.SUM, builder.field(0))) + SqlStdOperatorTable.SUM, builder.field(0))) .build(); - RelTraitSet relTraitSet = node.getInput(0).getTraitSet(); - assertTrue(relTraitSet.contains(EnumerableConvention.INSTANCE)); - assertTrue(relTraitSet.contains(RelDistributions.BROADCAST_DISTRIBUTED)); + final RelTraitSet expectedTraitSet = builder.getCluster().traitSet() + .replace(EnumerableConvention.INSTANCE) + .plus(RelDistributions.BROADCAST_DISTRIBUTED); + + assertThat(node.getInput(0).getTraitSet(), is(expectedTraitSet)); } /** Test case for * * [CALCITE-6340] RelBuilder drops set conventions when aggregating over duplicate - * projected fields.. + * projected fields */ @Test void testPruneProjectInputOfAggregatePreservesConvention() { final RelBuilder builder = createBuilder(config -> config.withPruneInputOfAggregate(true)); - final RelNode root = builder.scan("DEPT") + final RelNode node = builder.scan("DEPT") .adoptConvention(EnumerableConvention.INSTANCE) .project(builder.alias(builder.field(0), "a"), builder.alias(builder.field(0), "b")) .aggregate( builder.groupKey(0), builder.aggregateCall( - SqlStdOperatorTable.SUM, builder.field(0))).build(); + SqlStdOperatorTable.SUM, builder.field(0))).build(); - assertTrue(root.getInput(0).getTraitSet().contains(EnumerableConvention.INSTANCE)); + final RelTraitSet expectedTraitSet = builder.getCluster().traitSet() + .replace(EnumerableConvention.INSTANCE) + .replace(RelCollations.of(new RelFieldCollation(0))); + + if (Bug.CALCITE_6270_FIXED) { + assertThat(node.getInput(0).getTraitSet(), is(expectedTraitSet)); + } else { + assertThat(node.getInput(0).getTraitSet().get(0), is(expectedTraitSet.get(0))); + } } private RelNode buildRelWithDuplicateAggregates(