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

Example 91 with QueryShardContext

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

the class ValuesSourceConfigTests method testLong.

public void testLong() throws Exception {
    IndexService indexService = createIndex("index", Settings.EMPTY, "type", "long", "type=long");
    client().prepareIndex("index").setId("1").setSource("long", 42).setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE).get();
    try (Engine.Searcher searcher = indexService.getShard(0).acquireSearcher("test")) {
        QueryShardContext context = indexService.newQueryShardContext(0, searcher, () -> 42L, null);
        ValuesSourceConfig config = ValuesSourceConfig.resolve(context, null, "long", null, null, null, null, CoreValuesSourceType.BYTES);
        ValuesSource.Numeric valuesSource = (ValuesSource.Numeric) config.getValuesSource();
        LeafReaderContext ctx = searcher.getIndexReader().leaves().get(0);
        SortedNumericDocValues values = valuesSource.longValues(ctx);
        assertTrue(values.advanceExact(0));
        assertEquals(1, values.docValueCount());
        assertEquals(42, values.nextValue());
    }
}

Example 92 with QueryShardContext

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

the class ValuesSourceConfigTests method testKeyword.

public void testKeyword() throws Exception {
    IndexService indexService = createIndex("index", Settings.EMPTY, "type", "bytes", "type=keyword");
    client().prepareIndex("index").setId("1").setSource("bytes", "abc").setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE).get();
    try (Engine.Searcher searcher = indexService.getShard(0).acquireSearcher("test")) {
        QueryShardContext context = indexService.newQueryShardContext(0, searcher, () -> 42L, null);
        ValuesSourceConfig config = ValuesSourceConfig.resolve(context, null, "bytes", null, null, null, null, CoreValuesSourceType.BYTES);
        ValuesSource.Bytes valuesSource = (ValuesSource.Bytes) config.getValuesSource();
        LeafReaderContext ctx = searcher.getIndexReader().leaves().get(0);
        SortedBinaryDocValues values = valuesSource.bytesValues(ctx);
        assertTrue(values.advanceExact(0));
        assertEquals(1, values.docValueCount());
        assertEquals(new BytesRef("abc"), values.nextValue());
    }
}

Example 93 with QueryShardContext

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

the class QueryRescorerBuilderTests method testBuildRescoreSearchContext.

/**
 * test that build() outputs a {@link RescoreContext} that has the same properties
 * than the test builder
 */
public void testBuildRescoreSearchContext() throws OpenSearchParseException, IOException {
    final long nowInMillis = randomNonNegativeLong();
    Settings indexSettings = Settings.builder().put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT).build();
    IndexSettings idxSettings = IndexSettingsModule.newIndexSettings(randomAlphaOfLengthBetween(1, 10), indexSettings);
    // shard context will only need indicesQueriesRegistry for building Query objects nested in query rescorer
    QueryShardContext mockShardContext = new QueryShardContext(0, idxSettings, BigArrays.NON_RECYCLING_INSTANCE, null, null, null, null, null, xContentRegistry(), namedWriteableRegistry, null, null, () -> nowInMillis, null, null, () -> true, null) {

        @Override
        public MappedFieldType fieldMapper(String name) {
            TextFieldMapper.Builder builder = new TextFieldMapper.Builder(name, createDefaultIndexAnalyzers());
            return builder.build(new Mapper.BuilderContext(idxSettings.getSettings(), new ContentPath(1))).fieldType();
        }
    };
    for (int runs = 0; runs < NUMBER_OF_TESTBUILDERS; runs++) {
        QueryRescorerBuilder rescoreBuilder = randomRescoreBuilder();
        QueryRescoreContext rescoreContext = (QueryRescoreContext) rescoreBuilder.buildContext(mockShardContext);
        int expectedWindowSize = rescoreBuilder.windowSize() == null ? RescorerBuilder.DEFAULT_WINDOW_SIZE : rescoreBuilder.windowSize().intValue();
        assertEquals(expectedWindowSize, rescoreContext.getWindowSize());
        Query expectedQuery = Rewriteable.rewrite(rescoreBuilder.getRescoreQuery(), mockShardContext).toQuery(mockShardContext);
        assertEquals(expectedQuery, rescoreContext.query());
        assertEquals(rescoreBuilder.getQueryWeight(), rescoreContext.queryWeight(), Float.MIN_VALUE);
        assertEquals(rescoreBuilder.getRescoreQueryWeight(), rescoreContext.rescoreQueryWeight(), Float.MIN_VALUE);
        assertEquals(rescoreBuilder.getScoreMode(), rescoreContext.scoreMode());
    }
}

Example 94 with QueryShardContext

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

the class SliceBuilderTests method testToFilterRandom.

