review comments

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

@@ -361,27 +361,32 @@ public void testTooManyInputFiles() throws IOException
   }
 
   @Test
-  public void testUnionAllIsDisallowed()
+  public void testUnionAllUsingUnionDataSourceDisallowed()
   {
     final RowSignature rowSignature =
         RowSignature.builder().add("__time", ColumnType.LONG).build();
-    testIngestQuery()
+    // This plans the query using DruidUnionDataSourceRule since the DruidUnionDataSourceRule#isCompatible
+    // returns true (column names, types match, and it is a union on the table data sources).
+    // It gets planned correctly, however MSQ engine cannot plan the query correctly
+    testSelectQuery()
         .setSql("SELECT * FROM foo\n"
                 + "UNION ALL\n"
                 + "SELECT * FROM foo\n")
         .setExpectedRowSignature(rowSignature)
-        .setExpectedDataSource("foo1")
         .setExpectedMSQFault(QueryNotSupportedFault.instance())
         .verifyResults();
   }
 
   @Test
-  public void testUnionAllIsDisallowedWhilePlanning()
+  public void testUnionAllUsingTopLevelUnionDisallowedWhilePlanning()
   {
-    // 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.
+    // This results in a planning error however the planning error isn't an accurate representation of the actual error.
+    // Calcite tries to plan the query using DruidUnionRule, which passes with native, however fails with MSQ (due to engine
+    // feature ALLOW_TOP_LEVEL_UNION_ALL being absent)
+    // Calcite then tries to write it using DruidUnionDataSourceRule. However, the isCompatible check fails because column
+    // names mismatch. But it sets the planning error with this mismatch, which can be misleading since native queries can
+    // plan fine using the DruidUnionRule (that's disabled in MSQ)
+    // Once MSQ is able to support union datasources, we'd be able to execute this query fine in MSQ
     testIngestQuery()
         .setSql(
             "INSERT INTO druid.dst "

sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidRules.java

@@ -95,13 +95,16 @@ public static List<RelOptRule> rules(PlannerContext plannerContext)
             DruidOuterQueryRule.WHERE_FILTER,
             DruidOuterQueryRule.SELECT_PROJECT,
             DruidOuterQueryRule.SORT,
-            new DruidUnionRule(plannerContext),
             new DruidUnionDataSourceRule(plannerContext),
-            new DruidSortUnionRule(plannerContext),
             DruidJoinRule.instance(plannerContext)
         )
     );
 
+    if (plannerContext.featureAvailable(EngineFeature.ALLOW_TOP_LEVEL_UNION_ALL)) {
+      retVal.add(new DruidUnionRule(plannerContext));
+      retVal.add(DruidSortUnionRule.instance());
+    }
+
     if (plannerContext.featureAvailable(EngineFeature.WINDOW_FUNCTIONS)) {
       retVal.add(new DruidQueryRule<>(Window.class, PartialDruidQuery.Stage.WINDOW, PartialDruidQuery::withWindow));
       retVal.add(

sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidSortUnionRule.java

@@ -23,9 +23,7 @@
 import org.apache.calcite.plan.RelOptRuleCall;
 import org.apache.calcite.rel.core.Sort;
 import org.apache.calcite.rex.RexLiteral;
-import org.apache.druid.sql.calcite.planner.PlannerContext;
 import org.apache.druid.sql.calcite.rel.DruidUnionRel;
-import org.apache.druid.sql.calcite.run.EngineFeature;
 
 import java.util.Collections;
 
@@ -34,27 +32,22 @@
  */
 public class DruidSortUnionRule extends RelOptRule
 {
-  private final PlannerContext plannerContext;
 
-  public DruidSortUnionRule(PlannerContext plannerContext)
+  private static final DruidSortUnionRule INSTANCE = new DruidSortUnionRule();
+
+  private DruidSortUnionRule()
   {
     super(operand(Sort.class, operand(DruidUnionRel.class, any())));
-    this.plannerContext = plannerContext;
+  }
+
+  public static DruidSortUnionRule instance()
+  {
+    return INSTANCE;
   }
 
   @Override
   public boolean matches(final RelOptRuleCall call)
   {
-    // Defensive check. If the planner disallows top level union all, then the DruidUnionRule would have prevented
-    // creating the DruidUnionRel in the first place
-    if (!plannerContext.featureAvailable(EngineFeature.ALLOW_TOP_LEVEL_UNION_ALL)) {
-      plannerContext.setPlanningError(
-          "Top level 'UNION ALL' is unsupported by SQL engine [%s].",
-          plannerContext.getEngine().name()
-      );
-      return false;
-    }
-
     // LIMIT, no ORDER BY
     final Sort sort = call.rel(0);
     return sort.collation.getFieldCollations().isEmpty() && sort.fetch != null;

sql/src/main/java/org/apache/druid/sql/calcite/rule/DruidUnionRule.java

@@ -26,7 +26,6 @@
 import org.apache.druid.sql.calcite.planner.PlannerContext;
 import org.apache.druid.sql.calcite.rel.DruidRel;
 import org.apache.druid.sql.calcite.rel.DruidUnionRel;
-import org.apache.druid.sql.calcite.run.EngineFeature;
 
 import java.util.List;
 
@@ -52,13 +51,6 @@ public DruidUnionRule(PlannerContext plannerContext)
   @Override
   public boolean matches(RelOptRuleCall call)
   {
-    if (!plannerContext.featureAvailable(EngineFeature.ALLOW_TOP_LEVEL_UNION_ALL)) {
-      plannerContext.setPlanningError(
-          "Top level 'UNION ALL' is unsupported by SQL engine [%s].",
-          plannerContext.getEngine().name()
-      );
-      return false;
-    }
     // Make DruidUnionRule and DruidUnionDataSourceRule mutually exclusive.
     final Union unionRel = call.rel(0);
     final DruidRel<?> firstDruidRel = call.rel(1);