Examples with SearchHighlightContext org.opensearch.search.fetch.subphase.highlight.SearchHighlightContext used on opensource projects

Example 1 with SearchHighlightContext

use of org.opensearch.search.fetch.subphase.highlight.SearchHighlightContext in project OpenSearch by opensearch-project.

the class PercolatorHighlightSubFetchPhaseTests method testHitsExecutionNeeded.

public void testHitsExecutionNeeded() {
    PercolateQuery percolateQuery = new PercolateQuery("_name", ctx -> null, Collections.singletonList(new BytesArray("{}")), new MatchAllDocsQuery(), Mockito.mock(IndexSearcher.class), null, new MatchAllDocsQuery());
    PercolatorHighlightSubFetchPhase subFetchPhase = new PercolatorHighlightSubFetchPhase(emptyMap());
    FetchContext fetchContext = mock(FetchContext.class);
    Mockito.when(fetchContext.highlight()).thenReturn(new SearchHighlightContext(Collections.emptyList()));
    Mockito.when(fetchContext.query()).thenReturn(new MatchAllDocsQuery());
    assertNull(subFetchPhase.getProcessor(fetchContext));
    Mockito.when(fetchContext.query()).thenReturn(percolateQuery);
    assertNotNull(subFetchPhase.getProcessor(fetchContext));
}

Example 2 with SearchHighlightContext

use of org.opensearch.search.fetch.subphase.highlight.SearchHighlightContext in project OpenSearch by opensearch-project.

the class PercolatorHighlightSubFetchPhase method getProcessor.

@Override
public FetchSubPhaseProcessor getProcessor(FetchContext fetchContext) {
    if (fetchContext.highlight() == null) {
        return null;
    }
    List<PercolateQuery> percolateQueries = locatePercolatorQuery(fetchContext.query());
    if (percolateQueries.isEmpty()) {
        return null;
    }
    return new FetchSubPhaseProcessor() {

        LeafReaderContext ctx;

        @Override
        public void setNextReader(LeafReaderContext readerContext) {
            this.ctx = readerContext;
        }

        @Override
        public void process(HitContext hit) throws IOException {
            boolean singlePercolateQuery = percolateQueries.size() == 1;
            for (PercolateQuery percolateQuery : percolateQueries) {
                String fieldName = singlePercolateQuery ? PercolatorMatchedSlotSubFetchPhase.FIELD_NAME_PREFIX : PercolatorMatchedSlotSubFetchPhase.FIELD_NAME_PREFIX + "_" + percolateQuery.getName();
                IndexSearcher percolatorIndexSearcher = percolateQuery.getPercolatorIndexSearcher();
                PercolateQuery.QueryStore queryStore = percolateQuery.getQueryStore();
                LeafReaderContext percolatorLeafReaderContext = percolatorIndexSearcher.getIndexReader().leaves().get(0);
                final Query query = queryStore.getQueries(ctx).apply(hit.docId());
                if (query != null) {
                    DocumentField field = hit.hit().field(fieldName);
                    if (field == null) {
                        // so then continue highlighting with the next hit.
                        continue;
                    }
                    for (Object matchedSlot : field.getValues()) {
                        int slot = (int) matchedSlot;
                        BytesReference document = percolateQuery.getDocuments().get(slot);
                        HitContext subContext = new HitContext(new SearchHit(slot, "unknown", Collections.emptyMap(), Collections.emptyMap()), percolatorLeafReaderContext, slot, new SourceLookup());
                        subContext.sourceLookup().setSource(document);
                        // force source because MemoryIndex does not store fields
                        SearchHighlightContext highlight = new SearchHighlightContext(fetchContext.highlight().fields(), true);
                        FetchSubPhaseProcessor processor = highlightPhase.getProcessor(fetchContext, highlight, query);
                        processor.process(subContext);
                        for (Map.Entry<String, HighlightField> entry : subContext.hit().getHighlightFields().entrySet()) {
                            if (percolateQuery.getDocuments().size() == 1) {
                                String hlFieldName;
                                if (singlePercolateQuery) {
                                    hlFieldName = entry.getKey();
                                } else {
                                    hlFieldName = percolateQuery.getName() + "_" + entry.getKey();
                                }
                                hit.hit().getHighlightFields().put(hlFieldName, new HighlightField(hlFieldName, entry.getValue().fragments()));
                            } else {
                                // In case multiple documents are being percolated we need to identify to which document
                                // a highlight belongs to.
                                String hlFieldName;
                                if (singlePercolateQuery) {
                                    hlFieldName = slot + "_" + entry.getKey();
                                } else {
                                    hlFieldName = percolateQuery.getName() + "_" + slot + "_" + entry.getKey();
                                }
                                hit.hit().getHighlightFields().put(hlFieldName, new HighlightField(hlFieldName, entry.getValue().fragments()));
                            }
                        }
                    }
                }
            }
        }
    };
}