Search in sources :

Example 1 with MergeIterable

use of org.apache.druid.java.util.common.guava.MergeIterable in project druid by druid-io.

the class CachingClusteredClientTest method testQueryCaching.

@SuppressWarnings("unchecked")
public void testQueryCaching(final QueryRunner runner, final int numTimesToQuery, boolean expectBySegment, final Query query, // does this assume query intervals must be ordered?
Object... args) {
    final List<Interval> queryIntervals = Lists.newArrayListWithCapacity(args.length / 2);
    final List<List<Iterable<Result<Object>>>> expectedResults = Lists.newArrayListWithCapacity(queryIntervals.size());
    parseResults(queryIntervals, expectedResults, args);
    for (int i = 0; i < queryIntervals.size(); ++i) {
        List<Object> mocks = new ArrayList<>();
        mocks.add(serverView);
        final Interval actualQueryInterval = new Interval(queryIntervals.get(0).getStart(), queryIntervals.get(i).getEnd());
        final List<Map<DruidServer, ServerExpectations>> serverExpectationList = populateTimeline(queryIntervals, expectedResults, i, mocks);
        List<Capture> queryCaptures = new ArrayList<>();
        final Map<DruidServer, ServerExpectations> finalExpectation = serverExpectationList.get(serverExpectationList.size() - 1);
        for (Map.Entry<DruidServer, ServerExpectations> entry : finalExpectation.entrySet()) {
            DruidServer server = entry.getKey();
            ServerExpectations expectations = entry.getValue();
            EasyMock.expect(serverView.getQueryRunner(server)).andReturn(expectations.getQueryRunner()).once();
            final Capture<? extends QueryPlus> capture = Capture.newInstance();
            final Capture<? extends ResponseContext> context = Capture.newInstance();
            queryCaptures.add(capture);
            QueryRunner queryable = expectations.getQueryRunner();
            if (query instanceof TimeseriesQuery) {
                List<SegmentId> segmentIds = new ArrayList<>();
                List<Interval> intervals = new ArrayList<>();
                List<Iterable<Result<TimeseriesResultValue>>> results = new ArrayList<>();
                for (ServerExpectation expectation : expectations) {
                    segmentIds.add(expectation.getSegmentId());
                    intervals.add(expectation.getInterval());
                    results.add(expectation.getResults());
                }
                EasyMock.expect(queryable.run(EasyMock.capture(capture), EasyMock.capture(context))).andReturn(toQueryableTimeseriesResults(expectBySegment, segmentIds, intervals, results)).once();
            } else if (query instanceof TopNQuery) {
                List<SegmentId> segmentIds = new ArrayList<>();
                List<Interval> intervals = new ArrayList<>();
                List<Iterable<Result<TopNResultValue>>> results = new ArrayList<>();
                for (ServerExpectation expectation : expectations) {
                    segmentIds.add(expectation.getSegmentId());
                    intervals.add(expectation.getInterval());
                    results.add(expectation.getResults());
                }
                EasyMock.expect(queryable.run(EasyMock.capture(capture), EasyMock.capture(context))).andReturn(toQueryableTopNResults(segmentIds, intervals, results)).once();
            } else if (query instanceof SearchQuery) {
                List<SegmentId> segmentIds = new ArrayList<>();
                List<Interval> intervals = new ArrayList<>();
                List<Iterable<Result<SearchResultValue>>> results = new ArrayList<>();
                for (ServerExpectation expectation : expectations) {
                    segmentIds.add(expectation.getSegmentId());
                    intervals.add(expectation.getInterval());
                    results.add(expectation.getResults());
                }
                EasyMock.expect(queryable.run(EasyMock.capture(capture), EasyMock.capture(context))).andReturn(toQueryableSearchResults(segmentIds, intervals, results)).once();
            } else if (query instanceof GroupByQuery) {
                List<SegmentId> segmentIds = new ArrayList<>();
                List<Interval> intervals = new ArrayList<>();
                List<Iterable<ResultRow>> results = new ArrayList<>();
                for (ServerExpectation expectation : expectations) {
                    segmentIds.add(expectation.getSegmentId());
                    intervals.add(expectation.getInterval());
                    results.add(expectation.getResults());
                }
                EasyMock.expect(queryable.run(EasyMock.capture(capture), EasyMock.capture(context))).andReturn(toQueryableGroupByResults((GroupByQuery) query, segmentIds, intervals, results)).once();
            } else if (query instanceof TimeBoundaryQuery) {
                List<SegmentId> segmentIds = new ArrayList<>();
                List<Interval> intervals = new ArrayList<>();
                List<Iterable<Result<TimeBoundaryResultValue>>> results = new ArrayList<>();
                for (ServerExpectation expectation : expectations) {
                    segmentIds.add(expectation.getSegmentId());
                    intervals.add(expectation.getInterval());
                    results.add(expectation.getResults());
                }
                EasyMock.expect(queryable.run(EasyMock.capture(capture), EasyMock.capture(context))).andReturn(toQueryableTimeBoundaryResults(segmentIds, intervals, results)).once();
            } else {
                throw new ISE("Unknown query type[%s]", query.getClass());
            }
        }
        final int expectedResultsRangeStart;
        final int expectedResultsRangeEnd;
        if (query instanceof TimeBoundaryQuery) {
            expectedResultsRangeStart = i;
            expectedResultsRangeEnd = i + 1;
        } else {
            expectedResultsRangeStart = 0;
            expectedResultsRangeEnd = i + 1;
        }
        runWithMocks(new Runnable() {

            @Override
            public void run() {
                for (int i = 0; i < numTimesToQuery; ++i) {
                    TestHelper.assertExpectedResults(new MergeIterable(query instanceof GroupByQuery ? ((GroupByQuery) query).getResultOrdering() : Comparators.naturalNullsFirst(), FunctionalIterable.create(new RangeIterable(expectedResultsRangeStart, expectedResultsRangeEnd)).transformCat(new Function<Integer, Iterable<Iterable<Result<Object>>>>() {

                        @Override
                        public Iterable<Iterable<Result<Object>>> apply(@Nullable Integer input) {
                            List<Iterable<Result<Object>>> retVal = new ArrayList<>();
                            final Map<DruidServer, ServerExpectations> exps = serverExpectationList.get(input);
                            for (ServerExpectations expectations : exps.values()) {
                                for (ServerExpectation expectation : expectations) {
                                    retVal.add(expectation.getResults());
                                }
                            }
                            return retVal;
                        }
                    })), runner.run(QueryPlus.wrap(query.withQuerySegmentSpec(new MultipleIntervalSegmentSpec(ImmutableList.of(actualQueryInterval)))), initializeResponseContext()));
                    if (queryCompletedCallback != null) {
                        queryCompletedCallback.run();
                    }
                }
            }
        }, mocks.toArray());
        // make sure all the queries were sent down as 'bySegment'
        for (Capture queryCapture : queryCaptures) {
            QueryPlus capturedQueryPlus = (QueryPlus) queryCapture.getValue();
            Query capturedQuery = capturedQueryPlus.getQuery();
            if (expectBySegment) {
                Assert.assertEquals(true, capturedQuery.getContextValue(QueryContexts.BY_SEGMENT_KEY));
            } else {
                Assert.assertTrue(capturedQuery.getContextValue(QueryContexts.BY_SEGMENT_KEY) == null || capturedQuery.getContextValue(QueryContexts.BY_SEGMENT_KEY).equals(false));
            }
        }
    }
}
Also used : TimeseriesResultValue(org.apache.druid.query.timeseries.TimeseriesResultValue) TimeseriesQuery(org.apache.druid.query.timeseries.TimeseriesQuery) TopNQuery(org.apache.druid.query.topn.TopNQuery) Query(org.apache.druid.query.Query) TimeBoundaryQuery(org.apache.druid.query.timeboundary.TimeBoundaryQuery) SearchQuery(org.apache.druid.query.search.SearchQuery) GroupByQuery(org.apache.druid.query.groupby.GroupByQuery) ArrayList(java.util.ArrayList) Capture(org.easymock.Capture) Result(org.apache.druid.query.Result) GroupByQuery(org.apache.druid.query.groupby.GroupByQuery) MergeIterable(org.apache.druid.java.util.common.guava.MergeIterable) TimeBoundaryResultValue(org.apache.druid.query.timeboundary.TimeBoundaryResultValue) TopNQuery(org.apache.druid.query.topn.TopNQuery) ArrayList(java.util.ArrayList) List(java.util.List) ImmutableList(com.google.common.collect.ImmutableList) ResultRow(org.apache.druid.query.groupby.ResultRow) TimeseriesQuery(org.apache.druid.query.timeseries.TimeseriesQuery) QueryableDruidServer(org.apache.druid.client.selector.QueryableDruidServer) Map(java.util.Map) TreeMap(java.util.TreeMap) ImmutableMap(com.google.common.collect.ImmutableMap) HashMap(java.util.HashMap) Nullable(javax.annotation.Nullable) Interval(org.joda.time.Interval) MergeIterable(org.apache.druid.java.util.common.guava.MergeIterable) FunctionalIterable(org.apache.druid.java.util.common.guava.FunctionalIterable) MultipleIntervalSegmentSpec(org.apache.druid.query.spec.MultipleIntervalSegmentSpec) TimeBoundaryQuery(org.apache.druid.query.timeboundary.TimeBoundaryQuery) SearchResultValue(org.apache.druid.query.search.SearchResultValue) HashPartitionFunction(org.apache.druid.timeline.partition.HashPartitionFunction) Function(com.google.common.base.Function) HashFunction(com.google.common.hash.HashFunction) ISE(org.apache.druid.java.util.common.ISE) QueryPlus(org.apache.druid.query.QueryPlus) SearchQuery(org.apache.druid.query.search.SearchQuery) SegmentId(org.apache.druid.timeline.SegmentId) FinalizeResultsQueryRunner(org.apache.druid.query.FinalizeResultsQueryRunner) QueryRunner(org.apache.druid.query.QueryRunner)

