From d5b6b5c01ff17cac100a591a4cc4c9e83d4e1ace Mon Sep 17 00:00:00 2001
From: Ruben Quesada Lopez <rubenql@gmail.com>
Date: Thu, 21 Mar 2024 16:43:37 +0000
Subject: [PATCH] [CALCITE-6338] RelMdCollation#project can return an
 incomplete list of collations in the presence of aliasing

---
 .../calcite/rel/metadata/RelMdCollation.java  | 20 +++++++---
 .../apache/calcite/test/RelMetadataTest.java  | 40 +++++++++++++++++++
 core/src/test/resources/sql/sub-query.iq      |  5 +--
 3 files changed, 57 insertions(+), 8 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
index 7c107d39d654..9c728384f4a5 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
@@ -316,22 +316,32 @@ public static List<RelCollation> sort(RelCollation collation) {
         targetsWithMonotonicity.put(project.i, call.getOperator().getMonotonicity(binding));
       }
     }
-    final List<RelFieldCollation> fieldCollations = new ArrayList<>();
+    List<List<RelFieldCollation>> fieldCollationsList = new ArrayList<>();
   loop:
     for (RelCollation ic : inputCollations) {
       if (ic.getFieldCollations().isEmpty()) {
         continue;
       }
-      fieldCollations.clear();
+      fieldCollationsList.clear();
+      fieldCollationsList.add(new ArrayList<>());
       for (RelFieldCollation ifc : ic.getFieldCollations()) {
         final Collection<Integer> integers = targets.get(ifc.getFieldIndex());
         if (integers.isEmpty()) {
           continue loop; // cannot do this collation
         }
-        fieldCollations.add(ifc.withFieldIndex(integers.iterator().next()));
+        fieldCollationsList = fieldCollationsList.stream()
+            .flatMap(fieldCollations -> integers.stream()
+                .map(integer -> {
+                  List<RelFieldCollation> newFieldCollations = new ArrayList<>(fieldCollations);
+                  newFieldCollations.add(ifc.withFieldIndex(integer));
+                  return newFieldCollations;
+                })).collect(Collectors.toList());
+      }
+      assert !fieldCollationsList.isEmpty();
+      for (List<RelFieldCollation> fieldCollations : fieldCollationsList) {
+        assert !fieldCollations.isEmpty();
+        collations.add(RelCollations.of(fieldCollations));
       }
-      assert !fieldCollations.isEmpty();
-      collations.add(RelCollations.of(fieldCollations));
     }
 
     final List<RelFieldCollation> fieldCollationsForRexCalls =
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 db1f65c242ea..d9733486bd6a 100644
--- a/core/src/test/java/org/apache/calcite/test/RelMetadataTest.java
+++ b/core/src/test/java/org/apache/calcite/test/RelMetadataTest.java
@@ -1948,6 +1948,46 @@ public String colType(MyRelMetadataQuery myRelMetadataQuery, RelNode rel, int co
     assertEquals("[[0]]", collations.toString());
   }
 
+  /**
+   * Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-6338">[CALCITE-6338]
+   * RelMdCollation#project can return an incomplete list of collations
+   * in the presence of aliasing</a>.
+   */
+  @Test void testCollationProjectAliasing() {
+    final RelBuilder builder = RelBuilderTest.createBuilder();
+    final RelNode relNode1 = builder
+        .scan("EMP")
+        .sort(2, 3)
+        .project(builder.field(0), builder.field(2), builder.field(2), builder.field(3))
+        .build();
+    checkCollationProjectAliasing(relNode1, "[[1, 3], [2, 3]]");
+    final RelNode relNode2 = builder
+        .scan("EMP")
+        .sort(0, 1)
+        .project(builder.field(0), builder.field(0), builder.field(1), builder.field(1))
+        .build();
+    checkCollationProjectAliasing(relNode2, "[[0, 2], [0, 3], [1, 2], [1, 3]]");
+    final RelNode relNode3 = builder
+        .scan("EMP")
+        .sort(0, 1, 2)
+        .project(
+            builder.field(0), builder.field(0),
+            builder.field(1), builder.field(1), builder.field(1),
+            builder.field(2))
+        .build();
+    checkCollationProjectAliasing(relNode3,
+        "[[0, 2, 5], [0, 3, 5], [0, 4, 5], [1, 2, 5], [1, 3, 5], [1, 4, 5]]");
+  }
+
+  private void checkCollationProjectAliasing(RelNode relNode, String expectedCollation) {
+    final RelMetadataQuery mq = relNode.getCluster().getMetadataQuery();
+    assertThat(relNode, instanceOf(Project.class));
+    final ImmutableList<RelCollation> collations = mq.collations(relNode);
+    assertThat(collations, notNullValue());
+    assertEquals(expectedCollation, 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/sql/sub-query.iq b/core/src/test/resources/sql/sub-query.iq
index 4d76b3fb6b9e..16cf991f137d 100644
--- a/core/src/test/resources/sql/sub-query.iq
+++ b/core/src/test/resources/sql/sub-query.iq
@@ -2260,9 +2260,8 @@ EnumerableCalc(expr#0..5=[{inputs}], expr#6=[<>($t2, $t1)], expr#7=[1], expr#8=[
   EnumerableMergeJoin(condition=[=($0, $5)], joinType=[left])
     EnumerableCalc(expr#0..7=[{inputs}], EMPNO=[$t0])
       EnumerableTableScan(table=[[scott, EMP]])
-    EnumerableSort(sort0=[$4], dir0=[ASC])
-      EnumerableCalc(expr#0..7=[{inputs}], expr#8=[1:BIGINT], expr#9=[true], c=[$t8], d=[$t8], m=[$t0], trueLiteral=[$t9], EMPNO1=[$t0])
-        EnumerableTableScan(table=[[scott, EMP]])
+    EnumerableCalc(expr#0..7=[{inputs}], expr#8=[1:BIGINT], expr#9=[true], c=[$t8], d=[$t8], m=[$t0], trueLiteral=[$t9], EMPNO1=[$t0])
+      EnumerableTableScan(table=[[scott, EMP]])
 !plan
 +-------+
 | EMPNO |