[CALCITE-5903] RelMdCollation does not define collations for Enumerab…

…leLimit
rubenada · Aug 7, 2023 · e3b03aa · e3b03aa
1 parent 98f3048
commit e3b03aa
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 0 deletions.
diff --git a/core/src/main/java/org/apache/calcite/rel/metadata/RelMdCollation.java b/core/src/main/java/org/apache/calcite/rel/metadata/RelMdCollation.java
@@ -18,6 +18,7 @@
 
 import org.apache.calcite.adapter.enumerable.EnumerableCorrelate;
 import org.apache.calcite.adapter.enumerable.EnumerableHashJoin;
+import org.apache.calcite.adapter.enumerable.EnumerableLimit;
 import org.apache.calcite.adapter.enumerable.EnumerableMergeJoin;
 import org.apache.calcite.adapter.enumerable.EnumerableMergeUnion;
 import org.apache.calcite.adapter.enumerable.EnumerableNestedLoopJoin;
@@ -198,6 +199,11 @@ private RelMdCollation() {}
             join.getJoinType()));
   }
 
+  public @Nullable ImmutableList<RelCollation> collations(EnumerableLimit rel,
+      RelMetadataQuery mq) {
+    return mq.collations(rel.getInput());
+  }
+
   public @Nullable ImmutableList<RelCollation> collations(Sort sort,
       RelMetadataQuery mq) {
     return copyOf(

diff --git a/core/src/test/java/org/apache/calcite/test/RelMetadataTest.java b/core/src/test/java/org/apache/calcite/test/RelMetadataTest.java
@@ -16,6 +16,7 @@
  */
 package org.apache.calcite.test;
 import org.apache.calcite.adapter.enumerable.EnumerableConvention;
+import org.apache.calcite.adapter.enumerable.EnumerableLimit;
 import org.apache.calcite.adapter.enumerable.EnumerableMergeJoin;
 import org.apache.calcite.adapter.enumerable.EnumerableRules;
 import org.apache.calcite.config.CalciteSystemProperty;
@@ -147,6 +148,7 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.hasToString;
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
@@ -1566,6 +1568,40 @@ public String colType(MyRelMetadataQuery myRelMetadataQuery, RelNode rel, int co
     metadataConfig.applyMetadata(rel.getCluster());
   }
 
+  /**
+   * Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-5903">[CALCITE-5903]
+   * RelMdCollation does not define collations for EnumerableLimit</a>.
+   */
+  @Test void testCollationEnumerableLimit() {
+    final RelNode result = sql("select * from emp order by empno limit 10")
+        .withCluster(cluster -> {
+          final RelOptPlanner planner = new VolcanoPlanner();
+          planner.addRule(CoreRules.PROJECT_TO_CALC);
+          planner.addRule(EnumerableRules.ENUMERABLE_TABLE_SCAN_RULE);
+          planner.addRule(EnumerableRules.ENUMERABLE_CALC_RULE);
+          planner.addRule(EnumerableRules.ENUMERABLE_SORT_RULE);
+          planner.addRule(EnumerableRules.ENUMERABLE_LIMIT_RULE);
+          planner.addRelTraitDef(ConventionTraitDef.INSTANCE);
+          return RelOptCluster.create(planner, cluster.getRexBuilder());
+        })
+        .withRelTransform(rel -> {
+          final RelOptPlanner planner = rel.getCluster().getPlanner();
+          planner.setRoot(rel);
+          final RelTraitSet requiredOutputTraits =
+              rel.getCluster().traitSet().replace(EnumerableConvention.INSTANCE);
+          final RelNode rootRel = planner.changeTraits(rel, requiredOutputTraits);
+          planner.setRoot(rootRel);
+          return planner.findBestExp();
+        }).toRel();
+
+    assertThat(result, instanceOf(EnumerableLimit.class));
+    final RelMetadataQuery mq = result.getCluster().getMetadataQuery();
+    final ImmutableList<RelCollation> collations = mq.collations(result);
+    assertThat(collations, notNullValue());
+    assertEquals("[[0]]", collations.toString());
+  }
+
   /** Unit test for
    * {@link org.apache.calcite.rel.metadata.RelMdCollation#project}
    * and other helper functions for deducing collations. */

diff --git a/.../resources/org/apache/calcite/rel/metadata/janino/GeneratedMetadata_CollationHandler.java b/.../resources/org/apache/calcite/rel/metadata/janino/GeneratedMetadata_CollationHandler.java
@@ -64,6 +64,8 @@ private com.google.common.collect.ImmutableList collations_(
       return provider0.collations((org.apache.calcite.adapter.enumerable.EnumerableCorrelate) r, mq);
     } else if (r instanceof org.apache.calcite.adapter.enumerable.EnumerableHashJoin) {
       return provider0.collations((org.apache.calcite.adapter.enumerable.EnumerableHashJoin) r, mq);
+    } else if (r instanceof org.apache.calcite.adapter.enumerable.EnumerableLimit) {
+      return provider0.collations((org.apache.calcite.adapter.enumerable.EnumerableLimit) r, mq);
     } else if (r instanceof org.apache.calcite.adapter.enumerable.EnumerableMergeJoin) {
       return provider0.collations((org.apache.calcite.adapter.enumerable.EnumerableMergeJoin) r, mq);
     } else if (r instanceof org.apache.calcite.adapter.enumerable.EnumerableMergeUnion) {