UNION ALLs in MSQ

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/DataSourcePlan.java

@@ -51,6 +51,7 @@
 import org.apache.druid.query.QueryContext;
 import org.apache.druid.query.QueryDataSource;
 import org.apache.druid.query.TableDataSource;
+import org.apache.druid.query.UnionDataSource;
 import org.apache.druid.query.UnnestDataSource;
 import org.apache.druid.query.filter.DimFilter;
 import org.apache.druid.query.planning.DataSourceAnalysis;
@@ -170,6 +171,18 @@ public static DataSourcePlan forDataSource(
           minStageNumber,
           broadcast
       );
+    } else if (dataSource instanceof UnionDataSource) {
+      return forUnion(
+          queryKit,
+          queryId,
+          queryContext,
+          (UnionDataSource) dataSource,
+          querySegmentSpec,
+          filter,
+          maxWorkerCount,
+          minStageNumber,
+          broadcast
+      );
     } else if (dataSource instanceof JoinDataSource) {
       final JoinAlgorithm preferredJoinAlgorithm = PlannerContext.getJoinAlgorithm(queryContext);
       final JoinAlgorithm deducedJoinAlgorithm = deduceJoinAlgorithm(
@@ -458,6 +471,54 @@ private static DataSourcePlan forUnnest(
     );
   }
 
+  private static DataSourcePlan forUnion(
+      final QueryKit queryKit,
+      final String queryId,
+      final QueryContext queryContext,
+      final UnionDataSource unionDataSource,
+      final QuerySegmentSpec querySegmentSpec,
+      @Nullable DimFilter filter,
+      final int maxWorkerCount,
+      final int minStageNumber,
+      final boolean broadcast
+  )
+  {
+    // This is done to prevent loss of generality since MSQ can plan any type of DataSource.
+    List<DataSource> children = unionDataSource.getDataSources();
+
+    final QueryDefinitionBuilder subqueryDefBuilder = QueryDefinition.builder();
+    final List<DataSource> newChildren = new ArrayList<>();
+    final List<InputSpec> inputSpecs = new ArrayList<>();
+    final IntSet broadcastInputs = new IntOpenHashSet();
+
+    for (DataSource child : children) {
+      DataSourcePlan childDataSourcePlan = forDataSource(
+          queryKit,
+          queryId,
+          queryContext,
+          child,
+          querySegmentSpec,
+          filter,
+          maxWorkerCount,
+          Math.max(minStageNumber, subqueryDefBuilder.getNextStageNumber()),
+          broadcast
+      );
+
+      int shift = inputSpecs.size();
+
+      newChildren.add(shiftInputNumbers(childDataSourcePlan.getNewDataSource(), shift));
+      inputSpecs.addAll(childDataSourcePlan.getInputSpecs());
+      childDataSourcePlan.getSubQueryDefBuilder().ifPresent(subqueryDefBuilder::addAll);
+      childDataSourcePlan.getBroadcastInputs().forEach(inp -> broadcastInputs.add(inp + shift));
+    }
+    return new DataSourcePlan(
+        new UnionDataSource(newChildren),
+        inputSpecs,
+        broadcastInputs,
+        subqueryDefBuilder
+    );
+  }
+
   /**
    * Build a plan for broadcast hash-join.
    */

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/MSQTaskSqlEngine.java

@@ -114,6 +114,7 @@ public boolean featureAvailable(EngineFeature feature, PlannerContext plannerCon
       case TIME_BOUNDARY_QUERY:
       case GROUPING_SETS:
       case WINDOW_FUNCTIONS:
+      case ALLOW_TOP_LEVEL_UNION_ALL:
         return false;
       case UNNEST:
       case CAN_SELECT:

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

@@ -20,13 +20,16 @@
 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;
 import org.apache.druid.msq.indexing.error.InsertCannotAllocateSegmentFault;
 import org.apache.druid.msq.indexing.error.InsertCannotBeEmptyFault;
 import org.apache.druid.msq.indexing.error.InsertTimeNullFault;
 import org.apache.druid.msq.indexing.error.InsertTimeOutOfBoundsFault;
+import org.apache.druid.msq.indexing.error.QueryNotSupportedFault;
 import org.apache.druid.msq.indexing.error.TooManyClusteredByColumnsFault;
 import org.apache.druid.msq.indexing.error.TooManyColumnsFault;
 import org.apache.druid.msq.indexing.error.TooManyInputFilesFault;
@@ -356,4 +359,49 @@ public void testTooManyInputFiles() throws IOException
         .setExpectedMSQFault(new TooManyInputFilesFault(numFiles, Limits.MAX_INPUT_FILES_PER_WORKER, 2))
         .verifyResults();
   }
+
+  @Test
+  public void testUnionAllUsingUnionDataSourceDisallowed()
+  {
+    final RowSignature rowSignature =
+        RowSignature.builder().add("__time", ColumnType.LONG).build();
+    // 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)
+        .setExpectedMSQFault(QueryNotSupportedFault.instance())
+        .verifyResults();
+  }
+
+  @Test
+  public void testUnionAllUsingTopLevelUnionDisallowedWhilePlanning()
+  {
+    // 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 "
+            + "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,5 @@ public void testArrayAggQueryOnComplexDatatypes()
       );
     }
   }
