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

Example 1 with LeafSearchLookup

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

the class ScriptedMetricAggContextsTests method testReturnSource.

public void testReturnSource() throws IOException {
    ScriptedMetricAggContexts.MapScript.Factory factory = getEngine().compile("test", "state._source = params._source", ScriptedMetricAggContexts.MapScript.CONTEXT, Collections.emptyMap());
    Map<String, Object> params = new HashMap<>();
    Map<String, Object> state = new HashMap<>();
    MemoryIndex index = new MemoryIndex();
    // we don't need a real index, just need to construct a LeafReaderContext which cannot be mocked
    LeafReaderContext leafReaderContext = index.createSearcher().getIndexReader().leaves().get(0);
    SearchLookup lookup = mock(SearchLookup.class);
    LeafSearchLookup leafLookup = mock(LeafSearchLookup.class);
    when(lookup.getLeafSearchLookup(leafReaderContext)).thenReturn(leafLookup);
    SourceLookup sourceLookup = mock(SourceLookup.class);
    when(leafLookup.asMap()).thenReturn(Collections.singletonMap("_source", sourceLookup));
    when(sourceLookup.loadSourceIfNeeded()).thenReturn(Collections.singletonMap("test", 1));
    ScriptedMetricAggContexts.MapScript.LeafFactory leafFactory = factory.newFactory(params, state, lookup);
    ScriptedMetricAggContexts.MapScript script = leafFactory.newInstance(leafReaderContext);
    script.execute();
    assertTrue(state.containsKey("_source"));
    assertTrue(state.get("_source") instanceof Map && ((Map) state.get("_source")).containsKey("test"));
    assertEquals(1, ((Map) state.get("_source")).get("test"));
}

Example 2 with LeafSearchLookup

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

the class ScriptedMetricAggContextsTests method testMapSourceAccess.

public void testMapSourceAccess() throws IOException {
    ScriptedMetricAggContexts.MapScript.Factory factory = getEngine().compile("test", "state.testField = params._source.three", ScriptedMetricAggContexts.MapScript.CONTEXT, Collections.emptyMap());
    Map<String, Object> params = new HashMap<>();
    Map<String, Object> state = new HashMap<>();
    MemoryIndex index = new MemoryIndex();
    // we don't need a real index, just need to construct a LeafReaderContext which cannot be mocked
    LeafReaderContext leafReaderContext = index.createSearcher().getIndexReader().leaves().get(0);
    SearchLookup lookup = mock(SearchLookup.class);
    LeafSearchLookup leafLookup = mock(LeafSearchLookup.class);
    when(lookup.getLeafSearchLookup(leafReaderContext)).thenReturn(leafLookup);
    SourceLookup sourceLookup = mock(SourceLookup.class);
    when(leafLookup.asMap()).thenReturn(Collections.singletonMap("_source", sourceLookup));
    when(sourceLookup.loadSourceIfNeeded()).thenReturn(Collections.singletonMap("three", 3));
    ScriptedMetricAggContexts.MapScript.LeafFactory leafFactory = factory.newFactory(params, state, lookup);
    ScriptedMetricAggContexts.MapScript script = leafFactory.newInstance(leafReaderContext);
    script.execute();
    assertTrue(state.containsKey("testField"));
    assertEquals(3, state.get("testField"));
}

Example 3 with LeafSearchLookup

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

Example 4 with LeafSearchLookup

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

the class QueryShardContextTests method mockFieldType.

private static MappedFieldType mockFieldType(String fieldName, BiFunction<LeafSearchLookup, Integer, String> runtimeDocValues) {
    MappedFieldType fieldType = mock(MappedFieldType.class);
    when(fieldType.name()).thenReturn(fieldName);
    when(fieldType.fielddataBuilder(any(), any())).thenAnswer(builderInv -> {
        @SuppressWarnings("unchecked") Supplier<SearchLookup> searchLookup = ((Supplier<SearchLookup>) builderInv.getArguments()[1]);
        IndexFieldData<?> indexFieldData = mock(IndexFieldData.class);
        when(indexFieldData.load(any())).thenAnswer(loadArgs -> {
            LeafReaderContext leafReaderContext = (LeafReaderContext) loadArgs.getArguments()[0];
            LeafFieldData leafFieldData = mock(LeafFieldData.class);
            when(leafFieldData.getScriptValues()).thenAnswer(scriptValuesArgs -> new ScriptDocValues<String>() {

                String value;

                @Override
                public int size() {
                    return 1;
                }

                @Override
                public String get(int index) {
                    assert index == 0;
                    return value;
                }

                @Override
                public void setNextDocId(int docId) {
                    assert docId >= 0;
                    LeafSearchLookup leafLookup = searchLookup.get().getLeafSearchLookup(leafReaderContext);
                    leafLookup.setDocument(docId);
                    value = runtimeDocValues.apply(leafLookup, docId);
                }
            });
            return leafFieldData;
        });
        IndexFieldData.Builder builder = mock(IndexFieldData.Builder.class);
        when(builder.build(any(), any())).thenAnswer(buildInv -> indexFieldData);
        return builder;
    });
    return fieldType;
}