Add support for union queries

indexing-service/src/main/java/io/druid/indexing/overlord/ThreadPoolTaskRunner.java

@@ -22,6 +22,7 @@ package io.druid.indexing.overlord;
 import com.google.api.client.repackaged.com.google.common.base.Preconditions;
 import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
 import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.Futures;
@@ -153,13 +154,7 @@ public class ThreadPoolTaskRunner implements TaskRunner, QuerySegmentWalker
   private <T> QueryRunner<T> getQueryRunnerImpl(Query<T> query)
   {
     QueryRunner<T> queryRunner = null;
-    String queryDataSource;
+    final String queryDataSource = Iterables.getOnlyElement(query.getDataSource().getNames());
-    try {
-        queryDataSource = ((TableDataSource)query.getDataSource()).getName();
-    }
-    catch (ClassCastException e) {
-        throw new IllegalArgumentException("Subqueries are not welcome here");
-    }
 
     for (final ThreadPoolTaskRunnerWorkItem taskRunnerWorkItem : ImmutableList.copyOf(runningItems)) {
       final Task task = taskRunnerWorkItem.getTask();

processing/src/main/java/io/druid/query/DataSource.java

@@ -24,15 +24,20 @@ package io.druid.query;
 import com.fasterxml.jackson.annotation.JsonSubTypes;
 import com.fasterxml.jackson.annotation.JsonTypeInfo;
 
+import java.util.List;
+
 @JsonTypeInfo(use = JsonTypeInfo.Id.NAME,
               include = JsonTypeInfo.As.PROPERTY,
               property = "type",
               defaultImpl = LegacyDataSource.class)
 @JsonSubTypes({
                   @JsonSubTypes.Type(value = TableDataSource.class, name = "table"),
-                  @JsonSubTypes.Type(value = QueryDataSource.class, name = "query")
+                  @JsonSubTypes.Type(value = QueryDataSource.class, name = "query"),
+                  @JsonSubTypes.Type(value = UnionDataSource.class, name = "union")
               })
 public interface DataSource
 {
-  public String getName();
+  public Iterable<String> getNames();
+
+  public String getMetricName();
 }

processing/src/main/java/io/druid/query/Druids.java

@@ -537,7 +537,7 @@ public class Druids
     public SearchQueryBuilder copy(SearchQuery query)
     {
       return new SearchQueryBuilder()
-          .dataSource(((TableDataSource)query.getDataSource()).getName())
+          .dataSource(query.getDataSource())
           .intervals(query.getQuerySegmentSpec())
           .filters(query.getDimensionsFilter())
           .granularity(query.getGranularity())

processing/src/main/java/io/druid/query/Query.java

@@ -88,4 +88,6 @@ public interface Query<T>
   public Query<T> withId(String id);
 
   public String getId();
+
+  Query<T> withDataSource(DataSource dataSource);
 }

processing/src/main/java/io/druid/query/QueryDataSource.java

@@ -24,6 +24,9 @@ package io.druid.query;
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.annotation.JsonTypeName;
+import com.google.common.collect.Lists;
+
+import java.util.List;
 
 @JsonTypeName("query")
 public class QueryDataSource implements DataSource
@@ -38,9 +41,15 @@ public class QueryDataSource implements DataSource
   }
 
   @Override
-  public String getName()
+  public Iterable<String> getNames()
   {
-    return query.getDataSource().getName();
+    return query.getDataSource().getNames();
+  }
+
+  @Override
+  public String getMetricName()
+  {
+    return query.getDataSource().getMetricName();
   }
 
   @JsonProperty

processing/src/main/java/io/druid/query/TableDataSource.java

@@ -23,6 +23,9 @@ package io.druid.query;
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.annotation.JsonTypeName;
+import com.google.common.collect.Lists;
+
+import java.util.List;
 
 @JsonTypeName("table")
 public class TableDataSource implements DataSource
@@ -37,8 +40,18 @@ public class TableDataSource implements DataSource
   }
 
   @JsonProperty