+
 }

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

@@ -0,0 +1,177 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.msq.test;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.ImmutableList;
+import com.google.inject.Injector;
+import com.google.inject.Module;
+import org.apache.druid.common.config.NullHandling;
+import org.apache.druid.guice.DruidInjectorBuilder;
+import org.apache.druid.java.util.common.granularity.Granularities;
+import org.apache.druid.msq.exec.WorkerMemoryParameters;
+import org.apache.druid.msq.sql.MSQTaskSqlEngine;
+import org.apache.druid.query.QueryDataSource;
+import org.apache.druid.query.TableDataSource;
+import org.apache.druid.query.UnionDataSource;
+import org.apache.druid.query.aggregation.CountAggregatorFactory;
+import org.apache.druid.query.aggregation.LongSumAggregatorFactory;
+import org.apache.druid.query.dimension.DefaultDimensionSpec;
+import org.apache.druid.query.groupby.GroupByQuery;
+import org.apache.druid.query.groupby.TestGroupByBuffers;
+import org.apache.druid.server.QueryLifecycleFactory;
+import org.apache.druid.sql.calcite.BaseCalciteQueryTest;
+import org.apache.druid.sql.calcite.CalciteUnionQueryTest;
+import org.apache.druid.sql.calcite.QueryTestBuilder;
+import org.apache.druid.sql.calcite.filtration.Filtration;
+import org.apache.druid.sql.calcite.run.SqlEngine;
+import org.apache.druid.sql.calcite.util.CalciteTests;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+/**
+ * Runs {@link CalciteUnionQueryTest} but with MSQ engine
+ */
+public class CalciteUnionQueryMSQTest extends CalciteUnionQueryTest
+{
+  private TestGroupByBuffers groupByBuffers;
+
+  @Before
+  public void setup2()
+  {
+    groupByBuffers = TestGroupByBuffers.createDefault();
+  }
+
+  @After
+  public void teardown2()
+  {
+    groupByBuffers.close();
+  }
+
+  @Override
+  public void configureGuice(DruidInjectorBuilder builder)
+  {
+    super.configureGuice(builder);
+    builder.addModules(CalciteMSQTestsHelper.fetchModules(temporaryFolder, groupByBuffers).toArray(new Module[0]));
+  }
+
+
+  @Override
+  public SqlEngine createEngine(
+      QueryLifecycleFactory qlf,
+      ObjectMapper queryJsonMapper,
+      Injector injector
+  )
+  {
+    final WorkerMemoryParameters workerMemoryParameters =
+        WorkerMemoryParameters.createInstance(
+            WorkerMemoryParameters.PROCESSING_MINIMUM_BYTES * 50,
+            2,
+            10,
+            2,
+            0,
+            0
+        );
+    final MSQTestOverlordServiceClient indexingServiceClient = new MSQTestOverlordServiceClient(
+        queryJsonMapper,
+        injector,
+        new MSQTestTaskActionClient(queryJsonMapper),
+        workerMemoryParameters
+    );
+    return new MSQTaskSqlEngine(indexingServiceClient, queryJsonMapper);
+  }
+
+  @Override
+  protected QueryTestBuilder testBuilder()
+  {
+    return new QueryTestBuilder(new BaseCalciteQueryTest.CalciteTestConfig(true))
+        .addCustomRunner(new ExtractResultsFactory(() -> (MSQTestOverlordServiceClient) ((MSQTaskSqlEngine) queryFramework().engine()).overlordClient()))
+        .skipVectorize(true)
+        .verifyNativeQueries(new VerifyMSQSupportedNativeQueriesPredicate())
+        .msqCompatible(msqCompatible);
+  }
+
+  /**
+   * 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()
+  {
+
+  }
+
+  @Test
+  public void testUnionOnSubqueries()
+  {
+    testQuery(
+        "SELECT\n"
+        + "  SUM(cnt),\n"
+        + "  COUNT(*)\n"
+        + "FROM (\n"
+        + "  (SELECT dim2, SUM(cnt) AS cnt FROM druid.foo GROUP BY dim2)\n"
+        + "  UNION ALL\n"
+        + "  (SELECT dim2, SUM(cnt) AS cnt FROM druid.foo GROUP BY dim2)\n"
+        + ")",
+        ImmutableList.of(
+            GroupByQuery.builder()
+                        .setDataSource(
+                            new QueryDataSource(
+                                GroupByQuery.builder()
+                                            .setDataSource(
+                                                new UnionDataSource(
+                                                    ImmutableList.of(
+                                                        new TableDataSource(CalciteTests.DATASOURCE1),
+                                                        new TableDataSource(CalciteTests.DATASOURCE1)
+                                                    )
+                                                )
+                                            )
+                                            .setInterval(querySegmentSpec(Filtration.eternity()))
+                                            .setGranularity(Granularities.ALL)
+                                            .setDimensions(dimensions(new DefaultDimensionSpec("dim2", "d0")))
+                                            .setAggregatorSpecs(aggregators(new LongSumAggregatorFactory("a0", "cnt")))
+                                            .setContext(QUERY_CONTEXT_DEFAULT)
+                                            .build()
+                            )
+                        )
+                        .setInterval(querySegmentSpec(Filtration.eternity()))
+                        .setGranularity(Granularities.ALL)
+                        .setAggregatorSpecs(aggregators(
+                            new LongSumAggregatorFactory("_a0", "a0"),
+                            new CountAggregatorFactory("_a1")
+                        ))
+                        .setContext(QUERY_CONTEXT_DEFAULT)
+                        .build()
+        ),
+        NullHandling.replaceWithDefault() ?
+        ImmutableList.of(
+            new Object[]{12L, 3L}
+        ) :
+        ImmutableList.of(
+            new Object[]{12L, 4L}
+        )
+    );
+  }
+
+}

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

@@ -1390,9 +1390,11 @@ public Pair<MSQSpec, Pair<List<MSQResultsReport.ColumnAndType>, List<Object[]>>>
 
     public void verifyResults()
     {
-      Preconditions.checkArgument(expectedResultRows != null, "Result rows cannot be null");
-      Preconditions.checkArgument(expectedRowSignature != null, "Row signature cannot be null");
-      Preconditions.checkArgument(expectedMSQSpec != null, "MultiStageQuery Query spec cannot be null ");
+      if (expectedMSQFault == null) {
+        Preconditions.checkArgument(expectedResultRows != null, "Result rows cannot be null");
+        Preconditions.checkArgument(expectedRowSignature != null, "Row signature cannot be null");
+        Preconditions.checkArgument(expectedMSQSpec != null, "MultiStageQuery Query spec cannot be null ");
+      }
       Pair<MSQSpec, Pair<List<MSQResultsReport.ColumnAndType>, List<Object[]>>> specAndResults = runQueryWithResult();
 
       if (specAndResults == null) { // A fault was expected and the assertion has been done in the runQueryWithResult