Examples with SourceLookup org.opensearch.search.lookup.SourceLookup used on opensource projects

Example 6 with SourceLookup

use of org.opensearch.search.lookup.SourceLookup in project OpenSearch by opensearch-project.

the class PercolatorMatchedSlotSubFetchPhaseTests method testHitsExecute.

public void testHitsExecute() throws Exception {
    try (Directory directory = newDirectory()) {
        // Need a one doc index:
        try (RandomIndexWriter indexWriter = new RandomIndexWriter(random(), directory)) {
            Document document = new Document();
            indexWriter.addDocument(document);
        }
        PercolatorMatchedSlotSubFetchPhase phase = new PercolatorMatchedSlotSubFetchPhase();
        try (DirectoryReader reader = DirectoryReader.open(directory)) {
            LeafReaderContext context = reader.leaves().get(0);
            // A match:
            {
                HitContext hit = new HitContext(new SearchHit(0), context, 0, new SourceLookup());
                PercolateQuery.QueryStore queryStore = ctx -> docId -> new TermQuery(new Term("field", "value"));
                MemoryIndex memoryIndex = new MemoryIndex();
                memoryIndex.addField("field", "value", new WhitespaceAnalyzer());
                memoryIndex.addField(new NumericDocValuesField(SeqNoFieldMapper.PRIMARY_TERM_NAME, 0), null);
                PercolateQuery percolateQuery = new PercolateQuery("_name", queryStore, Collections.emptyList(), new MatchAllDocsQuery(), memoryIndex.createSearcher(), null, new MatchNoDocsQuery());
                FetchContext sc = mock(FetchContext.class);
                when(sc.query()).thenReturn(percolateQuery);
                FetchSubPhaseProcessor processor = phase.getProcessor(sc);
                assertNotNull(processor);
                processor.process(hit);
                assertNotNull(hit.hit().field(PercolatorMatchedSlotSubFetchPhase.FIELD_NAME_PREFIX));
                assertEquals(0, (int) hit.hit().field(PercolatorMatchedSlotSubFetchPhase.FIELD_NAME_PREFIX).getValue());
            }
            // No match:
            {
                HitContext hit = new HitContext(new SearchHit(0), context, 0, new SourceLookup());
                PercolateQuery.QueryStore queryStore = ctx -> docId -> new TermQuery(new Term("field", "value"));
                MemoryIndex memoryIndex = new MemoryIndex();
                memoryIndex.addField("field", "value1", new WhitespaceAnalyzer());
                memoryIndex.addField(new NumericDocValuesField(SeqNoFieldMapper.PRIMARY_TERM_NAME, 0), null);
                PercolateQuery percolateQuery = new PercolateQuery("_name", queryStore, Collections.emptyList(), new MatchAllDocsQuery(), memoryIndex.createSearcher(), null, new MatchNoDocsQuery());
                FetchContext sc = mock(FetchContext.class);
                when(sc.query()).thenReturn(percolateQuery);
                FetchSubPhaseProcessor processor = phase.getProcessor(sc);
                assertNotNull(processor);
                processor.process(hit);
                assertNull(hit.hit().field(PercolatorMatchedSlotSubFetchPhase.FIELD_NAME_PREFIX));
            }
            // No query:
            {
                HitContext hit = new HitContext(new SearchHit(0), context, 0, new SourceLookup());
                PercolateQuery.QueryStore queryStore = ctx -> docId -> null;
                MemoryIndex memoryIndex = new MemoryIndex();
                memoryIndex.addField("field", "value", new WhitespaceAnalyzer());
                memoryIndex.addField(new NumericDocValuesField(SeqNoFieldMapper.PRIMARY_TERM_NAME, 0), null);
                PercolateQuery percolateQuery = new PercolateQuery("_name", queryStore, Collections.emptyList(), new MatchAllDocsQuery(), memoryIndex.createSearcher(), null, new MatchNoDocsQuery());
                FetchContext sc = mock(FetchContext.class);
                when(sc.query()).thenReturn(percolateQuery);
                FetchSubPhaseProcessor processor = phase.getProcessor(sc);
                assertNotNull(processor);
                processor.process(hit);
                assertNull(hit.hit().field(PercolatorMatchedSlotSubFetchPhase.FIELD_NAME_PREFIX));
            }
        }
    }
}

Example 7 with SourceLookup

