Examples with LeafBucketCollector org.opensearch.search.aggregations.LeafBucketCollector used on opensource projects

Example 16 with LeafBucketCollector

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

the class DoubleValuesSource method getLeafCollector.

@Override
LeafBucketCollector getLeafCollector(LeafReaderContext context, LeafBucketCollector next) throws IOException {
    final SortedNumericDoubleValues dvs = docValuesFunc.apply(context);
    return new LeafBucketCollector() {

        @Override
        public void collect(int doc, long bucket) throws IOException {
            if (dvs.advanceExact(doc)) {
                int num = dvs.docValueCount();
                for (int i = 0; i < num; i++) {
                    currentValue = dvs.nextValue();
                    missingCurrentValue = false;
                    next.collect(doc, bucket);
                }
            } else if (missingBucket) {
                missingCurrentValue = true;
                next.collect(doc, bucket);
            }
        }
    };
}

Example 17 with LeafBucketCollector

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

the class ContextIndexSearcherTests method testIntersectScorerAndRoleBits.

public void testIntersectScorerAndRoleBits() throws Exception {
    final Directory directory = newDirectory();
    IndexWriter iw = new IndexWriter(directory, new IndexWriterConfig(new StandardAnalyzer()).setMergePolicy(NoMergePolicy.INSTANCE));
    Document document = new Document();
    document.add(new StringField("field1", "value1", Field.Store.NO));
    document.add(new StringField("field2", "value1", Field.Store.NO));
    iw.addDocument(document);
    document = new Document();
    document.add(new StringField("field1", "value2", Field.Store.NO));
    document.add(new StringField("field2", "value1", Field.Store.NO));
    iw.addDocument(document);
    document = new Document();
    document.add(new StringField("field1", "value3", Field.Store.NO));
    document.add(new StringField("field2", "value1", Field.Store.NO));
    iw.addDocument(document);
    document = new Document();
    document.add(new StringField("field1", "value4", Field.Store.NO));
    document.add(new StringField("field2", "value1", Field.Store.NO));
    iw.addDocument(document);
    iw.commit();
    iw.deleteDocuments(new Term("field1", "value3"));
    iw.close();
    DirectoryReader directoryReader = DirectoryReader.open(directory);
    IndexSearcher searcher = new IndexSearcher(directoryReader);
    Weight weight = searcher.createWeight(new BoostQuery(new ConstantScoreQuery(new TermQuery(new Term("field2", "value1"))), 3f), ScoreMode.COMPLETE, 1f);
    LeafReaderContext leaf = directoryReader.leaves().get(0);
    CombinedBitSet bitSet = new CombinedBitSet(query(leaf, "field1", "value1"), leaf.reader().getLiveDocs());
    LeafCollector leafCollector = new LeafBucketCollector() {

        Scorable scorer;

        @Override
        public void setScorer(Scorable scorer) throws IOException {
            this.scorer = scorer;
        }

        @Override
        public void collect(int doc, long bucket) throws IOException {
            assertThat(doc, equalTo(0));
            assertThat(scorer.score(), equalTo(3f));
        }
    };
    intersectScorerAndBitSet(weight.scorer(leaf), bitSet, leafCollector, () -> {
    });
    bitSet = new CombinedBitSet(query(leaf, "field1", "value2"), leaf.reader().getLiveDocs());
    leafCollector = new LeafBucketCollector() {

        @Override
        public void collect(int doc, long bucket) throws IOException {
            assertThat(doc, equalTo(1));
        }
    };
    intersectScorerAndBitSet(weight.scorer(leaf), bitSet, leafCollector, () -> {
    });
    bitSet = new CombinedBitSet(query(leaf, "field1", "value3"), leaf.reader().getLiveDocs());
    leafCollector = new LeafBucketCollector() {

        @Override
        public void collect(int doc, long bucket) throws IOException {
            fail("docId [" + doc + "] should have been deleted");
        }
    };
    intersectScorerAndBitSet(weight.scorer(leaf), bitSet, leafCollector, () -> {
    });
    bitSet = new CombinedBitSet(query(leaf, "field1", "value4"), leaf.reader().getLiveDocs());
    leafCollector = new LeafBucketCollector() {

        @Override
        public void collect(int doc, long bucket) throws IOException {
            assertThat(doc, equalTo(3));
        }
    };
    intersectScorerAndBitSet(weight.scorer(leaf), bitSet, leafCollector, () -> {
    });
    directoryReader.close();
    directory.close();
}

