add CoreRules.UNION_TO_DISTINCT

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQFaultsTest.java

@@ -20,6 +20,8 @@
 package org.apache.druid.msq.exec;
 
 import com.google.common.collect.ImmutableMap;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.error.DruidExceptionMatcher;
 import org.apache.druid.indexing.common.actions.SegmentAllocateAction;
 import org.apache.druid.java.util.common.Intervals;
 import org.apache.druid.java.util.common.StringUtils;
@@ -372,4 +374,29 @@ public void testUnionAllIsDisallowed()
         .setExpectedMSQFault(QueryNotSupportedFault.instance())
         .verifyResults();
   }
+
+  @Test
+  public void testUnionAllIsDisallowedWhilePlanning()
+  {
+    // This results in a planning error however the planning error isn't an accurate representation of the actual error
+    // because Calcite rewrites it using CoreRules.UNION_TO_DISTINCT, which plans it using Union Datasource.
+    // However, this fails since the column names mismatch. Once MSQ is able to support Union datasources, the planning
+    // error would become an accurate representation of the error.
+    testIngestQuery()
+        .setSql(
+            "INSERT INTO druid.dst "
+            + "SELECT dim2, dim1, m1 FROM foo2 "
+            + "UNION ALL "
+            + "SELECT dim1, dim2, m1 FROM foo "
+            + "PARTITIONED BY ALL TIME")
+        .setExpectedValidationErrorMatcher(
+            new DruidExceptionMatcher(
+                DruidException.Persona.ADMIN,
+                DruidException.Category.INVALID_INPUT,
+                "general"
+            ).expectMessageIs("Query planning failed for unknown reason, our best guess is this "
+                              + "[SQL requires union between two tables and column names queried for each table are different "
+                              + "Left: [dim2, dim1, m1], Right: [dim1, dim2, m1].]"))
+        .verifyPlanningErrors();
+  }
 }

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/CalciteSelectQueryMSQTest.java

@@ -174,4 +174,17 @@ public void testArrayAggQueryOnComplexDatatypes()
       );
     }
   }
+
+  /**
+   * Doesn't pass through Druid however the planning error is different as it rewrites to a union datasource.
+   * This test is disabled because MSQ wants to support union datasources, and it makes little sense to add highly
+   * conditional planning error for the same. Planning errors are merely hints, and this is one of those times
+   * when the hint is incorrect till MSQ starts supporting the union datasource.
+   */
+  @Test
+  @Override
+  public void testUnionIsUnplannable()
+  {
+
+  }
 }

sql/src/main/java/org/apache/druid/sql/calcite/planner/CalciteRulesManager.java

@@ -177,6 +177,7 @@ public class CalciteRulesManager
       ImmutableList.of(
           AbstractConverter.ExpandConversionRule.INSTANCE,
           CoreRules.AGGREGATE_REMOVE,
+          CoreRules.UNION_TO_DISTINCT,
           CoreRules.PROJECT_REMOVE,
           CoreRules.AGGREGATE_JOIN_TRANSPOSE,
           CoreRules.AGGREGATE_PROJECT_MERGE,