Examples with MergeIterable io.druid.java.util.common.guava.MergeIterable used on opensource projects

Example 1 with MergeIterable

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

the class ChainedExecutionQueryRunner method run.

@Override
public Sequence<T> run(final Query<T> query, final Map<String, Object> responseContext) {
    final int priority = BaseQuery.getContextPriority(query, 0);
    final Ordering ordering = query.getResultOrdering();
    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)
            ListenableFuture<List<Iterable<T>>> futures = Futures.allAsList(Lists.newArrayList(Iterables.transform(queryables, new Function<QueryRunner<T>, ListenableFuture<Iterable<T>>>() {

                @Override
                public ListenableFuture<Iterable<T>> apply(final QueryRunner<T> input) {
                    if (input == null) {
                        throw new ISE("Null queryRunner! Looks to be some segment unmapping action happening");
                    }
                    return exec.submit(new AbstractPrioritizedCallable<Iterable<T>>(priority) {

                        @Override
                        public Iterable<T> call() throws Exception {
                            try {
                                Sequence<T> result = input.run(query, responseContext);
                                if (result == null) {
                                    throw new ISE("Got a null result! Segments are missing!");
                                }
                                List<T> retVal = Sequences.toList(result, Lists.<T>newArrayList());
                                if (retVal == null) {
                                    throw new ISE("Got a null list of results! WTF?!");
                                }
                                return retVal;
                            } catch (QueryInterruptedException e) {
                                throw Throwables.propagate(e);
                            } catch (Exception e) {
                                log.error(e, "Exception with one of the sequences!");
                                throw Throwables.propagate(e);
                            }
                        }
                    });
                }
            })));
            queryWatcher.registerQuery(query, futures);
            try {
                final Number timeout = query.getContextValue(QueryContextKeys.TIMEOUT, (Number) null);
                return new MergeIterable<>(ordering.nullsFirst(), timeout == null ? futures.get() : futures.get(timeout.longValue(), TimeUnit.MILLISECONDS)).iterator();
            } catch (InterruptedException e) {
                log.warn(e, "Query interrupted, cancelling pending results, query id [%s]", query.getId());
                futures.cancel(true);
                throw new QueryInterruptedException(e);
            } catch (CancellationException e) {
                throw new QueryInterruptedException(e);
            } catch (TimeoutException e) {
                log.info("Query timeout, cancelling pending results for query id [%s]", query.getId());
                futures.cancel(true);
                throw new QueryInterruptedException(e);
            } catch (ExecutionException e) {
                throw Throwables.propagate(e.getCause());
            }
        }

        @Override
        public void cleanup(Iterator<T> tIterator) {
        }
    });
}

Example 2 with MergeIterable

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

the class IndexMerger method merge.

public File merge(List<IndexableAdapter> indexes, final boolean rollup, final AggregatorFactory[] metricAggs, File outDir, IndexSpec indexSpec, ProgressIndicator progress) throws IOException {
    FileUtils.deleteDirectory(outDir);
    FileUtils.forceMkdir(outDir);
    final List<String> mergedDimensions = getMergedDimensions(indexes);
    final List<String> mergedMetrics = Lists.transform(mergeIndexed(Lists.newArrayList(FunctionalIterable.create(indexes).transform(new Function<IndexableAdapter, Iterable<String>>() {

        @Override
        public Iterable<String> apply(@Nullable IndexableAdapter input) {
            return input.getMetricNames();
        }
    }))), new Function<String, String>() {

        @Override
        public String apply(@Nullable String input) {
            return input;
        }
    });
    final AggregatorFactory[] sortedMetricAggs = new AggregatorFactory[mergedMetrics.size()];
    for (int i = 0; i < metricAggs.length; i++) {
        AggregatorFactory metricAgg = metricAggs[i];
        int metricIndex = mergedMetrics.indexOf(metricAgg.getName());
        /*
        If metricIndex is negative, one of the metricAggs was not present in the union of metrics from the indices
        we are merging
       */
        if (metricIndex > -1) {
            sortedMetricAggs[metricIndex] = metricAgg;
        }
    }
    /*
      If there is nothing at sortedMetricAggs[i], then we did not have a metricAgg whose name matched the name
      of the ith element of mergedMetrics. I.e. There was a metric in the indices to merge that we did not ask for.
     */
    for (int i = 0; i < sortedMetricAggs.length; i++) {
        if (sortedMetricAggs[i] == null) {
            throw new IAE("Indices to merge contained metric[%s], but requested metrics did not", mergedMetrics.get(i));
        }
    }
    for (int i = 0; i < mergedMetrics.size(); i++) {
        if (!sortedMetricAggs[i].getName().equals(mergedMetrics.get(i))) {
            throw new IAE("Metric mismatch, index[%d] [%s] != [%s]", i, sortedMetricAggs[i].getName(), mergedMetrics.get(i));
        }
    }
    Function<ArrayList<Iterable<Rowboat>>, Iterable<Rowboat>> rowMergerFn = new Function<ArrayList<Iterable<Rowboat>>, Iterable<Rowboat>>() {

        @Override
        public Iterable<Rowboat> apply(@Nullable ArrayList<Iterable<Rowboat>> boats) {
            if (rollup) {
                return CombiningIterable.create(new MergeIterable<Rowboat>(Ordering.<Rowboat>natural().nullsFirst(), boats), Ordering.<Rowboat>natural().nullsFirst(), new RowboatMergeFunction(sortedMetricAggs));
            } else {
                return new MergeIterable<Rowboat>(new Ordering<Rowboat>() {

                    @Override
                    public int compare(Rowboat left, Rowboat right) {
                        return Longs.compare(left.getTimestamp(), right.getTimestamp());
                    }
                }.nullsFirst(), boats);
            }
        }
    };
    return makeIndexFiles(indexes, sortedMetricAggs, outDir, progress, mergedDimensions, mergedMetrics, rowMergerFn, indexSpec);
}

