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

Example 26 with QueryShardContext

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

the class FunctionScoreQueryBuilderTests method testSingleScriptFunction.

/**
 * Please see https://github.com/elastic/elasticsearch/issues/35123 for context.
 */
public void testSingleScriptFunction() throws IOException {
    QueryBuilder queryBuilder = termQuery(KEYWORD_FIELD_NAME, "value");
    ScoreFunctionBuilder<ScriptScoreFunctionBuilder> functionBuilder = new ScriptScoreFunctionBuilder(new Script(ScriptType.INLINE, MockScriptEngine.NAME, "1", Collections.emptyMap()));
    FunctionScoreQueryBuilder builder = functionScoreQuery(queryBuilder, functionBuilder);
    if (randomBoolean()) {
        builder.boostMode(randomFrom(CombineFunction.values()));
    }
    QueryShardContext shardContext = createShardContext();
    Query query = builder.rewrite(shardContext).toQuery(shardContext);
    assertThat(query, instanceOf(FunctionScoreQuery.class));
    CombineFunction expectedBoostMode = builder.boostMode() != null ? builder.boostMode() : FunctionScoreQueryBuilder.DEFAULT_BOOST_MODE;
    CombineFunction actualBoostMode = ((FunctionScoreQuery) query).getCombineFunction();
    assertEquals(expectedBoostMode, actualBoostMode);
}

Example 27 with QueryShardContext

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

the class FunctionScoreQueryBuilderTests method testMustRewrite.

@Override
public void testMustRewrite() throws IOException {
    QueryShardContext context = createShardContext();
    context.setAllowUnmappedFields(true);
    TermQueryBuilder termQueryBuilder = new TermQueryBuilder("unmapped_field", "foo");
    // main query needs rewriting
    FunctionScoreQueryBuilder functionQueryBuilder1 = new FunctionScoreQueryBuilder(termQueryBuilder);
    functionQueryBuilder1.setMinScore(1);
    IllegalStateException e = expectThrows(IllegalStateException.class, () -> functionQueryBuilder1.toQuery(context));
    assertEquals("Rewrite first", e.getMessage());
    // filter needs rewriting
    FunctionScoreQueryBuilder functionQueryBuilder2 = new FunctionScoreQueryBuilder(new MatchAllQueryBuilder(), new FilterFunctionBuilder[] { new FilterFunctionBuilder(termQueryBuilder, new RandomScoreFunctionBuilder()) });
    e = expectThrows(IllegalStateException.class, () -> functionQueryBuilder2.toQuery(context));
    assertEquals("Rewrite first", e.getMessage());
}

Example 28 with QueryShardContext

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

the class FunctionScoreQueryBuilderTests method testCacheability.

/**
 * Check that this query is generally cacheable except for builders using {@link ScriptScoreFunctionBuilder} or
 * {@link RandomScoreFunctionBuilder} without a seed
 */
@Override
public void testCacheability() throws IOException {
    FunctionScoreQueryBuilder queryBuilder = createTestQueryBuilder();
    boolean isCacheable = isCacheable(queryBuilder);
    QueryShardContext context = createShardContext();
    QueryBuilder rewriteQuery = rewriteQuery(queryBuilder, new QueryShardContext(context));
    assertNotNull(rewriteQuery.toQuery(context));
    // we occasionally need to update the expected "isCacheable" flag after rewrite for MatchNoneQueryBuilder
    if (rewriteQuery instanceof MatchNoneQueryBuilder) {
        isCacheable = true;
    }
    assertEquals("query should " + (isCacheable ? "" : "not") + " be cacheable: " + queryBuilder.toString(), isCacheable, context.isCacheable());
    ScoreFunctionBuilder<?> scriptScoreFunction = new ScriptScoreFunctionBuilder(new Script(ScriptType.INLINE, MockScriptEngine.NAME, "1", Collections.emptyMap()));
    queryBuilder = new FunctionScoreQueryBuilder(new FilterFunctionBuilder[] { new FilterFunctionBuilder(RandomQueryBuilder.createQuery(random()), scriptScoreFunction) });
    context = createShardContext();
    rewriteQuery = rewriteQuery(queryBuilder, new QueryShardContext(context));
    assertNotNull(rewriteQuery.toQuery(context));
    assertTrue("function script query should be cacheable" + queryBuilder.toString(), context.isCacheable());
    RandomScoreFunctionBuilder randomScoreFunctionBuilder = new RandomScoreFunctionBuilderWithFixedSeed();
    queryBuilder = new FunctionScoreQueryBuilder(new FilterFunctionBuilder[] { new FilterFunctionBuilder(RandomQueryBuilder.createQuery(random()), randomScoreFunctionBuilder) });
    context = createShardContext();
    rewriteQuery = rewriteQuery(queryBuilder, new QueryShardContext(context));
    assertNotNull(rewriteQuery.toQuery(context));
    assertFalse("function random query should not be cacheable: " + queryBuilder.toString(), context.isCacheable());
}

Example 29 with QueryShardContext

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

the class ScoreFunctionBuilderTests method testRandomScoreFunctionWithSeed.

public void testRandomScoreFunctionWithSeed() throws Exception {
    RandomScoreFunctionBuilder builder = new RandomScoreFunctionBuilder();
    builder.seed(42);
    QueryShardContext context = Mockito.mock(QueryShardContext.class);
    Settings indexSettings = Settings.builder().put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT).put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1).put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1).build();
    IndexSettings settings = new IndexSettings(IndexMetadata.builder("index").settings(indexSettings).build(), Settings.EMPTY);
    Mockito.when(context.index()).thenReturn(settings.getIndex());
    Mockito.when(context.getShardId()).thenReturn(0);
    Mockito.when(context.getIndexSettings()).thenReturn(settings);
    MapperService mapperService = Mockito.mock(MapperService.class);
    MappedFieldType ft = new NumberFieldMapper.NumberFieldType("foo", NumberType.LONG);
    Mockito.when(mapperService.fieldType(Mockito.anyString())).thenReturn(ft);
    Mockito.when(context.getMapperService()).thenReturn(mapperService);
    builder.toFunction(context);
    assertWarnings("OpenSearch requires that a [field] parameter is provided when a [seed] is set");
}

Example 30 with QueryShardContext

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

the class MultiMatchQueryTests method testCrossFieldMultiMatchQuery.

public void testCrossFieldMultiMatchQuery() throws IOException {
    QueryShardContext queryShardContext = indexService.newQueryShardContext(randomInt(20), null, () -> {
        throw new UnsupportedOperationException();
    }, null);
    queryShardContext.setAllowUnmappedFields(true);
    for (float tieBreaker : new float[] { 0.0f, 0.5f }) {
        Query parsedQuery = multiMatchQuery("banon").field("name.first", 2).field("name.last", 3).field("foobar").type(MultiMatchQueryBuilder.Type.CROSS_FIELDS).tieBreaker(tieBreaker).toQuery(queryShardContext);
        try (Engine.Searcher searcher = indexService.getShard(0).acquireSearcher("test")) {
            Query rewrittenQuery = searcher.rewrite(parsedQuery);
            Query tq1 = new BoostQuery(new TermQuery(new Term("name.last", "banon")), 3);
            Query tq2 = new BoostQuery(new TermQuery(new Term("name.first", "banon")), 2);
            Query expected = new DisjunctionMaxQuery(Arrays.asList(tq2, tq1), tieBreaker);
            assertEquals(expected, rewrittenQuery);
        }
    }
}