public void testToFilterRandom() throws IOException {
    Directory dir = new ByteBuffersDirectory();
    try (IndexWriter writer = new IndexWriter(dir, newIndexWriterConfig(new MockAnalyzer(random())))) {
        writer.commit();
    }
    try (IndexReader reader = DirectoryReader.open(dir)) {
        QueryShardContext context = createShardContext(Version.CURRENT, reader, "field", DocValuesType.SORTED_NUMERIC, 1, 0);
        SliceBuilder builder = new SliceBuilder("field", 5, 10);
        Query query = builder.toFilter(null, createRequest(0), context, Version.CURRENT);
        assertThat(query, instanceOf(DocValuesSliceQuery.class));
        assertThat(builder.toFilter(null, createRequest(0), context, Version.CURRENT), equalTo(query));
        try (IndexReader newReader = DirectoryReader.open(dir)) {
            when(context.getIndexReader()).thenReturn(newReader);
            assertThat(builder.toFilter(null, createRequest(0), context, Version.CURRENT), equalTo(query));
        }
        // numSlices > numShards
        int numSlices = randomIntBetween(10, 100);
        int numShards = randomIntBetween(1, 9);
        Map<Integer, AtomicInteger> numSliceMap = new HashMap<>();
        for (int i = 0; i < numSlices; i++) {
            for (int j = 0; j < numShards; j++) {
                SliceBuilder slice = new SliceBuilder("_id", i, numSlices);
                context = createShardContext(Version.CURRENT, reader, "_id", DocValuesType.SORTED, numShards, j);
                Query q = slice.toFilter(null, createRequest(j), context, Version.CURRENT);
                if (q instanceof TermsSliceQuery || q instanceof MatchAllDocsQuery) {
                    AtomicInteger count = numSliceMap.get(j);
                    if (count == null) {
                        count = new AtomicInteger(0);
                        numSliceMap.put(j, count);
                    }
                    count.incrementAndGet();
                    if (q instanceof MatchAllDocsQuery) {
                        assertThat(count.get(), equalTo(1));
                    }
                } else {
                    assertThat(q, instanceOf(MatchNoDocsQuery.class));
                }
            }
        }
        int total = 0;
        for (Map.Entry<Integer, AtomicInteger> e : numSliceMap.entrySet()) {
            total += e.getValue().get();
        }
        assertThat(total, equalTo(numSlices));
        // numShards > numSlices
        numShards = randomIntBetween(4, 100);
        numSlices = randomIntBetween(2, numShards - 1);
        List<Integer> targetShards = new ArrayList<>();
        for (int i = 0; i < numSlices; i++) {
            for (int j = 0; j < numShards; j++) {
                SliceBuilder slice = new SliceBuilder("_id", i, numSlices);
                context = createShardContext(Version.CURRENT, reader, "_id", DocValuesType.SORTED, numShards, j);
                Query q = slice.toFilter(null, createRequest(j), context, Version.CURRENT);
                if (q instanceof MatchNoDocsQuery == false) {
                    assertThat(q, instanceOf(MatchAllDocsQuery.class));
                    targetShards.add(j);
                }
            }
        }
        assertThat(targetShards.size(), equalTo(numShards));
        assertThat(new HashSet<>(targetShards).size(), equalTo(numShards));
        // numShards == numSlices
        numShards = randomIntBetween(2, 10);
        numSlices = numShards;
        for (int i = 0; i < numSlices; i++) {
            for (int j = 0; j < numShards; j++) {
                SliceBuilder slice = new SliceBuilder("_id", i, numSlices);
                context = createShardContext(Version.CURRENT, reader, "_id", DocValuesType.SORTED, numShards, j);
                Query q = slice.toFilter(null, createRequest(j), context, Version.CURRENT);
                if (i == j) {
                    assertThat(q, instanceOf(MatchAllDocsQuery.class));
                } else {
                    assertThat(q, instanceOf(MatchNoDocsQuery.class));
                }
            }
        }
    }
}

Example 95 with QueryShardContext

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

the class SliceBuilderTests method testToFilterSimple.

public void testToFilterSimple() throws IOException {
    Directory dir = new ByteBuffersDirectory();
    try (IndexWriter writer = new IndexWriter(dir, newIndexWriterConfig(new MockAnalyzer(random())))) {
        writer.commit();
    }
    try (IndexReader reader = DirectoryReader.open(dir)) {
        QueryShardContext context = createShardContext(Version.CURRENT, reader, "_id", DocValuesType.SORTED_NUMERIC, 1, 0);
        SliceBuilder builder = new SliceBuilder(5, 10);
        Query query = builder.toFilter(null, createRequest(0), context, Version.CURRENT);
        assertThat(query, instanceOf(TermsSliceQuery.class));
        assertThat(builder.toFilter(null, createRequest(0), context, Version.CURRENT), equalTo(query));
        try (IndexReader newReader = DirectoryReader.open(dir)) {
            when(context.getIndexReader()).thenReturn(newReader);
            assertThat(builder.toFilter(null, createRequest(0), context, Version.CURRENT), equalTo(query));
        }
    }
}