diff --git a/processing/src/main/java/io/druid/query/BaseQuery.java b/processing/src/main/java/io/druid/query/BaseQuery.java index 86587d47ae8..5d1db267156 100644 --- a/processing/src/main/java/io/druid/query/BaseQuery.java +++ b/processing/src/main/java/io/druid/query/BaseQuery.java @@ -20,6 +20,7 @@ package io.druid.query; import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableMap; import com.google.common.collect.Maps; @@ -109,7 +110,21 @@ public abstract class BaseQuery> implements Query @Override public QueryRunner getRunner(QuerySegmentWalker walker) { - return querySegmentSpec.lookup(this, walker); + return getQuerySegmentSpecForLookUp(this).lookup(this, walker); + } + + @VisibleForTesting + public static QuerySegmentSpec getQuerySegmentSpecForLookUp(BaseQuery query) + { + if (query.getDataSource() instanceof QueryDataSource) { + QueryDataSource ds = (QueryDataSource) query.getDataSource(); + Query subquery = ds.getQuery(); + if (subquery instanceof BaseQuery) { + return getQuerySegmentSpecForLookUp((BaseQuery) subquery); + } + throw new IllegalStateException("Invalid subquery type " + subquery.getClass()); + } + return query.getQuerySegmentSpec(); } @Override diff --git a/processing/src/test/java/io/druid/query/groupby/GroupByQueryTest.java b/processing/src/test/java/io/druid/query/groupby/GroupByQueryTest.java index 13a996f948d..7ffae74f0dc 100644 --- a/processing/src/test/java/io/druid/query/groupby/GroupByQueryTest.java +++ b/processing/src/test/java/io/druid/query/groupby/GroupByQueryTest.java @@ -26,7 +26,9 @@ import com.google.common.collect.Lists; import com.google.common.collect.Ordering; import io.druid.data.input.MapBasedRow; import io.druid.data.input.Row; +import io.druid.java.util.common.Intervals; import io.druid.java.util.common.granularity.Granularities; +import io.druid.query.BaseQuery; import io.druid.query.Query; import io.druid.query.QueryRunnerTestHelper; import io.druid.query.aggregation.AggregatorFactory; @@ -38,6 +40,8 @@ import io.druid.query.dimension.DimensionSpec; import io.druid.query.groupby.orderby.DefaultLimitSpec; import io.druid.query.groupby.orderby.OrderByColumnSpec; import io.druid.query.ordering.StringComparators; +import io.druid.query.spec.MultipleIntervalSegmentSpec; +import io.druid.query.spec.QuerySegmentSpec; import io.druid.segment.TestHelper; import io.druid.segment.column.ValueType; import org.junit.Assert; @@ -45,6 +49,7 @@ import org.junit.Test; import java.io.IOException; import java.util.Arrays; +import java.util.List; public class GroupByQueryTest { @@ -103,4 +108,28 @@ public class GroupByQueryTest ); Assert.assertEquals(-1, compare); } + + @Test + public void testSegmentLookUpForNestedQueries() + { + QuerySegmentSpec innerQuerySegmentSpec = new MultipleIntervalSegmentSpec(Lists.newArrayList(Intervals.of( + "2011-11-07/2011-11-08"))); + QuerySegmentSpec outerQuerySegmentSpec = new MultipleIntervalSegmentSpec(Lists.newArrayList((Intervals.of( + "2011-11-04/2011-11-08")))); + List aggs = Lists.newArrayList(QueryRunnerTestHelper.rowsCount); + final GroupByQuery innerQuery = GroupByQuery.builder() + .setDataSource("blah") + .setInterval(innerQuerySegmentSpec) + .setGranularity(Granularities.DAY) + .setAggregatorSpecs(aggs) + .build(); + final GroupByQuery query = GroupByQuery + .builder() + .setDataSource(innerQuery) + .setInterval(outerQuerySegmentSpec) + .setAggregatorSpecs(aggs) + .setGranularity(Granularities.DAY) + .build(); + Assert.assertEquals(innerQuerySegmentSpec, BaseQuery.getQuerySegmentSpecForLookUp(query)); + } }