Support Union in Decoupled planning

benchmarks/src/test/java/org/apache/druid/benchmark/query/CachingClusteredClientBenchmark.java

@@ -258,39 +258,37 @@ public int getNumThreads()
       }
     };
 
-    conglomerate = new DefaultQueryRunnerFactoryConglomerate(
-        ImmutableMap.<Class<? extends Query>, QueryRunnerFactory>builder()
-            .put(
-                TimeseriesQuery.class,
-                new TimeseriesQueryRunnerFactory(
-                    new TimeseriesQueryQueryToolChest(),
-                    new TimeseriesQueryEngine(),
-                    QueryRunnerTestHelper.NOOP_QUERYWATCHER
-                )
+    conglomerate = DefaultQueryRunnerFactoryConglomerate.buildFromQueryRunnerFactories(ImmutableMap.<Class<? extends Query>, QueryRunnerFactory>builder()
+        .put(
+            TimeseriesQuery.class,
+            new TimeseriesQueryRunnerFactory(
+                new TimeseriesQueryQueryToolChest(),
+                new TimeseriesQueryEngine(),
+                QueryRunnerTestHelper.NOOP_QUERYWATCHER
             )
-            .put(
-                TopNQuery.class,
-                new TopNQueryRunnerFactory(
-                    new StupidPool<>(
-                        "TopNQueryRunnerFactory-bufferPool",
-                        () -> ByteBuffer.allocate(PROCESSING_BUFFER_SIZE)
-                    ),
-                    new TopNQueryQueryToolChest(new TopNQueryConfig()),
-                    QueryRunnerTestHelper.NOOP_QUERYWATCHER
-                )
+        )
+        .put(
+            TopNQuery.class,
+            new TopNQueryRunnerFactory(
+                new StupidPool<>(
+                    "TopNQueryRunnerFactory-bufferPool",
+                    () -> ByteBuffer.allocate(PROCESSING_BUFFER_SIZE)
+                ),
+                new TopNQueryQueryToolChest(new TopNQueryConfig()),
+                QueryRunnerTestHelper.NOOP_QUERYWATCHER
             )
-            .put(
-                GroupByQuery.class,
-                makeGroupByQueryRunnerFactory(
-                    GroupByQueryRunnerTest.DEFAULT_MAPPER,
-                    new GroupByQueryConfig()
-                    {
-                    },
-                    processingConfig
-                )
+        )
+        .put(
+            GroupByQuery.class,
+            makeGroupByQueryRunnerFactory(
+                GroupByQueryRunnerTest.DEFAULT_MAPPER,
+                new GroupByQueryConfig()
+                {
+                },
+                processingConfig
             )
-            .build()
-    );
+        )
+        .build());
 
     toolChestWarehouse = new QueryToolChestWarehouse()
     {

extensions-contrib/moving-average-query/src/test/java/org/apache/druid/query/movingaverage/MovingAverageQueryTest.java

@@ -56,8 +56,8 @@
 import org.apache.druid.query.Query;
 import org.apache.druid.query.QueryPlus;
 import org.apache.druid.query.QueryRunner;
+import org.apache.druid.query.QueryRunnerFactoryConglomerate;
 import org.apache.druid.query.QuerySegmentWalker;
-import org.apache.druid.query.QueryToolChestWarehouse;
 import org.apache.druid.query.Result;
 import org.apache.druid.query.RetryQueryRunnerConfig;
 import org.apache.druid.query.SegmentDescriptor;
@@ -104,7 +104,7 @@
 public class MovingAverageQueryTest extends InitializedNullHandlingTest
 {
   private final ObjectMapper jsonMapper;
-  private final QueryToolChestWarehouse warehouse;
+  private final QueryRunnerFactoryConglomerate conglomerate;
   private final RetryQueryRunnerConfig retryConfig;
   private final ServerConfig serverConfig;
 
@@ -167,7 +167,7 @@ public <T> QueryRunner<T> getQueryRunnerForSegments(Query<T> query, Iterable<Seg
     Injector injector = Initialization.makeInjectorWithModules(baseInjector, modules);
 
     jsonMapper = injector.getInstance(ObjectMapper.class);
-    warehouse = injector.getInstance(QueryToolChestWarehouse.class);
+    conglomerate = injector.getInstance(QueryRunnerFactoryConglomerate.class);
     retryConfig = injector.getInstance(RetryQueryRunnerConfig.class);
     serverConfig = injector.getInstance(ServerConfig.class);
 
@@ -321,7 +321,7 @@ public long getMaxQueuedBytes()
     };
 
     CachingClusteredClient baseClient = new CachingClusteredClient(
-        warehouse,
+        conglomerate,
         new TimelineServerView()
         {
           @Override
@@ -375,7 +375,7 @@ public void registerServerRemovedCallback(Executor exec, ServerRemovedCallback c
         new NoopServiceEmitter(),
         baseClient,
         null /* local client; unused in this test, so pass in null */,
-        warehouse,
+        conglomerate,
         new MapJoinableFactory(ImmutableSet.of(), ImmutableMap.of()),
         retryConfig,
         jsonMapper,
@@ -392,7 +392,7 @@ public void registerServerRemovedCallback(Executor exec, ServerRemovedCallback c
     final Sequence<?> res = query.getRunner(walker).run(queryPlus);
 
     List actualResults = new ArrayList();
-    actualResults = (List<MapBasedRow>) res.accumulate(actualResults, Accumulators.list());
+    actualResults = res.accumulate(actualResults, Accumulators.list());
 
     expectedResults = consistentTypeCasting(expectedResults);
     actualResults = consistentTypeCasting(actualResults);

extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/indexing/kafka/KafkaIndexTaskTest.java

@@ -2891,28 +2891,26 @@ private static DataSchema cloneDataSchema(final DataSchema dataSchema)
   @Override
   protected QueryRunnerFactoryConglomerate makeQueryRunnerConglomerate()
   {
-    return new DefaultQueryRunnerFactoryConglomerate(
-        ImmutableMap.<Class<? extends Query>, QueryRunnerFactory>builder()
-                    .put(
-                        TimeseriesQuery.class,
-                        new TimeseriesQueryRunnerFactory(
-                            new TimeseriesQueryQueryToolChest(),
-                            new TimeseriesQueryEngine(),
-                            (query, future) -> {
-                              // do nothing
-                            }
-                        )
+    return DefaultQueryRunnerFactoryConglomerate.buildFromQueryRunnerFactories(ImmutableMap.<Class<? extends Query>, QueryRunnerFactory>builder()
+                .put(
+                    TimeseriesQuery.class,
+                    new TimeseriesQueryRunnerFactory(
+                        new TimeseriesQueryQueryToolChest(),
+                        new TimeseriesQueryEngine(),
+                        (query, future) -> {
+                          // do nothing
+                        }
                     )
-                    .put(
-                        ScanQuery.class,
-                        new ScanQueryRunnerFactory(
-                            new ScanQueryQueryToolChest(DefaultGenericQueryMetricsFactory.instance()),
-                            new ScanQueryEngine(),
-                            new ScanQueryConfig()
-                        )
+                )
+                .put(
+                    ScanQuery.class,
+                    new ScanQueryRunnerFactory(
+                        new ScanQueryQueryToolChest(DefaultGenericQueryMetricsFactory.instance()),
+                        new ScanQueryEngine(),
+                        new ScanQueryConfig()
                     )
-                    .build()
-    );
+                )
+                .build());
   }
 
   private void makeToolboxFactory() throws IOException

extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/KinesisIndexTaskTest.java

@@ -2418,18 +2418,16 @@ private static DataSchema cloneDataSchema(final DataSchema dataSchema)
   @Override
   protected QueryRunnerFactoryConglomerate makeQueryRunnerConglomerate()
   {
-    return new DefaultQueryRunnerFactoryConglomerate(
-        ImmutableMap.of(
-            TimeseriesQuery.class,
-            new TimeseriesQueryRunnerFactory(
-                new TimeseriesQueryQueryToolChest(),
-                new TimeseriesQueryEngine(),
-                (query, future) -> {
-                  // do nothing
-                }
-            )
+    return DefaultQueryRunnerFactoryConglomerate.buildFromQueryRunnerFactories(ImmutableMap.of(
+        TimeseriesQuery.class,
+        new TimeseriesQueryRunnerFactory(
+            new TimeseriesQueryQueryToolChest(),
+            new TimeseriesQueryEngine(),
+            (query, future) -> {
+              // do nothing
+            }
         )
-    );
+    ));
   }
 
   private void makeToolboxFactory() throws IOException

processing/src/main/java/org/apache/druid/error/DruidException.java

@@ -524,4 +524,8 @@ public String getErrorCode()
     protected abstract DruidException makeException(DruidExceptionBuilder bob);
   }
 
+  public static DruidException methodNotSupported()
+  {
+    return defensive("Method Not supported.");
+  }
 }

processing/src/main/java/org/apache/druid/frame/allocation/ArenaMemoryAllocatorFactory.java

@@ -24,6 +24,8 @@
  */
 public class ArenaMemoryAllocatorFactory implements MemoryAllocatorFactory
 {
+  private static final int FRAME_SIZE = 8_000_000;
+
   private final int capacity;
 
   public ArenaMemoryAllocatorFactory(final int capacity)
@@ -42,4 +44,9 @@ public long allocatorCapacity()
   {
     return capacity;
   }
+
+  public static MemoryAllocatorFactory makeDefault()
+  {
+    return new ArenaMemoryAllocatorFactory(FRAME_SIZE);
+  }
 }

processing/src/main/java/org/apache/druid/java/util/common/guava/Sequences.java

@@ -22,7 +22,6 @@
 import com.google.common.base.Function;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Predicate;
-
 import java.io.Closeable;
 import java.util.Arrays;
 import java.util.Collections;
@@ -147,4 +146,9 @@ public <OutType> Yielder<OutType> toYielder(OutType initValue, YieldingAccumulat
       return Yielders.done(initValue, null);
     }
   }
+
+  public static <T> Sequence<T> of(T... values)
+  {
+    return simple(Arrays.asList(values));
+  }
 }

processing/src/main/java/org/apache/druid/query/DefaultQueryRunnerFactoryConglomerate.java

@@ -19,29 +19,52 @@
 
 package org.apache.druid.query;
 
+import com.google.common.collect.Maps;
 import com.google.inject.Inject;
-
 import java.util.IdentityHashMap;
 import java.util.Map;
 
-/**
-*/
 public class DefaultQueryRunnerFactoryConglomerate implements QueryRunnerFactoryConglomerate
 {
   private final Map<Class<? extends Query>, QueryRunnerFactory> factories;
+  private final Map<Class<? extends Query>, QueryToolChest> toolchests;
+
+  public static DefaultQueryRunnerFactoryConglomerate buildFromQueryRunnerFactories(
+      Map<Class<? extends Query>, QueryRunnerFactory> factories)
+  {
+    return new DefaultQueryRunnerFactoryConglomerate(factories, Maps.transformValues(factories, f -> f.getToolchest()));
+  }
 
   @Inject
-  public DefaultQueryRunnerFactoryConglomerate(Map<Class<? extends Query>, QueryRunnerFactory> factories)
+  public DefaultQueryRunnerFactoryConglomerate(Map<Class<? extends Query>, QueryRunnerFactory> factories,
+      Map<Class<? extends Query>, QueryToolChest> toolchests)
   {
-    // Accesses to IdentityHashMap should be faster than to HashMap or ImmutableMap.
-    // Class doesn't override Object.equals().
     this.factories = new IdentityHashMap<>(factories);
+    this.toolchests = new IdentityHashMap<>(toolchests);
   }
 
   @Override
   @SuppressWarnings("unchecked")
   public <T, QueryType extends Query<T>> QueryRunnerFactory<T, QueryType> findFactory(QueryType query)
   {
-    return (QueryRunnerFactory<T, QueryType>) factories.get(query.getClass());
+    return factories.get(query.getClass());
+  }
+
+  @Override
+  @SuppressWarnings("unchecked")
+  public <T, QueryType extends Query<T>> QueryToolChest<T, QueryType> getToolChest(QueryType query)
+  {
+    return toolchests.get(query.getClass());
+  }
+
+  @Override
+  @SuppressWarnings("unchecked")
+  public <T, QueryType extends Query<T>> QueryLogic<T> getQueryExecutor(QueryType query)
+  {
+    QueryToolChest<T, QueryType> toolchest = getToolChest(query);
+    if (toolchest instanceof QueryLogic) {
+      return (QueryLogic<T>) toolchest;
+    }
+    return null;
   }
 }

processing/src/main/java/org/apache/druid/query/FluentQueryRunner.java

@@ -64,7 +64,7 @@ public FluentQueryRunner<T> applyPostMergeDecoration()
 
   public FluentQueryRunner<T> applyPreMergeDecoration()
   {
-    return from(new UnionQueryRunner<>(toolChest.preMergeQueryDecoration(baseRunner)));
+    return from(new UnionDataSourceQueryRunner<>(toolChest.preMergeQueryDecoration(baseRunner)));
   }
 
   public FluentQueryRunner<T> emitCPUTimeMetric(ServiceEmitter emitter)

processing/src/main/java/org/apache/druid/query/Query.java

@@ -33,15 +33,18 @@
 import org.apache.druid.query.groupby.GroupByQuery;
 import org.apache.druid.query.metadata.metadata.SegmentMetadataQuery;
 import org.apache.druid.query.operator.WindowOperatorQuery;
+import org.apache.druid.query.planning.DataSourceAnalysis;
 import org.apache.druid.query.scan.ScanQuery;
 import org.apache.druid.query.search.SearchQuery;
 import org.apache.druid.query.select.SelectQuery;
 import org.apache.druid.query.spec.QuerySegmentSpec;
 import org.apache.druid.query.timeboundary.TimeBoundaryQuery;
 import org.apache.druid.query.timeseries.TimeseriesQuery;
 import org.apache.druid.query.topn.TopNQuery;
+import org.apache.druid.query.union.UnionQuery;
 import org.apache.druid.segment.Segment;
 import org.apache.druid.segment.VirtualColumns;
+import org.apache.druid.segment.column.RowSignature;
 import org.apache.druid.utils.CollectionUtils;
 import org.joda.time.DateTimeZone;
 import org.joda.time.Duration;
@@ -66,6 +69,7 @@
     @JsonSubTypes.Type(name = Query.TIMESERIES, value = TimeseriesQuery.class),
     @JsonSubTypes.Type(name = Query.TOPN, value = TopNQuery.class),
     @JsonSubTypes.Type(name = Query.WINDOW_OPERATOR, value = WindowOperatorQuery.class),
+    @JsonSubTypes.Type(name = Query.UNION_QUERY, value = UnionQuery.class),
 })
 public interface Query<T>
 {
@@ -79,6 +83,7 @@ public interface Query<T>
   String TIMESERIES = "timeseries";
   String TOPN = "topN";
   String WINDOW_OPERATOR = "windowOperator";
+  String UNION_QUERY = "union";
 
   DataSource getDataSource();
 
@@ -285,4 +290,14 @@ default Interval getSingleInterval()
         )
     );
   }
+
+  default DataSourceAnalysis getDataSourceAnalysis()
+  {
+    return getDataSource().getAnalysis().maybeWithBaseQuery(this);
+  }
+
+  default RowSignature getResultRowSignature(RowSignature.Finalization finalization)
+  {
+    return null;
+  }
 }