Examples with SliceBuilder org.opensearch.search.slice.SliceBuilder used on opensource projects

Example 1 with SliceBuilder

use of org.opensearch.search.slice.SliceBuilder in project OpenSearch by opensearch-project.

the class DefaultSearchContext method buildFilteredQuery.

@Override
public Query buildFilteredQuery(Query query) {
    List<Query> filters = new ArrayList<>();
    if (mapperService().hasNested() && new NestedHelper(mapperService()).mightMatchNestedDocs(query) && (aliasFilter == null || new NestedHelper(mapperService()).mightMatchNestedDocs(aliasFilter))) {
        filters.add(Queries.newNonNestedFilter(mapperService().getIndexSettings().getIndexVersionCreated()));
    }
    if (aliasFilter != null) {
        filters.add(aliasFilter);
    }
    if (sliceBuilder != null) {
        Query slicedQuery = sliceBuilder.toFilter(clusterService, request, queryShardContext, minNodeVersion);
        if (slicedQuery instanceof MatchNoDocsQuery) {
            return slicedQuery;
        } else {
            filters.add(slicedQuery);
        }
    }
    if (filters.isEmpty()) {
        return query;
    } else {
        BooleanQuery.Builder builder = new BooleanQuery.Builder();
        builder.add(query, Occur.MUST);
        for (Query filter : filters) {
            builder.add(filter, Occur.FILTER);
        }
        return builder.build();
    }
}

Example 2 with SliceBuilder

use of org.opensearch.search.slice.SliceBuilder in project OpenSearch by opensearch-project.

the class BulkByScrollParallelizationHelper method sliceIntoSubRequests.

/**
 * Slice a search request into {@code times} separate search requests slicing on {@code field}. Note that the slices are *shallow*
 * copies of this request so don't change them.
 */
static SearchRequest[] sliceIntoSubRequests(SearchRequest request, String field, int times) {
    SearchRequest[] slices = new SearchRequest[times];
    for (int slice = 0; slice < times; slice++) {
        SliceBuilder sliceBuilder = new SliceBuilder(field, slice, times);
        SearchSourceBuilder slicedSource;
        if (request.source() == null) {
            slicedSource = new SearchSourceBuilder().slice(sliceBuilder);
        } else {
            if (request.source().slice() != null) {
                throw new IllegalStateException("Can't slice a request that already has a slice configuration");
            }
            slicedSource = request.source().shallowCopy().slice(sliceBuilder);
        }
        SearchRequest searchRequest = new SearchRequest(request);
        searchRequest.source(slicedSource);
        slices[slice] = searchRequest;
    }
    return slices;
}

Example 3 with SliceBuilder

use of org.opensearch.search.slice.SliceBuilder in project OpenSearch by opensearch-project.

the class BulkByScrollParallelizationHelper method setWorkerCount.

private static <Request extends AbstractBulkByScrollRequest<Request>> void setWorkerCount(Request request, BulkByScrollTask task, int slices) {
    if (slices > 1) {
        task.setWorkerCount(slices);
    } else {
        SliceBuilder sliceBuilder = request.getSearchRequest().source().slice();
        Integer sliceId = sliceBuilder == null ? null : sliceBuilder.getId();
        task.setWorker(request.getRequestsPerSecond(), sliceId);
    }
}

Example 4 with SliceBuilder

use of org.opensearch.search.slice.SliceBuilder in project OpenSearch by opensearch-project.

the class RandomSearchRequestGenerator method randomSearchSourceBuilder.