+  public String getName(){
+    return name;
+  }
+
   @Override
-  public String getName()
+  public List<String> getNames()
+  {
+    return Lists.newArrayList(name);
+  }
+
+  @Override
+  public String getMetricName()
   {
     return name;
   }

processing/src/main/java/io/druid/query/UnionDataSource.java

@@ -0,0 +1,101 @@
+/*
+ * Druid - a distributed column store.
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * as published by the Free Software Foundation; either version 2
+ * of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+ *
+ * This file Copyright (C) 2014 N3TWORK, Inc. and contributed to the Druid project
+ * under the Druid Corporate Contributor License Agreement.
+ */
+
+package io.druid.query;
+
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.base.Function;
+import com.google.common.base.Preconditions;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Lists;
+
+import javax.annotation.Nullable;
+import java.util.List;
+import java.util.SortedSet;
+import java.util.TreeSet;
+
+public class UnionDataSource implements DataSource
+{
+  @JsonProperty
+  private final List<DataSource> dataSources;
+
+  @JsonCreator
+  public UnionDataSource(@JsonProperty("dataSources") List<DataSource> dataSources)
+  {
+    Preconditions.checkNotNull(dataSources, "datasources cannot be null for uniondatasource");
+    this.dataSources = dataSources;
+  }
+
+  @Override
+  public Iterable<String> getNames()
+  {
+   return Iterables.concat(Iterables.transform(dataSources, new Function<DataSource, Iterable<String>>()
+   {
+     @Override
+     public Iterable<String> apply(DataSource input)
+     {
+       return input.getNames();
+     }
+   }));
+  }
+
+  @Override
+  public String getMetricName()
+  {
+    SortedSet<String> str = new TreeSet<>();
+    for(DataSource ds : dataSources){
+      str.add(ds.getMetricName());
+    }
+    return str.toString();
+  }
+
+  @JsonProperty
+  public List<DataSource> getDataSources(){
+    return dataSources;
+  }
+
+  @Override
+  public boolean equals(Object o)
+  {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+
+    UnionDataSource that = (UnionDataSource) o;
+
+    if (!dataSources.equals(that.dataSources)) {
+      return false;
+    }
+
+    return true;
+  }
+
+  @Override
+  public int hashCode()
+  {
+    return dataSources.hashCode();
+  }
+}

processing/src/main/java/io/druid/query/UnionQueryRunner.java

@@ -0,0 +1,65 @@
+/*
+ * Druid - a distributed column store.
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * as published by the Free Software Foundation; either version 2
+ * of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+ *
+ * This file Copyright (C) 2014 N3TWORK, Inc. and contributed to the Druid project
+ * under the Druid Corporate Contributor License Agreement.
+ */
+
+package io.druid.query;
+
+import com.google.common.base.Function;
+import com.google.common.collect.Iterables;
+import com.metamx.common.guava.Sequence;
+import com.metamx.common.guava.Sequences;
+
+public class UnionQueryRunner<T> implements QueryRunner<T>
+{
+  private final QueryRunner<T> baseRunner;
+
+  public UnionQueryRunner(QueryRunner<T> baseRunner)
+  {
+    this.baseRunner = baseRunner;
+  }
+
+  @Override
+  public Sequence<T> run(final Query<T> query)
+  {
+    DataSource dataSource = query.getDataSource();
+    if (dataSource instanceof UnionDataSource) {
+      System.out.println("Breaking into Single Source" + dataSource.getMetricName());
+      return Sequences.concat(
+          Iterables.transform(
+              ((UnionDataSource) dataSource).getDataSources(),
+              new Function<DataSource, Sequence<T>>()
+              {
+                @Override
+                public Sequence<T> apply(DataSource singleSource)
+                {
+                  System.out.println("Running with Single Source" + singleSource.getNames());
+                  return baseRunner.run(
+                      query.withDataSource(singleSource)
+                  );
+                }
+              }
+          )
+      );
+    } else {
+      return baseRunner.run(query);
+    }
+  }
+
+}