Example 18 with LeafBucketCollector

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

the class MergingBucketsDeferringCollectorTests method testBucketMergeAndDeleteLastEntry.

@AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/60021")
public void testBucketMergeAndDeleteLastEntry() throws Exception {
    testCase((deferringCollector, delegate) -> new LeafBucketCollector() {

        @Override
        public void collect(int doc, long owningBucketOrd) throws IOException {
            // Only collects at top level
            assert owningBucketOrd == 0;
            delegate.collect(doc, doc);
            if (doc == 7) {
                deferringCollector.mergeBuckets(oldBucket -> oldBucket <= 3 ? 0 : -1);
            }
        }
    }, (deferringCollector, finalCollector) -> {
        deferringCollector.prepareSelectedBuckets(0, 8, 9);
        assertThat(finalCollector.collection, equalTo(org.opensearch.common.collect.Map.of(0L, org.opensearch.common.collect.List.of(0, 1, 2, 3), 1L, org.opensearch.common.collect.List.of(8), 2L, org.opensearch.common.collect.List.of(9))));
    });
}

Example 19 with LeafBucketCollector

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

the class MergingBucketsDeferringCollectorTests method testBucketMergeNoDelete.

public void testBucketMergeNoDelete() throws Exception {
    testCase((deferringCollector, delegate) -> new LeafBucketCollector() {

        @Override
        public void collect(int doc, long owningBucketOrd) throws IOException {
            // Only collects at top level
            assert owningBucketOrd == 0;
            delegate.collect(doc, doc);
            if (doc == 7) {
                deferringCollector.mergeBuckets(oldBucket -> 0);
            }
        }
    }, (deferringCollector, finalCollector) -> {
        deferringCollector.prepareSelectedBuckets(0, 8, 9);
        equalTo(org.opensearch.common.collect.Map.of(0L, org.opensearch.common.collect.List.of(0, 1, 2, 3, 4, 5, 6, 7), 1L, org.opensearch.common.collect.List.of(8), 2L, org.opensearch.common.collect.List.of(9)));
    });
}

Example 20 with LeafBucketCollector

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

the class MergingBucketsDeferringCollectorTests method testCase.

private void testCase(BiFunction<MergingBucketsDeferringCollector, LeafBucketCollector, LeafBucketCollector> leafCollector, CheckedBiConsumer<MergingBucketsDeferringCollector, CollectingBucketCollector, IOException> verify) throws IOException {
    try (Directory directory = newDirectory()) {
        try (IndexWriter indexWriter = new IndexWriter(directory, new IndexWriterConfig())) {
            for (int i = 0; i < 10; i++) {
                indexWriter.addDocument(new Document());
            }
        }
        try (IndexReader indexReader = DirectoryReader.open(directory)) {
            IndexSearcher indexSearcher = new IndexSearcher(indexReader);
            Query query = new MatchAllDocsQuery();
            SearchContext searchContext = createSearchContext(indexSearcher, createIndexSettings(), query, null);
            when(searchContext.query()).thenReturn(query);
            MergingBucketsDeferringCollector deferringCollector = new MergingBucketsDeferringCollector(searchContext, false);
            CollectingBucketCollector finalCollector = new CollectingBucketCollector();
            deferringCollector.setDeferredCollector(Collections.singleton(finalCollector));
            deferringCollector.preCollection();
            indexSearcher.search(query, new BucketCollector() {

                @Override
                public ScoreMode scoreMode() {
                    return ScoreMode.COMPLETE_NO_SCORES;
                }

                @Override
                public void preCollection() throws IOException {
                }

                @Override
                public void postCollection() throws IOException {
                }

                @Override
                public LeafBucketCollector getLeafCollector(LeafReaderContext ctx) throws IOException {
                    LeafBucketCollector delegate = deferringCollector.getLeafCollector(ctx);
                    return leafCollector.apply(deferringCollector, delegate);
                }
            });
            deferringCollector.postCollection();
            verify.accept(deferringCollector, finalCollector);
        }
    }
}