Examples with QueryShardContext org.opensearch.index.query.QueryShardContext used on opensource projects

Example 1 with QueryShardContext

use of org.opensearch.index.query.QueryShardContext in project OpenSearch by opensearch-project.

the class NeedsScoreTests method testNeedsScores.

public void testNeedsScores() {
    IndexService index = createIndex("test", Settings.EMPTY, "type", "d", "type=double");
    Map<ScriptContext<?>, List<Whitelist>> contexts = new HashMap<>();
    contexts.put(NumberSortScript.CONTEXT, Whitelist.BASE_WHITELISTS);
    PainlessScriptEngine service = new PainlessScriptEngine(Settings.EMPTY, contexts);
    QueryShardContext shardContext = index.newQueryShardContext(0, null, () -> 0, null);
    NumberSortScript.Factory factory = service.compile(null, "1.2", NumberSortScript.CONTEXT, Collections.emptyMap());
    NumberSortScript.LeafFactory ss = factory.newFactory(Collections.emptyMap(), shardContext.lookup());
    assertFalse(ss.needs_score());
    factory = service.compile(null, "doc['d'].value", NumberSortScript.CONTEXT, Collections.emptyMap());
    ss = factory.newFactory(Collections.emptyMap(), shardContext.lookup());
    assertFalse(ss.needs_score());
    factory = service.compile(null, "1/_score", NumberSortScript.CONTEXT, Collections.emptyMap());
    ss = factory.newFactory(Collections.emptyMap(), shardContext.lookup());
    assertTrue(ss.needs_score());
    factory = service.compile(null, "doc['d'].value * _score", NumberSortScript.CONTEXT, Collections.emptyMap());
    ss = factory.newFactory(Collections.emptyMap(), shardContext.lookup());
    assertTrue(ss.needs_score());
}

Example 2 with QueryShardContext

use of org.opensearch.index.query.QueryShardContext in project OpenSearch by opensearch-project.

the class CandidateQueryTests method testDuel2.

public void testDuel2() throws Exception {
    List<String> stringValues = new ArrayList<>();
    stringValues.add("value1");
    stringValues.add("value2");
    stringValues.add("value3");
    MappedFieldType intFieldType = mapperService.fieldType("int_field");
    List<int[]> ranges = new ArrayList<>();
    ranges.add(new int[] { -5, 5 });
    ranges.add(new int[] { 0, 10 });
    ranges.add(new int[] { 15, 50 });
    QueryShardContext context = createSearchContext(indexService).getQueryShardContext();
    List<ParseContext.Document> documents = new ArrayList<>();
    {
        addQuery(new TermQuery(new Term("string_field", randomFrom(stringValues))), documents);
    }
    {
        addQuery(new PhraseQuery(0, "string_field", stringValues.toArray(new String[0])), documents);
    }
    {
        int[] range = randomFrom(ranges);
        Query rangeQuery = intFieldType.rangeQuery(range[0], range[1], true, true, null, null, null, context);
        addQuery(rangeQuery, documents);
    }
    {
        int numBooleanQueries = randomIntBetween(1, 5);
        for (int i = 0; i < numBooleanQueries; i++) {
            Query randomBQ = randomBQ(1, stringValues, ranges, intFieldType, context);
            addQuery(randomBQ, documents);
        }
    }
    {
        addQuery(new MatchNoDocsQuery(), documents);
    }
    {
        addQuery(new MatchAllDocsQuery(), documents);
    }
    indexWriter.addDocuments(documents);
    indexWriter.close();
    directoryReader = DirectoryReader.open(directory);
    IndexSearcher shardSearcher = newSearcher(directoryReader);
    // Disable query cache, because ControlQuery cannot be cached...
    shardSearcher.setQueryCache(null);
    Document document = new Document();
    for (String value : stringValues) {
        document.add(new TextField("string_field", value, Field.Store.NO));
        logger.info("Test with document: {}" + document);
        MemoryIndex memoryIndex = MemoryIndex.fromDocument(document, new WhitespaceAnalyzer());
        duelRun(queryStore, memoryIndex, shardSearcher);
    }
    for (int[] range : ranges) {
        List<Field> numberFields = NumberFieldMapper.NumberType.INTEGER.createFields("int_field", between(range[0], range[1]), true, true, false);
        for (Field numberField : numberFields) {
            document.add(numberField);
        }
        logger.info("Test with document: {}" + document);
        MemoryIndex memoryIndex = MemoryIndex.fromDocument(document, new WhitespaceAnalyzer());
        duelRun(queryStore, memoryIndex, shardSearcher);
    }
}

Example 3 with QueryShardContext

use of org.opensearch.index.query.QueryShardContext in project OpenSearch by opensearch-project.

the class PercolatorFieldMapperTests method testQueryWithRewrite.