processing/src/main/java/io/druid/query/groupby/GroupByQuery.java

@@ -257,6 +257,24 @@ public class GroupByQuery extends BaseQuery<Row>
     );
   }
 
+  @Override
+  public Query<Row> withDataSource(DataSource dataSource)
+  {
+    return new GroupByQuery(
+        dataSource,
+        getQuerySegmentSpec(),
+        dimFilter,
+        granularity,
+        dimensions,
+        aggregatorSpecs,
+        postAggregatorSpecs,
+        havingSpec,
+        limitSpec,
+        orderByLimitFn,
+        getContext()
+    );
+  }
+
   public static class Builder
   {
     private DataSource dataSource;

processing/src/main/java/io/druid/query/groupby/GroupByQueryQueryToolChest.java

@@ -41,6 +41,7 @@ import io.druid.query.QueryDataSource;
 import io.druid.query.QueryRunner;
 import io.druid.query.QueryToolChest;
 import io.druid.query.SubqueryQueryRunner;
+import io.druid.query.UnionQueryRunner;
 import io.druid.query.aggregation.AggregatorFactory;
 import io.druid.query.aggregation.MetricManipulationFn;
 import io.druid.segment.incremental.IncrementalIndex;
@@ -163,7 +164,7 @@ public class GroupByQueryQueryToolChest extends QueryToolChest<Row, GroupByQuery
     }
 
     return new ServiceMetricEvent.Builder()
-        .setUser2(query.getDataSource().toString())
+        .setUser2(query.getDataSource().getMetricName())
         .setUser3(String.format("%,d dims", query.getDimensions().size()))
         .setUser4("groupBy")
         .setUser5(Joiner.on(",").join(query.getIntervals()))
@@ -202,7 +203,10 @@ public class GroupByQueryQueryToolChest extends QueryToolChest<Row, GroupByQuery
   @Override
   public QueryRunner<Row> preMergeQueryDecoration(QueryRunner<Row> runner)
   {
-    return new SubqueryQueryRunner<Row>(
+    return new UnionQueryRunner<Row>(
-        new IntervalChunkingQueryRunner<Row>(runner, configSupplier.get().getChunkPeriod()));
+        new SubqueryQueryRunner<Row>(
+            new IntervalChunkingQueryRunner<Row>(runner, configSupplier.get().getChunkPeriod())
+        )
+    );
   }
 }

processing/src/main/java/io/druid/query/metadata/SegmentMetadataQueryQueryToolChest.java

@@ -147,7 +147,7 @@ public class SegmentMetadataQueryQueryToolChest extends QueryToolChest<SegmentAn
     }
 
     return new ServiceMetricEvent.Builder()
-        .setUser2(query.getDataSource().toString())
+        .setUser2(query.getDataSource().getMetricName())
         .setUser4(query.getType())
         .setUser5(Joiner.on(",").join(query.getIntervals()))
         .setUser6(String.valueOf(query.hasFilters()))

processing/src/main/java/io/druid/query/metadata/metadata/SegmentMetadataQuery.java

@@ -92,6 +92,17 @@ public class SegmentMetadataQuery extends BaseQuery<SegmentAnalysis>
         spec, toInclude, merge, getContext());
   }
 
+  @Override
+  public Query<SegmentAnalysis> withDataSource(DataSource dataSource)
+  {
+    return new SegmentMetadataQuery(
+        dataSource,
+        getQuerySegmentSpec(),
+        toInclude,
+        merge,
+        getContext());
+  }
+
   @Override
   public boolean equals(Object o)
   {

processing/src/main/java/io/druid/query/search/SearchQueryQueryToolChest.java

@@ -45,6 +45,7 @@ import io.druid.query.QueryToolChest;
 import io.druid.query.Result;
 import io.druid.query.ResultGranularTimestampComparator;
 import io.druid.query.ResultMergeQueryRunner;
+import io.druid.query.UnionQueryRunner;
 import io.druid.query.aggregation.MetricManipulationFn;
 import io.druid.query.filter.DimFilter;
 import io.druid.query.search.search.SearchHit;
@@ -121,7 +122,7 @@ public class SearchQueryQueryToolChest extends QueryToolChest<Result<SearchResul
     }
 
     return new ServiceMetricEvent.Builder()