use of org.opensearch.search.lookup.SourceLookup 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()));
                            }
                        }
                    }
                }
            }
        }
    };
}

Example 8 with SourceLookup

use of org.opensearch.search.lookup.SourceLookup in project OpenSearch by opensearch-project.

the class UpdateHelper method extractGetResult.

/**
 * Applies {@link UpdateRequest#fetchSource()} to the _source of the updated document to be returned in a update response.
 */
public static GetResult extractGetResult(final UpdateRequest request, String concreteIndex, long seqNo, long primaryTerm, long version, final Map<String, Object> source, XContentType sourceContentType, @Nullable final BytesReference sourceAsBytes) {
    if (request.fetchSource() == null || request.fetchSource().fetchSource() == false) {
        return null;
    }
    BytesReference sourceFilteredAsBytes = sourceAsBytes;
    if (request.fetchSource().includes().length > 0 || request.fetchSource().excludes().length > 0) {
        SourceLookup sourceLookup = new SourceLookup();
        sourceLookup.setSource(source);
        Object value = sourceLookup.filter(request.fetchSource());
        try {
            final int initialCapacity = Math.min(1024, sourceAsBytes.length());
            BytesStreamOutput streamOutput = new BytesStreamOutput(initialCapacity);
            try (XContentBuilder builder = new XContentBuilder(sourceContentType.xContent(), streamOutput)) {
                builder.value(value);
                sourceFilteredAsBytes = BytesReference.bytes(builder);
            }
        } catch (IOException e) {
            throw new OpenSearchException("Error filtering source", e);
        }
    }
    // TODO when using delete/none, we can still return the source as bytes by generating it (using the sourceContentType)
    return new GetResult(concreteIndex, request.id(), seqNo, primaryTerm, version, true, sourceFilteredAsBytes, Collections.emptyMap(), Collections.emptyMap());
}

Example 9 with SourceLookup

use of org.opensearch.search.lookup.SourceLookup in project OpenSearch by opensearch-project.

the class FetchFieldsPhase method getProcessor.

@Override
public FetchSubPhaseProcessor getProcessor(FetchContext fetchContext) {
    FetchFieldsContext fetchFieldsContext = fetchContext.fetchFieldsContext();
    if (fetchFieldsContext == null) {
        return null;
    }
    SearchLookup searchLookup = fetchContext.searchLookup();
    if (fetchContext.mapperService().documentMapper().sourceMapper().enabled() == false) {
        throw new IllegalArgumentException("Unable to retrieve the requested [fields] since _source is disabled " + "in the mappings for index [" + fetchContext.getIndexName() + "]");
    }
    FieldFetcher fieldFetcher = FieldFetcher.create(fetchContext.getQueryShardContext(), searchLookup, fetchFieldsContext.fields());
    return new FetchSubPhaseProcessor() {

        @Override
        public void setNextReader(LeafReaderContext readerContext) {
            fieldFetcher.setNextReader(readerContext);
        }

        @Override
        public void process(HitContext hitContext) throws IOException {
            SearchHit hit = hitContext.hit();
            SourceLookup sourceLookup = hitContext.sourceLookup();
            Set<String> ignoredFields = getIgnoredFields(hit);
            Map<String, DocumentField> documentFields = fieldFetcher.fetch(sourceLookup, ignoredFields);
            for (Map.Entry<String, DocumentField> entry : documentFields.entrySet()) {
                hit.setDocumentField(entry.getKey(), entry.getValue());
            }
        }
    };
}

Example 10 with SourceLookup

use of org.opensearch.search.lookup.SourceLookup in project OpenSearch by opensearch-project.

the class InnerHitsPhase method getProcessor.

@Override
public FetchSubPhaseProcessor getProcessor(FetchContext searchContext) {
    if (searchContext.innerHits() == null) {
        return null;
    }
    Map<String, InnerHitsContext.InnerHitSubContext> innerHits = searchContext.innerHits().getInnerHits();
    return new FetchSubPhaseProcessor() {

        @Override
        public void setNextReader(LeafReaderContext readerContext) {
        }

        @Override
        public void process(HitContext hitContext) throws IOException {
            SearchHit hit = hitContext.hit();
            SourceLookup rootLookup = searchContext.getRootSourceLookup(hitContext);
            hitExecute(innerHits, hit, rootLookup);
        }
    };
}