Example 3 with MergeIterable

use of io.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 = Lists.newArrayList();
        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 = Lists.newArrayList();
        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 Query> capture = new Capture();
            final Capture<? extends Map> context = new Capture();
            queryCaptures.add(capture);
            QueryRunner queryable = expectations.getQueryRunner();
            if (query instanceof TimeseriesQuery) {
                List<String> segmentIds = Lists.newArrayList();
                List<Interval> intervals = Lists.newArrayList();
                List<Iterable<Result<TimeseriesResultValue>>> results = Lists.newArrayList();
                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<String> segmentIds = Lists.newArrayList();
                List<Interval> intervals = Lists.newArrayList();
                List<Iterable<Result<TopNResultValue>>> results = Lists.newArrayList();
                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<String> segmentIds = Lists.newArrayList();
                List<Interval> intervals = Lists.newArrayList();
                List<Iterable<Result<SearchResultValue>>> results = Lists.newArrayList();
                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 SelectQuery) {
                List<String> segmentIds = Lists.newArrayList();
                List<Interval> intervals = Lists.newArrayList();
                List<Iterable<Result<SelectResultValue>>> results = Lists.newArrayList();
                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(toQueryableSelectResults(segmentIds, intervals, results)).once();
            } else if (query instanceof GroupByQuery) {
                List<String> segmentIds = Lists.newArrayList();
                List<Interval> intervals = Lists.newArrayList();
                List<Iterable<Row>> results = Lists.newArrayList();
                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(segmentIds, intervals, results)).once();
            } else if (query instanceof TimeBoundaryQuery) {
                List<String> segmentIds = Lists.newArrayList();
                List<Interval> intervals = Lists.newArrayList();
                List<Iterable<Result<TimeBoundaryResultValue>>> results = Lists.newArrayList();
                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() {
                HashMap<String, List> context = new HashMap<String, List>();
                for (int i = 0; i < numTimesToQuery; ++i) {
                    TestHelper.assertExpectedResults(new MergeIterable<>(Ordering.<Result<Object>>natural().nullsFirst(), 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 = Lists.newArrayList();
                            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(query.withQuerySegmentSpec(new MultipleIntervalSegmentSpec(ImmutableList.of(actualQueryInterval))), context));
                    if (queryCompletedCallback != null) {
                        queryCompletedCallback.run();
                    }
                }
            }
        }, mocks.toArray());
        // make sure all the queries were sent down as 'bySegment'
        for (Capture queryCapture : queryCaptures) {
            Query capturedQuery = (Query) queryCapture.getValue();
            if (expectBySegment) {
                Assert.assertEquals(true, capturedQuery.getContextValue("bySegment"));
            } else {
                Assert.assertTrue(capturedQuery.getContextValue("bySegment") == null || capturedQuery.getContextValue("bySegment").equals(false));
            }
        }
    }
}