Examples with ResponseContext org.apache.druid.query.context.ResponseContext used on opensource projects

Example 11 with ResponseContext

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

the class GroupByQueryRunnerTest method testMergeResultsWithLimitPushDownSortByAgg.

@Test
public void testMergeResultsWithLimitPushDownSortByAgg() {
    if (!config.getDefaultStrategy().equals(GroupByStrategySelector.STRATEGY_V2)) {
        return;
    }
    GroupByQuery.Builder builder = makeQueryBuilder().setDataSource(QueryRunnerTestHelper.DATA_SOURCE).setInterval("2011-04-02/2011-04-04").setDimensions(new DefaultDimensionSpec("quality", "alias")).setAggregatorSpecs(QueryRunnerTestHelper.ROWS_COUNT, new LongSumAggregatorFactory("idx", "index")).setLimitSpec(new DefaultLimitSpec(Collections.singletonList(new OrderByColumnSpec("idx", OrderByColumnSpec.Direction.DESCENDING)), 5)).overrideContext(ImmutableMap.of(GroupByQueryConfig.CTX_KEY_FORCE_LIMIT_PUSH_DOWN, true)).setGranularity(Granularities.ALL);
    final GroupByQuery allGranQuery = builder.build();
    QueryRunner mergedRunner = factory.getToolchest().mergeResults(new QueryRunner<ResultRow>() {

        @Override
        public Sequence<ResultRow> run(QueryPlus<ResultRow> queryPlus, ResponseContext responseContext) {
            // simulate two daily segments
            final QueryPlus<ResultRow> queryPlus1 = queryPlus.withQuery(queryPlus.getQuery().withQuerySegmentSpec(new MultipleIntervalSegmentSpec(Collections.singletonList(Intervals.of("2011-04-02/2011-04-03")))));
            final QueryPlus<ResultRow> queryPlus2 = queryPlus.withQuery(queryPlus.getQuery().withQuerySegmentSpec(new MultipleIntervalSegmentSpec(Collections.singletonList(Intervals.of("2011-04-03/2011-04-04")))));
            return factory.getToolchest().mergeResults((queryPlus3, responseContext1) -> new MergeSequence<>(queryPlus3.getQuery().getResultOrdering(), Sequences.simple(Arrays.asList(runner.run(queryPlus1, responseContext1), runner.run(queryPlus2, responseContext1))))).run(queryPlus, responseContext);
        }
    });
    List<ResultRow> allGranExpectedResults = Arrays.asList(makeRow(allGranQuery, "2011-04-02", "alias", "mezzanine", "rows", 6L, "idx", 4420L), makeRow(allGranQuery, "2011-04-02", "alias", "premium", "rows", 6L, "idx", 4416L), makeRow(allGranQuery, "2011-04-02", "alias", "entertainment", "rows", 2L, "idx", 319L), makeRow(allGranQuery, "2011-04-02", "alias", "automotive", "rows", 2L, "idx", 269L), makeRow(allGranQuery, "2011-04-02", "alias", "travel", "rows", 2L, "idx", 243L));
    Iterable<ResultRow> results = mergedRunner.run(QueryPlus.wrap(allGranQuery)).toList();
    TestHelper.assertExpectedObjects(allGranExpectedResults, results, "merged");
}

Example 12 with ResponseContext

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

the class GroupByQueryRunnerTest method testMergedPostAggHavingSpec.

