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

Example 6 with SearchLookup

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

the class ExpressionFieldScriptTests method setUp.

@Override
public void setUp() throws Exception {
    super.setUp();
    NumberFieldMapper.NumberFieldType fieldType = new NumberFieldMapper.NumberFieldType("field", NumberFieldMapper.NumberType.DOUBLE);
    MapperService mapperService = mock(MapperService.class);
    when(mapperService.fieldType("field")).thenReturn(fieldType);
    when(mapperService.fieldType("alias")).thenReturn(fieldType);
    SortedNumericDoubleValues doubleValues = mock(SortedNumericDoubleValues.class);
    when(doubleValues.advanceExact(anyInt())).thenReturn(true);
    when(doubleValues.nextValue()).thenReturn(2.718);
    LeafNumericFieldData atomicFieldData = mock(LeafNumericFieldData.class);
    when(atomicFieldData.getDoubleValues()).thenReturn(doubleValues);
    IndexNumericFieldData fieldData = mock(IndexNumericFieldData.class);
    when(fieldData.getFieldName()).thenReturn("field");
    when(fieldData.load(any())).thenReturn(atomicFieldData);
    service = new ExpressionScriptEngine();
    lookup = new SearchLookup(mapperService, (ignored, lookup) -> fieldData);
}

Example 7 with SearchLookup

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

the class ExpressionTermsSetQueryTests method setUp.

@Override
public void setUp() throws Exception {
    super.setUp();
    NumberFieldType fieldType = new NumberFieldType("field", NumberType.DOUBLE);
    MapperService mapperService = mock(MapperService.class);
    when(mapperService.fieldType("field")).thenReturn(fieldType);
    when(mapperService.fieldType("alias")).thenReturn(fieldType);
    SortedNumericDoubleValues doubleValues = mock(SortedNumericDoubleValues.class);
    when(doubleValues.advanceExact(anyInt())).thenReturn(true);
    when(doubleValues.nextValue()).thenReturn(2.718);
    LeafNumericFieldData atomicFieldData = mock(LeafNumericFieldData.class);
    when(atomicFieldData.getDoubleValues()).thenReturn(doubleValues);
    IndexNumericFieldData fieldData = mock(IndexNumericFieldData.class);
    when(fieldData.getFieldName()).thenReturn("field");
    when(fieldData.load(any())).thenReturn(atomicFieldData);
    service = new ExpressionScriptEngine();
    lookup = new SearchLookup(mapperService, (ignored, lookup) -> fieldData);
}

Example 8 with SearchLookup

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

the class SearchService method parseSource.