-        .setUser2(query.getDataSource().toString())
+        .setUser2(query.getDataSource().getMetricName())
         .setUser4("search")
         .setUser5(COMMA_JOIN.join(query.getIntervals()))
         .setUser6(String.valueOf(query.hasFilters()))
@@ -258,7 +259,9 @@ public class SearchQueryQueryToolChest extends QueryToolChest<Result<SearchResul
   public QueryRunner<Result<SearchResultValue>> preMergeQueryDecoration(QueryRunner<Result<SearchResultValue>> runner)
   {
     return new SearchThresholdAdjustingQueryRunner(
-        new IntervalChunkingQueryRunner<Result<SearchResultValue>>(runner, config.getChunkPeriod()),
+        new UnionQueryRunner<Result<SearchResultValue>>(
+            new IntervalChunkingQueryRunner<Result<SearchResultValue>>(runner, config.getChunkPeriod())
+        ),
         config
     );
   }

processing/src/main/java/io/druid/query/search/search/SearchQuery.java

@@ -111,6 +111,22 @@ public class SearchQuery extends BaseQuery<Result<SearchResultValue>>
     );
   }
 
+  @Override
+  public Query<Result<SearchResultValue>> withDataSource(DataSource dataSource)
+  {
+    return new SearchQuery(
+        dataSource,
+        dimFilter,
+        granularity,
+        limit,
+        getQuerySegmentSpec(),
+        dimensions,
+        querySpec,
+        sortSpec,
+        getContext()
+    );
+  }
+
   @Override
   public SearchQuery withOverriddenContext(Map<String, Object> contextOverrides)
   {

processing/src/main/java/io/druid/query/select/SelectQuery.java

@@ -120,6 +120,21 @@ public class SelectQuery extends BaseQuery<Result<SelectResultValue>>
     );
   }
 
+  @Override
+  public Query<Result<SelectResultValue>> withDataSource(DataSource dataSource)
+  {
+    return new SelectQuery(
+        dataSource,
+        getQuerySegmentSpec(),
+        dimFilter,
+        granularity,
+        dimensions,
+        metrics,
+        pagingSpec,
+        getContext()
+    );
+  }
+
   public SelectQuery withOverriddenContext(Map<String, Object> contextOverrides)
   {
     return new SelectQuery(

processing/src/main/java/io/druid/query/select/SelectQueryQueryToolChest.java

@@ -42,6 +42,7 @@ import io.druid.query.QueryToolChest;
 import io.druid.query.Result;
 import io.druid.query.ResultGranularTimestampComparator;
 import io.druid.query.ResultMergeQueryRunner;
+import io.druid.query.UnionQueryRunner;
 import io.druid.query.aggregation.MetricManipulationFn;
 import io.druid.query.filter.DimFilter;
 import org.joda.time.DateTime;
@@ -123,7 +124,7 @@ public class SelectQueryQueryToolChest extends QueryToolChest<Result<SelectResul
     }
 
     return new ServiceMetricEvent.Builder()
-        .setUser2(query.getDataSource().toString())
+        .setUser2(query.getDataSource().getMetricName())
         .setUser4("Select")
         .setUser5(COMMA_JOIN.join(query.getIntervals()))
         .setUser6(String.valueOf(query.hasFilters()))
@@ -277,7 +278,12 @@ public class SelectQueryQueryToolChest extends QueryToolChest<Result<SelectResul
   @Override
   public QueryRunner<Result<SelectResultValue>> preMergeQueryDecoration(QueryRunner<Result<SelectResultValue>> runner)
   {
-    return new IntervalChunkingQueryRunner<Result<SelectResultValue>>(runner, config.getChunkPeriod());
+    return new UnionQueryRunner<Result<SelectResultValue>>(
+        new IntervalChunkingQueryRunner<Result<SelectResultValue>>(
+            runner,
+            config.getChunkPeriod()
+        )
+    );
   }
 
   public Ordering<Result<SelectResultValue>> getOrdering()

processing/src/main/java/io/druid/query/timeboundary/TimeBoundaryQuery.java

@@ -97,6 +97,16 @@ public class TimeBoundaryQuery extends BaseQuery<Result<TimeBoundaryResultValue>
     );
   }
 