public void testQueryWithRewrite() throws Exception {
    addQueryFieldMappings();
    client().prepareIndex("remote").setId("1").setSource("field", "value").get();
    QueryBuilder queryBuilder = termsLookupQuery("field", new TermsLookup("remote", "1", "field"));
    ParsedDocument doc = mapperService.documentMapper("doc").parse(new SourceToParse("test", "doc", "1", BytesReference.bytes(XContentFactory.jsonBuilder().startObject().field(fieldName, queryBuilder).endObject()), XContentType.JSON));
    BytesRef qbSource = doc.rootDoc().getFields(fieldType.queryBuilderField.name())[0].binaryValue();
    QueryShardContext shardContext = indexService.newQueryShardContext(randomInt(20), null, () -> {
        throw new UnsupportedOperationException();
    }, null);
    PlainActionFuture<QueryBuilder> future = new PlainActionFuture<>();
    Rewriteable.rewriteAndFetch(queryBuilder, shardContext, future);
    assertQueryBuilder(qbSource, future.get());
}

Example 4 with QueryShardContext

use of org.opensearch.index.query.QueryShardContext in project OpenSearch by opensearch-project.

the class QueryBuilderStoreTests method testStoringQueryBuilders.

public void testStoringQueryBuilders() throws IOException {
    try (Directory directory = newDirectory()) {
        TermQueryBuilder[] queryBuilders = new TermQueryBuilder[randomIntBetween(1, 16)];
        IndexWriterConfig config = new IndexWriterConfig(new WhitespaceAnalyzer());
        config.setMergePolicy(NoMergePolicy.INSTANCE);
        Settings settings = Settings.builder().put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT).build();
        BinaryFieldMapper fieldMapper = PercolatorFieldMapper.Builder.createQueryBuilderFieldBuilder(new Mapper.BuilderContext(settings, new ContentPath(0)));
        Version version = Version.CURRENT;
        try (IndexWriter indexWriter = new IndexWriter(directory, config)) {
            for (int i = 0; i < queryBuilders.length; i++) {
                queryBuilders[i] = new TermQueryBuilder(randomAlphaOfLength(4), randomAlphaOfLength(8));
                ParseContext parseContext = mock(ParseContext.class);
                ParseContext.Document document = new ParseContext.Document();
                when(parseContext.doc()).thenReturn(document);
                PercolatorFieldMapper.createQueryBuilderField(version, fieldMapper, queryBuilders[i], parseContext);
                indexWriter.addDocument(document);
            }
        }
        QueryShardContext queryShardContext = mock(QueryShardContext.class);
        when(queryShardContext.indexVersionCreated()).thenReturn(version);
        when(queryShardContext.getWriteableRegistry()).thenReturn(writableRegistry());
        when(queryShardContext.getXContentRegistry()).thenReturn(xContentRegistry());
        when(queryShardContext.getForField(fieldMapper.fieldType())).thenReturn(new BytesBinaryIndexFieldData(fieldMapper.name(), CoreValuesSourceType.BYTES));
        when(queryShardContext.fieldMapper(Mockito.anyString())).thenAnswer(invocation -> {
            final String fieldName = (String) invocation.getArguments()[0];
            return new KeywordFieldMapper.KeywordFieldType(fieldName);
        });
        PercolateQuery.QueryStore queryStore = PercolateQueryBuilder.createStore(fieldMapper.fieldType(), queryShardContext);
        try (IndexReader indexReader = DirectoryReader.open(directory)) {
            LeafReaderContext leafContext = indexReader.leaves().get(0);
            CheckedFunction<Integer, Query, IOException> queries = queryStore.getQueries(leafContext);
            assertEquals(queryBuilders.length, leafContext.reader().numDocs());
            for (int i = 0; i < queryBuilders.length; i++) {
                TermQuery query = (TermQuery) queries.apply(i);
                assertEquals(queryBuilders[i].fieldName(), query.getTerm().field());
                assertEquals(queryBuilders[i].value(), query.getTerm().text());
            }
        }
    }
}

Example 5 with QueryShardContext

use of org.opensearch.index.query.QueryShardContext in project OpenSearch by opensearch-project.

the class PercolateQueryBuilderTests method testFromJsonWithType.

public void testFromJsonWithType() throws IOException {
    indexedDocumentIndex = randomAlphaOfLength(4);
    indexedDocumentId = randomAlphaOfLength(4);
    indexedDocumentVersion = Versions.MATCH_ANY;
    documentSource = Collections.singletonList(randomSource(new HashSet<>()));
    QueryShardContext queryShardContext = createShardContext();
    QueryBuilder queryBuilder = parseQuery("{\"percolate\" : { \"index\": \"" + indexedDocumentIndex + "\", \"type\": \"_doc\", \"id\": \"" + indexedDocumentId + "\", \"field\":\"" + queryField + "\"}}");
    rewriteAndFetch(queryBuilder, queryShardContext).toQuery(queryShardContext);
}