Examples with QueryDataSource org.apache.druid.query.QueryDataSource used on opensource projects

Example 51 with QueryDataSource

use of org.apache.druid.query.QueryDataSource in project druid by druid-io.

the class HllSketchSqlAggregatorTest method testAvgDailyCountDistinctHllSketch.

@Test
public void testAvgDailyCountDistinctHllSketch() throws Exception {
    // Can't vectorize due to outer query, which runs on an inline datasource.
    cannotVectorize();
    final List<Object[]> expectedResults = ImmutableList.of(new Object[] { 1L });
    testQuery("SELECT\n" + "  AVG(u)\n" + "FROM (" + "  SELECT FLOOR(__time TO DAY), APPROX_COUNT_DISTINCT_DS_HLL(cnt) AS u\n" + "  FROM druid.foo\n" + "  GROUP BY 1\n" + ")", ImmutableList.of(GroupByQuery.builder().setDataSource(new QueryDataSource(Druids.newTimeseriesQueryBuilder().dataSource(CalciteTests.DATASOURCE1).intervals(new MultipleIntervalSegmentSpec(ImmutableList.of(Filtration.eternity()))).granularity(new PeriodGranularity(Period.days(1), null, DateTimeZone.UTC)).aggregators(Collections.singletonList(new HllSketchBuildAggregatorFactory("a0:a", "cnt", null, null, ROUND))).postAggregators(ImmutableList.of(new FinalizingFieldAccessPostAggregator("a0", "a0:a"))).context(QUERY_CONTEXT_DEFAULT).build().withOverriddenContext(BaseCalciteQueryTest.getTimeseriesContextWithFloorTime(ImmutableMap.of(TimeseriesQuery.SKIP_EMPTY_BUCKETS, true, BaseQuery.SQL_QUERY_ID, "dummy"), "d0")))).setInterval(new MultipleIntervalSegmentSpec(ImmutableList.of(Filtration.eternity()))).setGranularity(Granularities.ALL).setAggregatorSpecs(NullHandling.replaceWithDefault() ? Arrays.asList(new LongSumAggregatorFactory("_a0:sum", "a0"), new CountAggregatorFactory("_a0:count")) : Arrays.asList(new LongSumAggregatorFactory("_a0:sum", "a0"), new FilteredAggregatorFactory(new CountAggregatorFactory("_a0:count"), BaseCalciteQueryTest.not(BaseCalciteQueryTest.selector("a0", null, null))))).setPostAggregatorSpecs(ImmutableList.of(new ArithmeticPostAggregator("_a0", "quotient", ImmutableList.of(new FieldAccessPostAggregator(null, "_a0:sum"), new FieldAccessPostAggregator(null, "_a0:count"))))).setContext(QUERY_CONTEXT_DEFAULT).build()), expectedResults);
}

Example 52 with QueryDataSource

use of org.apache.druid.query.QueryDataSource in project druid by druid-io.

the class GroupByStrategyV2 method numMergeBuffersNeededForSubtotalsSpec.

private static int numMergeBuffersNeededForSubtotalsSpec(GroupByQuery query) {
    List<List<String>> subtotalSpecs = query.getSubtotalsSpec();
    final DataSource dataSource = query.getDataSource();
    int numMergeBuffersNeededForSubQuerySubtotal = 0;
    if (dataSource instanceof QueryDataSource) {
        Query<?> subQuery = ((QueryDataSource) dataSource).getQuery();
        if (subQuery instanceof GroupByQuery) {
            numMergeBuffersNeededForSubQuerySubtotal = numMergeBuffersNeededForSubtotalsSpec((GroupByQuery) subQuery);
        }
    }
    if (subtotalSpecs == null || subtotalSpecs.size() == 0) {
        return numMergeBuffersNeededForSubQuerySubtotal;
    }
    List<String> queryDimOutputNames = query.getDimensions().stream().map(DimensionSpec::getOutputName).collect(Collectors.toList());
    for (List<String> subtotalSpec : subtotalSpecs) {
        if (!Utils.isPrefix(subtotalSpec, queryDimOutputNames)) {
            return 2;
        }
    }
    return Math.max(1, numMergeBuffersNeededForSubQuerySubtotal);
}

Example 53 with QueryDataSource

use of org.apache.druid.query.QueryDataSource in project druid by druid-io.

the class GroupByQueryMergeBufferTest method testDoubleNestedGroupBy.