+  @Override
+  public Query<Result<TimeBoundaryResultValue>> withDataSource(DataSource dataSource)
+  {
+    return new TimeBoundaryQuery(
+        dataSource,
+        getQuerySegmentSpec(),
+        getContext()
+    );
+  }
+
   public byte[] getCacheKey()
   {
     return ByteBuffer.allocate(1)

processing/src/main/java/io/druid/query/timeboundary/TimeBoundaryQueryQueryToolChest.java

@@ -117,7 +117,7 @@ public class TimeBoundaryQueryQueryToolChest
   public ServiceMetricEvent.Builder makeMetricBuilder(TimeBoundaryQuery query)
   {
     return new ServiceMetricEvent.Builder()
-        .setUser2(query.getDataSource().toString())
+        .setUser2(query.getDataSource().getMetricName())
         .setUser4(query.getType())
         .setUser6("false");
   }

processing/src/main/java/io/druid/query/timeseries/TimeseriesQuery.java

@@ -116,6 +116,20 @@ public class TimeseriesQuery extends BaseQuery<Result<TimeseriesResultValue>>
     );
   }
 
+  @Override
+  public Query<Result<TimeseriesResultValue>> withDataSource(DataSource dataSource)
+  {
+    return new TimeseriesQuery(
+        dataSource,
+        getQuerySegmentSpec(),
+        dimFilter,
+        granularity,
+        aggregatorSpecs,
+        postAggregatorSpecs,
+        getContext()
+    );
+  }
+
   public TimeseriesQuery withOverriddenContext(Map<String, Object> contextOverrides)
   {
     return new TimeseriesQuery(

processing/src/main/java/io/druid/query/timeseries/TimeseriesQueryQueryToolChest.java

@@ -42,6 +42,7 @@ import io.druid.query.QueryToolChest;
 import io.druid.query.Result;
 import io.druid.query.ResultGranularTimestampComparator;
 import io.druid.query.ResultMergeQueryRunner;
+import io.druid.query.UnionQueryRunner;
 import io.druid.query.aggregation.AggregatorFactory;
 import io.druid.query.aggregation.MetricManipulationFn;
 import io.druid.query.aggregation.PostAggregator;
@@ -122,7 +123,7 @@ public class TimeseriesQueryQueryToolChest extends QueryToolChest<Result<Timeser
     }
 
     return new ServiceMetricEvent.Builder()
-        .setUser2(query.getDataSource().toString())
+        .setUser2(query.getDataSource().getMetricName())
         .setUser4("timeseries")
         .setUser5(COMMA_JOIN.join(query.getIntervals()))
         .setUser6(String.valueOf(query.hasFilters()))
@@ -229,7 +230,12 @@ public class TimeseriesQueryQueryToolChest extends QueryToolChest<Result<Timeser
   @Override
   public QueryRunner<Result<TimeseriesResultValue>> preMergeQueryDecoration(QueryRunner<Result<TimeseriesResultValue>> runner)
   {
-    return new IntervalChunkingQueryRunner<Result<TimeseriesResultValue>>(runner, config.getChunkPeriod());
+    return new UnionQueryRunner<Result<TimeseriesResultValue>>(
+        new IntervalChunkingQueryRunner<Result<TimeseriesResultValue>>(
+            runner,
+            config.getChunkPeriod()
+        )
+    );
   }
 
   public Ordering<Result<TimeseriesResultValue>> getOrdering()

processing/src/main/java/io/druid/query/topn/TopNQuery.java

@@ -27,6 +27,7 @@ import io.druid.granularity.QueryGranularity;
 import io.druid.query.BaseQuery;
 import io.druid.query.DataSource;
 import io.druid.query.Queries;
+import io.druid.query.Query;
 import io.druid.query.Result;
 import io.druid.query.aggregation.AggregatorFactory;
 import io.druid.query.aggregation.PostAggregator;
@@ -162,6 +163,23 @@ public class TopNQuery extends BaseQuery<Result<TopNResultValue>>
     );
   }
 
