From e3b03aa47aafa8d9b97102f71532652c658c56c3 Mon Sep 17 00:00:00 2001
From: rubenada <rubenql@gmail.com>
Date: Mon, 7 Aug 2023 15:14:32 +0100
Subject: [PATCH] [CALCITE-5903] RelMdCollation does not define collations for
 EnumerableLimit

---
 .../calcite/rel/metadata/RelMdCollation.java  |  6 ++++
 .../apache/calcite/test/RelMetadataTest.java  | 36 +++++++++++++++++++
 .../GeneratedMetadata_CollationHandler.java   |  2 ++
 3 files changed, 44 insertions(+)

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
index 815968d03102..7c107d39d654 100644
--- 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
index 19c5ebc0e638..57a58a39836c 100644
--- 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/core/src/test/resources/org/apache/calcite/rel/metadata/janino/GeneratedMetadata_CollationHandler.java b/core/src/test/resources/org/apache/calcite/rel/metadata/janino/GeneratedMetadata_CollationHandler.java
index 0dec36c46aac..912ce05dfe30 100644
--- a/core/src/test/resources/org/apache/calcite/rel/metadata/janino/GeneratedMetadata_CollationHandler.java
+++ b/core/src/test/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) {