public static SearchSourceBuilder randomSearchSourceBuilder(Supplier<HighlightBuilder> randomHighlightBuilder, Supplier<SuggestBuilder> randomSuggestBuilder, Supplier<RescorerBuilder<?>> randomRescoreBuilder, Supplier<List<SearchExtBuilder>> randomExtBuilders, Supplier<CollapseBuilder> randomCollapseBuilder) {
    SearchSourceBuilder builder = new SearchSourceBuilder();
    if (randomBoolean()) {
        builder.from(randomIntBetween(0, 10000));
    }
    if (randomBoolean()) {
        builder.size(randomIntBetween(0, 10000));
    }
    if (randomBoolean()) {
        builder.explain(randomBoolean());
    }
    if (randomBoolean()) {
        builder.version(randomBoolean());
    }
    if (randomBoolean()) {
        builder.seqNoAndPrimaryTerm(randomBoolean());
    }
    if (randomBoolean()) {
        builder.trackScores(randomBoolean());
    }
    if (randomBoolean()) {
        builder.minScore(randomFloat() * 1000);
    }
    if (randomBoolean()) {
        builder.timeout(TimeValue.parseTimeValue(randomTimeValue(), null, "timeout"));
    }
    if (randomBoolean()) {
        builder.terminateAfter(randomIntBetween(1, 100000));
    }
    if (randomBoolean()) {
        if (randomBoolean()) {
            builder.trackTotalHits(randomBoolean());
        } else {
            builder.trackTotalHitsUpTo(randomIntBetween(SearchContext.TRACK_TOTAL_HITS_DISABLED, SearchContext.TRACK_TOTAL_HITS_ACCURATE));
        }
    }
    switch(randomInt(2)) {
        case 0:
            builder.storedFields();
            break;
        case 1:
            builder.storedField("_none_");
            break;
        case 2:
            int fieldsSize = randomInt(25);
            List<String> fields = new ArrayList<>(fieldsSize);
            for (int i = 0; i < fieldsSize; i++) {
                fields.add(randomAlphaOfLengthBetween(5, 50));
            }
            builder.storedFields(fields);
            break;
        default:
            throw new IllegalStateException();
    }
    if (randomBoolean()) {
        int numFields = randomInt(5);
        for (int i = 0; i < numFields; i++) {
            builder.fetchField(randomAlphaOfLengthBetween(5, 10));
        }
    }
    if (randomBoolean()) {
        int scriptFieldsSize = randomInt(25);
        for (int i = 0; i < scriptFieldsSize; i++) {
            if (randomBoolean()) {
                builder.scriptField(randomAlphaOfLengthBetween(5, 50), mockScript("foo"), randomBoolean());
            } else {
                builder.scriptField(randomAlphaOfLengthBetween(5, 50), mockScript("foo"));
            }
        }
    }
    if (randomBoolean()) {
        FetchSourceContext fetchSourceContext;
        int branch = randomInt(5);
        String[] includes = new String[randomIntBetween(0, 20)];
        for (int i = 0; i < includes.length; i++) {
            includes[i] = randomAlphaOfLengthBetween(5, 20);
        }
        String[] excludes = new String[randomIntBetween(0, 20)];
        for (int i = 0; i < excludes.length; i++) {
            excludes[i] = randomAlphaOfLengthBetween(5, 20);
        }
        switch(branch) {
            case 0:
                fetchSourceContext = new FetchSourceContext(randomBoolean());
                break;
            case 1:
                fetchSourceContext = new FetchSourceContext(true, includes, excludes);
                break;
            case 2:
                fetchSourceContext = new FetchSourceContext(true, new String[] { randomAlphaOfLengthBetween(5, 20) }, new String[] { randomAlphaOfLengthBetween(5, 20) });
                break;
            case 3:
                fetchSourceContext = new FetchSourceContext(true, includes, excludes);
                break;
            case 4:
                fetchSourceContext = new FetchSourceContext(true, includes, null);
                break;
            case 5:
                fetchSourceContext = new FetchSourceContext(true, new String[] { randomAlphaOfLengthBetween(5, 20) }, null);
                break;
            default:
                throw new IllegalStateException();
        }
        builder.fetchSource(fetchSourceContext);
    }
    if (randomBoolean()) {
        int size = randomIntBetween(0, 20);
        List<String> statsGroups = new ArrayList<>(size);
        for (int i = 0; i < size; i++) {
            statsGroups.add(randomAlphaOfLengthBetween(5, 20));
        }
        builder.stats(statsGroups);
    }
    if (randomBoolean()) {
        int indexBoostSize = randomIntBetween(1, 10);
        for (int i = 0; i < indexBoostSize; i++) {
            builder.indexBoost(randomAlphaOfLengthBetween(5, 20), randomFloat() * 10);
        }
    }
    if (randomBoolean()) {
        builder.query(QueryBuilders.termQuery(randomAlphaOfLengthBetween(5, 20), randomAlphaOfLengthBetween(5, 20)));
    }
    if (randomBoolean()) {
        builder.postFilter(QueryBuilders.termQuery(randomAlphaOfLengthBetween(5, 20), randomAlphaOfLengthBetween(5, 20)));
    }
    if (randomBoolean()) {
        int numSorts = randomIntBetween(1, 5);
        for (int i = 0; i < numSorts; i++) {
            int branch = randomInt(5);
            switch(branch) {
                case 0:
                    builder.sort(SortBuilders.fieldSort(randomAlphaOfLengthBetween(5, 20)).order(randomFrom(SortOrder.values())));
                    break;
                case 1:
                    builder.sort(SortBuilders.geoDistanceSort(randomAlphaOfLengthBetween(5, 20), AbstractQueryTestCase.randomGeohash(1, 12)).order(randomFrom(SortOrder.values())));
                    break;
                case 2:
                    builder.sort(SortBuilders.scoreSort().order(randomFrom(SortOrder.values())));
                    break;
                case 3:
                    builder.sort(SortBuilders.scriptSort(new Script(ScriptType.INLINE, Script.DEFAULT_SCRIPT_LANG, "foo", emptyMap()), ScriptSortBuilder.ScriptSortType.NUMBER).order(randomFrom(SortOrder.values())));
                    break;
                case 4:
                    builder.sort(randomAlphaOfLengthBetween(5, 20));
                    break;
                case 5:
                    builder.sort(randomAlphaOfLengthBetween(5, 20), randomFrom(SortOrder.values()));
                    break;
            }
        }
    }
    if (randomBoolean()) {
        int numSearchFrom = randomIntBetween(1, 5);
        try {
            // We build a json version of the search_from first in order to
            // ensure that every number type remain the same before/after xcontent (de)serialization.
            // This is not a problem because the final type of each field value is extracted from associated sort field.
            // This little trick ensure that equals and hashcode are the same when using the xcontent serialization.
            XContentBuilder jsonBuilder = XContentFactory.jsonBuilder();
            jsonBuilder.startObject();
            jsonBuilder.startArray("search_from");
            for (int i = 0; i < numSearchFrom; i++) {
                int branch = randomInt(8);
                switch(branch) {
                    case 0:
                        jsonBuilder.value(randomInt());
                        break;
                    case 1:
                        jsonBuilder.value(randomFloat());
                        break;
                    case 2:
                        jsonBuilder.value(randomLong());
                        break;
                    case 3:
                        jsonBuilder.value(randomDouble());
                        break;
                    case 4:
                        jsonBuilder.value(randomAlphaOfLengthBetween(5, 20));
                        break;
                    case 5:
                        jsonBuilder.value(randomBoolean());
                        break;
                    case 6:
                        jsonBuilder.value(randomByte());
                        break;
                    case 7:
                        jsonBuilder.value(randomShort());
                        break;
                    case 8:
                        jsonBuilder.value(new Text(randomAlphaOfLengthBetween(5, 20)));
                        break;
                }
            }
            jsonBuilder.endArray();
            jsonBuilder.endObject();
            XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(NamedXContentRegistry.EMPTY, DeprecationHandler.THROW_UNSUPPORTED_OPERATION, BytesReference.bytes(jsonBuilder).streamInput());
            parser.nextToken();
            parser.nextToken();
            parser.nextToken();
            builder.searchAfter(SearchAfterBuilder.fromXContent(parser).getSortValues());
        } catch (IOException e) {
            throw new RuntimeException("Error building search_from", e);
        }
    }
    if (randomBoolean()) {
        builder.highlighter(randomHighlightBuilder.get());
    }
    if (randomBoolean()) {
        builder.suggest(randomSuggestBuilder.get());
    }
    if (randomBoolean()) {
        int numRescores = randomIntBetween(1, 5);
        for (int i = 0; i < numRescores; i++) {
            builder.addRescorer(randomRescoreBuilder.get());
        }
    }
    if (randomBoolean()) {
        builder.aggregation(AggregationBuilders.avg(randomAlphaOfLengthBetween(5, 20)).field("foo"));
    }
    if (randomBoolean()) {
        builder.ext(randomExtBuilders.get());
    }
    if (randomBoolean()) {
        String field = randomBoolean() ? null : randomAlphaOfLengthBetween(5, 20);
        int max = between(2, 1000);
        int id = randomInt(max - 1);
        if (field == null) {
            builder.slice(new SliceBuilder(id, max));
        } else {
            builder.slice(new SliceBuilder(field, id, max));
        }
    }
    if (randomBoolean()) {
        builder.collapse(randomCollapseBuilder.get());
    }
    if (randomBoolean()) {
        PointInTimeBuilder pit = new PointInTimeBuilder(randomAlphaOfLengthBetween(3, 10));
        if (randomBoolean()) {
            pit.setKeepAlive(TimeValue.timeValueMinutes(randomIntBetween(1, 60)));
        }
        builder.pointInTimeBuilder(pit);
    }
    return builder;
}

Example 5 with SliceBuilder

use of org.opensearch.search.slice.SliceBuilder in project OpenSearch by opensearch-project.

the class ReindexRequestTests method testNoSliceBuilderSetWithSlicedRequest.

public void testNoSliceBuilderSetWithSlicedRequest() {
    ReindexRequest reindex = newRequest();
    reindex.getSearchRequest().source().slice(new SliceBuilder(0, 4));
    reindex.setSlices(between(2, Integer.MAX_VALUE));
    ActionRequestValidationException e = reindex.validate();
    assertEquals("Validation Failed: 1: can't specify both manual and automatic slicing at the same time;", e.getMessage());
}