@Test
public void testDoubleNestedGroupBy() {
    final GroupByQuery query = GroupByQuery.builder().setDataSource(new QueryDataSource(GroupByQuery.builder().setDataSource(GroupByQuery.builder().setDataSource(QueryRunnerTestHelper.DATA_SOURCE).setInterval(QueryRunnerTestHelper.FIRST_TO_THIRD).setGranularity(Granularities.ALL).setDimensions(new DefaultDimensionSpec("quality", "alias"), new DefaultDimensionSpec("market", null)).setAggregatorSpecs(Collections.singletonList(QueryRunnerTestHelper.ROWS_COUNT)).build()).setInterval(QueryRunnerTestHelper.FIRST_TO_THIRD).setGranularity(Granularities.ALL).setDimensions(new DefaultDimensionSpec("quality", "alias")).setAggregatorSpecs(Collections.singletonList(QueryRunnerTestHelper.ROWS_COUNT)).build())).setGranularity(Granularities.ALL).setInterval(QueryRunnerTestHelper.FIRST_TO_THIRD).setAggregatorSpecs(new LongSumAggregatorFactory("rows", "rows")).setContext(ImmutableMap.of(QueryContexts.TIMEOUT_KEY, TIMEOUT)).build();
    Assert.assertEquals(2, GroupByStrategyV2.countRequiredMergeBufferNum(query));
    GroupByQueryRunnerTestHelper.runQuery(FACTORY, runner, query);
    // This should be 1 because the broker needs 2 buffers and the queryable node needs one.
    Assert.assertEquals(1, MERGE_BUFFER_POOL.getMinRemainBufferNum());
    Assert.assertEquals(4, MERGE_BUFFER_POOL.getPoolSize());
}

Example 54 with QueryDataSource

use of org.apache.druid.query.QueryDataSource in project druid by druid-io.

the class GroupByQueryMergeBufferTest method testNestedGroupBy.

@Test
public void testNestedGroupBy() {
    final GroupByQuery query = GroupByQuery.builder().setDataSource(new QueryDataSource(GroupByQuery.builder().setDataSource(QueryRunnerTestHelper.DATA_SOURCE).setInterval(QueryRunnerTestHelper.FIRST_TO_THIRD).setGranularity(Granularities.ALL).setDimensions(new DefaultDimensionSpec("quality", "alias")).setAggregatorSpecs(Collections.singletonList(QueryRunnerTestHelper.ROWS_COUNT)).build())).setGranularity(Granularities.ALL).setInterval(QueryRunnerTestHelper.FIRST_TO_THIRD).setAggregatorSpecs(new LongSumAggregatorFactory("rows", "rows")).setContext(ImmutableMap.of(QueryContexts.TIMEOUT_KEY, TIMEOUT)).build();
    Assert.assertEquals(1, GroupByStrategyV2.countRequiredMergeBufferNum(query));
    GroupByQueryRunnerTestHelper.runQuery(FACTORY, runner, query);
    Assert.assertEquals(2, MERGE_BUFFER_POOL.getMinRemainBufferNum());
    Assert.assertEquals(4, MERGE_BUFFER_POOL.getPoolSize());
}

Example 55 with QueryDataSource

use of org.apache.druid.query.QueryDataSource in project druid by druid-io.

the class GroupByQueryMergeBufferTest method testNestedGroupByWithSubtotals.

@Test
public void testNestedGroupByWithSubtotals() {
    final GroupByQuery query = GroupByQuery.builder().setDataSource(new QueryDataSource(GroupByQuery.builder().setDataSource(QueryRunnerTestHelper.DATA_SOURCE).setInterval(QueryRunnerTestHelper.FIRST_TO_THIRD).setGranularity(Granularities.ALL).setDimensions(Arrays.asList(DefaultDimensionSpec.of("quality"), DefaultDimensionSpec.of("market"), DefaultDimensionSpec.of("placement"))).setAggregatorSpecs(Collections.singletonList(QueryRunnerTestHelper.ROWS_COUNT)).build())).setGranularity(Granularities.ALL).setInterval(QueryRunnerTestHelper.FIRST_TO_THIRD).setDimensions(Arrays.asList(DefaultDimensionSpec.of("quality"), DefaultDimensionSpec.of("market"))).setSubtotalsSpec(Arrays.asList(Collections.singletonList("quality"), Collections.singletonList("market"))).setAggregatorSpecs(new LongSumAggregatorFactory("rows", "rows")).setContext(ImmutableMap.of(QueryContexts.TIMEOUT_KEY, TIMEOUT)).build();
    Assert.assertEquals(3, GroupByStrategyV2.countRequiredMergeBufferNum(query));
    GroupByQueryRunnerTestHelper.runQuery(FACTORY, runner, query);
    // 2 for subtotal, 1 for nested group by and 1 for GroupByQueryRunnerFactory#mergeRunners
    Assert.assertEquals(0, MERGE_BUFFER_POOL.getMinRemainBufferNum());
    Assert.assertEquals(4, MERGE_BUFFER_POOL.getPoolSize());
}