@Test
public void testMergedPostAggHavingSpec() {
    GroupByQuery.Builder builder = makeQueryBuilder().setDataSource(QueryRunnerTestHelper.DATA_SOURCE).setInterval("2011-04-02/2011-04-04").setDimensions(new DefaultDimensionSpec("quality", "alias")).setAggregatorSpecs(QueryRunnerTestHelper.ROWS_COUNT, new LongSumAggregatorFactory("idx", "index")).setPostAggregatorSpecs(Collections.singletonList(new ArithmeticPostAggregator("rows_times_10", "*", Arrays.asList(new FieldAccessPostAggregator("rows", "rows"), new ConstantPostAggregator("const", 10L))))).setGranularity(new PeriodGranularity(new Period("P1M"), null, null)).setHavingSpec(new OrHavingSpec(ImmutableList.of(new GreaterThanHavingSpec("rows_times_10", 20L), new EqualToHavingSpec("idx", 217L))));
    GroupByQuery query = builder.build();
    // Same query, but with expressions instead of arithmetic.
    final GroupByQuery expressionQuery = query.withPostAggregatorSpecs(Collections.singletonList(new ExpressionPostAggregator("rows_times_10", "rows * 10.0", null, TestExprMacroTable.INSTANCE)));
    List<ResultRow> expectedResults = Arrays.asList(makeRow(query, "2011-04-01", "alias", "business", "rows", 2L, "idx", 217L, "rows_times_10", 20.0), makeRow(query, "2011-04-01", "alias", "mezzanine", "rows", 6L, "idx", 4420L, "rows_times_10", 60.0), makeRow(query, "2011-04-01", "alias", "premium", "rows", 6L, "idx", 4416L, "rows_times_10", 60.0));
    QueryRunner mergedRunner = factory.getToolchest().mergeResults(new QueryRunner<ResultRow>() {

        @Override
        public Sequence<ResultRow> run(QueryPlus<ResultRow> queryPlus, ResponseContext responseContext) {
            // simulate two daily segments
            final QueryPlus queryPlus1 = queryPlus.withQuery(queryPlus.getQuery().withQuerySegmentSpec(new MultipleIntervalSegmentSpec(Collections.singletonList(Intervals.of("2011-04-02/2011-04-03")))));
            final QueryPlus queryPlus2 = queryPlus.withQuery(queryPlus.getQuery().withQuerySegmentSpec(new MultipleIntervalSegmentSpec(Collections.singletonList(Intervals.of("2011-04-03/2011-04-04")))));
            return new MergeSequence(queryPlus.getQuery().getResultOrdering(), Sequences.simple(Arrays.asList(runner.run(queryPlus1, responseContext), runner.run(queryPlus2, responseContext))));
        }
    });
    ResponseContext context = ResponseContext.createEmpty();
    // add an extra layer of merging, simulate broker forwarding query to historical
    TestHelper.assertExpectedObjects(expectedResults, factory.getToolchest().postMergeQueryDecoration(factory.getToolchest().mergeResults(factory.getToolchest().preMergeQueryDecoration(mergedRunner))).run(QueryPlus.wrap(query)), "merged");
    TestHelper.assertExpectedObjects(expectedResults, factory.getToolchest().postMergeQueryDecoration(factory.getToolchest().mergeResults(factory.getToolchest().preMergeQueryDecoration(mergedRunner))).run(QueryPlus.wrap(expressionQuery)), "merged");
}

Example 13 with ResponseContext

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

the class GroupByQueryRunnerTest method testMergeResultsWithLimitPushDownSortByDimAggDim.