+  @Override
+  public Query<Result<TopNResultValue>> withDataSource(DataSource dataSource)
+  {
+    return new TopNQuery(
+        getDataSource(),
+        dimensionSpec,
+        topNMetricSpec,
+        threshold,
+        getQuerySegmentSpec(),
+        dimFilter,
+        granularity,
+        aggregatorSpecs,
+        postAggregatorSpecs,
+        getContext()
+    );
+  }
+
   public TopNQuery withThreshold(int threshold)
   {
     return new TopNQuery(

processing/src/main/java/io/druid/query/topn/TopNQueryQueryToolChest.java

@@ -45,6 +45,7 @@ import io.druid.query.QueryToolChest;
 import io.druid.query.Result;
 import io.druid.query.ResultGranularTimestampComparator;
 import io.druid.query.ResultMergeQueryRunner;
+import io.druid.query.UnionQueryRunner;
 import io.druid.query.aggregation.AggregatorFactory;
 import io.druid.query.aggregation.AggregatorUtil;
 import io.druid.query.aggregation.MetricManipulationFn;
@@ -131,7 +132,7 @@ public class TopNQueryQueryToolChest extends QueryToolChest<Result<TopNResultVal
     }
 
     return new ServiceMetricEvent.Builder()
-        .setUser2(query.getDataSource().toString())
+        .setUser2(query.getDataSource().getMetricName())
         .setUser4(String.format("topN/%s/%s", query.getThreshold(), query.getDimensionSpec().getDimension()))
         .setUser5(COMMA_JOIN.join(query.getIntervals()))
         .setUser6(String.valueOf(query.hasFilters()))
@@ -370,7 +371,13 @@ public class TopNQueryQueryToolChest extends QueryToolChest<Result<TopNResultVal
   @Override
   public QueryRunner<Result<TopNResultValue>> preMergeQueryDecoration(QueryRunner<Result<TopNResultValue>> runner)
   {
-    return new IntervalChunkingQueryRunner<Result<TopNResultValue>>(runner, config.getChunkPeriod());
+    return new UnionQueryRunner<Result<TopNResultValue>>(
+        new IntervalChunkingQueryRunner<Result<TopNResultValue>>(
+            runner,
+            config
+                .getChunkPeriod()
+        )
+    );
   }
 
   @Override

processing/src/test/java/io/druid/query/DataSourceTest.java

@@ -22,7 +22,10 @@
 package io.druid.query;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.base.Function;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
+import com.metamx.common.guava.Sequence;
 import io.druid.jackson.DefaultObjectMapper;
 import io.druid.query.aggregation.AggregatorFactory;
 import io.druid.query.aggregation.LongSumAggregatorFactory;
@@ -85,4 +88,26 @@ public class DataSourceTest
     Assert.assertEquals(new QueryDataSource(query), dataSource);
   }
 
