[CALCITE-6460] SortRemoveConstantKeysRule fails with AssertionError d…

…ue to mismatched collation on resulting Sort
asolimando · Jul 11, 2024 · a98508f · a98508f
1 parent dad9073
commit a98508f
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/core/src/main/java/org/apache/calcite/rel/rules/SortRemoveConstantKeysRule.java b/core/src/main/java/org/apache/calcite/rel/rules/SortRemoveConstantKeysRule.java
@@ -19,6 +19,7 @@
 import org.apache.calcite.plan.RelOptPredicateList;
 import org.apache.calcite.plan.RelOptRuleCall;
 import org.apache.calcite.plan.RelRule;
+import org.apache.calcite.rel.RelCollation;
 import org.apache.calcite.rel.RelCollationTraitDef;
 import org.apache.calcite.rel.RelCollations;
 import org.apache.calcite.rel.RelFieldCollation;
@@ -77,8 +78,13 @@ protected SortRemoveConstantKeysRule(Config config) {
       return;
     }
 
+    final RelCollation collation =
+        RelCollationTraitDef.INSTANCE.canonize(RelCollations.of(collationsList));
     final Sort result =
-        sort.copy(sort.getTraitSet(), input, RelCollations.of(collationsList));
+        sort.copy(
+            sort.getTraitSet().replaceIf(RelCollationTraitDef.INSTANCE, () -> collation),
+            input,
+            collation);
     call.transformTo(result);
     call.getPlanner().prune(sort);
   }

diff --git a/core/src/test/java/org/apache/calcite/test/RelOptRulesTest.java b/core/src/test/java/org/apache/calcite/test/RelOptRulesTest.java
@@ -1383,6 +1383,11 @@ private void checkSemiOrAntiJoinProjectTranspose(JoinRelType type) {
         + "group by deptno, sal\n"
         + "order by deptno, sal desc nulls first";
     sql(sql)
+        // Use VolcanoPlanner with collation to ensure the correct creation of the new Sort
+        .withVolcanoPlanner(false,  p -> {
+          p.addRelTraitDef(RelCollationTraitDef.INSTANCE);
+          RelOptUtil.registerDefaultRules(p, false, false);
+        })
         .withRule(CoreRules.SORT_REMOVE_CONSTANT_KEYS)
         .check();
   }