@Test
public void testMergeResultsWithLimitPushDownSortByDimAggDim() {
    if (!config.getDefaultStrategy().equals(GroupByStrategySelector.STRATEGY_V2)) {
        return;
    }
    GroupByQuery.Builder builder = makeQueryBuilder().setDataSource(QueryRunnerTestHelper.DATA_SOURCE).setInterval("2011-04-02/2011-04-04").setDimensions(new DefaultDimensionSpec("quality", "alias"), new DefaultDimensionSpec("market", "market")).setAggregatorSpecs(QueryRunnerTestHelper.ROWS_COUNT, new LongSumAggregatorFactory("idx", "index")).setLimitSpec(new DefaultLimitSpec(Lists.newArrayList(new OrderByColumnSpec("alias", OrderByColumnSpec.Direction.DESCENDING), new OrderByColumnSpec("idx", OrderByColumnSpec.Direction.DESCENDING), new OrderByColumnSpec("market", OrderByColumnSpec.Direction.DESCENDING)), 5)).overrideContext(ImmutableMap.of(GroupByQueryConfig.CTX_KEY_FORCE_LIMIT_PUSH_DOWN, true)).setGranularity(Granularities.ALL);
    final GroupByQuery allGranQuery = builder.build();
    QueryRunner mergedRunner = factory.getToolchest().mergeResults(new QueryRunner<ResultRow>() {

        @Override
        public Sequence<ResultRow> run(QueryPlus<ResultRow> queryPlus, ResponseContext responseContext) {
            // simulate two daily segments
            final QueryPlus<ResultRow> queryPlus1 = queryPlus.withQuery(queryPlus.getQuery().withQuerySegmentSpec(new MultipleIntervalSegmentSpec(Collections.singletonList(Intervals.of("2011-04-02/2011-04-03")))));
            final QueryPlus<ResultRow> queryPlus2 = queryPlus.withQuery(queryPlus.getQuery().withQuerySegmentSpec(new MultipleIntervalSegmentSpec(Collections.singletonList(Intervals.of("2011-04-03/2011-04-04")))));
            return factory.getToolchest().mergeResults((queryPlus3, responseContext1) -> new MergeSequence<>(queryPlus3.getQuery().getResultOrdering(), Sequences.simple(Arrays.asList(runner.run(queryPlus1, responseContext1), runner.run(queryPlus2, responseContext1))))).run(queryPlus, responseContext);
        }
    });
    List<ResultRow> allGranExpectedResults = Arrays.asList(makeRow(allGranQuery, "2011-04-02", "alias", "travel", "market", "spot", "rows", 2L, "idx", 243L), makeRow(allGranQuery, "2011-04-02", "alias", "technology", "market", "spot", "rows", 2L, "idx", 177L), makeRow(allGranQuery, "2011-04-02", "alias", "premium", "market", "total_market", "rows", 2L, "idx", 2342L), makeRow(allGranQuery, "2011-04-02", "alias", "premium", "market", "upfront", "rows", 2L, "idx", 1817L), makeRow(allGranQuery, "2011-04-02", "alias", "premium", "market", "spot", "rows", 2L, "idx", 257L));
    Iterable<ResultRow> results = mergedRunner.run(QueryPlus.wrap(allGranQuery)).toList();
    TestHelper.assertExpectedObjects(allGranExpectedResults, results, "merged");
}

Example 14 with ResponseContext

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

the class GroupByQueryRunnerTest method doTestMergeResultsWithOrderBy.

private void doTestMergeResultsWithOrderBy(GroupByQuery baseQuery, LimitSpec limitSpec, List<ResultRow> expectedResults) {
    QueryRunner mergedRunner = factory.getToolchest().mergeResults(new QueryRunner<ResultRow>() {

        @Override
        public Sequence<ResultRow> run(QueryPlus<ResultRow> queryPlus, ResponseContext responseContext) {
            // simulate two daily segments
            final QueryPlus queryPlus1 = queryPlus.withQuery(queryPlus.getQuery().withQuerySegmentSpec(new MultipleIntervalSegmentSpec(Collections.singletonList(Intervals.of("2011-04-02/2011-04-03")))));
            final QueryPlus queryPlus2 = queryPlus.withQuery(queryPlus.getQuery().withQuerySegmentSpec(new MultipleIntervalSegmentSpec(Collections.singletonList(Intervals.of("2011-04-03/2011-04-04")))));
            return new MergeSequence(queryPlus.getQuery().getResultOrdering(), Sequences.simple(Arrays.asList(runner.run(queryPlus1, responseContext), runner.run(queryPlus2, responseContext))));
        }
    });
    final GroupByQuery query = baseQuery.withLimitSpec(limitSpec);
    TestHelper.assertExpectedObjects(expectedResults, mergedRunner.run(QueryPlus.wrap(query)), "merged");
}

Example 15 with ResponseContext

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