+  @Test
+  public void testUnionDataSource() throws Exception
+  {
+    DataSource dataSource = jsonMapper.readValue(
+        "{\"type\":\"union\", \"dataSources\":[\"ds1\", \"ds2\"]}",
+        DataSource.class
+    );
+    Assert.assertTrue(dataSource instanceof UnionDataSource);
+    Assert.assertEquals(
+        Lists.newArrayList(new TableDataSource("ds1"), new TableDataSource("ds2")),
+        Lists.newArrayList(((UnionDataSource) dataSource).getDataSources())
+    );
+    Assert.assertEquals(
+        Lists.newArrayList("ds1", "ds2"),
+        Lists.newArrayList(dataSource.getNames())
+    );
+    Assert.assertEquals(Lists.newArrayList("ds1", "ds2").toString(), dataSource.getMetricName());
+
+    final DataSource serde = jsonMapper.readValue(jsonMapper.writeValueAsString(dataSource), DataSource.class);
+    Assert.assertEquals(dataSource, serde);
+  }
+
 }

processing/src/test/java/io/druid/query/metadata/SegmentMetadataQueryTest.java

@@ -20,6 +20,7 @@
 package io.druid.query.metadata;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.Iterables;
 import io.druid.jackson.DefaultObjectMapper;
 import io.druid.query.Query;
 import io.druid.query.metadata.metadata.SegmentMetadataQuery;
@@ -41,7 +42,7 @@ public class SegmentMetadataQueryTest
                       + "}";
     Query query = mapper.readValue(queryStr, Query.class);
     Assert.assertTrue(query instanceof SegmentMetadataQuery);
-    Assert.assertEquals("test_ds", query.getDataSource().getName());
+    Assert.assertEquals("test_ds", Iterables.getOnlyElement(query.getDataSource().getNames()));
     Assert.assertEquals(new Interval("2013-12-04T00:00:00.000Z/2013-12-05T00:00:00.000Z"), query.getIntervals().get(0));
 
     // test serialize and deserialize

server/src/main/java/io/druid/client/BrokerServerView.java

@@ -20,6 +20,7 @@
 package io.druid.client;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Ordering;
 import com.google.inject.Inject;
