diff --git a/processing/src/main/java/io/druid/query/groupby/GroupByQueryQueryToolChest.java b/processing/src/main/java/io/druid/query/groupby/GroupByQueryQueryToolChest.java index 2f0a6bf92d1..2a379886518 100644 --- a/processing/src/main/java/io/druid/query/groupby/GroupByQueryQueryToolChest.java +++ b/processing/src/main/java/io/druid/query/groupby/GroupByQueryQueryToolChest.java @@ -199,35 +199,42 @@ public class GroupByQueryQueryToolChest extends QueryToolChest( - outerQuery.applyLimit( - Sequences.concat( - Sequences.map( - Sequences.simple(outerQuery.getIntervals()), - new Function>() - { - @Override - public Sequence apply(Interval interval) - { - return engine.process( - outerQuery.withQuerySegmentSpec( - new MultipleIntervalSegmentSpec(ImmutableList.of(interval)) - ), - new IncrementalIndexStorageAdapter(index) - ); - } - } - ) + final IncrementalIndex outerQueryResultIndex = makeIncrementalIndex( + outerQuery, + Sequences.concat( + Sequences.map( + Sequences.simple(outerQuery.getIntervals()), + new Function>() + { + @Override + public Sequence apply(Interval interval) + { + return engine.process( + outerQuery.withQuerySegmentSpec( + new MultipleIntervalSegmentSpec(ImmutableList.of(interval)) + ), + new IncrementalIndexStorageAdapter(innerQueryResultIndex) + ); + } + } ) - ), - index + ) ); + + innerQueryResultIndex.close(); + + return new ResourceClosingSequence<>( + outerQuery.applyLimit(postAggregate(query, outerQueryResultIndex)), + outerQueryResultIndex + ); + } else { final IncrementalIndex index = makeIncrementalIndex( query, runner.run( diff --git a/processing/src/test/java/io/druid/query/groupby/GroupByQueryRunnerTest.java b/processing/src/test/java/io/druid/query/groupby/GroupByQueryRunnerTest.java index f2f38cdf78a..9a4f2712d85 100644 --- a/processing/src/test/java/io/druid/query/groupby/GroupByQueryRunnerTest.java +++ b/processing/src/test/java/io/druid/query/groupby/GroupByQueryRunnerTest.java @@ -2809,6 +2809,63 @@ public class GroupByQueryRunnerTest TestHelper.assertExpectedObjects(expectedResults, results, ""); } + @Test + public void testSubqueryWithExtractionFnInOuterQuery() + { + //https://github.com/druid-io/druid/issues/2556 + + GroupByQuery subquery = GroupByQuery + .builder() + .setDataSource(QueryRunnerTestHelper.dataSource) + .setQuerySegmentSpec(QueryRunnerTestHelper.firstToThird) + .setDimensions(Lists.newArrayList(new DefaultDimensionSpec("quality", "alias"))) + .setDimFilter(new JavaScriptDimFilter("quality", "function(dim){ return true; }")) + .setAggregatorSpecs( + Arrays.asList( + QueryRunnerTestHelper.rowsCount, + new LongSumAggregatorFactory("idx", "index") + ) + ) + .setGranularity(QueryRunnerTestHelper.dayGran) + .build(); + + GroupByQuery query = GroupByQuery + .builder() + .setDataSource(subquery) + .setQuerySegmentSpec( + new MultipleIntervalSegmentSpec( + ImmutableList.of( + new Interval("2011-04-01T00:00:00.000Z/2011-04-03T00:00:00.000Z") + ) + ) + ) + .setDimensions(Lists.newArrayList( + new ExtractionDimensionSpec( + "alias", + "alias", + new RegexDimExtractionFn("(a).*", true, "a") + ) + ) + ) + .setAggregatorSpecs( + Arrays.asList( + new LongSumAggregatorFactory("rows", "rows"), + new LongSumAggregatorFactory("idx", "idx") + ) + ) + .setGranularity(QueryRunnerTestHelper.dayGran) + .build(); + + List expectedResults = Arrays.asList( + GroupByQueryRunnerTestHelper.createExpectedRow("2011-04-01", "alias", "a", "rows", 13L, "idx", 6619L), + GroupByQueryRunnerTestHelper.createExpectedRow("2011-04-02", "alias", "a", "rows", 13L, "idx", 5827L) + ); + + // Subqueries are handled by the ToolChest + Iterable results = GroupByQueryRunnerTestHelper.runQuery(factory, runner, query); + TestHelper.assertExpectedObjects(expectedResults, results, ""); + } + @Test public void testDifferentGroupingSubquery() {