Examples with MultiBucketConsumer org.opensearch.search.aggregations.MultiBucketConsumerService.MultiBucketConsumer used on opensource projects

Example 1 with MultiBucketConsumer

use of org.opensearch.search.aggregations.MultiBucketConsumerService.MultiBucketConsumer in project OpenSearch by opensearch-project.

the class InternalAggregationTestCase method testReduceRandom.

public void testReduceRandom() throws IOException {
    String name = randomAlphaOfLength(5);
    int size = between(1, 200);
    List<T> inputs = randomResultsToReduce(name, size);
    assertThat(inputs, hasSize(size));
    List<InternalAggregation> toReduce = new ArrayList<>();
    toReduce.addAll(inputs);
    // Sort aggs so that unmapped come last. This mimicks the behavior of InternalAggregations.reduce()
    inputs.sort(INTERNAL_AGG_COMPARATOR);
    ScriptService mockScriptService = mockScriptService();
    MockBigArrays bigArrays = new MockBigArrays(new MockPageCacheRecycler(Settings.EMPTY), new NoneCircuitBreakerService());
    if (randomBoolean() && toReduce.size() > 1) {
        // sometimes do a partial reduce
        Collections.shuffle(toReduce, random());
        int r = randomIntBetween(1, inputs.size());
        List<InternalAggregation> toPartialReduce = toReduce.subList(0, r);
        // Sort aggs so that unmapped come last. This mimicks the behavior of InternalAggregations.reduce()
        toPartialReduce.sort(INTERNAL_AGG_COMPARATOR);
        InternalAggregation.ReduceContext context = InternalAggregation.ReduceContext.forPartialReduction(bigArrays, mockScriptService, () -> PipelineAggregator.PipelineTree.EMPTY);
        @SuppressWarnings("unchecked") T reduced = (T) toPartialReduce.get(0).reduce(toPartialReduce, context);
        int initialBucketCount = 0;
        for (InternalAggregation internalAggregation : toPartialReduce) {
            initialBucketCount += countInnerBucket(internalAggregation);
        }
        int reducedBucketCount = countInnerBucket(reduced);
        // check that non final reduction never adds buckets
        assertThat(reducedBucketCount, lessThanOrEqualTo(initialBucketCount));
        /*
             * Sometimes serializing and deserializing the partially reduced
             * result to simulate the compaction that we attempt after a
             * partial reduce. And to simulate cross cluster search.
             */
        if (randomBoolean()) {
            reduced = copyNamedWriteable(reduced, getNamedWriteableRegistry(), categoryClass());
        }
        toReduce = new ArrayList<>(toReduce.subList(r, inputs.size()));
        toReduce.add(reduced);
    }
    MultiBucketConsumer bucketConsumer = new MultiBucketConsumer(DEFAULT_MAX_BUCKETS, new NoneCircuitBreakerService().getBreaker(CircuitBreaker.REQUEST));
    InternalAggregation.ReduceContext context = InternalAggregation.ReduceContext.forFinalReduction(bigArrays, mockScriptService, bucketConsumer, PipelineTree.EMPTY);
    @SuppressWarnings("unchecked") T reduced = (T) inputs.get(0).reduce(toReduce, context);
    doAssertReducedMultiBucketConsumer(reduced, bucketConsumer);
    assertReduced(reduced, inputs);
}

Example 2 with MultiBucketConsumer

use of org.opensearch.search.aggregations.MultiBucketConsumerService.MultiBucketConsumer in project OpenSearch by opensearch-project.

the class AggregatorTestCase method searchAndReduce.

/**
 * Collects all documents that match the provided query {@link Query} and
 * returns the reduced {@link InternalAggregation}.
 * <p>
 * Half the time it aggregates each leaf individually and reduces all
 * results together. The other half the time it aggregates across the entire
 * index at once and runs a final reduction on the single resulting agg.
 */
protected <A extends InternalAggregation, C extends Aggregator> A searchAndReduce(IndexSettings indexSettings, IndexSearcher searcher, Query query, AggregationBuilder builder, int maxBucket, MappedFieldType... fieldTypes) throws IOException {
    final IndexReaderContext ctx = searcher.getTopReaderContext();
    final PipelineTree pipelines = builder.buildPipelineTree();
    List<InternalAggregation> aggs = new ArrayList<>();
    Query rewritten = searcher.rewrite(query);
    MultiBucketConsumer bucketConsumer = new MultiBucketConsumer(maxBucket, new NoneCircuitBreakerService().getBreaker(CircuitBreaker.REQUEST));
    C root = createAggregator(query, builder, searcher, bucketConsumer, fieldTypes);
    if (randomBoolean() && searcher.getIndexReader().leaves().size() > 0) {
        assertThat(ctx, instanceOf(CompositeReaderContext.class));
        final CompositeReaderContext compCTX = (CompositeReaderContext) ctx;
        final int size = compCTX.leaves().size();
        final ShardSearcher[] subSearchers = new ShardSearcher[size];
        for (int searcherIDX = 0; searcherIDX < subSearchers.length; searcherIDX++) {
            final LeafReaderContext leave = compCTX.leaves().get(searcherIDX);
            subSearchers[searcherIDX] = new ShardSearcher(leave, compCTX);
        }
        for (ShardSearcher subSearcher : subSearchers) {
            MultiBucketConsumer shardBucketConsumer = new MultiBucketConsumer(maxBucket, new NoneCircuitBreakerService().getBreaker(CircuitBreaker.REQUEST));
            C a = createAggregator(query, builder, subSearcher, indexSettings, shardBucketConsumer, fieldTypes);
            a.preCollection();
            Weight weight = subSearcher.createWeight(rewritten, ScoreMode.COMPLETE, 1f);
            subSearcher.search(weight, a);
            a.postCollection();
            aggs.add(a.buildTopLevel());
        }
    } else {
        root.preCollection();
        searcher.search(rewritten, root);
        root.postCollection();
        aggs.add(root.buildTopLevel());
    }
    if (randomBoolean() && aggs.size() > 1) {
        // sometimes do an incremental reduce
        int toReduceSize = aggs.size();
        Collections.shuffle(aggs, random());
        int r = randomIntBetween(1, toReduceSize);
        List<InternalAggregation> toReduce = aggs.subList(0, r);
        InternalAggregation.ReduceContext context = InternalAggregation.ReduceContext.forPartialReduction(root.context().bigArrays(), getMockScriptService(), () -> PipelineAggregator.PipelineTree.EMPTY);
        A reduced = (A) aggs.get(0).reduce(toReduce, context);
        aggs = new ArrayList<>(aggs.subList(r, toReduceSize));
        aggs.add(reduced);
    }
    // now do the final reduce
    MultiBucketConsumer reduceBucketConsumer = new MultiBucketConsumer(maxBucket, new NoneCircuitBreakerService().getBreaker(CircuitBreaker.REQUEST));
    InternalAggregation.ReduceContext context = InternalAggregation.ReduceContext.forFinalReduction(root.context().bigArrays(), getMockScriptService(), reduceBucketConsumer, pipelines);
    @SuppressWarnings("unchecked") A internalAgg = (A) aggs.get(0).reduce(aggs, context);
    // materialize any parent pipelines
    internalAgg = (A) internalAgg.reducePipelines(internalAgg, context, pipelines);
    // materialize any sibling pipelines at top level
    for (PipelineAggregator pipelineAggregator : pipelines.aggregators()) {
        internalAgg = (A) pipelineAggregator.reduce(internalAgg, context);
    }
    doAssertReducedMultiBucketConsumer(internalAgg, reduceBucketConsumer);
    return internalAgg;
}