Example 2 with MergeIterable

use of org.apache.druid.java.util.common.guava.MergeIterable in project druid by druid-io.

the class ChainedExecutionQueryRunner method run.

@Override
public Sequence<T> run(final QueryPlus<T> queryPlus, final ResponseContext responseContext) {
    Query<T> query = queryPlus.getQuery();
    final int priority = QueryContexts.getPriority(query);
    final Ordering ordering = query.getResultOrdering();
    final QueryPlus<T> threadSafeQueryPlus = queryPlus.withoutThreadUnsafeState();
    return new BaseSequence<T, Iterator<T>>(new BaseSequence.IteratorMaker<T, Iterator<T>>() {

        @Override
        public Iterator<T> make() {
            // Make it a List<> to materialize all of the values (so that it will submit everything to the executor)
            List<ListenableFuture<Iterable<T>>> futures = Lists.newArrayList(Iterables.transform(queryables, input -> {
                if (input == null) {
                    throw new ISE("Null queryRunner! Looks to be some segment unmapping action happening");
                }
                return queryProcessingPool.submitRunnerTask(new AbstractPrioritizedQueryRunnerCallable<Iterable<T>, T>(priority, input) {

                    @Override
                    public Iterable<T> call() {
                        try {
                            Sequence<T> result = input.run(threadSafeQueryPlus, responseContext);
                            if (result == null) {
                                throw new ISE("Got a null result! Segments are missing!");
                            }
                            List<T> retVal = result.toList();
                            if (retVal == null) {
                                throw new ISE("Got a null list of results");
                            }
                            return retVal;
                        } catch (QueryInterruptedException e) {
                            throw new RuntimeException(e);
                        } catch (QueryTimeoutException e) {
                            throw e;
                        } catch (Exception e) {
                            log.noStackTrace().error(e, "Exception with one of the sequences!");
                            Throwables.propagateIfPossible(e);
                            throw new RuntimeException(e);
                        }
                    }
                });
            }));
            ListenableFuture<List<Iterable<T>>> future = Futures.allAsList(futures);
            queryWatcher.registerQueryFuture(query, future);
            try {
                return new MergeIterable<>(ordering.nullsFirst(), QueryContexts.hasTimeout(query) ? future.get(QueryContexts.getTimeout(query), TimeUnit.MILLISECONDS) : future.get()).iterator();
            } catch (InterruptedException e) {
                log.noStackTrace().warn(e, "Query interrupted, cancelling pending results, query id [%s]", query.getId());
                // Note: canceling combinedFuture first so that it can complete with INTERRUPTED as its final state. See ChainedExecutionQueryRunnerTest.testQueryTimeout()
                GuavaUtils.cancelAll(true, future, futures);
                throw new QueryInterruptedException(e);
            } catch (CancellationException e) {
                throw new QueryInterruptedException(e);
            } catch (TimeoutException e) {
                log.warn("Query timeout, cancelling pending results for query id [%s]", query.getId());
                GuavaUtils.cancelAll(true, future, futures);
                throw new QueryTimeoutException(StringUtils.nonStrictFormat("Query [%s] timed out", query.getId()));
            } catch (ExecutionException e) {
                GuavaUtils.cancelAll(true, future, futures);
                Throwables.propagateIfPossible(e.getCause());
                throw new RuntimeException(e.getCause());
            }
        }

        @Override
        public void cleanup(Iterator<T> tIterator) {
        }
    });
}
Also used : MergeIterable(org.apache.druid.java.util.common.guava.MergeIterable) MergeIterable(org.apache.druid.java.util.common.guava.MergeIterable) Ordering(com.google.common.collect.Ordering) Iterator(java.util.Iterator) List(java.util.List) ISE(org.apache.druid.java.util.common.ISE) ExecutionException(java.util.concurrent.ExecutionException) TimeoutException(java.util.concurrent.TimeoutException) Sequence(org.apache.druid.java.util.common.guava.Sequence) BaseSequence(org.apache.druid.java.util.common.guava.BaseSequence) BaseSequence(org.apache.druid.java.util.common.guava.BaseSequence) CancellationException(java.util.concurrent.CancellationException) TimeoutException(java.util.concurrent.TimeoutException) ExecutionException(java.util.concurrent.ExecutionException) CancellationException(java.util.concurrent.CancellationException) ListenableFuture(com.google.common.util.concurrent.ListenableFuture)

Aggregations

List (java.util.List)2 ISE (org.apache.druid.java.util.common.ISE)2 MergeIterable (org.apache.druid.java.util.common.guava.MergeIterable)2 Function (com.google.common.base.Function)1 ImmutableList (com.google.common.collect.ImmutableList)1 ImmutableMap (com.google.common.collect.ImmutableMap)1 Ordering (com.google.common.collect.Ordering)1 HashFunction (com.google.common.hash.HashFunction)1 ListenableFuture (com.google.common.util.concurrent.ListenableFuture)1 ArrayList (java.util.ArrayList)1 HashMap (java.util.HashMap)1 Iterator (java.util.Iterator)1 Map (java.util.Map)1 TreeMap (java.util.TreeMap)1 CancellationException (java.util.concurrent.CancellationException)1 ExecutionException (java.util.concurrent.ExecutionException)1 TimeoutException (java.util.concurrent.TimeoutException)1 Nullable (javax.annotation.Nullable)1 QueryableDruidServer (org.apache.druid.client.selector.QueryableDruidServer)1 BaseSequence (org.apache.druid.java.util.common.guava.BaseSequence)1