@@ -237,17 +238,7 @@ public class BrokerServerView implements TimelineServerView
   @Override
   public VersionedIntervalTimeline<String, ServerSelector> getTimeline(DataSource dataSource)
   {
-    String table;
+    String table = Iterables.getOnlyElement(dataSource.getNames());
-    while (dataSource instanceof QueryDataSource) {
-      dataSource = ((QueryDataSource) dataSource).getQuery().getDataSource();
-    }
-
-    if (dataSource instanceof TableDataSource) {
-      table = ((TableDataSource) dataSource).getName();
-    } else {
-      throw new UnsupportedOperationException("Unsupported data source type: " + dataSource.getClass().getSimpleName());
-    }
-
     synchronized (lock) {
       return timelines.get(table);
     }

server/src/main/java/io/druid/segment/realtime/RealtimeManager.java

@@ -21,6 +21,7 @@ package io.druid.segment.realtime;
 
 import com.google.common.base.Preconditions;
 import com.google.common.base.Throwables;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Maps;
 import com.google.common.io.Closeables;
 import com.google.inject.Inject;
@@ -126,18 +127,7 @@ public class RealtimeManager implements QuerySegmentWalker
 
   private <T> String getDataSourceName(Query<T> query)
   {
-    DataSource dataSource = query.getDataSource();
+    return Iterables.getOnlyElement(query.getDataSource().getNames());
-    if (!(dataSource instanceof TableDataSource)) {
-      throw new UnsupportedOperationException("data source type '" + dataSource.getClass().getName() + "' unsupported");
-    }
-
-    String dataSourceName;
-    try {
-      dataSourceName = ((TableDataSource) query.getDataSource()).getName();
-    } catch (ClassCastException e) {
-      throw new UnsupportedOperationException("Subqueries are only supported in the broker");
-    }
-    return dataSourceName;
   }
 
 

server/src/main/java/io/druid/server/AsyncQueryForwardingServlet.java

@@ -282,7 +282,7 @@ public class AsyncQueryForwardingServlet extends HttpServlet
 
           emitter.emit(
               new ServiceMetricEvent.Builder()
-                  .setUser2(theQuery.getDataSource().getName())
+                  .setUser2(theQuery.getDataSource().getMetricName())
                   .setUser4(theQuery.getType())
                   .setUser5(COMMA_JOIN.join(theQuery.getIntervals()))
                   .setUser6(String.valueOf(theQuery.hasFilters()))

server/src/main/java/io/druid/server/QueryResource.java

@@ -131,7 +131,7 @@ public class QueryResource
 
       emitter.emit(
           new ServiceMetricEvent.Builder()
-              .setUser2(query.getDataSource().toString())
+              .setUser2(query.getDataSource().getMetricName())
               .setUser4(query.getType())
               .setUser5(COMMA_JOIN.join(query.getIntervals()))
               .setUser6(String.valueOf(query.hasFilters()))

server/src/main/java/io/druid/server/coordination/ServerManager.java

@@ -21,7 +21,9 @@ package io.druid.server.coordination;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.base.Function;
+import com.google.common.base.Preconditions;
 import com.google.common.base.Predicates;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Ordering;
 import com.google.inject.Inject;
 import com.metamx.common.ISE;
@@ -252,14 +254,7 @@ public class ServerManager implements QuerySegmentWalker
     if (!(dataSource instanceof TableDataSource)) {
       throw new UnsupportedOperationException("data source type '" + dataSource.getClass().getName() + "' unsupported");
     }
-
+    String dataSourceName = getDataSourceName(dataSource);
-    String dataSourceName;
-    try {
-      dataSourceName = ((TableDataSource) query.getDataSource()).getName();
-    }
-    catch (ClassCastException e) {
-      throw new UnsupportedOperationException("Subqueries are only supported in the broker");
-    }
 
     final VersionedIntervalTimeline<String, ReferenceCountingSegment> timeline = dataSources.get(dataSourceName);
 
@@ -325,6 +320,16 @@ public class ServerManager implements QuerySegmentWalker
     return new FinalizeResultsQueryRunner<T>(toolChest.mergeResults(factory.mergeRunners(exec, adapters)), toolChest);
   }
 
+  private String getDataSourceName(DataSource dataSource)
+  {
+    Preconditions.checkArgument(
+        dataSource instanceof TableDataSource,
+        "Subqueries and Unions are only supported in the broker"
+    );
+
+    return Iterables.getOnlyElement(dataSource.getNames());
+  }
+
   @Override
   public <T> QueryRunner<T> getQueryRunnerForSegments(Query<T> query, Iterable<SegmentDescriptor> specs)
   {
@@ -338,13 +343,7 @@ public class ServerManager implements QuerySegmentWalker
 
     final QueryToolChest<T, Query<T>> toolChest = factory.getToolchest();
 
-    String dataSourceName;
+    String dataSourceName = getDataSourceName(query.getDataSource());
-    try {
-      dataSourceName = ((TableDataSource) query.getDataSource()).getName();
-    }
-    catch (ClassCastException e) {
-      throw new UnsupportedOperationException("Subqueries are only supported in the broker");
-    }
 
     final VersionedIntervalTimeline<String, ReferenceCountingSegment> timeline = dataSources.get(dataSourceName);
 

server/src/main/java/io/druid/server/router/TieredBrokerHostSelector.java

@@ -137,7 +137,8 @@ public class TieredBrokerHostSelector<T> implements HostSelector<T>
     }
 
     if (brokerServiceName == null) {
-      List<Rule> rules = ruleManager.getRulesWithDefault((query.getDataSource()).getName());
+      // For Union Queries tier will be selected on the rules for first dataSource.
+      List<Rule> rules = ruleManager.getRulesWithDefault(Iterables.getFirst(query.getDataSource().getNames(), null));
 
       // find the rule that can apply to the entire set of intervals
       DateTime now = new DateTime();