the class GroupByTimeseriesQueryRunnerTest method constructorFeeder.

@SuppressWarnings("unchecked")
@Parameterized.Parameters(name = "{0}, vectorize = {1}")
public static Iterable<Object[]> constructorFeeder() {
    GroupByQueryConfig config = new GroupByQueryConfig();
    config.setMaxIntermediateRows(10000);
    final Pair<GroupByQueryRunnerFactory, Closer> factoryAndCloser = GroupByQueryRunnerTest.makeQueryRunnerFactory(config);
    final GroupByQueryRunnerFactory factory = factoryAndCloser.lhs;
    RESOURCE_CLOSER.register(factoryAndCloser.rhs);
    final List<Object[]> constructors = new ArrayList<>();
    for (QueryRunner<ResultRow> runner : QueryRunnerTestHelper.makeQueryRunners(factory)) {
        final QueryRunner modifiedRunner = new QueryRunner() {

            @Override
            public Sequence run(QueryPlus queryPlus, ResponseContext responseContext) {
                TimeseriesQuery tsQuery = (TimeseriesQuery) queryPlus.getQuery();
                QueryRunner<ResultRow> newRunner = factory.mergeRunners(Execs.directExecutor(), ImmutableList.of(runner));
                QueryToolChest toolChest = factory.getToolchest();
                newRunner = new FinalizeResultsQueryRunner<>(toolChest.mergeResults(toolChest.preMergeQueryDecoration(newRunner)), toolChest);
                final String timeDimension = tsQuery.getTimestampResultField();
                final List<VirtualColumn> virtualColumns = new ArrayList<>(Arrays.asList(tsQuery.getVirtualColumns().getVirtualColumns()));
                Map<String, Object> theContext = tsQuery.getContext();
                if (timeDimension != null) {
                    theContext = new HashMap<>(tsQuery.getContext());
                    final PeriodGranularity granularity = (PeriodGranularity) tsQuery.getGranularity();
                    virtualColumns.add(new ExpressionVirtualColumn("v0", StringUtils.format("timestamp_floor(__time, '%s')", granularity.getPeriod()), ColumnType.LONG, TestExprMacroTable.INSTANCE));
                    theContext.put(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD, timeDimension);
                    theContext.put(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD_GRANULARITY, granularity);
                    theContext.put(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD_INDEX, 0);
                }
                GroupByQuery newQuery = GroupByQuery.builder().setDataSource(tsQuery.getDataSource()).setQuerySegmentSpec(tsQuery.getQuerySegmentSpec()).setGranularity(tsQuery.getGranularity()).setDimFilter(tsQuery.getDimensionsFilter()).setDimensions(timeDimension == null ? ImmutableList.of() : ImmutableList.of(new DefaultDimensionSpec("v0", timeDimension, ColumnType.LONG))).setAggregatorSpecs(tsQuery.getAggregatorSpecs()).setPostAggregatorSpecs(tsQuery.getPostAggregatorSpecs()).setVirtualColumns(VirtualColumns.create(virtualColumns)).setContext(theContext).build();
                return Sequences.map(newRunner.run(queryPlus.withQuery(newQuery), responseContext), new Function<ResultRow, Result<TimeseriesResultValue>>() {

                    @Override
                    public Result<TimeseriesResultValue> apply(final ResultRow input) {
                        final MapBasedRow mapBasedRow = input.toMapBasedRow(newQuery);
                        return new Result<>(mapBasedRow.getTimestamp(), new TimeseriesResultValue(mapBasedRow.getEvent()));
                    }
                });
            }

            @Override
            public String toString() {
                return runner.toString();
            }
        };
        for (boolean vectorize : ImmutableList.of(false, true)) {
            // Add vectorization tests for any indexes that support it.
            if (!vectorize || QueryRunnerTestHelper.isTestRunnerVectorizable(runner)) {
                constructors.add(new Object[] { modifiedRunner, vectorize });
            }
        }
    }
    return constructors;
}