private void parseSource(DefaultSearchContext context, SearchSourceBuilder source, boolean includeAggregations) {
    // nothing to parse...
    if (source == null) {
        return;
    }
    SearchShardTarget shardTarget = context.shardTarget();
    QueryShardContext queryShardContext = context.getQueryShardContext();
    context.from(source.from());
    context.size(source.size());
    Map<String, InnerHitContextBuilder> innerHitBuilders = new HashMap<>();
    if (source.query() != null) {
        InnerHitContextBuilder.extractInnerHits(source.query(), innerHitBuilders);
        context.parsedQuery(queryShardContext.toQuery(source.query()));
    }
    if (source.postFilter() != null) {
        InnerHitContextBuilder.extractInnerHits(source.postFilter(), innerHitBuilders);
        context.parsedPostFilter(queryShardContext.toQuery(source.postFilter()));
    }
    if (innerHitBuilders.size() > 0) {
        for (Map.Entry<String, InnerHitContextBuilder> entry : innerHitBuilders.entrySet()) {
            try {
                entry.getValue().build(context, context.innerHits());
            } catch (IOException e) {
                throw new SearchException(shardTarget, "failed to build inner_hits", e);
            }
        }
    }
    if (source.sorts() != null) {
        try {
            Optional<SortAndFormats> optionalSort = SortBuilder.buildSort(source.sorts(), context.getQueryShardContext());
            if (optionalSort.isPresent()) {
                context.sort(optionalSort.get());
            }
        } catch (IOException e) {
            throw new SearchException(shardTarget, "failed to create sort elements", e);
        }
    }
    context.trackScores(source.trackScores());
    if (source.trackTotalHitsUpTo() != null && source.trackTotalHitsUpTo() != SearchContext.TRACK_TOTAL_HITS_ACCURATE && context.scrollContext() != null) {
        throw new SearchException(shardTarget, "disabling [track_total_hits] is not allowed in a scroll context");
    }
    if (source.trackTotalHitsUpTo() != null) {
        context.trackTotalHitsUpTo(source.trackTotalHitsUpTo());
    }
    if (source.minScore() != null) {
        context.minimumScore(source.minScore());
    }
    if (source.profile()) {
        context.setProfilers(new Profilers(context.searcher()));
    }
    if (source.timeout() != null) {
        context.timeout(source.timeout());
    }
    context.terminateAfter(source.terminateAfter());
    if (source.aggregations() != null && includeAggregations) {
        try {
            AggregatorFactories factories = source.aggregations().build(queryShardContext, null);
            context.aggregations(new SearchContextAggregations(factories, multiBucketConsumerService.create()));
        } catch (IOException e) {
            throw new AggregationInitializationException("Failed to create aggregators", e);
        }
    }
    if (source.suggest() != null) {
        try {
            context.suggest(source.suggest().build(queryShardContext));
        } catch (IOException e) {
            throw new SearchException(shardTarget, "failed to create SuggestionSearchContext", e);
        }
    }
    if (source.rescores() != null) {
        try {
            for (RescorerBuilder<?> rescore : source.rescores()) {
                context.addRescore(rescore.buildContext(queryShardContext));
            }
        } catch (IOException e) {
            throw new SearchException(shardTarget, "failed to create RescoreSearchContext", e);
        }
    }
    if (source.explain() != null) {
        context.explain(source.explain());
    }
    if (source.fetchSource() != null) {
        context.fetchSourceContext(source.fetchSource());
    }
    if (source.docValueFields() != null) {
        FetchDocValuesContext docValuesContext = FetchDocValuesContext.create(context.mapperService(), source.docValueFields());
        context.docValuesContext(docValuesContext);
    }
    if (source.fetchFields() != null) {
        FetchFieldsContext fetchFieldsContext = new FetchFieldsContext(source.fetchFields());
        context.fetchFieldsContext(fetchFieldsContext);
    }
    if (source.highlighter() != null) {
        HighlightBuilder highlightBuilder = source.highlighter();
        try {
            context.highlight(highlightBuilder.build(queryShardContext));
        } catch (IOException e) {
            throw new SearchException(shardTarget, "failed to create SearchContextHighlighter", e);
        }
    }
    if (source.scriptFields() != null && source.size() != 0) {
        int maxAllowedScriptFields = context.mapperService().getIndexSettings().getMaxScriptFields();
        if (source.scriptFields().size() > maxAllowedScriptFields) {
            throw new IllegalArgumentException("Trying to retrieve too many script_fields. Must be less than or equal to: [" + maxAllowedScriptFields + "] but was [" + source.scriptFields().size() + "]. This limit can be set by changing the [" + IndexSettings.MAX_SCRIPT_FIELDS_SETTING.getKey() + "] index level setting.");
        }
        for (org.opensearch.search.builder.SearchSourceBuilder.ScriptField field : source.scriptFields()) {
            FieldScript.Factory factory = scriptService.compile(field.script(), FieldScript.CONTEXT);
            SearchLookup lookup = context.getQueryShardContext().lookup();
            FieldScript.LeafFactory searchScript = factory.newFactory(field.script().getParams(), lookup);
            context.scriptFields().add(new ScriptField(field.fieldName(), searchScript, field.ignoreFailure()));
        }
    }
    if (source.ext() != null) {
        for (SearchExtBuilder searchExtBuilder : source.ext()) {
            context.addSearchExt(searchExtBuilder);
        }
    }
    if (source.version() != null) {
        context.version(source.version());
    }
    if (source.seqNoAndPrimaryTerm() != null) {
        context.seqNoAndPrimaryTerm(source.seqNoAndPrimaryTerm());
    }
    if (source.stats() != null) {
        context.groupStats(source.stats());
    }
    if (CollectionUtils.isEmpty(source.searchAfter()) == false) {
        if (context.scrollContext() != null) {
            throw new SearchException(shardTarget, "`search_after` cannot be used in a scroll context.");
        }
        if (context.from() > 0) {
            throw new SearchException(shardTarget, "`from` parameter must be set to 0 when `search_after` is used.");
        }
        FieldDoc fieldDoc = SearchAfterBuilder.buildFieldDoc(context.sort(), source.searchAfter());
        context.searchAfter(fieldDoc);
    }
    if (source.slice() != null) {
        if (context.scrollContext() == null) {
            throw new SearchException(shardTarget, "`slice` cannot be used outside of a scroll context");
        }
        context.sliceBuilder(source.slice());
    }
    if (source.storedFields() != null) {
        if (source.storedFields().fetchFields() == false) {
            if (context.sourceRequested()) {
                throw new SearchException(shardTarget, "[stored_fields] cannot be disabled if [_source] is requested");
            }
            if (context.fetchFieldsContext() != null) {
                throw new SearchException(shardTarget, "[stored_fields] cannot be disabled when using the [fields] option");
            }
        }
        context.storedFieldsContext(source.storedFields());
    }
    if (source.collapse() != null) {
        if (context.scrollContext() != null) {
            throw new SearchException(shardTarget, "cannot use `collapse` in a scroll context");
        }
        if (context.searchAfter() != null) {
            throw new SearchException(shardTarget, "cannot use `collapse` in conjunction with `search_after`");
        }
        if (context.rescore() != null && context.rescore().isEmpty() == false) {
            throw new SearchException(shardTarget, "cannot use `collapse` in conjunction with `rescore`");
        }
        final CollapseContext collapseContext = source.collapse().build(queryShardContext);
        context.collapse(collapseContext);
    }
}

Example 9 with SearchLookup

use of org.opensearch.search.lookup.SearchLookup 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 SearchLookup

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

the class ScriptScoreQueryTests method newFactory.

private ScoreScript.LeafFactory newFactory(Script script, boolean needsScore, Function<ScoreScript.ExplanationHolder, Double> function) {
    SearchLookup lookup = mock(SearchLookup.class);
    LeafSearchLookup leafLookup = mock(LeafSearchLookup.class);
    when(lookup.getLeafSearchLookup(any())).thenReturn(leafLookup);
    return new ScoreScript.LeafFactory() {

        @Override
        public boolean needs_score() {
            return needsScore;
        }

        @Override
        public ScoreScript newInstance(LeafReaderContext ctx) throws IOException {
            return new ScoreScript(script.getParams(), lookup, leafReaderContext) {

                @Override
                public double execute(ExplanationHolder explanation) {
                    return function.apply(explanation);
                }
            };
        }
    };
}